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/actions/test-llamacpp-build/action.yml b/.github/actions/test-llamacpp-build/action.yml new file mode 100644 index 000000000000..4602cb30da5f --- /dev/null +++ b/.github/actions/test-llamacpp-build/action.yml @@ -0,0 +1,238 @@ +# Same as lemonade-sdk/llamacpp-rocm (MIT). +name: 'Test Llama.cpp Build' +description: 'Download and test llama.cpp build artifacts' + +inputs: + os_type: + description: 'Operating system type (Windows or Linux)' + required: true + gfx_target: + description: 'GPU target (gfx1151)' + required: true + artifact_name: + description: 'Full artifact name to download' + required: true + +runs: + using: "composite" + steps: + - name: Download build artifacts + uses: actions/download-artifact@v4 + with: + name: ${{ inputs.artifact_name }} + path: llama-binaries + + - name: Download test model (Windows) + if: inputs.os_type == 'Windows' + shell: powershell + run: | + $modelUrl = "https://huggingface.co/unsloth/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_0.gguf?download=true" + $modelPath = "Qwen3-0.6B-Q4_0.gguf" + + Write-Host "Downloading test model from: $modelUrl" + curl.exe -L -o $modelPath $modelUrl + + # Verify download + if (Test-Path $modelPath) { + $fileSize = (Get-Item $modelPath).Length + Write-Host "Model downloaded successfully. Size: $fileSize bytes" + } else { + Write-Error "Failed to download model" + exit 1 + } + + - name: Download test model (Linux) + if: inputs.os_type == 'Linux' + shell: bash + run: | + model_url="https://huggingface.co/unsloth/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_0.gguf?download=true" + model_path="Qwen3-0.6B-Q4_0.gguf" + + echo "Downloading test model from: $model_url" + curl -L -o "$model_path" "$model_url" + + # Verify download + if [ -f "$model_path" ]; then + file_size=$(stat -c%s "$model_path") + echo "Model downloaded successfully. Size: $file_size bytes" + else + echo "Failed to download model" + exit 1 + fi + + - name: Set up library path (Linux) + if: inputs.os_type == 'Linux' + shell: bash + run: | + # Add the llama-binaries directory to LD_LIBRARY_PATH + echo "LD_LIBRARY_PATH=$(pwd)/llama-binaries:$LD_LIBRARY_PATH" >> $GITHUB_ENV + + - name: Run llama-cli test (Windows) + if: inputs.os_type == 'Windows' + shell: powershell + run: | + $llamaCliPath = "llama-binaries\llama-cli.exe" + $modelPath = "Qwen3-0.6B-Q4_0.gguf" + + # Verify llama-cli exists + if (-not (Test-Path $llamaCliPath)) { + Write-Error "llama-cli.exe not found at: $llamaCliPath" + Write-Host "Available files in llama-binaries:" + Get-ChildItem -Recurse "llama-binaries" + exit 1 + } + + Write-Host "Running llama-cli test for ${{ inputs.gfx_target }}..." + Write-Host "Command: $llamaCliPath -m `"$modelPath`" -ngl 99 -p `"I believe the meaning of life is`" -st -v" + + # Run the command and capture output + try { + $process = Start-Process -FilePath $llamaCliPath -ArgumentList @("-m", "`"$modelPath`"", "-ngl", "99", "-p", "`"I believe the meaning of life is`"", "-st", "-v") -NoNewWindow -Wait -PassThru -RedirectStandardOutput "stdout.txt" -RedirectStandardError "stderr.txt" + + # Wait for the process to complete + $process.WaitForExit() + + # Read both stdout and stderr + $stdout = Get-Content "stdout.txt" -Raw -ErrorAction SilentlyContinue + $stderr = Get-Content "stderr.txt" -Raw -ErrorAction SilentlyContinue + + # Always print the output (both stdout and stderr) + Write-Host "=== LLAMA-CLI STDOUT ===" + if ($stdout) { + Write-Host $stdout + } else { + Write-Host "(empty)" + } + Write-Host "=== LLAMA-CLI STDERR ===" + if ($stderr) { + Write-Host $stderr + } else { + Write-Host "(empty)" + } + Write-Host "=== END OUTPUT ===" + Write-Host "Process exit code: $($process.ExitCode)" + + # Combine output for pattern matching (stderr often contains important CUDA/GPU info) + $output = $stdout + $stderr + + # Clean up temp files + Remove-Item "stdout.txt" -ErrorAction SilentlyContinue + Remove-Item "stderr.txt" -ErrorAction SilentlyContinue + } + catch { + Write-Error "Failed to start llama-cli process: $_" + Write-Error "This usually indicates a missing dependency (DLL) or permission issue" + # Try to get more details about dependencies + Write-Host "Checking for dependency issues..." + Write-Host "Files in llama-binaries directory:" + Get-ChildItem "llama-binaries" | Format-Table Name, Length + exit 1 + } + + # Check if the process exited successfully + if ($process.ExitCode -ne 0) { + Write-Host "❌ llama-cli exited with error code: $($process.ExitCode)" + Write-Host "This error code typically indicates:" + Write-Host " - Exit code 127: Command or library not found (missing DLL)" + Write-Host " - Exit code 1: General error" + Write-Host " - Other codes: Application-specific errors" + Write-Host "" + Write-Host "Please check the STDERR output above for specific error messages." + exit 1 + } + + # Check for expected outputs + $foundGpuOffload = $output -match "offloaded 29/29 layers to GPU" + $foundThinkTag = $output -match "" + + Write-Host "=== TEST RESULTS ===" + Write-Host "Expected 'offloaded 29/29 layers to GPU': $(if ($foundGpuOffload) { 'FOUND' } else { 'NOT FOUND' })" + Write-Host "Expected '' tag: $(if ($foundThinkTag) { 'FOUND' } else { 'NOT FOUND' })" + + if ($foundGpuOffload -and $foundThinkTag) { + Write-Host "✅ Test PASSED - Both expected outputs found" + } else { + Write-Host "❌ Test FAILED - Missing expected outputs" + exit 1 + } + + - name: Run llama-cli test (Linux) + if: inputs.os_type == 'Linux' + shell: bash + run: | + llama_cli_path="./llama-binaries/llama-cli" + model_path="Qwen3-0.6B-Q4_0.gguf" + output_file="llama_output.txt" + + # Make llama-cli executable + chmod +x "$llama_cli_path" + + # Verify llama-cli exists + if [ ! -f "$llama_cli_path" ]; then + echo "llama-cli not found at: $llama_cli_path" + echo "Available files in llama-binaries:" + find llama-binaries -type f + exit 1 + fi + + echo "Running llama-cli test for ${{ inputs.gfx_target }}..." + echo "Command: $llama_cli_path -m \"$model_path\" -ngl 99 -p \"I believe the meaning of life is\" -st -v" + + # Run the command and capture output to file (more robust than variable) + set +e # Don't exit on error, we want to capture the output and check exit code manually + "$llama_cli_path" -m "$model_path" -ngl 99 -p "I believe the meaning of life is" -st -v > "$output_file" 2>&1 + exit_code=$? + set -e # Re-enable exit on error + + echo "=== LLAMA-CLI OUTPUT ===" + if [ -f "$output_file" ] && [ -s "$output_file" ]; then + cat "$output_file" + else + echo "(empty)" + fi + echo "=== END OUTPUT ===" + echo "Process exit code: $exit_code" + + # Check if the process exited successfully + if [ $exit_code -ne 0 ]; then + echo "❌ llama-cli exited with error code: $exit_code" + echo "This error code typically indicates:" + echo " - Exit code 127: Command or library not found (missing shared library)" + echo " - Exit code 1: General error" + echo " - Other codes: Application-specific errors" + echo "" + echo "Checking for missing library dependencies..." + ldd "$llama_cli_path" || echo "ldd command failed" + echo "" + echo "Please check the output above for specific error messages." + rm -f "$output_file" + exit 1 + fi + + # Check for expected outputs using grep -F (fixed strings, no regex interpretation) + # This handles special characters like angle brackets more reliably + found_gpu_offload=false + found_think_tag=false + + if grep -F "offloaded 29/29 layers to GPU" "$output_file" > /dev/null 2>&1; then + found_gpu_offload=true + fi + + if grep -F "" "$output_file" > /dev/null 2>&1; then + found_think_tag=true + fi + + echo "=== TEST RESULTS ===" + echo "Expected 'offloaded 29/29 layers to GPU': $(if [ "$found_gpu_offload" = true ]; then echo 'FOUND'; else echo 'NOT FOUND'; fi)" + echo "Expected '' tag: $(if [ "$found_think_tag" = true ]; then echo 'FOUND'; else echo 'NOT FOUND'; fi)" + + # Clean up output file + rm -f "$output_file" + + if [ "$found_gpu_offload" = true ] && [ "$found_think_tag" = true ]; then + echo "✅ Test PASSED - Both expected outputs found" + else + echo "❌ Test FAILED - Missing expected outputs" + exit 1 + fi + 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/bench-gfx1151.yml b/.github/workflows/bench-gfx1151.yml new file mode 100644 index 000000000000..5cc398bc6adf --- /dev/null +++ b/.github/workflows/bench-gfx1151.yml @@ -0,0 +1,464 @@ +# On-demand gfx1151 correctness and benchmark lane for the AMD DevLab Dispatch +# ephemeral runner pool. Dispatch it from the Actions tab with the branch, tag or +# SHA to measure. It then: +# +# 1. checks out that ref +# 2. builds llama-bench and test-backend-ops for gfx1151 (HIP by default) +# 3. runs test-backend-ops in `test` mode, then in `perf` mode, on the GPU device +# 4. downloads the fitting quant of each requested model from Hugging Face and +# runs llama-bench over the requested prompt / generation / depth / batch matrix +# 5. uploads every log plus the raw jsonl as one artifact +# +# Dispatch-only on purpose: it holds a GPU for tens of minutes per model and the +# DevLab pool queue is measured in hours, so it must never become a required check. +# +# This lane measures ONE build. It is not the Correctness Flow in AGENTS.md: there +# is no merge-base baseline, no corpus or logit comparison and no paired oracle. +# Use it to get numbers for a branch, then do the paired run separately. +# +# A failing test-backend-ops case does not stop the benchmark: the perf and +# llama-bench steps run whenever the build succeeded, so one marginal numerical +# failure cannot hide the throughput data. The run still fails overall, so the +# correctness signal is never lost - per AGENTS.md a speedup from a build that +# fails correctness must not be retained, only measured and reported as such. +# +# test-backend-ops perf runs with GGML_CUDA_DISABLE_GRAPHS=1 because SOLVE_TRI +# aborts under HIP graph capture (hipBLAS CUBLAS_STATUS_INTERNAL_ERROR from +# solve_tri_f32_cublas). Those perf numbers are therefore graphs-off; the +# llama-bench numbers keep graphs on. +# Runner labels: `devlab-dispatch` is the opt-in gate that reaches the AMD Ryzen AI +# Dev Lab ephemeral pool at all, and `strix-halo` narrows it to Strix Halo silicon +# inside that pool. `strix-halo` on its own reaches nothing. Do not use +# `runs-on: group:` - the pool matches on labels only and does not join runner +# groups. A wrong label does not fail, it queues silently, so keep both. +# +# Model assets, public Hugging Face, byte sizes read from the HF API on 2026-10-07: +# +# flash-next ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF IQ3_XXS 2 shards 75839998528 +# smaller fallbacks in the same repo: IQ2_XS 68026093024, Q2_0 66423878624 +# qwen38-27b unsloth/Qwen3.8-27B-GGUF UD-Q4_K_XL 1 file 17559178144 +# ornith-35b-a3b ornith-ai/Ornith-1.5-35B-A3B-GGUF Q4_K_M 1 file 21713463040 +# +# The defaults are about 115 GB of download. Each model is deleted before the next +# one is fetched, but the largest single model plus its KV cache must fit in RAM: +# n_ctx here is 2048 + 256 + 30000 = 32304 tokens. The pre-flight check below is a +# coarse filter on the file size plus a tunable margin, not a real KV budget - a +# model that slips through fails loudly at context creation instead. +# +# Naming: the request said "Ornith 1.5 30B-A3B". The published model is +# Ornith-1.5-35B-A3B (41 blocks, 256 experts, 8 used, arch qwen35moe). +# +# HF_TOKEN_CI is not configured as a secret on this repository, so the downloads +# below run anonymously. That works for these public repos but is rate limited; +# add the secret to get the authenticated path. +# +# Execution history: first run 2026-10-07 (pwilkin/llama.cpp run 37682939495). +# Checkout, the HIP build, GPU device detection, test-backend-ops and the artifact +# upload all worked, and a DevLab runner was assigned with exactly the labels +# self-hosted,Linux,devlab-dispatch,strix-halo. test-backend-ops reported +# 31003/31004 with one marginal MUL_MAT_ID q4_1 failure (ERR 0.000530464 against +# a 0.000500000 tolerance), which skipped perf and llama-bench on that run. No +# throughput number has been measured yet. The runner labels were confirmed by +# the AMD DevLab operator, and the model repos and byte sizes come from the +# public Hugging Face API. + +name: Benchmark gfx1151 (on-demand) + +on: + workflow_dispatch: + inputs: + ref: + description: 'Branch, tag or SHA to build and benchmark' + required: false + default: 'master' + backend: + description: 'GPU backend to build and test' + required: true + default: 'hip' + type: choice + options: [hip, vulkan] + models: + description: 'Comma-separated subset of flash-next,qwen38-27b,ornith-35b-a3b' + required: false + default: 'flash-next,qwen38-27b,ornith-35b-a3b' + reps: + description: 'llama-bench repetitions per test (-r)' + required: false + default: '5' + mem_margin_gib: + description: 'GiB held back for KV cache and compute buffers when deciding if a model fits' + required: false + default: '16' + skip_ops: + description: 'Skip test-backend-ops' + required: false + default: false + type: boolean + skip_bench: + description: 'Skip llama-bench' + required: false + default: false + type: boolean + +permissions: + contents: read + +# No cancel-in-progress: the pool queues silently for hours and cancelling kills +# runs that were never assigned a host. Serialise on the requested ref instead. +concurrency: + group: ${{ github.workflow }}-${{ inputs.ref }} + cancel-in-progress: false + +env: + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 0 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + bench: + name: ${{ inputs.backend }} / ${{ inputs.ref }} + runs-on: [self-hosted, Linux, devlab-dispatch, strix-halo] + timeout-minutes: 720 + + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + ref: ${{ inputs.ref }} + persist-credentials: false + + - name: Report environment + env: + INPUT_REF: ${{ inputs.ref }} + INPUT_BACKEND: ${{ inputs.backend }} + INPUT_MODELS: ${{ inputs.models }} + INPUT_REPS: ${{ inputs.reps }} + run: | + # diagnostics only, so never fail the job on a missing probe + set +e + mkdir -p "$RUNNER_TEMP/logs" + { + git log --oneline -1 + uname -a + nproc + free -g + df -h "$RUNNER_TEMP" . + rocminfo | grep -E '^ Name|Marketing Name' | head -6 + rocm-smi --showproductname | head -8 + cmake --version | head -1 + hipconfig --version + ccache --version | head -1 + } 2>&1 | tee "$RUNNER_TEMP/logs/environment.log" + + { + echo "### Benchmark gfx1151 (on-demand)" + echo + echo "- ref: \`$INPUT_REF\`" + echo "- commit: \`$(git rev-parse HEAD)\`" + echo "- backend: \`$INPUT_BACKEND\`" + echo "- models: \`$INPUT_MODELS\`" + echo "- repetitions: \`$INPUT_REPS\`" + } >> "$GITHUB_STEP_SUMMARY" + + - name: Build + id: build + env: + INPUT_BACKEND: ${{ inputs.backend }} + run: | + set -euo pipefail + mkdir -p "$RUNNER_TEMP/logs" + + backend="$INPUT_BACKEND" + case "$backend" in + hip) + # same shape as .github/workflows/hip-quality-check.yml, which is + # this repository's own gfx1151 recipe + cmake_extra=( + -DGGML_HIP=ON -DGGML_VULKAN=OFF -DGGML_CUDA=OFF + -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" + -DGPU_TARGETS=gfx1151 + -DGGML_HIP_EXPORT_METRICS=Off + ) + ;; + vulkan) + cmake_extra=(-DGGML_VULKAN=ON -DGGML_HIP=OFF -DGGML_CUDA=OFF) + ;; + *) + echo "::error::unknown backend '$backend'" >&2 + exit 2 + ;; + esac + + cmake -B build -S . \ + -DCMAKE_BUILD_TYPE=Release \ + "${cmake_extra[@]}" \ + -DLLAMA_BUILD_COMMON=ON \ + -DLLAMA_BUILD_TOOLS=ON \ + -DLLAMA_BUILD_TESTS=ON \ + -DLLAMA_BUILD_EXAMPLES=OFF \ + -DLLAMA_BUILD_SERVER=OFF \ + 2>&1 | tee "$RUNNER_TEMP/logs/build-cmake.log" + + cmake --build build --config Release -j "$(nproc)" \ + --target llama-bench test-backend-ops \ + 2>&1 | tee "$RUNNER_TEMP/logs/build-make.log" + + - name: Detect GPU device + id: device + run: | + # the runner starts bash with -e; drop it so the failure path below runs + set +e + set -uo pipefail + mkdir -p "$RUNNER_TEMP/logs" + ./build/bin/llama-bench --list-devices > "$RUNNER_TEMP/devices.txt" + rc=$? + cp "$RUNNER_TEMP/devices.txt" "$RUNNER_TEMP/logs/devices.txt" + cat "$RUNNER_TEMP/logs/devices.txt" + if [ "$rc" -ne 0 ]; then + echo "::error::llama-bench --list-devices failed (exit $rc)" >&2 + exit "$rc" + fi + dev=$(grep -oE '^ (ROCm|Vulkan)[0-9]+:' "$RUNNER_TEMP/devices.txt" | head -1 | tr -d ' :') || dev="" + if [ -z "$dev" ]; then + echo "::error::no GPU device registered; the requested backend did not build in" >&2 + exit 1 + fi + echo "OPS_DEVICE=$dev" >> "$GITHUB_ENV" + echo "device: $dev" + + - name: test-backend-ops test + if: ${{ !inputs.skip_ops }} + run: | + # the runner starts bash with -e; drop it so the guards below always run + set +e + set -uo pipefail + mkdir -p "$RUNNER_TEMP/logs" + log="$RUNNER_TEMP/logs/ops-test.log" + ./build/bin/test-backend-ops test -b "$OPS_DEVICE" 2>&1 | tee "$log" + rc=${PIPESTATUS[0]} + + # `-b` with a name that matches no device exits 0 and still prints + # "N/N backends passed", so prove the device really ran. + awk -v dev="$OPS_DEVICE" ' + $0 ~ "^Backend [0-9]+/[0-9]+: " dev "$" { + seen = 1 + if ((getline nxt) > 0 && nxt !~ /^ Skipping/) ran = 1 + } + END { exit !(seen && ran) } + ' "$log" || { echo "::error::$OPS_DEVICE never ran; the -b filter matched nothing" >&2; exit 1; } + + summary=$(grep -oE '[0-9]+/[0-9]+ tests passed' "$log" | tail -1) + total=$(printf '%s\n' "$summary" | cut -d/ -f2 | awk '{print $1}') + echo "test-backend-ops summary: ${summary:-none}" + if [ -z "$total" ] || [ "$total" -le 0 ]; then + echo "::error::no test-backend-ops cases ran" >&2 + exit 1 + fi + + exit "$rc" + + - name: test-backend-ops perf + if: ${{ !cancelled() && steps.build.outcome == 'success' && !inputs.skip_ops }} + env: + # SOLVE_TRI aborts under HIP graph capture: hipBLAS returns + # CUBLAS_STATUS_INTERNAL_ERROR from hipblasStrsmBatched inside + # solve_tri_f32_cublas (ggml-cuda/solve_tri.cu:71) and ggml aborts. + # Reproduced at 4977288b3 on ROCm 10.0 locally and on ROCm 7.2.1 on + # the DevLab host; test mode passes 24/24, so it is specific to the + # perf path under capture. Graphs off lets every case run, so the + # perf numbers below are a graphs-off measurement. Remove this once + # the capture interaction is fixed. + GGML_CUDA_DISABLE_GRAPHS: "1" + run: | + set +e + set -uo pipefail + mkdir -p "$RUNNER_TEMP/logs" + log="$RUNNER_TEMP/logs/ops-perf.log" + ./build/bin/test-backend-ops perf -b "$OPS_DEVICE" 2>&1 | tee "$log" + rc=${PIPESTATUS[0]} + + awk -v dev="$OPS_DEVICE" ' + $0 ~ "^Backend [0-9]+/[0-9]+: " dev "$" { + seen = 1 + if ((getline nxt) > 0 && nxt !~ /^ Skipping/) ran = 1 + } + END { exit !(seen && ran) } + ' "$log" || { echo "::error::$OPS_DEVICE never ran; the -b filter matched nothing" >&2; exit 1; } + + # perf mode returns 0 even when nothing was measured + if ! grep -q 'us/run' "$log"; then + echo "::error::no perf rows were produced" >&2 + exit 1 + fi + + exit "$rc" + + - name: llama-bench + if: ${{ !cancelled() && steps.build.outcome == 'success' && !inputs.skip_bench }} + env: + MODELS: ${{ inputs.models }} + REPS: ${{ inputs.reps }} + MARGIN_GIB: ${{ inputs.mem_margin_gib }} + run: | + set +e + set -uo pipefail + + models="$MODELS" + reps="$REPS" + margin_gib="$MARGIN_GIB" + + case "$reps" in + '' | *[!0-9]*) + echo "::error::reps must be a positive integer, got '$reps'" >&2 + exit 2 + ;; + esac + if [ "${#reps}" -gt 6 ] || [ "$(( 10#$reps ))" -le 0 ]; then + echo "::error::reps must be between 1 and 999999, got '$reps'" >&2 + exit 2 + fi + reps=$(( 10#$reps )) + + case "$margin_gib" in + '' | *[!0-9]*) + echo "::error::mem_margin_gib must be a nonnegative integer, got '$margin_gib'" >&2 + exit 2 + ;; + esac + if [ "${#margin_gib}" -gt 5 ] || [ "$(( 10#$margin_gib ))" -gt 4096 ]; then + echo "::error::mem_margin_gib must be between 0 and 4096, got '$margin_gib'" >&2 + exit 2 + fi + margin_gib=$(( 10#$margin_gib )) + + model_root="$RUNNER_TEMP/models" + logdir="$RUNNER_TEMP/logs" + mkdir -p "$model_root" "$logdir" + + margin_bytes=$(( margin_gib * 1024 * 1024 * 1024 )) + avail_bytes=$(awk '/^MemAvailable:/ {print $2 * 1024; exit}' /proc/meminfo) + + # repo|space separated files|expected total bytes + model_spec() { + case "$1" in + flash-next) + echo 'ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF|IQ3_XXS/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00001-of-00002.gguf IQ3_XXS/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS-00002-of-00002.gguf|75839998528' + ;; + qwen38-27b) + echo 'unsloth/Qwen3.8-27B-GGUF|Qwen3.8-27B-UD-Q4_K_XL.gguf|17559178144' + ;; + ornith-35b-a3b) + echo 'ornith-ai/Ornith-1.5-35B-A3B-GGUF|Ornith-1.5-35B-Q4_K_M.gguf|21713463040' + ;; + *) return 1 ;; + esac + } + + # reject a bad name before anything is downloaded + IFS=',' read -ra want <<< "$models" + for name in "${want[@]}"; do + name=$(printf '%s' "$name" | tr -d '[:space:]') + [ -n "$name" ] || continue + if ! model_spec "$name" > /dev/null; then + echo "::error::unknown model '$name'; choose from flash-next,qwen38-27b,ornith-35b-a3b" >&2 + exit 1 + fi + done + + ran=0 + for name in "${want[@]}"; do + name=$(printf '%s' "$name" | tr -d '[:space:]') + [ -n "$name" ] || continue + + spec=$(model_spec "$name") + repo=${spec%%|*} + rest=${spec#*|} + files=${rest%%|*} + bytes=${rest##*|} + + need=$(( bytes + margin_bytes )) + if [ "$avail_bytes" -gt 0 ] && [ "$need" -gt "$avail_bytes" ]; then + echo "::warning::SKIP $name: needs $(( need / 1073741824 )) GiB, $(( avail_bytes / 1073741824 )) GiB available" + continue + fi + + dest="$model_root/$name" + mkdir -p "$dest" + free_bytes=$(df -B1 --output=avail "$dest" | tail -1) + if [ "$free_bytes" -lt $(( bytes + bytes / 5 )) ]; then + echo "::warning::SKIP $name: needs $(( (bytes + bytes / 5) / 1073741824 )) GiB of disk, $(( free_bytes / 1073741824 )) GiB free" + continue + fi + + for f in $files; do + out="$dest/$(basename "$f")" + if [ -s "$out" ]; then + echo "already present: $out" + continue + fi + url="https://huggingface.co/$repo/resolve/main/$f" + echo "fetching $url" + if [ -n "${HF_TOKEN:-}" ]; then + curl -fL --retry 5 --retry-delay 5 --retry-all-errors -C - \ + -H "Authorization: Bearer $HF_TOKEN" -o "$out" "$url" \ + || { echo "::error::download failed: $url" >&2; exit 1; } + else + curl -fL --retry 5 --retry-delay 5 --retry-all-errors -C - \ + -o "$out" "$url" \ + || { echo "::error::download failed: $url" >&2; exit 1; } + fi + done + + # llama.cpp resolves the sibling shards from the first shard path + model="$dest/$(basename "${files%% *}")" + log="$logdir/$name.log" + jsonl="$logdir/$name.jsonl" + echo "=== llama-bench $name: $model ===" + + # n_ctx is n_prompt + n_gen + n_depth = 2048 + 256 + 30000. + # -fa on and -ctk/-ctv f16 are spelled out so the numbers do not move + # with the auto defaults. HIP_LAUNCH_BLOCKING is deliberately NOT set: + # the gpu-rocm job sets it to work around an async-execution + # correctness issue, but it serialises every kernel launch and would + # invalidate every throughput number measured here. + timeout 7200 ./build/bin/llama-bench \ + -m "$model" \ + -p 2048 -n 256 \ + -b 2048,4096,8192 -ub 2048,4096,8192 \ + -d 0,30000 \ + -lzm on-direct \ + -ngl 99 -fa on -ctk f16 -ctv f16 \ + -r "$reps" -o jsonl --progress \ + > "$jsonl" 2> "$log" + rc=$? + cat "$log" + if [ "$rc" -ne 0 ]; then + echo "::error::llama-bench failed for $name (exit $rc)" >&2 + exit "$rc" + fi + if [ ! -s "$jsonl" ]; then + echo "::error::llama-bench produced no jsonl rows for $name" >&2 + exit 1 + fi + echo "jsonl rows: $(wc -l < "$jsonl")" + ran=$(( ran + 1 )) + + rm -rf "$dest" + done + + if [ "$ran" -eq 0 ]; then + echo "::error::no model was benchmarked" >&2 + exit 1 + fi + + - name: Upload logs + if: ${{ always() }} + uses: actions/upload-artifact@v4 + with: + name: bench-gfx1151-${{ github.run_id }} + path: ${{ runner.temp }}/logs + retention-days: 30 + if-no-files-found: warn diff --git a/.github/workflows/build-and-test-snapdragon.yml b/.github/workflows/build-and-test-snapdragon.yml index 296780acb4cc..7058dafee1fc 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 90960a7f61e8..9b0953427df0 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 f87f1346354f..f39f0616ee4b 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 cde01952fc6f..da71b32e9216 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 68b6c01900bf..0b8160b096b6 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: registry.mthreads.com/mcconline/musa_sdk:5.2.0-devel-ubuntu22.04-s5000 diff --git a/.github/workflows/build-cuda-windows.yml b/.github/workflows/build-cuda-windows.yml index e08553e6cdfc..f10aef6c0ee1 100644 --- a/.github/workflows/build-cuda-windows.yml +++ b/.github/workflows/build-cuda-windows.yml @@ -94,9 +94,9 @@ jobs: strategy: matrix: include: - # sync with release.yml + # sync with release.yml; this fork only supports Strix Halo (gfx1151) - name: "radeon" - gpu_targets: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032" + gpu_targets: "gfx1151" steps: - name: Clone 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-llamacpp-rocm.yml b/.github/workflows/build-llamacpp-rocm.yml new file mode 100644 index 000000000000..f6c39725c6e9 --- /dev/null +++ b/.github/workflows/build-llamacpp-rocm.yml @@ -0,0 +1,1120 @@ +# Nightly ROCm release builds of this fork, same layout as lemonade-sdk/llamacpp-rocm +# (https://github.com/lemonade-sdk/llamacpp-rocm, MIT) so the zips drop into the same tooling. +# Differences: builds this repository instead of cloning upstream llama.cpp, releases on every merge to +# master instead of nightly, builds Linux only by default (Windows still selectable by dispatch), uses the +# job token when no GH_TOKEN secret exists, and runs the self-hosted Strix Halo tests only when the runners +# are registered. +name: Build Llama.cpp + ROCm +on: + workflow_dispatch: + inputs: + operating_systems: + description: 'Operating systems to build for (comma-separated: windows,ubuntu)' + required: false + default: 'ubuntu' + gfx_target: + description: 'AMD GPU targets (comma-separated; this fork supports Strix Halo / gfx1151 only)' + required: false + default: 'gfx1151' + rocm_version: + description: 'ROCm version to use (e.g., 7.13.0a20260318) or "latest" to auto-detect' + required: false + default: 'latest' + llamacpp_version: + description: 'Branch or tag of this repository to build, or "latest" for master' + required: false + default: 'latest' + vulkan: + description: 'Also build the Linux Vulkan zip' + required: false + default: true + type: boolean + create_release: + description: 'Create a GitHub release after successful build' + required: false + default: true + type: boolean + + # Trigger on pushes to pull requests that change the release pipeline + pull_request: + types: [opened, synchronize, reopened] + paths: + - '.github/workflows/build-llamacpp-rocm.yml' + - '.github/actions/test-llamacpp-build/**' + - 'scripts/rocm-release/**' + + # Release every merge to master + push: + branches: [master] + +# a newer merge supersedes a release that is still building +concurrency: + group: ${{ github.workflow }}-${{ github.event_name }}-${{ github.ref }} + cancel-in-progress: true + +env: + OPERATING_SYSTEMS: ${{ github.event.inputs.operating_systems || 'ubuntu' }} + # this fork only supports Strix Halo (gfx1151) + GFX_TARGETS: ${{ github.event.inputs.gfx_target || 'gfx1151' }} + ROCM_VERSION: ${{ github.event.inputs.rocm_version || 'latest' }} + LLAMACPP_VERSION: ${{ github.event.inputs.llamacpp_version || 'latest' }} + # on a push, build exactly the merged commit + BUILD_SHA: ${{ github.event_name == 'push' && github.sha || '' }} + +jobs: + prepare-matrix: + runs-on: ubuntu-22.04 + outputs: + windows_matrix: ${{ steps.set-matrix.outputs.windows_matrix }} + ubuntu_matrix: ${{ steps.set-matrix.outputs.ubuntu_matrix }} + should_build_windows: ${{ steps.set-matrix.outputs.should_build_windows }} + should_build_ubuntu: ${{ steps.set-matrix.outputs.should_build_ubuntu }} + steps: + - name: Set matrix + id: set-matrix + run: | + targets="${{ env.GFX_TARGETS }}" + operating_systems="${{ env.OPERATING_SYSTEMS }}" + + echo "Input targets: $targets" + echo "Input operating systems: $operating_systems" + + # Convert targets to JSON array + matrix_targets=$(echo "$targets" \ + | tr ',' '\n' \ + | sed 's/^ *//;s/ *$//' \ + | sed 's/^"//;s/"$//' \ + | jq -R . \ + | jq -s '{gfx_target: .}' \ + | jq -c) + + # Check which operating systems to build + should_build_windows="false" + should_build_ubuntu="false" + + if [[ "$operating_systems" == *"windows"* ]]; then + should_build_windows="true" + echo "windows_matrix=$matrix_targets" >> $GITHUB_OUTPUT + fi + + if [[ "$operating_systems" == *"ubuntu"* ]]; then + should_build_ubuntu="true" + echo "ubuntu_matrix=$matrix_targets" >> $GITHUB_OUTPUT + fi + + echo "should_build_windows=$should_build_windows" >> $GITHUB_OUTPUT + echo "should_build_ubuntu=$should_build_ubuntu" >> $GITHUB_OUTPUT + + echo "Windows build: $should_build_windows" + echo "Ubuntu build: $should_build_ubuntu" + echo "Generated matrix: $matrix_targets" + + build-windows: + # Pinned to windows-2022 (VS 2022 / MSVC 14.4x). The windows-latest image moved to + # VS 2026 (MSVC 14.51) on 2026-06-08, whose declares isgreater/isless/etc. as + # constexpr (implicitly __host__ __device__), which collides with ROCm clang's + # __device__-only HIP math wrappers and breaks the HIP build. See llama.cpp#22570. + runs-on: windows-2022 + needs: prepare-matrix + if: needs.prepare-matrix.outputs.should_build_windows == 'true' + strategy: + matrix: ${{fromJson(needs.prepare-matrix.outputs.windows_matrix)}} + fail-fast: false + outputs: + rocm_version: ${{ steps.set-outputs.outputs.rocm_version }} + llamacpp_commit_hash: ${{ steps.set-outputs.outputs.llamacpp_commit_hash }} + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Clean up existing directories (safety precaution) + run: | + # Remove existing llama.cpp directory if it exists + if (Test-Path "llama.cpp") { + Write-Host "Removing existing llama.cpp directory..." + Remove-Item -Recurse -Force "llama.cpp" + } + + # Remove existing C:\opt\rocm directory if it exists + if (Test-Path "C:\opt\rocm") { + Write-Host "Removing existing C:\opt\rocm directory..." + Remove-Item -Recurse -Force "C:\opt\rocm" + } + + # Remove any existing ROCm tarball + if (Test-Path "rocm.tar.gz") { + Write-Host "Removing existing rocm.tar.gz..." + Remove-Item -Force "rocm.tar.gz" + } + + Write-Host "Cleanup completed successfully" + + - name: Install Visual Studio Build Tools + run: | + # Download and install Visual Studio Build Tools + $vsInstallerUrl = "https://aka.ms/vs/17/release/vs_buildtools.exe" + $vsInstallerPath = "$env:TEMP\vs_buildtools.exe" + + Write-Host "Downloading Visual Studio Build Tools..." + Invoke-WebRequest -Uri $vsInstallerUrl -OutFile $vsInstallerPath + + Write-Host "Installing Visual Studio Build Tools..." + Start-Process -FilePath $vsInstallerPath -ArgumentList "--quiet", "--wait", "--norestart", "--add", "Microsoft.VisualStudio.Component.VC.Tools.x86.x64", "--add", "Microsoft.VisualStudio.Component.VC.CMake.Project", "--add", "Microsoft.VisualStudio.Component.VC.ATL", "--add", "Microsoft.VisualStudio.Component.Windows11SDK.22621" -Wait + + # Clean up installer + Remove-Item $vsInstallerPath -Force + + - name: Install build dependencies + run: | + Write-Host "Installing build dependencies using manual methods..." + + # Install Ninja + Write-Host "Installing Ninja..." + $ninjaUrl = "https://github.com/ninja-build/ninja/releases/download/v1.11.1/ninja-win.zip" + $ninjaPath = "$env:TEMP\ninja-win.zip" + $ninjaDir = "C:\ninja" + New-Item -ItemType Directory -Force -Path $ninjaDir + Invoke-WebRequest -Uri $ninjaUrl -OutFile $ninjaPath + Expand-Archive -Path $ninjaPath -DestinationPath $ninjaDir -Force + + # Install Strawberry Perl + Write-Host "Installing Strawberry Perl..." + $perlUrl = "http://strawberryperl.com/download/5.32.1.1/strawberry-perl-5.32.1.1-64bit.msi" + $perlPath = "$env:TEMP\strawberry-perl-5.32.1.1-64bit.msi" + Invoke-WebRequest -Uri $perlUrl -OutFile $perlPath + Start-Process msiexec.exe -ArgumentList "/i $perlPath /quiet /norestart" -Wait + + # Verify installations + $env:PATH = "C:\ninja;C:\Strawberry\perl\bin;$env:PATH" + Write-Host "Verifying installations..." + ninja --version + perl --version + + Write-Host "Manual installation of build dependencies completed" + + - name: Download ROCm nightly tarball + run: | + # Determine ROCm version to use + $rocmVersion = "${{ env.ROCM_VERSION }}" + $currentTarget = "${{ matrix.gfx_target }}" + + $s3Target = $currentTarget + + # TheRock now publishes nightlies to the ROCm core tarball index. That + # index is a static HTML page with an embedded JSON `files` array (name + + # mtime), so we parse that instead of the retired S3 XML bucket listing. + $baseUrl = "https://nightly.repo.amd.com/rocm/core/tarball" + if ($rocmVersion -eq "latest") { + Write-Host "Auto-detecting latest ROCm version for target: $currentTarget" + $indexHtml = (Invoke-WebRequest "$baseUrl/" -UseBasicParsing).Content + $filesMatch = [regex]::Match($indexHtml, 'const files = (\[.*?\]);', [System.Text.RegularExpressions.RegexOptions]::Singleline) + if (-not $filesMatch.Success) { + Write-Error "Failed to parse file index from $baseUrl/" + exit 1 + } + + # Pick the newest by the YYYYMMDD build date embedded in the filename. + # Assign before piping so the parsed array enumerates element-by-element. + $allFiles = $filesMatch.Groups[1].Value | ConvertFrom-Json + $prefix = "therock-dist-windows-$s3Target-" + # Anchor the match so the version comes immediately after the prefix. + # This excludes variants like "...-all-tests-.tar.gz" that + # TheRock also publishes and share the same prefix/date. + $versionPattern = "^$([regex]::Escape($prefix))\d+\.\d+\.\d+(a|rc)\d+\.tar\.gz$" + $latest = $allFiles | + Where-Object { $_.name -match $versionPattern } | + Sort-Object { [regex]::Match($_.name, '(\d{8})\.tar\.gz$').Groups[1].Value } | + Select-Object -Last 1 + if (-not $latest) { + Write-Error "No tarball found for prefix '$prefix' at $baseUrl/" + exit 1 + } + $latestFile = $latest.name + Write-Host "Found latest file: $latestFile" + + # Extract version from the filename for environment variable + if ($latestFile -match "therock-dist-windows-$s3Target-(\d+\.\d+\.\d+(?:a|rc)\d+)\.tar\.gz") { + $rocmVersion = $matches[1] + Write-Host "Detected latest ROCm version: $rocmVersion" + } else { + Write-Error "Failed to extract ROCm version from latest file: $latestFile" + Write-Error "Expected pattern: therock-dist-windows-$s3Target-.tar.gz" + exit 1 + } + + $rocmUrl = "$baseUrl/$latestFile" + } else { + # For specific versions, construct the URL directly + $rocmUrl = "$baseUrl/therock-dist-windows-$s3Target-$rocmVersion.tar.gz" + } + + # Store the version for use in other steps + echo "DETECTED_ROCM_VERSION=$rocmVersion" >> $env:GITHUB_ENV + + Write-Host "Downloading ROCm from: $rocmUrl" + Invoke-WebRequest -Uri $rocmUrl -OutFile "rocm.tar.gz" + + - name: Extract ROCm to C:\opt\rocm + run: | + # Create directory if it doesn't exist + New-Item -ItemType Directory -Force -Path "C:\opt\rocm" + + # Extract the tarball + tar -xzf rocm.tar.gz -C C:\opt\rocm --strip-components=1 + + - name: Clone llama.cpp + run: | + $llamacppVersion = "${{ env.LLAMACPP_VERSION }}" + if ($llamacppVersion -eq "latest") { + Write-Host "Cloning master branch (latest) with blobless clone" + git clone --filter=blob:none --single-branch --branch master https://github.com/${{ github.repository }}.git llama.cpp + } else { + Write-Host "Cloning llama.cpp version: $llamacppVersion with blobless clone" + git clone --filter=blob:none --single-branch --branch $llamacppVersion https://github.com/${{ github.repository }}.git llama.cpp + } + cd llama.cpp + if ("${{ env.BUILD_SHA }}" -ne "") { git checkout --detach "${{ env.BUILD_SHA }}" } + + # Get commit hash (5 digits) and store it as environment variable + $commitHash = git rev-parse --short=5 HEAD + echo "LLAMACPP_COMMIT_HASH=$commitHash" >> $env:GITHUB_ENV + Write-Host "llama.cpp commit hash (5 digits): $commitHash" + + # Show current commit info + Write-Host "Current llama.cpp commit:" + git log --oneline -1 + + - name: Build Llama.cpp + ROCm + shell: cmd + run: | + + set "current_target=${{ matrix.gfx_target }}" + echo GPU target: %current_target% + + REM Set up environment variables and PATH + set HIP_PATH=C:\opt\rocm + set HIP_PLATFORM=amd + set PATH=%HIP_PATH%\lib\llvm\bin;%HIP_PATH%\bin;%PATH% + + REM Set up x64 Native Tools Command Prompt environment + call "C:\Program Files (x86)\Microsoft Visual Studio\Installer\vswhere.exe" -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath > vs_path.txt + set /p VS_PATH=> $env:GITHUB_OUTPUT + echo "llamacpp_commit_hash=$env:LLAMACPP_COMMIT_HASH" >> $env:GITHUB_OUTPUT + + Write-Host "Final rocm_version: $rocmVersion" + Write-Host "Final llamacpp_commit_hash: $env:LLAMACPP_COMMIT_HASH" + + - name: Create release summary + run: | + $summary = @" + # Llama.cpp + ROCm Build Summary (Windows - ${{ matrix.gfx_target }}) + + ## Build Configuration + - **GPU Target**: ${{ matrix.gfx_target }} + - **ROCm Version**: ${{ env.DETECTED_ROCM_VERSION || env.ROCM_VERSION }} + - **llama.cpp Commit Hash**: ${{ env.LLAMACPP_COMMIT_HASH }} + - **Build Type**: Release + - **Platform**: Windows + + ## Build Artifacts + The compiled binaries are available in the build artifacts. + + "@ + + Set-Content -Path "build-summary-windows-${{ matrix.gfx_target }}.md" -Value $summary + + - name: Upload build summary + id: upload-summary + uses: actions/upload-artifact@v4 + with: + name: build-summary-windows-${{ matrix.gfx_target }} + path: build-summary-windows-${{ matrix.gfx_target }}.md + retention-days: 30 + + - name: Clean up build artifacts and temporary files + if: always() + run: | + Write-Host "Starting cleanup process..." + + # Remove ROCm tarball + if (Test-Path "rocm.tar.gz") { + Write-Host "Removing rocm.tar.gz..." + Remove-Item -Force "rocm.tar.gz" + } + + # Remove ROCm installation directory + if (Test-Path "C:\opt\rocm") { + Write-Host "Removing C:\opt\rocm directory..." + Remove-Item -Recurse -Force "C:\opt\rocm" + } + # Remove llama.cpp source directory (keeping only the build artifacts) + if (Test-Path "llama.cpp") { + Write-Host "Removing llama.cpp source directory..." + Remove-Item -Recurse -Force "llama.cpp" + } + + # Clean up any temporary build files + if (Test-Path "build-summary-windows-${{ matrix.gfx_target }}.md") { + Write-Host "Removing build-summary-windows-${{ matrix.gfx_target }}.md..." + Remove-Item -Force "build-summary-windows-${{ matrix.gfx_target }}.md" + } + + Write-Host "Cleanup completed successfully" + + build-ubuntu: + runs-on: ubuntu-22.04 + needs: prepare-matrix + if: needs.prepare-matrix.outputs.should_build_ubuntu == 'true' + strategy: + matrix: ${{fromJson(needs.prepare-matrix.outputs.ubuntu_matrix)}} + fail-fast: false + outputs: + rocm_version: ${{ steps.set-outputs.outputs.rocm_version }} + llamacpp_commit_hash: ${{ steps.set-outputs.outputs.llamacpp_commit_hash }} + + + steps: + - name: Free disk space + # Remove unused runner files to free up disk space + run: curl -fsSL https://raw.githubusercontent.com/kou/arrow/e49d8ae15583ceff03237571569099a6ad62be32/ci/scripts/util_free_space.sh | bash + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Clean up existing directories (safety precaution) + run: | + # Remove existing llama.cpp directory if it exists + if [ -d "llama.cpp" ]; then + echo "Removing existing llama.cpp directory..." + rm -rf llama.cpp + fi + + # Remove existing /opt/rocm directory if it exists + if [ -d "/opt/rocm" ]; then + echo "Removing existing /opt/rocm directory..." + sudo rm -rf /opt/rocm + fi + + # Remove any existing ROCm tarball + if [ -f "rocm.tar.gz" ]; then + echo "Removing existing rocm.tar.gz..." + rm -f rocm.tar.gz + fi + + echo "Cleanup completed successfully" + + - name: Install build dependencies + run: | + echo "Installing build dependencies..." + sudo apt update + sudo apt install -y cmake ninja-build unzip curl + + # Verify installations + echo "Verifying installations..." + cmake --version + ninja --version + echo "Build dependencies installation completed" + + - name: Download and extract ROCm directly to /opt/rocm + run: | + # Determine ROCm version to use + rocm_version="${{ env.ROCM_VERSION }}" + current_target="${{ matrix.gfx_target }}" + + s3_target="$current_target" + + # TheRock now publishes nightlies to the ROCm core tarball index. That + # index is a static HTML page with an embedded JSON `files` array (name + + # mtime), so we parse that instead of the retired S3 XML bucket listing. + base_url="https://nightly.repo.amd.com/rocm/core/tarball" + if [ "$rocm_version" = "latest" ]; then + echo "Auto-detecting latest ROCm version for target: $current_target" + prefix="therock-dist-linux-${s3_target}-" + + files_json=$(curl -s "$base_url/" | tr '\n' ' ' | grep -oP 'const files = \K\[.*?\](?=\s*;)') + if [ -z "$files_json" ]; then + echo "Failed to parse file index from $base_url/" + exit 1 + fi + + # Pick the newest by the YYYYMMDD build date embedded in the filename. + # Anchor the match so the version comes immediately after the prefix. + # This excludes variants like "...-all-tests-.tar.gz" that + # TheRock also publishes and share the same prefix/date. + latest_file=$(echo "$files_json" | jq -r --arg p "$prefix" \ + '[.[] | select(.name | test("^" + $p + "[0-9]+\\.[0-9]+\\.[0-9]+(a|rc)[0-9]+\\.tar\\.gz$"))] | sort_by(.name | capture("(?[0-9]{8})\\.tar\\.gz$").d) | last | .name // empty') + if [ -z "$latest_file" ]; then + echo "No tarball found for prefix '$prefix' at $base_url/" + exit 1 + fi + echo "Found latest file: $latest_file" + + # Extract version from the filename for environment variable + if [[ "$latest_file" =~ therock-dist-linux-${s3_target}-([0-9]+\.[0-9]+\.[0-9]+(a|rc)[0-9]+)\.tar\.gz ]]; then + rocm_version="${BASH_REMATCH[1]}" + echo "Detected latest ROCm version: $rocm_version" + else + echo "Failed to extract ROCm version from latest file: $latest_file" + echo "Expected pattern: therock-dist-linux-${s3_target}-.tar.gz" + exit 1 + fi + + rocm_url="$base_url/$latest_file" + else + # For specific versions, construct the URL directly + rocm_url="$base_url/therock-dist-linux-${s3_target}-${rocm_version}.tar.gz" + fi + + # Store the version for use in other steps + echo "DETECTED_ROCM_VERSION=$rocm_version" >> $GITHUB_ENV + + echo "Streaming ROCm from: $rocm_url directly to extraction" + + # Create directory if it doesn't exist + sudo mkdir -p /opt/rocm + + # Stream download directly into tar extraction (no intermediate file) + curl -sL "$rocm_url" | sudo tar --use-compress-program=gzip -xf - -C /opt/rocm --strip-components=1 + + - name: Set ROCm environment variables + run: | + echo "Setting ROCm environment variables..." + + # Set environment variables for this step and subsequent steps + echo "HIP_PATH=/opt/rocm" >> $GITHUB_ENV + echo "ROCM_PATH=/opt/rocm" >> $GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $GITHUB_ENV + echo "HIP_CLANG_PATH=/opt/rocm/llvm/bin" >> $GITHUB_ENV + echo "HIP_INCLUDE_PATH=/opt/rocm/include" >> $GITHUB_ENV + echo "HIP_LIB_PATH=/opt/rocm/lib" >> $GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode" >> $GITHUB_ENV + + # Update PATH + echo "/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH" >> $GITHUB_PATH + + # Set library paths + echo "LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib64:/opt/rocm/llvm/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV + echo "LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib64:${LIBRARY_PATH:-}" >> $GITHUB_ENV + echo "CPATH=/opt/rocm/include:${CPATH:-}" >> $GITHUB_ENV + echo "PKG_CONFIG_PATH=/opt/rocm/lib/pkgconfig:${PKG_CONFIG_PATH:-}" >> $GITHUB_ENV + + echo "ROCm environment variables set successfully" + + + - name: Clone llama.cpp + run: | + llamacpp_version="${{ env.LLAMACPP_VERSION }}" + if [ "$llamacpp_version" = "latest" ]; then + echo "Cloning master branch (latest) with blobless clone" + git clone --filter=blob:none --single-branch --branch master https://github.com/${{ github.repository }}.git llama.cpp + else + echo "Cloning llama.cpp version: $llamacpp_version with blobless clone" + git clone --filter=blob:none --single-branch --branch "$llamacpp_version" https://github.com/${{ github.repository }}.git llama.cpp + fi + cd llama.cpp + if [ -n "${{ env.BUILD_SHA }}" ]; then git checkout --detach "${{ env.BUILD_SHA }}"; fi + + # Get commit hash (5 digits) and store it as environment variable + commit_hash=$(git rev-parse --short=5 HEAD) + echo "LLAMACPP_COMMIT_HASH=$commit_hash" >> $GITHUB_ENV + echo "llama.cpp commit hash (5 digits): $commit_hash" + + # Show current commit info + echo "Current llama.cpp commit:" + git log --oneline -1 + + + - name: Build Llama.cpp + ROCm + run: | + current_target="${{ matrix.gfx_target }}" + echo "GPU target: $current_target" + + # Create build directory + cd llama.cpp + mkdir build + cd build + + # Configure the project + cmake .. -G Ninja \ + -DCMAKE_C_COMPILER=/opt/rocm/llvm/bin/clang \ + -DCMAKE_CXX_COMPILER=/opt/rocm/llvm/bin/clang++ \ + -DCMAKE_CXX_FLAGS="-I/opt/rocm/include" \ + -DCMAKE_CROSSCOMPILING=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DGPU_TARGETS="$current_target" \ + -DBUILD_SHARED_LIBS=ON \ + -DLLAMA_BUILD_TESTS=OFF \ + -DGGML_HIP=ON \ + -DGGML_OPENMP=OFF \ + -DGGML_CUDA_FORCE_CUBLAS=OFF \ + -DGGML_RPC=ON \ + -DGGML_HIP_ROCWMMA_FATTN=OFF \ + -DLLAMA_BUILD_BORINGSSL=ON \ + -DGGML_NATIVE=OFF \ + -DGGML_STATIC=OFF \ + -DCMAKE_SYSTEM_NAME=Linux + + # Build the project + cmake --build . -j $(nproc) + + - name: Copy ROCm core libs to build directory + run: | + build_bin_path="llama.cpp/build/bin" + rocm_bin_path="/opt/rocm/bin" + + # Copy the rocblas/library folder and all its contents + rocblas_lib_path="/opt/rocm/lib/rocblas/library" + if [ -d "$rocblas_lib_path" ]; then + echo "Copying rocblas/library folder and all contents..." + dest_rocblas_path="$build_bin_path/rocblas/library" + mkdir -p "$(dirname "$dest_rocblas_path")" + cp -r "$rocblas_lib_path" "$(dirname "$dest_rocblas_path")/" + echo "Copied: rocblas/library folder with all contents" + + # List the contents of the copied rocblas/library folder + echo "Contents of rocblas/library:" + find "$dest_rocblas_path" -type f -exec ls -la {} \; | head -20 + else + echo "Warning: rocblas/library folder not found at: $rocblas_lib_path" + fi + + # Copy the hipblaslt/library folder and all its contents + hipblaslt_lib_path="/opt/rocm/lib/hipblaslt/library" + if [ -d "$hipblaslt_lib_path" ]; then + echo "Copying hipblaslt/library folder and all contents..." + dest_hipblaslt_path="$build_bin_path/hipblaslt/library" + mkdir -p "$(dirname "$dest_hipblaslt_path")" + cp -r "$hipblaslt_lib_path" "$(dirname "$dest_hipblaslt_path")/" + echo "Copied: hipblaslt/library folder with all contents" + + # List the contents of the copied hipblaslt/library folder + echo "Contents of hipblaslt/library:" + find "$dest_hipblaslt_path" -type f -exec ls -la {} \; | head -20 + else + echo "Warning: hipblaslt/library folder not found at: $hipblaslt_lib_path" + fi + + # Copy required ROCm libraries to build directory + # If artifacts from ROCm or Llama.cpp change, you may need to update this list + # To get a new list of all libraries, run: + # scripts/rocm-release/gather_required_libs.py --rocm-dir /opt/rocm --dest-dir llama.cpp/build/bin + echo "Copying required ROCm libraries to build directory..." + cp -v /opt/rocm/lib/libhipblas.so* "$build_bin_path/" 2>/dev/null || echo "libhipblas.so* not found" + cp -v /opt/rocm/lib/librocblas.so* "$build_bin_path/" 2>/dev/null || echo "librocblas.so* not found" + cp -v /opt/rocm/lib/libamdhip64.so* "$build_bin_path/" 2>/dev/null || echo "libamdhip64.so* not found" + cp -v /opt/rocm/lib/librocsolver.so* "$build_bin_path/" 2>/dev/null || echo "librocsolver.so* not found" + cp -v /opt/rocm/lib/libroctx64.so* "$build_bin_path/" 2>/dev/null || echo "libroctx64.so* not found" + cp -v /opt/rocm/lib/libhipblaslt.so* "$build_bin_path/" 2>/dev/null || echo "libhipblaslt.so* not found" + cp -v /opt/rocm/lib/libtensilelite-host.so* "$build_bin_path/" 2>/dev/null || echo "libtensilelite-host.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_liblzma.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_liblzma.so* not found" + cp -v /opt/rocm/lib/librocprofiler-register.so* "$build_bin_path/" 2>/dev/null || echo "librocprofiler-register.so* not found" + cp -v /opt/rocm/lib/libamd_comgr.so* "$build_bin_path/" 2>/dev/null || echo "libamd_comgr.so* not found" + cp -v /opt/rocm/lib/libamd_comgr_loader.so* "$build_bin_path/" 2>/dev/null || echo "libamd_comgr_loader.so* not found" + cp -v /opt/rocm/lib/libhsa-runtime64.so* "$build_bin_path/" 2>/dev/null || echo "libhsa-runtime64.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_numa.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_numa.so* not found" + cp -v /opt/rocm/lib/librocroller.so* "$build_bin_path/" 2>/dev/null || echo "librocroller.so* not found" + cp -v /opt/rocm/lib/liborigami.so* "$build_bin_path/" 2>/dev/null || echo "liborigami.so* not found" + cp -v /opt/rocm/lib/librocm_kpack.so* "$build_bin_path/" 2>/dev/null || echo "librocm_kpack.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_z.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_z.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_zstd.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_zstd.so* not found" + cp -v /opt/rocm/lib/llvm/lib/libLLVM.so* "$build_bin_path/" 2>/dev/null || echo "libLLVM.so* not found" + cp -v /opt/rocm/lib/llvm/lib/libclang-cpp.so* "$build_bin_path/" 2>/dev/null || echo "libclang-cpp.so* not found" + cp -v /opt/rocm/lib/llvm/lib/libomp.so* "$build_bin_path/" 2>/dev/null || echo "libomp.so* not found" + + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_elf.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_elf.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_drm.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_drm.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_drm_amdgpu.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_drm_amdgpu.so* not found" + cp -v /opt/rocm/lib/rocm_sysdeps/lib/librocm_sysdeps_bz2.so* "$build_bin_path/" 2>/dev/null || echo "librocm_sysdeps_bz2.so* not found" + + echo "Finished copying required ROCm libraries" + + - name: Set RPATH for portable distribution + run: | + sudo apt-get install -y patchelf + cd llama.cpp/build/bin + # Set RPATH to $ORIGIN so all libraries (including the comgr stub loader) find deps locally + for file in *.so* llama-*; do + [ -f "$file" ] && [ ! -L "$file" ] && patchelf --set-rpath '$ORIGIN' "$file" 2>/dev/null || true + done + + - name: List build artifacts (including ROCm files) + run: | + cd llama.cpp/build/bin + echo "Final build artifacts (including ROCm library files):" + ls -la + + - name: Upload build artifacts + id: upload-artifacts + uses: actions/upload-artifact@v4 + with: + name: llama-ubuntu-rocm-${{ matrix.gfx_target }}-x64 + path: llama.cpp/build/bin/ + retention-days: 30 + + - name: Set job outputs + id: set-outputs + run: | + echo "Setting job outputs..." + echo "Current target: ${{ matrix.gfx_target }}" + echo "DETECTED_ROCM_VERSION: ${DETECTED_ROCM_VERSION}" + echo "ROCM_VERSION: ${{ env.ROCM_VERSION }}" + echo "LLAMACPP_COMMIT_HASH: ${LLAMACPP_COMMIT_HASH}" + + rocm_version="${DETECTED_ROCM_VERSION:-${{ env.ROCM_VERSION }}}" + echo "rocm_version=$rocm_version" >> $GITHUB_OUTPUT + echo "llamacpp_commit_hash=${LLAMACPP_COMMIT_HASH}" >> $GITHUB_OUTPUT + + echo "Final rocm_version: $rocm_version" + echo "Final llamacpp_commit_hash: ${LLAMACPP_COMMIT_HASH}" + + - name: Create release summary + run: | + cat > "build-summary-ubuntu-${{ matrix.gfx_target }}.md" << 'EOF' + # Llama.cpp + ROCm Build Summary (Ubuntu - ${{ matrix.gfx_target }}) + + ## Build Configuration + - **GPU Target**: ${{ matrix.gfx_target }} + - **ROCm Version**: ${{ env.DETECTED_ROCM_VERSION || env.ROCM_VERSION }} + - **llama.cpp Commit Hash**: ${{ env.LLAMACPP_COMMIT_HASH }} + - **Build Type**: Release + - **Platform**: Ubuntu + + ## Build Artifacts + The compiled binaries are available in the build artifacts. + EOF + + - name: Upload build summary + id: upload-summary + uses: actions/upload-artifact@v4 + with: + name: build-summary-ubuntu-${{ matrix.gfx_target }} + path: build-summary-ubuntu-${{ matrix.gfx_target }}.md + retention-days: 30 + + - name: Clean up build artifacts and temporary files + if: always() + run: | + echo "Starting cleanup process..." + + # Remove ROCm tarball + if [ -f "rocm.tar.gz" ]; then + echo "Removing rocm.tar.gz..." + rm -f rocm.tar.gz + fi + + # Remove ROCm installation directory + if [ -d "/opt/rocm" ]; then + echo "Removing /opt/rocm directory..." + sudo rm -rf /opt/rocm + fi + + # Remove llama.cpp source directory (keeping only the build artifacts) + if [ -d "llama.cpp" ]; then + echo "Removing llama.cpp source directory..." + rm -rf llama.cpp + fi + + # Clean up any temporary build files + if [ -f "build-summary-ubuntu-${{ matrix.gfx_target }}.md" ]; then + echo "Removing build-summary-ubuntu-${{ matrix.gfx_target }}.md..." + rm -f "build-summary-ubuntu-${{ matrix.gfx_target }}.md" + fi + + echo "Cleanup completed successfully" + + build-ubuntu-vulkan: + # Linux Vulkan build of the same commit, packaged next to the ROCm zips in the same release. + # Portable: CPU variants and backends are loaded at runtime, RPATH is $ORIGIN, and the system + # Vulkan loader and driver (Mesa RADV) are used. + runs-on: ubuntu-22.04 + if: github.event.inputs.vulkan != 'false' + outputs: + llamacpp_commit_hash: ${{ steps.set-outputs.outputs.llamacpp_commit_hash }} + + steps: + - name: Checkout source + uses: actions/checkout@v4 + with: + ref: ${{ env.BUILD_SHA != '' && env.BUILD_SHA || (env.LLAMACPP_VERSION == 'latest' && 'master' || env.LLAMACPP_VERSION) }} + path: llama.cpp + + - name: Install build dependencies + run: | + wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo apt-key add - + sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-jammy.list https://packages.lunarg.com/vulkan/lunarg-vulkan-jammy.list + sudo apt-get update -y + sudo apt-get install -y build-essential cmake ninja-build vulkan-sdk zip + + - name: Build Llama.cpp + Vulkan + run: | + cd llama.cpp + echo "LLAMACPP_COMMIT_HASH=$(git rev-parse --short=5 HEAD)" >> $GITHUB_ENV + git log --oneline -1 + cmake -B build -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_RPATH='$ORIGIN' \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DBUILD_SHARED_LIBS=ON \ + -DLLAMA_BUILD_TESTS=OFF \ + -DGGML_VULKAN=ON \ + -DGGML_RPC=ON \ + -DGGML_NATIVE=OFF \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DLLAMA_BUILD_BORINGSSL=ON + cmake --build build -j $(nproc) + cp LICENSE build/bin/ + ls -la build/bin + + - name: Upload build artifacts + uses: actions/upload-artifact@v4 + with: + name: llama-ubuntu-vulkan-x64 + path: llama.cpp/build/bin/ + retention-days: 30 + + - name: Set job outputs + id: set-outputs + run: echo "llamacpp_commit_hash=${LLAMACPP_COMMIT_HASH}" >> $GITHUB_OUTPUT + + test-stx-halo: + runs-on: ${{ matrix.runner }} + needs: [prepare-matrix, build-windows, build-ubuntu] + if: | + vars.STX_HALO_RUNNERS == 'true' && + (needs.build-windows.result == 'success' || needs.build-ubuntu.result == 'success') && + contains(github.event.inputs.gfx_target || 'gfx1151', 'gfx1151') + strategy: + matrix: + include: + - os: Windows + gfx_target: gfx1151 + artifact_prefix: windows + runner: ["stx-halo", "Windows"] + - os: Linux + gfx_target: gfx1151 + artifact_prefix: ubuntu + runner: ["stx-halo", "Linux"] + fail-fast: false + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Test llama.cpp build + uses: ./.github/actions/test-llamacpp-build + with: + os_type: ${{ matrix.os }} + gfx_target: ${{ matrix.gfx_target }} + artifact_name: llama-${{ matrix.artifact_prefix }}-rocm-${{ matrix.gfx_target }}-x64 + + create-release: + needs: [prepare-matrix, build-windows, build-ubuntu, build-ubuntu-vulkan, test-stx-halo] + runs-on: ubuntu-22.04 + permissions: + contents: write + pull-requests: write + if: | + always() && + (needs.build-windows.result == 'success' || needs.build-windows.result == 'skipped') && + (needs.build-ubuntu.result == 'success' || needs.build-ubuntu.result == 'skipped') && + (needs.build-ubuntu-vulkan.result == 'success' || needs.build-ubuntu-vulkan.result == 'skipped') && + (needs.build-windows.result == 'success' || needs.build-ubuntu.result == 'success') && + (needs.test-stx-halo.result == 'success' || needs.test-stx-halo.result == 'skipped') && + github.event_name != 'pull_request' && + (github.event_name == 'workflow_dispatch' && + (github.event.inputs.create_release == 'true' || github.event.inputs.create_release == null) || + github.event_name == 'push') + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Download all build artifacts + uses: actions/download-artifact@v4 + with: + path: ./all-artifacts + + - name: Generate release tag + id: generate-tag + env: + GITHUB_TOKEN: ${{ secrets.GH_TOKEN || github.token }} + run: | + # Get all existing release tags that match the pattern b#### + existing_tags=$(gh release list --limit 1000 --json tagName --jq '.[].tagName' | grep -E '^b[0-9]{4}$' | sort -V || echo "") + + if [ -z "$existing_tags" ]; then + # No existing sequential tags found, start with b1000 + next_number=1000 + echo "No existing sequential release tags found, starting with b1000" + else + # Find the highest existing number + highest_tag=$(echo "$existing_tags" | tail -n 1) + highest_number=$(echo "$highest_tag" | sed 's/^b//') + next_number=$((highest_number + 1)) + echo "Highest existing tag: $highest_tag (number: $highest_number)" + echo "Next number will be: $next_number" + fi + + # Format as b#### (4-digit number with leading zeros) + TAG=$(printf "b%04d" $next_number) + + echo "tag=${TAG}" >> $GITHUB_OUTPUT + echo "Generated release tag: ${TAG}" + + - name: Check if tag already exists + id: check-tag + run: | + TAG="${{ steps.generate-tag.outputs.tag }}" + + # Check if the tag already exists + if git ls-remote --tags origin "$TAG" | grep -q "$TAG"; then + echo "Tag $TAG already exists, skipping release creation" + echo "tag_exists=true" >> $GITHUB_OUTPUT + else + echo "Tag $TAG does not exist, proceeding with release creation" + echo "tag_exists=false" >> $GITHUB_OUTPUT + fi + + - name: Create archives for all target artifacts + if: steps.check-tag.outputs.tag_exists == 'false' + run: | + # Parse targets and operating systems from environment + targets="${{ env.GFX_TARGETS }}" + operating_systems="${{ env.OPERATING_SYSTEMS }}" + TAG="${{ steps.generate-tag.outputs.tag }}" + + echo "Processing targets: $targets" + echo "Processing operating systems: $operating_systems" + echo "Using release tag: $TAG" + + # Create individual archives for each target and OS combination + IFS=',' read -ra TARGET_ARRAY <<< "$targets" + IFS=',' read -ra OS_ARRAY <<< "$operating_systems" + + for os in "${OS_ARRAY[@]}"; do + os=$(echo "$os" | xargs) # trim whitespace + for target in "${TARGET_ARRAY[@]}"; do + target=$(echo "$target" | xargs) # trim whitespace + + echo "Processing OS: $os, target: $target" + + # Use artifact name to find the directory + artifact_name="llama-${os}-rocm-${target}-x64" + artifact_dir="./all-artifacts/${artifact_name}" + + # Create final archive with release tag + final_archive_name="llama-${TAG}-${os}-rocm-${target}-x64" + + if [ -d "$artifact_dir" ]; then + echo "Creating archive: ${final_archive_name}.zip" + cd "$artifact_dir" + zip -r "../../${final_archive_name}.zip" * + cd ../../ + else + echo "Warning: Artifact directory not found: $artifact_dir" + ls -la ./all-artifacts/ + fi + done + done + + if [ -d "./all-artifacts/llama-ubuntu-vulkan-x64" ]; then + echo "Creating archive: llama-${TAG}-ubuntu-vulkan-x64.zip" + (cd ./all-artifacts/llama-ubuntu-vulkan-x64 && zip -r "../../llama-${TAG}-ubuntu-vulkan-x64.zip" *) + fi + + echo "Created archives:" + ls -la *.zip + + - name: Create Release + if: steps.check-tag.outputs.tag_exists == 'false' + env: + GITHUB_TOKEN: ${{ secrets.GH_TOKEN || github.token }} + run: | + TAG="${{ steps.generate-tag.outputs.tag }}" + # Get ROCm version from either Windows or Ubuntu build (whichever succeeded) + ROCM_VERSION="${{ needs.build-windows.outputs.rocm_version || needs.build-ubuntu.outputs.rocm_version }}" + # Get llama.cpp commit hash from either Windows or Ubuntu build (whichever succeeded) + LLAMACPP_COMMIT_HASH="${{ needs.build-windows.outputs.llamacpp_commit_hash || needs.build-ubuntu.outputs.llamacpp_commit_hash }}" + targets="${{ env.GFX_TARGETS }}" + operating_systems="${{ env.OPERATING_SYSTEMS }}" + + echo "Creating release with tag: $TAG" + echo "GPU Targets: $targets" + echo "Operating Systems: $operating_systems" + echo "ROCm Version: $ROCM_VERSION" + echo "Llama.cpp Commit: $LLAMACPP_COMMIT_HASH" + + # Verify archives exist + ls -la *.zip + + # Prepare upload files list + upload_files="" + IFS=',' read -ra TARGET_ARRAY <<< "$targets" + IFS=',' read -ra OS_ARRAY <<< "$operating_systems" + + for os in "${OS_ARRAY[@]}"; do + os=$(echo "$os" | xargs) # trim whitespace + for target in "${TARGET_ARRAY[@]}"; do + target=$(echo "$target" | xargs) # trim whitespace + + final_archive_name="llama-${TAG}-${os}-rocm-${target}-x64" + if [ -f "${final_archive_name}.zip" ]; then + upload_files="${upload_files} ${final_archive_name}.zip" + fi + done + done + + vulkan_line="" + if [ -f "llama-${TAG}-ubuntu-vulkan-x64.zip" ]; then + upload_files="${upload_files} llama-${TAG}-ubuntu-vulkan-x64.zip" + vulkan_line="**Vulkan**: ubuntu x64" + fi + + echo "Files to upload: $upload_files" + + # Create release with GitHub CLI + gh release create "$TAG" \ + --title "$TAG" \ + --notes "**Build Number**: $TAG + **Operating System(s)**: $operating_systems + **GPU Target(s)**: $targets + **ROCm Version**: $ROCM_VERSION + $vulkan_line + **Llama.cpp Commit Hash**: $LLAMACPP_COMMIT_HASH + **Build Date**: $(date -u '+%Y-%m-%d %H:%M:%S UTC') + + This release includes compiled llama.cpp binaries with ROCm support for AMD Strix Halo (gfx1151), with all essential ROCm runtime libraries included." \ + $upload_files 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 daa08b1bf949..1a6d11858061 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-sycl.yml b/.github/workflows/build-sycl.yml index ddc4e7f2b1a1..d26635dd7c7b 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/.github/workflows/ci-self-hosted-cpu.yml b/.github/workflows/ci-self-hosted-cpu.yml index 18bc9f144205..61b1400109b5 100644 --- a/.github/workflows/ci-self-hosted-cpu.yml +++ b/.github/workflows/ci-self-hosted-cpu.yml @@ -17,6 +17,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-cpu.yml', diff --git a/.github/workflows/ci-self-hosted-cuda.yml b/.github/workflows/ci-self-hosted-cuda.yml index 5951b52d3781..9caada9811b1 100644 --- a/.github/workflows/ci-self-hosted-cuda.yml +++ b/.github/workflows/ci-self-hosted-cuda.yml @@ -19,6 +19,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-cuda.yml', diff --git a/.github/workflows/ci-self-hosted-kleidiai.yml b/.github/workflows/ci-self-hosted-kleidiai.yml index c955aa0022d4..44ee6bec6d36 100644 --- a/.github/workflows/ci-self-hosted-kleidiai.yml +++ b/.github/workflows/ci-self-hosted-kleidiai.yml @@ -17,6 +17,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-kleidiai.yml', diff --git a/.github/workflows/ci-self-hosted-metal.yml b/.github/workflows/ci-self-hosted-metal.yml index 7af30e294a93..3033099f036d 100644 --- a/.github/workflows/ci-self-hosted-metal.yml +++ b/.github/workflows/ci-self-hosted-metal.yml @@ -20,6 +20,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-metal.yml', diff --git a/.github/workflows/ci-self-hosted-openvino.yml b/.github/workflows/ci-self-hosted-openvino.yml index e0947c46e045..532adb6b3b32 100644 --- a/.github/workflows/ci-self-hosted-openvino.yml +++ b/.github/workflows/ci-self-hosted-openvino.yml @@ -17,6 +17,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-openvino.yml', diff --git a/.github/workflows/ci-self-hosted-vulkan.yml b/.github/workflows/ci-self-hosted-vulkan.yml index ffed4b09b0c8..95400116532d 100644 --- a/.github/workflows/ci-self-hosted-vulkan.yml +++ b/.github/workflows/ci-self-hosted-vulkan.yml @@ -19,6 +19,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-vulkan.yml', diff --git a/.github/workflows/ci-self-hosted-webgpu.yml b/.github/workflows/ci-self-hosted-webgpu.yml index a6a36a8ebab3..8835428a2eb3 100644 --- a/.github/workflows/ci-self-hosted-webgpu.yml +++ b/.github/workflows/ci-self-hosted-webgpu.yml @@ -18,6 +18,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/ci-self-hosted-webgpu.yml', diff --git a/.github/workflows/fusion.yml b/.github/workflows/fusion.yml index 7c8596467aab..ae60a33e8294 100644 --- a/.github/workflows/fusion.yml +++ b/.github/workflows/fusion.yml @@ -14,7 +14,10 @@ on: 'src/models/**' ] + # disabled for pull requests: the metal job runs on a self-hosted macOS runner and + # hangs for the full 24h timeout when the runner is unavailable, blocking PRs pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/fusion.yml', diff --git a/.github/workflows/hip-quality-check.yml b/.github/workflows/hip-quality-check.yml index ee4e746f2f10..22019f04b264 100644 --- a/.github/workflows/hip-quality-check.yml +++ b/.github/workflows/hip-quality-check.yml @@ -71,7 +71,7 @@ jobs: run: | cmake -B build -S . \ -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \ - -DGPU_TARGETS=gfx942 \ + -DGPU_TARGETS=gfx1151 \ -DGGML_HIP=ON \ -DGGML_HIP_EXPORT_METRICS=Off \ -DCMAKE_HIP_FLAGS="-Werror -Wno-tautological-compare" \ @@ -84,7 +84,7 @@ jobs: run: | cmake -B build -S . \ -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \ - -DGPU_TARGETS=gfx908 \ + -DGPU_TARGETS=gfx1151 \ -DGGML_HIP=ON \ -DGGML_HIP_EXPORT_METRICS=On \ -DCMAKE_HIP_FLAGS="" \ diff --git a/.github/workflows/nightly-failure-alert.yml b/.github/workflows/nightly-failure-alert.yml new file mode 100644 index 000000000000..e26dc76a3ed3 --- /dev/null +++ b/.github/workflows/nightly-failure-alert.yml @@ -0,0 +1,66 @@ +# Same as lemonade-sdk/llamacpp-rocm (MIT). +name: Nightly Failure Alert +on: + workflow_run: + workflows: ["Build Llama.cpp + ROCm"] + types: [completed] + +env: + # GitHub user(s) to alert, space-separated + ALERT_USERS: dzannotti + ALERT_LABEL: nightly-failure + +jobs: + alert: + if: | + github.event.workflow_run.event == 'push' && + github.event.workflow_run.conclusion != 'success' + runs-on: ubuntu-22.04 + permissions: + issues: write + actions: read + steps: + - name: Open or update failure issue + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GH_REPO: ${{ github.repository }} + RUN_ID: ${{ github.event.workflow_run.id }} + RUN_URL: ${{ github.event.workflow_run.html_url }} + CONCLUSION: ${{ github.event.workflow_run.conclusion }} + RUN_DATE: ${{ github.event.workflow_run.run_started_at }} + run: | + failed_jobs=$(gh api --paginate "repos/$GH_REPO/actions/runs/$RUN_ID/jobs?per_page=100" \ + --jq '.jobs[] | select(.conclusion != "success" and .conclusion != "skipped") | "- [\(.name)](\(.html_url)) — \(.conclusion)"') + [ -z "$failed_jobs" ] && failed_jobs="- (no individual job failures reported)" + + mentions=$(for u in $ALERT_USERS; do printf '@%s ' "$u"; done) + assignees=$(echo $ALERT_USERS | tr ' ' ',') + + body=$(cat <> $GITHUB_OUTPUT + echo "Final release tag: $release_tag" + + - name: Set matrix + id: set-matrix + run: | + targets="${{ env.GFX_TARGETS }}" + operating_systems="${{ env.OPERATING_SYSTEMS }}" + + echo "Input targets: $targets" + echo "Input operating systems: $operating_systems" + # Convert targets to JSON array + matrix_targets=$(echo "$targets" \ + | tr ',' '\n' \ + | sed 's/^ *//;s/ *$//' \ + | sed 's/^"//;s/"$//' \ + | jq -R . \ + | jq -s '{gfx_target: .}' \ + | jq -c) + # Check which operating systems to test + should_test_windows="false" + should_test_ubuntu="false" + + if [[ "$operating_systems" == *"windows"* ]]; then + should_test_windows="true" + echo "windows_matrix=$matrix_targets" >> $GITHUB_OUTPUT + fi + + if [[ "$operating_systems" == *"ubuntu"* ]]; then + should_test_ubuntu="true" + echo "ubuntu_matrix=$matrix_targets" >> $GITHUB_OUTPUT + fi + + echo "should_test_windows=$should_test_windows" >> $GITHUB_OUTPUT + echo "should_test_ubuntu=$should_test_ubuntu" >> $GITHUB_OUTPUT + + echo "Windows test: $should_test_windows" + echo "Ubuntu test: $should_test_ubuntu" + echo "Generated matrix: $matrix_targets" + test-windows: + runs-on: [stx-halo, Windows] + needs: prepare-matrix + if: vars.STX_HALO_RUNNERS == 'true' && needs.prepare-matrix.outputs.should_test_windows == 'true' + strategy: + matrix: ${{fromJson(needs.prepare-matrix.outputs.windows_matrix)}} + fail-fast: false + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Download release artifacts + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + $releaseTag = "${{ needs.prepare-matrix.outputs.latest_release_tag }}" + $target = "${{ matrix.gfx_target }}" + $assetName = "llama-${releaseTag}-windows-rocm-${target}-x64.zip" + $repo = "${{ github.repository }}" + + Write-Host "Downloading release asset: $assetName" + Write-Host "From release: $releaseTag" + + # Get release information using GitHub REST API + $headers = @{ "Authorization" = "token $env:GITHUB_TOKEN" } + $releaseUrl = "https://api.github.com/repos/$repo/releases/tags/$releaseTag" + + try { + $release = Invoke-RestMethod -Uri $releaseUrl -Headers $headers + $asset = $release.assets | Where-Object { $_.name -eq $assetName } + + if (-not $asset) { + Write-Error "Asset '$assetName' not found in release '$releaseTag'" + exit 1 + } + + Write-Host "Found asset: $($asset.name) (Size: $($asset.size) bytes)" + + # Download the specific asset + Invoke-WebRequest -Uri $asset.browser_download_url -OutFile $assetName -Headers $headers + Write-Host "Downloaded $assetName successfully" + } + catch { + Write-Error "Failed to download release asset: $_" + exit 1 + } + + # Extract the downloaded zip file + Write-Host "Extracting $assetName..." + Expand-Archive -Path $assetName -DestinationPath "llama-binaries" -Force + + # List contents + Write-Host "Contents of llama-binaries:" + Get-ChildItem -Recurse "llama-binaries" | Format-Table Name, Length + + - name: Download test model + run: | + $modelUrl = "https://huggingface.co/unsloth/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_0.gguf?download=true" + $modelPath = "Qwen3-0.6B-Q4_0.gguf" + + Write-Host "Downloading test model from: $modelUrl" + curl.exe -L -o $modelPath $modelUrl + + # Verify download + if (Test-Path $modelPath) { + $fileSize = (Get-Item $modelPath).Length + Write-Host "Model downloaded successfully. Size: $fileSize bytes" + } else { + Write-Error "Failed to download model" + exit 1 + } + + - name: Run llama-cli test + run: | + $llamaCliPath = "llama-binaries\llama-cli.exe" + $modelPath = "Qwen3-0.6B-Q4_0.gguf" + + # Verify llama-cli exists + if (-not (Test-Path $llamaCliPath)) { + Write-Error "llama-cli.exe not found at: $llamaCliPath" + Write-Host "Available files in llama-binaries:" + Get-ChildItem -Recurse "llama-binaries" + exit 1 + } + + Write-Host "Running llama-cli test..." + Write-Host "Command: $llamaCliPath -m `"$modelPath`" -ngl 99 -p `"I believe the meaning of life is`" -st" + + # Run the command and capture output + try { + $process = Start-Process -FilePath $llamaCliPath -ArgumentList @("-m", "`"$modelPath`"", "-ngl", "99", "-p", "`"I believe the meaning of life is`"", "-st") -NoNewWindow -Wait -PassThru -RedirectStandardOutput "stdout.txt" -RedirectStandardError "stderr.txt" + + # Wait for the process to complete + $process.WaitForExit() + + # Read both stdout and stderr + $stdout = Get-Content "stdout.txt" -Raw -ErrorAction SilentlyContinue + $stderr = Get-Content "stderr.txt" -Raw -ErrorAction SilentlyContinue + + # Combine output (stderr often contains important CUDA/GPU info) + $output = $stdout + $stderr + + # Clean up temp files + Remove-Item "stdout.txt" -ErrorAction SilentlyContinue + Remove-Item "stderr.txt" -ErrorAction SilentlyContinue + + Write-Host "Process exit code: $($process.ExitCode)" + } + catch { + Write-Error "Failed to run llama-cli: $_" + exit 1 + } + + Write-Host "=== LLAMA-CLI OUTPUT ===" + Write-Host $output + Write-Host "=== END OUTPUT ===" + + # Check if the process exited successfully + if ($process.ExitCode -ne 0) { + Write-Host "❌ llama-cli exited with error code: $($process.ExitCode)" + exit 1 + } + + # Check for expected outputs + $foundGpuOffload = $output -match "offloaded 29/29 layers to GPU" + $foundThinkTag = $output -match "" + + Write-Host "=== TEST RESULTS ===" + Write-Host "Expected 'offloaded 29/29 layers to GPU': $(if ($foundGpuOffload) { 'FOUND' } else { 'NOT FOUND' })" + Write-Host "Expected '' tag: $(if ($foundThinkTag) { 'FOUND' } else { 'NOT FOUND' })" + + if ($foundGpuOffload -and $foundThinkTag) { + Write-Host "✅ Test PASSED - Both expected outputs found" + } else { + Write-Host "❌ Test FAILED - Missing expected outputs" + exit 1 + } + test-ubuntu: + runs-on: [stx-halo, Linux] + needs: prepare-matrix + if: vars.STX_HALO_RUNNERS == 'true' && needs.prepare-matrix.outputs.should_test_ubuntu == 'true' + strategy: + matrix: ${{fromJson(needs.prepare-matrix.outputs.ubuntu_matrix)}} + fail-fast: false + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Install jq + run: | + if ! command -v jq &> /dev/null; then + echo "Installing jq..." + # Download and install jq binary directly + mkdir -p ~/bin + curl -L -o ~/bin/jq https://github.com/jqlang/jq/releases/download/jq-1.7.1/jq-linux-amd64 + chmod +x ~/bin/jq + export PATH="$HOME/bin:$PATH" + echo "PATH=$HOME/bin:$PATH" >> $GITHUB_ENV + echo "jq installed to ~/bin/jq" + else + echo "jq is already installed" + fi + + - name: Download release artifacts + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + release_tag="${{ needs.prepare-matrix.outputs.latest_release_tag }}" + target="${{ matrix.gfx_target }}" + asset_name="llama-${release_tag}-ubuntu-rocm-${target}-x64.zip" + repo="${{ github.repository }}" + + echo "Downloading release asset: $asset_name" + echo "From release: $release_tag" + + # Get release information using GitHub REST API + release_url="https://api.github.com/repos/$repo/releases/tags/$release_tag" + + # Get release data and extract asset download URL + release_data=$(curl -s -H "Authorization: token $GITHUB_TOKEN" "$release_url") + + if [ $? -ne 0 ]; then + echo "Error: Failed to fetch release information" + exit 1 + fi + + # Extract the browser_download_url for the specific asset + download_url=$(echo "$release_data" | jq -r ".assets[] | select(.name == \"$asset_name\") | .browser_download_url") + + if [ "$download_url" = "null" ] || [ -z "$download_url" ]; then + echo "Error: Asset '$asset_name' not found in release '$release_tag'" + echo "Available assets:" + echo "$release_data" | jq -r '.assets[].name' + exit 1 + fi + + asset_size=$(echo "$release_data" | jq -r ".assets[] | select(.name == \"$asset_name\") | .size") + echo "Found asset: $asset_name (Size: $asset_size bytes)" + + # Download the specific asset + curl -L -H "Authorization: token $GITHUB_TOKEN" -o "$asset_name" "$download_url" + + if [ $? -eq 0 ]; then + echo "Downloaded $asset_name successfully" + else + echo "Error: Failed to download $asset_name" + exit 1 + fi + + # Extract the downloaded zip file + echo "Extracting $asset_name..." + unzip -q "$asset_name" -d llama-binaries + + # List contents + echo "Contents of llama-binaries:" + find llama-binaries -type f -exec ls -la {} \; + + - name: Download test model + run: | + model_url="https://huggingface.co/unsloth/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_0.gguf?download=true" + model_path="Qwen3-0.6B-Q4_0.gguf" + + echo "Downloading test model from: $model_url" + curl -L -o "$model_path" "$model_url" + + # Verify download + if [ -f "$model_path" ]; then + file_size=$(stat -c%s "$model_path") + echo "Model downloaded successfully. Size: $file_size bytes" + else + echo "Failed to download model" + exit 1 + fi + + - name: Set up library path + run: | + # Add the llama-binaries directory to LD_LIBRARY_PATH + echo "LD_LIBRARY_PATH=$(pwd)/llama-binaries:$LD_LIBRARY_PATH" >> $GITHUB_ENV + + - name: Run llama-cli test + run: | + llama_cli_path="./llama-binaries/llama-cli" + model_path="Qwen3-0.6B-Q4_0.gguf" + + # Make llama-cli executable + chmod +x "$llama_cli_path" + + # Verify llama-cli exists + if [ ! -f "$llama_cli_path" ]; then + echo "llama-cli not found at: $llama_cli_path" + echo "Available files in llama-binaries:" + find llama-binaries -type f + exit 1 + fi + + echo "Running llama-cli test..." + echo "Command: $llama_cli_path -m \"$model_path\" -ngl 99 -p \"I believe the meaning of life is\" -st" + + # Run the command and capture output + output=$("$llama_cli_path" -m "$model_path" -ngl 99 -p "I believe the meaning of life is" -st 2>&1) + + echo "=== LLAMA-CLI OUTPUT ===" + echo "$output" + echo "=== END OUTPUT ===" + + # Check for expected outputs + if echo "$output" | grep -q "offloaded 29/29 layers to GPU"; then + found_gpu_offload=true + else + found_gpu_offload=false + fi + + if echo "$output" | grep -q ""; then + found_think_tag=true + else + found_think_tag=false + fi + + echo "=== TEST RESULTS ===" + echo "Expected 'offloaded 29/29 layers to GPU': $(if [ "$found_gpu_offload" = true ]; then echo 'FOUND'; else echo 'NOT FOUND'; fi)" + echo "Expected '' tag: $(if [ "$found_think_tag" = true ]; then echo 'FOUND'; else echo 'NOT FOUND'; fi)" + + if [ "$found_gpu_offload" = true ] && [ "$found_think_tag" = true ]; then + echo "✅ Test PASSED - Both expected outputs found" + else + echo "❌ Test FAILED - Missing expected outputs" + exit 1 + fi 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..dfc78d252e12 100644 --- a/README.md +++ b/README.md @@ -1,91 +1,124 @@ -# 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 for strix around. And help the community with a single fast +llama.cpp fork instead of many competing ones. + +## 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) -- 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 +cmake -B build -DGGML_VULKAN=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build --config Release -j +``` -Once installed: +**ROCm / HIP** (needs ROCm installed; build for `gfx1151` explicitly) ```sh -# Download and run a model directly from Hugging Face -llama cli -hf ggml-org/Qwen3.5-0.8B-GGUF +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 +``` + +Then: -# Launch OpenAI-compatible API server -llama serve -hf ggml-org/Qwen3.5-0.8B-GGUF +```sh +# chat, pulling the model straight from Hugging Face +./build/bin/llama-cli -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: + +**Strix Halo only** (lives here, measured on `gfx1151`) + +| 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` | +| 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 +133,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 b70da5e9e6e0..cb344f684579 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); } @@ -2707,15 +2730,24 @@ 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"); } } ).set_env("LLAMA_ARG_LAZY_MODE")); + add_opt(common_arg( + {"--ple"}, "FNAME", + "GGUF with per-layer embedding (PLE) tables to use instead of the tables in the model file; each tensor in it replaces the lookup table of the same name in the model and must have the same shape (default: unused)", + [](common_params & params, const std::string & value) { + params.model.ple = value; + } + ).set_env("LLAMA_ARG_PLE")); add_opt(common_arg( {"--numa"}, "TYPE", "attempt optimizations that help on some NUMA systems\n" @@ -2753,6 +2785,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", @@ -3713,6 +3752,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)", @@ -3721,13 +3804,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"}, @@ -4143,6 +4281,30 @@ 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-max-slots"}, "N", + string_format("draft only while at most N slots are generating; with more, the batch of all slots already " + "keeps the GPU busy and verify rows cost more than they save (default: %d, 0 = no limit)", + params.speculative.draft.max_slots), + [](common_params & params, int value) { + params.speculative.draft.max_slots = std::max(0, value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_MAX_SLOTS")); + 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}).set_env("LLAMA_ARG_SPEC_DRAFT_MTP_VOCAB")); 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), @@ -4150,6 +4312,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_min = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); + add_opt(common_arg( + {"--spec-n-rs-seq"}, "N", + "cap the recurrent-state rollback slots for models that support them (default: -1, follow --spec-draft-n-max).\n" + "the recurrent cache holds (1 + N) copies of each sequence's state, so on a large linear-attention\n" + "model each slot costs hundreds of MiB. a rollback deeper than N stays correct, it just uses the\n" + "slower host checkpoint. 0 disables the fast path entirely", + [](common_params & params, int value) { + params.speculative.n_rs_seq_max = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_N_RS_SEQ")); add_opt(common_arg( {"--spec-synth-len"}, "L", "target mean synthetic acceptance length, including the target token (benchmarking only)", @@ -4242,6 +4414,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", @@ -4756,6 +5033,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 203d1b4e11b8..d89e73d59860 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 @@ -140,6 +141,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 6d57b542177e..7c8f69583120 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1219,7 +1219,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; @@ -1608,11 +1609,11 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; mparams.lazy_mode = params.lazy_mode; + mparams.path_ple = params.model.ple.empty() ? nullptr : params.model.ple.c_str(); mparams.tensor_split = params.tensor_split; 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 0a09fed7f60a..eaac0edaf6fe 100644 --- a/common/common.h +++ b/common/common.h @@ -288,12 +288,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; @@ -307,6 +326,7 @@ struct common_params_model { std::string hf_repo = ""; // HF repo std::string hf_file = ""; // HF file std::string docker_repo = ""; // Docker repo + std::string ple = ""; // GGUF with per-layer embedding tables used instead of the model's own std::string get_name() const { if (!hf_repo.empty()) { @@ -343,6 +363,14 @@ 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; + // server: no new drafts while more than this many slots are generating (0 = no limit) + int32_t max_slots = 0; + + // 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; @@ -369,6 +397,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 }; @@ -385,6 +427,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(); } @@ -393,12 +437,28 @@ struct common_params_speculative { return synth_len != -1.0 || !synth_rates.empty(); } + bool has_prefill() const { + return prefill.enabled && !prefill.model.empty(); + } + + // cap on the rollback snapshot slots, -1 = follow draft.n_max (see need_n_rs_seq) + int32_t n_rs_seq_max = -1; + 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; }); - return needs_rs_seq ? draft.n_max : 0u; + if (!needs_rs_seq) { + return 0u; + } + + // the recurrent cache holds (1 + n_rs_seq) copies of every sequence's state, which on a + // large linear-attention model is GiBs. Capping trades speed for memory: a rollback + // deeper than n_rs_seq is still correct, it just falls back to the host checkpoint. + const int32_t n = n_rs_seq_max >= 0 ? std::min(n_rs_seq_max, draft.n_max) : draft.n_max; + + return (uint32_t) std::max(0, n); } }; @@ -583,7 +643,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 208f551f1b2b..41d545bd1a5f 100644 --- a/common/parsers/qwen3-coder.cpp +++ b/common/parsers/qwen3-coder.cpp @@ -134,6 +134,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 27a752cc8df6..2c2a573a1623 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 82e9e92238b9..ecbe365e71cc 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; @@ -201,6 +276,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { zeros.assign(llama_model_n_embd_inp(llama_get_model(ctx_dft)), 0.0f); 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, @@ -348,6 +424,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; @@ -377,7 +455,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--; @@ -488,6 +566,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; @@ -792,6 +871,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); @@ -823,7 +904,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; @@ -1007,6 +1088,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"); @@ -1078,6 +1160,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; } @@ -1190,6 +1274,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) { @@ -1200,7 +1322,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.size(); @@ -1378,6 +1500,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, @@ -1417,7 +1540,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) { @@ -1455,6 +1582,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; @@ -1648,6 +1777,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]); @@ -1677,7 +1810,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; @@ -2526,7 +2659,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 @@ -2908,6 +3042,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 5561481e7716..de8517a2f95c 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 64d606176eb2..a968b71ed7f1 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..41b54b5221b9 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Iterable, cast +from typing import Callable, Iterable, cast import torch from torch import Tensor @@ -26,8 +26,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 +61,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 +105,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/build.md b/docs/build.md index e46c9c9383d7..c380e8e3fac2 100644 --- a/docs/build.md +++ b/docs/build.md @@ -374,10 +374,10 @@ This provides GPU acceleration on HIP-supported AMD GPUs. Make sure to have ROCm installed. You can download it from your Linux distro's package manager or from here: [ROCm Quick Start (Linux)](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/tutorial/quick-start.html#rocm-install-quick). -- Using `CMake` for Linux (assuming a gfx1030-compatible AMD GPU): +- Using `CMake` for Linux (Strix Halo, `gfx1151`): ```bash HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \ - cmake -S . -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1030 -DCMAKE_BUILD_TYPE=Release \ + cmake -S . -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1151 -DCMAKE_BUILD_TYPE=Release \ && cmake --build build --config Release -- -j 16 ``` @@ -394,22 +394,20 @@ You can download it from your Linux distro's package manager or from here: [ROCm ```bash HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -p)" \ HIP_DEVICE_LIB_PATH= \ - cmake -S . -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1030 -DCMAKE_BUILD_TYPE=Release \ + cmake -S . -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1151 -DCMAKE_BUILD_TYPE=Release \ && cmake --build build -- -j 16 ``` -- Using `CMake` for Windows (using x64 Native Tools Command Prompt for VS, and assuming a gfx1100-compatible AMD GPU): +- Using `CMake` for Windows (using x64 Native Tools Command Prompt for VS, Strix Halo, `gfx1151`): ```bash set PATH=%HIP_PATH%\bin;%PATH% - cmake -S . -B build -G Ninja -DGPU_TARGETS=gfx1100 -DGGML_HIP=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release + cmake -S . -B build -G Ninja -DGPU_TARGETS=gfx1151 -DGGML_HIP=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release cmake --build build ``` - If necessary, adapt `GPU_TARGETS` to the GPU arch you want to compile for. The above example uses `gfx1100` that corresponds to Radeon RX 7900XTX/XT/GRE. You can find a list of targets [here](https://llvm.org/docs/AMDGPUUsage.html#processors) - Find your gpu version string by matching the most significant version information from `rocminfo | grep gfx | head -1 | awk '{print $2}'` with the list of processors, e.g. `gfx1035` maps to `gfx1030`. + `gfx1151` is the Strix Halo (Ryzen AI Max+) integrated GPU, the only target this fork supports. You can confirm it with `rocminfo | grep gfx | head -1 | awk '{print $2}'`. The environment variable [`HIP_VISIBLE_DEVICES`](https://rocm.docs.amd.com/en/latest/understand/gpu_isolation.html#hip-visible-devices) can be used to specify which GPU(s) will be used. -If your GPU is not officially supported you can use the environment variable [`HSA_OVERRIDE_GFX_VERSION`] set to a similar GPU, for example 10.3.0 on RDNA2 (e.g. gfx1030, gfx1031, or gfx1035) or 11.0.0 on RDNA3. Note that [`HSA_OVERRIDE_GFX_VERSION`] is [not supported on Windows](https://github.com/ROCm/ROCm/issues/2654) ### Unified Memory 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..83ee3c3821f7 --- /dev/null +++ b/docs/development/vulkan-sparse-fa-deferred.md @@ -0,0 +1,106 @@ +# 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. + +--- + +## 2026-09-29 sync: what changed in this directory + +This sync did not revert `ggml/src/ggml-vulkan/` wholesale. Each side was measured on a Radeon 8060S +(RADV Mesa 26.0.8), and the faster side was kept per area: + +- Matrix multiplication is upstream's (#25773 tile map, #27471, #27952 int8 coopmat1, #28415 IQ4_XS, + #28822 IQ3_S, #25483). Strix's legacy pipeline arrays and its env-gated mul_mat_id experiments + (`GGML_VK_MMID_*`, `GGML_VK_DENSE_F16B`, `GGML_VK_MMID_SCALE_EPILOGUE`) are gone. test-backend-ops perf, same + session: MUL_MAT geomean 1.10x, MUL_MAT_ID geomean 1.14x in upstream's favour. Strix was faster only for f16 + expert weights at n=128-256 and for iq3_xxs MoE; these are not ported yet. +- ROCmFPx takes upstream's per-type (LUT) matmul path and the q8_1 MMQ list. +- Everything else keeps the strix kernels: flash attention (dynamic KV, dequant-once, contiguize, gather/union, + DeepSeek V4 sparse), transposed CONCAT, lightning indexer, mat-vec column chunking, submit bounding. Per-op + timings at pp2048 showed upstream CONCAT 12-25x slower on the delta-net models, and FA/SSM_CONV/GDN 2-3x + slower on Signal-3.8-27B. +- The DeepSeek V4 hyper-connection ops are upstream's. The merged qwen4exp graph emits the gated and identity + forms, which the strix kernels could not run. +- `f1e44dcc1` (NV queue-submit workaround) is back. It is inert on RADV. + +Upstream sparse Flash Attention (#28105) is still deferred for the reasons above. `USE_SPARSE` / bit 16 is +not reintroduced, and `flash_attn_sparse_compact.comp` stays deleted. + +--- + +## 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 | +| `sync/halo-master-2026-09-29` | the 2026-09-29 sync (this section) | + +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/rocm-releases.md b/docs/rocm-releases.md new file mode 100644 index 000000000000..651fea73beab --- /dev/null +++ b/docs/rocm-releases.md @@ -0,0 +1,48 @@ +# ROCm and Vulkan release builds + +`.github/workflows/build-llamacpp-rocm.yml` publishes prebuilt ROCm binaries of this fork. It copies the +pipeline of [lemonade-sdk/llamacpp-rocm](https://github.com/lemonade-sdk/llamacpp-rocm) (MIT) and uses the same +artifact layout, so tools that consume their releases can consume these ones too. + +## What a release contains + +- One zip per GPU target: `llama-bNNNN-ubuntu-rocm-{target}-x64.zip`. Windows zips + (`llama-bNNNN-windows-rocm-{target}-x64.zip`) are built only on a manual run that asks for them. +- Target: `gfx1151` (Strix Halo) only. +- One Linux Vulkan zip per release: `llama-bNNNN-ubuntu-vulkan-x64.zip`. It is built from the same commit with + `GGML_VULKAN=ON`, loads the CPU variants and backends at runtime (`GGML_BACKEND_DL`, + `GGML_CPU_ALL_VARIANTS`), sets RPATH to `$ORIGIN`, and uses the system Vulkan loader and driver + (Mesa RADV on Strix Halo). A manual run can turn it off with the `vulkan` input. +- Each ROCm zip holds the llama.cpp binaries and shared libraries, built with `GGML_HIP=ON`, `GGML_RPC=ON` and + `BUILD_SHARED_LIBS=ON`. +- Each zip also holds the ROCm runtime libraries it needs (hipBLAS, rocBLAS, hipBLASLt and their kernel + libraries). On Linux, RPATH is set to `$ORIGIN`, so the zip runs without a system ROCm install. +- ROCm comes from the newest TheRock nightly tarball on `nightly.repo.amd.com` unless one is pinned. +- Releases are tagged `b1000`, `b1001` and so on, counting up. This numbering does not collide with the `b1xxxx` + tags carried over from upstream llama.cpp. +- The release notes record the build number, targets, ROCm version, source commit (5 characters) and build date. + +## When it runs + +- On every merge to `master`. This builds the merged commit and publishes a release. A newer merge cancels a + release that is still building. +- On demand with `workflow_dispatch`. Inputs: OS list (`ubuntu` by default, `windows,ubuntu` for both), targets, + ROCm version, branch or tag, and whether to publish. +- On pull requests that change the pipeline itself. These runs build but do not publish. + +If a release build for a merge fails, `nightly-failure-alert.yml` opens or updates an issue labelled `nightly-failure`. + +## Hardware tests + +`test-stx-halo` runs the built zip on real hardware. It is skipped unless the repository variable +`STX_HALO_RUNNERS` is set to `true`, because jobs whose self-hosted runners are missing sit in the queue until +they time out. To enable it: + +1. Register runners with the labels `stx-halo` + `Windows` / `Linux`. +2. Set `STX_HALO_RUNNERS` to `true`. + +`test-llamacpp-rocm.yml` tests an already published release on the same runners. + +## Tokens + +The workflow uses the `GH_TOKEN` secret when it exists. Otherwise it uses the job token with `contents: write`. 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 224bdef927ac..13fa8bc05296 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/rocmfp4/rocmfp4_hip_scale.cuh b/ggml/rocmfp4/rocmfp4_hip_scale.cuh new file mode 100644 index 000000000000..b69ac0d87f85 --- /dev/null +++ b/ggml/rocmfp4/rocmfp4_hip_scale.cuh @@ -0,0 +1,168 @@ +// 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 + +#ifndef GGML_ROCMFP4_USE_SCALE_LUT +#define GGML_ROCMFP4_USE_SCALE_LUT 0 +#endif + +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static __device__ __constant__ const float rocmfp4_scale_ue4m3_half_lut[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 +#endif + +static __device__ __forceinline__ float rocmfp4_u32_as_f32(uint32_t bits) { +#if defined(GGML_USE_HIP) + return __uint_as_float(bits); +#else + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +#endif +} + +// Decode invalid scale bytes (> 0x7e) as 0, matching the CPU decoder in +// rocmfp4.c. Loader validation of these bytes only runs with --check-tensors, +// so kernels must not trust them. +static __device__ __forceinline__ float rocmfp4_ue4m3_to_fp32_half_finite(uint8_t x) { + if (x > 0x7e) { + return 0.0f; + } +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT + return rocmfp4_scale_ue4m3_half_lut[x]; +#else + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +#endif +} + +static __device__ __forceinline__ float rocmfpx_ue4m3_to_fp32_finite(uint8_t x) { + if (x > 0x7e) { + return 0.0f; + } + + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +} + +static __device__ __forceinline__ uint8_t rocmfpx_nearest_scale_ue4m3_cuda(float target_scale) { + if (!(target_scale > 0.0f) || !isfinite(target_scale)) { + return 0; + } + + uint8_t lo = 1; + uint8_t hi = 0x7e; + while (lo < hi) { + const uint8_t mid = lo + (hi - lo) / 2; + if (rocmfpx_ue4m3_to_fp32_finite(mid) < target_scale) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfpx_ue4m3_to_fp32_finite(lo); + const float lo_scale = rocmfpx_ue4m3_to_fp32_finite((uint8_t) (lo - 1)); + return (target_scale - lo_scale <= hi_scale - target_scale) ? (uint8_t) (lo - 1) : lo; +} + +static __device__ __forceinline__ int8_t rocmfp4_decode_i8(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} diff --git a/ggml/rocmfpx/README.md b/ggml/rocmfpx/README.md new file mode 100644 index 000000000000..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..778589080771 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -93,7 +93,7 @@ enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_te // dynamic tensor allocator -#define GGML_VBUFFER_MAX_CHUNKS 16 +#define GGML_VBUFFER_INITIAL_CHUNKS 16 // relative memory address within an allocation that can be split into multiple buffers (chunks) struct buffer_address { @@ -121,8 +121,9 @@ struct tallocr_chunk { struct ggml_dyn_tallocr { size_t alignment; size_t max_chunk_size; - struct tallocr_chunk * chunks[GGML_VBUFFER_MAX_CHUNKS]; + struct tallocr_chunk ** chunks; // [cap_chunks], grown on demand int n_chunks; + int cap_chunks; #ifdef GGML_ALLOCATOR_DEBUG struct { @@ -158,20 +159,23 @@ static void ggml_dyn_tallocr_remove_block(struct tallocr_chunk * chunk, int idx) } static int ggml_dyn_tallocr_new_chunk(struct ggml_dyn_tallocr * alloc, size_t min_size) { - if (alloc->n_chunks >= GGML_VBUFFER_MAX_CHUNKS) { - return -1; + // The chunk list grows as needed. It used to be a fixed array of 16 whose last chunk was unbounded, so a + // compute buffer of more than ~16 chunks (16 GiB with Vulkan's 1 GiB chunks) packed everything left into + // one oversized chunk that the backend then refused (past RADV's 4 GiB buffer limit, issue 22). + if (alloc->n_chunks == alloc->cap_chunks) { + const int cap = alloc->cap_chunks ? 2*alloc->cap_chunks : GGML_VBUFFER_INITIAL_CHUNKS; + struct tallocr_chunk ** chunks = (struct tallocr_chunk **) realloc(alloc->chunks, cap*sizeof(*chunks)); + if (chunks == NULL) { + return -1; + } + alloc->chunks = chunks; + alloc->cap_chunks = cap; } struct tallocr_chunk * chunk = calloc(1, sizeof(struct tallocr_chunk)); chunk->n_free_blocks = 1; chunk->free_blocks[0].offset = 0; - // available space in a chunk is limited to max_chunk_size, but can be higher if: - // 1. a single tensor exceeds the maximum, and cannot fit any other way - // 2. we are running out of chunks - // backends will either manage to allocate the larger size, or report an error. + // available space in a chunk is limited to max_chunk_size, unless a single tensor exceeds it chunk->free_blocks[0].size = MAX(min_size, alloc->max_chunk_size); - if (alloc->n_chunks == GGML_VBUFFER_MAX_CHUNKS - 1) { - chunk->free_blocks[0].size = SIZE_MAX/2; - } alloc->chunks[alloc->n_chunks] = chunk; alloc->n_chunks++; return alloc->n_chunks - 1; @@ -252,7 +256,7 @@ static struct buffer_address ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * al best_fit_block = 0; } if (best_fit_chunk == -1) { - // since the last chunk always has virtually endless memory, this should never happen + // only reached if the chunk list itself cannot grow GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n", __func__, size, max_avail); GGML_ABORT("graph allocation: failed to reserve memory"); @@ -350,7 +354,7 @@ static void ggml_dyn_tallocr_free_bytes(struct ggml_dyn_tallocr * alloc, struct } static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; i++) { + for (int i = 0; i < alloc->n_chunks; i++) { free(alloc->chunks[i]); alloc->chunks[i] = NULL; } @@ -369,8 +373,9 @@ static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment, size_t m *alloc = (struct ggml_dyn_tallocr) { /*.alignment = */ alignment, /*.max_chunk_size = */ MIN(max_buffer_size, SIZE_MAX/2), // clamp to avoid overflows - /*.chunks = */ {NULL}, + /*.chunks = */ NULL, /*.n_chunks = */ 0, + /*.cap_chunks = */ 0, #ifdef GGML_ALLOCATOR_DEBUG /*.allocated_tensors = */ {{0}}, #endif @@ -385,6 +390,7 @@ static void ggml_dyn_tallocr_free(struct ggml_dyn_tallocr * alloc) { for (int i = 0; i < alloc->n_chunks; ++i) { free(alloc->chunks[i]); } + free(alloc->chunks); free(alloc); } @@ -396,26 +402,45 @@ static size_t ggml_dyn_tallocr_max_size(struct ggml_dyn_tallocr * alloc, int chu // virtual buffer with contiguous memory range, split into multiple backend buffers (chunks) struct vbuffer { - ggml_backend_buffer_t chunks[GGML_VBUFFER_MAX_CHUNKS]; + ggml_backend_buffer_t * chunks; // [n_chunks] + int n_chunks; }; static void ggml_vbuffer_free(struct vbuffer * buf) { if (buf == NULL) { return; } - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; ++i) { + for (int i = 0; i < buf->n_chunks; ++i) { ggml_backend_buffer_free(buf->chunks[i]); } + free(buf->chunks); free(buf); } +// make room for n chunks (new slots are empty) +static bool ggml_vbuffer_reserve_chunks(struct vbuffer * buf, int n) { + if (n <= buf->n_chunks) { + return true; + } + ggml_backend_buffer_t * chunks = (ggml_backend_buffer_t *) realloc(buf->chunks, n*sizeof(*chunks)); + if (chunks == NULL) { + return false; + } + for (int i = buf->n_chunks; i < n; ++i) { + chunks[i] = NULL; + } + buf->chunks = chunks; + buf->n_chunks = n; + return true; +} + static size_t ggml_vbuffer_chunk_size(struct vbuffer * buf, int chunk) { - return buf->chunks[chunk] ? ggml_backend_buffer_get_size(buf->chunks[chunk]) : 0; + return chunk < buf->n_chunks && buf->chunks[chunk] ? ggml_backend_buffer_get_size(buf->chunks[chunk]) : 0; } static size_t ggml_vbuffer_size(struct vbuffer * buf) { size_t size = 0; - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + for (int i = 0; i < buf->n_chunks && buf->chunks[i]; ++i) { size += ggml_backend_buffer_get_size(buf->chunks[i]); } return size; @@ -426,6 +451,10 @@ static struct vbuffer * ggml_vbuffer_alloc(ggml_backend_buffer_type_t buft, cons if (buf == NULL) { return NULL; } + if (!ggml_vbuffer_reserve_chunks(buf, talloc->n_chunks)) { + free(buf); + return NULL; + } for (int n = 0; n < talloc->n_chunks; n++) { size_t chunk_size = talloc->chunks[n]->max_size; @@ -446,7 +475,7 @@ static void ggml_vbuffer_tensor_alloc(struct vbuffer * buf, struct ggml_tensor * } static void ggml_vbuffer_reset(struct vbuffer * buf) { - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + for (int i = 0; i < buf->n_chunks && buf->chunks[i]; ++i) { ggml_backend_buffer_reset(buf->chunks[i]); } } @@ -493,6 +522,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 +602,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 +625,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 +687,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 +738,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 +764,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]; @@ -1124,6 +1198,20 @@ static void free_buffers(ggml_backend_buffer_t ** buffers, const size_t * n_buff free(*buffers); } +static bool alloc_ctx_tensors_finalize_views(struct ggml_context * ctx) { + for (struct ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->view_src != NULL && t->buffer == NULL) { + enum ggml_status status = ggml_backend_view_init(t); + if (status != GGML_STATUS_SUCCESS) { + GGML_LOG_ERROR("%s: failed to initialize view tensor %s\n", __func__, t->name); + return false; + } + } + } + + return true; +} + static bool alloc_tensor_range(struct ggml_context * ctx, struct ggml_tensor * first, struct ggml_tensor * last, ggml_backend_buffer_type_t buft, size_t size, @@ -1142,24 +1230,14 @@ static bool alloc_tensor_range(struct ggml_context * ctx, struct ggml_tallocr tallocr = ggml_tallocr_new(buffer); for (struct ggml_tensor * t = first; t != last; t = ggml_get_next_tensor(ctx, t)) { - enum ggml_status status = GGML_STATUS_SUCCESS; - if (t->data == NULL) { - if (t->view_src == NULL) { - status = ggml_tallocr_alloc(&tallocr, t); - } else if (t->buffer == NULL) { - status = ggml_backend_view_init(t); - } - } else { - if (t->view_src != NULL && t->buffer == NULL) { - // view of a pre-allocated tensor - status = ggml_backend_view_init(t); + if (t->view_src == NULL && t->data == NULL) { + enum ggml_status status = ggml_tallocr_alloc(&tallocr, t); + if (status != GGML_STATUS_SUCCESS) { + GGML_LOG_ERROR("%s: failed to allocate tensor %s\n", __func__, t->name); + free_buffers(buffers, n_buffers); + return false; } } - if (status != GGML_STATUS_SUCCESS) { - GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t->name); - free_buffers(buffers, n_buffers); - return false; - } } return true; @@ -1217,6 +1295,11 @@ static ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft_impl( return NULL; } + if (!alloc_ctx_tensors_finalize_views(ctx)) { + free_buffers(&buffers, &n_buffers); + return NULL; + } + ggml_backend_buffer_t buffer; if (n_buffers == 1) { buffer = buffers[0]; 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 20bf965017e3..accb58fbf721 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])); + } } } @@ -1668,6 +1818,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]); @@ -1782,7 +1934,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]); @@ -1798,6 +1957,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 { @@ -1820,6 +1980,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; } @@ -1842,6 +2003,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; } @@ -1869,7 +2034,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) @@ -1896,6 +2091,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]; @@ -1934,6 +2131,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); @@ -1989,13 +2187,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); @@ -2003,6 +2296,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; @@ -2034,6 +2335,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, @@ -2110,6 +2413,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 24c47569c0a8..8f08c95298e6 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -16,6 +16,10 @@ #include "common.h" #include "tiled/tiled.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, @@ -2993,6 +3193,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 a07e1f96373c..37798b920a46 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])); #if defined(GGML_SIMD) && !defined(__x86_64__) && !defined(_M_X64) #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 24115da09960..6d0ad66ef4c6 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 5363dd001cd0..ce70293aa466 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -21,6 +21,7 @@ #endif #endif #include "ggml-common.h" +#include "../../rocmfp4/rocmfp4.h" #include #include @@ -40,6 +41,9 @@ #include "vendors/cuda.h" #endif // defined(GGML_USE_HIP) +// device-side UE4M3 scale decoder for ROCmFP4 (needs the runtime header above) +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" + #define STRINGIZE_IMPL(...) #__VA_ARGS__ #define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__) @@ -89,6 +93,13 @@ #define GGML_CUDA_CC_IS_RDNA2(cc) (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3) #define GGML_CUDA_CC_IS_RDNA3_0(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA3_5) #define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4) + +// max selected-key depth the QSA prefill/decode kernels serve. The selected list addresses whole blocks of 4 keys, +// and the union stores block ids as uint32 with 0xFFFFFFFF reserved as the padding sentinel, so the space is not +// the limit: this bound matches the largest context the Q3.8-Flash-Next card supports with RoPE/YaRN scaling. +// The model-side maskless gate (qwen4exp.cpp) and both kernel support checks must agree with it exactly, so that +// any n_kv above the bound takes the masked route instead of reaching a kernel that refuses the op (GGML_ABORT). +#define GGML_QSA_MAX_KEYS 1048576 #define GGML_CUDA_CC_IS_RDNA3(cc) (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) #define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4) #define GGML_CUDA_CC_IS_GCN_APU(cc) ((cc) == GGML_CUDA_CC_GFX909 || (cc) == GGML_CUDA_CC_GFX90C) @@ -115,6 +126,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 @@ -708,6 +727,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__) @@ -1051,6 +1088,20 @@ struct ggml_cuda_type_traits { static constexpr int bs = sizeof(block_nvfp4); }; +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; @@ -1443,6 +1494,8 @@ struct ggml_cuda_stream_context { } }; +struct ggml_cuda_mmb_context; + struct ggml_backend_cuda_context { int device; std::string name; @@ -1455,6 +1508,15 @@ 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_small_batch = false; // set by ggml_backend_cuda_set_mmb_small_batch, 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) @@ -1571,6 +1633,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; @@ -1580,6 +1646,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 0619f4760e67..7749e9bfc5c7 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -243,6 +243,20 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_ dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x); } +template +static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_rocmfp4(vx, i, yy + i*QK_K, threadIdx.x); +} + +template +static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const int64_t i = blockIdx.x; + + dequantize_rocmfp4_fast(vx, i, yy + i*QK_K, threadIdx.x); +} + template static void dequantize_block_cuda(const void * vx, dst_t * y, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, @@ -374,6 +388,18 @@ static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t dequantize_block_mxfp4<<>>(vx, y); } +template +static void dequantize_row_rocmfp4_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4<<>>(vx, y); +} + +template +static void dequantize_row_rocmfp4_fast_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4_fast<<>>(vx, y); +} + template static __global__ void dequantize_block_nvfp4( const void * __restrict__ vx, @@ -473,6 +499,25 @@ 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) { if (k % 4 == 0 && @@ -484,6 +529,19 @@ static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, return; } + 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); } @@ -533,6 +591,10 @@ to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) { return dequantize_row_iq3_s_cuda; case GGML_TYPE_MXFP4: return dequantize_row_mxfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; case GGML_TYPE_F32: @@ -593,6 +655,10 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { return dequantize_row_iq3_s_cuda; case GGML_TYPE_MXFP4: return dequantize_row_mxfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; case GGML_TYPE_F32: @@ -650,6 +716,10 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { return dequantize_row_iq3_s_cuda; case GGML_TYPE_MXFP4: return dequantize_row_mxfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; case GGML_TYPE_F16: 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..6489e65442f4 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,112 @@ 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); +} + +template +static __device__ __forceinline__ void dequantize_rocmfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_rocmfp4 * x = (const block_rocmfp4 *) vx + ibs*(QK_K/QK_ROCMFP4); + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + // dual UE4M3 scales: low nibbles are weights j (e[0]), high nibbles are weights j+16 (e[1]) + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = ggml_cuda_cast(d0 * kvalues_rocmfp4[q4[j] & 0xf]); + y[j+16] = ggml_cuda_cast(d1 * kvalues_rocmfp4[q4[j] >> 4]); + } +} + +template +static __device__ __forceinline__ void dequantize_rocmfp4_fast(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { + + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx + ibs*(QK_K/QK_ROCMFP4); + + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = ggml_cuda_cast(d * kvalues_rocmfp4[q4[j] & 0xf]); + y[j+16] = ggml_cuda_cast(d * kvalues_rocmfp4[q4[j] >> 4]); + } +} 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 6d1ce52dbc77..c6b67334ab59 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1202,6 +1202,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 083d3228abdd..1a9d2388991f 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -157,8 +157,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); @@ -1781,7 +1781,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(); } } @@ -1868,7 +1870,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 b5ea915e132e..6c9a4c8f0460 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 ad4d1df2cc75..6d48040e8fa3 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) // one list per group of ncols1 queries: a column is selected if any query of the group can see it @@ -538,6 +542,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); @@ -670,6 +688,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; @@ -697,11 +719,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)) { @@ -736,6 +766,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; @@ -757,6 +793,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..117ffe98679b 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,8 +1,58 @@ #include "gated_delta_net.cuh" #include "ggml-cuda/common.cuh" +#include "unary.cuh" + +#include +#include + + +// 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 +140,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 +154,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 +168,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 +183,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 +216,890 @@ 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]; + } + } + } +} + + +// --------------------------------------------------------------------------------------------- +// Chunked Gated DeltaNet prefill (RDNA3.5, S_v = 128, non-KDA, final state only, n_seqs == 1). +// Per chunk of CL tokens (G = in-chunk cumsum of g, D[r][s] = exp(G_r - G_s)): +// A[r][s] = beta_r (k_r.k_s) D[r][s] (s < r), T' = (I + A)^-1 diag(beta) +// U = T' (V - diag(e^G) K S0) (all delta-rule corrections, one product) +// O = scale diag(e^G) Q S0 + P U, P[r][s] = scale (q_r.k_s) D[r][s] (s <= r) +// S1 = e^{G_L} S0 + K^T diag(e^{G_L - G}) U +// gdn_chunk_prep_kernel: fully parallel, one block per (chunk, k-head) for its H/Hk value heads. Produces the chunk's +// q/k (optionally through the fused causal conv + SiLU and the fused L2 norm) and v as fp16 rows, KK^T / QK^T on +// WMMA, T by a barrier-free wave-level forward substitution, and writes T', P and the gate factors. +// gdn_chunk_scan_kernel: one block (8 waves) per value head, sequential over chunks; wave w owns 16 state columns +// in fp32 WMMA accumulators. The chunk's q/k/T'/P are staged once per block in LDS (K^T transposed on the way in) +// and the next chunk is prefetched into registers. Optionally applies the gated RMS norm of the output (fused). +// WMMA operands are fp16 with fp32 accumulation: not bit-identical to the sequential kernels (NMSE ~2e-7). +// GGML_GDN_CHUNK=0 disables the path (and every fusion that depends on it), GGML_GDN_CHUNK_MIN sets the token threshold. +#if defined(GGML_USE_HIP) && defined(RDNA3_5) +#define GDNC_ENABLED 1 +#endif +#ifndef GDNC_SCAN_WAVES +#define GDNC_SCAN_WAVES 6 +#endif +typedef _Float16 gdnc_v16h __attribute__((ext_vector_type(16))); +typedef float gdnc_v8f __attribute__((ext_vector_type(8))); + +#if defined(GDNC_ENABLED) +#define GDNC_WMMA(a, b, c) __builtin_amdgcn_wmma_f32_16x16x16_f16_w32((a), (b), (c)) +// A/B operand fragment from a row-major fp16 matrix: lane holds row (lane & 15), 16 consecutive k +static __device__ __forceinline__ gdnc_v16h gdnc_frag(const _Float16 * base, const int stride, const int row0, const int k0, const int lane) { + const _Float16 * p = base + (row0 + (lane & 15)) * stride + k0; + gdnc_v16h f; + const uint4 a = *(const uint4 *) p, b = *(const uint4 *) (p + 8); + __builtin_memcpy(&f, &a, 16); __builtin_memcpy(((char *) &f) + 16, &b, 16); + return f; +} +// same, from 16 consecutive fp32 values +static __device__ __forceinline__ gdnc_v16h gdnc_frag_f32(const float * p) { + gdnc_v16h f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float4 x = ((const float4 *) p)[j]; + f[4 * j + 0] = (_Float16) x.x; f[4 * j + 1] = (_Float16) x.y; + f[4 * j + 2] = (_Float16) x.z; f[4 * j + 3] = (_Float16) x.w; + } + return f; +} +// accumulator tile (lane: column lane&15, rows 2e + (lane>>4)) -> B fragment (lane: column lane&15, rows 0..15) +static __device__ __forceinline__ gdnc_v16h gdnc_acc_to_frag(const gdnc_v8f & acc, const int lane) { + gdnc_v16h f; + const bool hi = lane >= 16; +#pragma unroll + for (int e = 0; e < 8; ++e) { + const float own = acc[e]; + const float oth = __int_as_float(__builtin_amdgcn_permlanex16(__float_as_int(own), __float_as_int(own), 0x76543210, 0xFEDCBA98, true, false)); + f[2 * e] = (_Float16) (hi ? oth : own); + f[2 * e + 1] = (_Float16) (hi ? own : oth); + } + return f; +} +#endif + +// inputs of the chunk kernels: the q/k/v views of the GDN op, or (fused conv) the raw conv input x [T, C] with the +// conv history [C, 3] and taps [C, 4]: q/k/v element (t, c) is then silu(sum_j w[c][j] * in(t - 3 + j, c)), the +// ssm_conv_f32 + SiLU arithmetic of gdn_conv_direct_kernel +struct gdnc_in { + const float * q; const float * k; const float * v; + int64_t sq1, sq2, sv1, sv2; + const float * x; const float * st; const float * cw; // x == nullptr: read the views + int64_t C, qc, kc, vc; // x row stride, channel of element 0 of the q / k / v views +}; + +#if defined(GDNC_ENABLED) +// channels c..c+3 of token t through the causal conv + SiLU; wt[ch][j] = tap j of channel c + ch +static __device__ __forceinline__ float4 gdnc_conv4(const gdnc_in & in, const float (&wt)[4][4], const int64_t t, const int64_t c) { + float s[4] = { 0.0f, 0.0f, 0.0f, 0.0f }; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int64_t u = t - 3 + j; + float xv[4]; + if (u >= 0) { + const float4 v4 = *(const float4 *) (in.x + u * in.C + c); + xv[0] = v4.x; xv[1] = v4.y; xv[2] = v4.z; xv[3] = v4.w; + } else { +#pragma unroll + for (int ch = 0; ch < 4; ++ch) xv[ch] = in.st[(c + ch) * 3 + (u + 3)]; + } +#pragma unroll + for (int ch = 0; ch < 4; ++ch) s[ch] = fmaf(xv[ch], wt[ch][j], s[ch]); + } + float4 o; + o.x = ggml_cuda_op_silu_single(s[0] + 0.0f); o.y = ggml_cuda_op_silu_single(s[1] + 0.0f); + o.z = ggml_cuda_op_silu_single(s[2] + 0.0f); o.w = ggml_cuda_op_silu_single(s[3] + 0.0f); + return o; +} +static __device__ __forceinline__ void gdnc_load_taps(const gdnc_in & in, const int64_t c, float (&wt)[4][4]) { +#pragma unroll + for (int ch = 0; ch < 4; ++ch) { + const float4 w4 = *(const float4 *) (in.cw + (c + ch) * 4); + wt[ch][0] = w4.x; wt[ch][1] = w4.y; wt[ch][2] = w4.z; wt[ch][3] = w4.w; + } +} +static __device__ __forceinline__ uint2 gdnc_pack4h(const float4 v) { + const _Float16 h[4] = { (_Float16) v.x, (_Float16) v.y, (_Float16) v.z, (_Float16) v.w }; + uint2 u; + __builtin_memcpy(&u, h, 8); + return u; +} +#endif + +// K1: one block per (chunk, k-head), handling the H/Hk value heads that share it. Produces the (normalized) q/k and v +// of the chunk as fp16 rows (Qh/Kh [Hk, T, 128], Vh [H, T, 128]) for the scan, plus T', P and the gate factors. +template +__global__ void __launch_bounds__(256) +gdn_chunk_prep_kernel(const gdnc_in in, const float * __restrict__ g, + const float * __restrict__ beta, _Float16 * __restrict__ Tb, _Float16 * __restrict__ Pb, float2 * __restrict__ Gb, + _Float16 * __restrict__ Qh, _Float16 * __restrict__ Kh, _Float16 * __restrict__ Vh, + const int H, const int Hk, const int n_tokens, const int nch, + const int64_t sb1, const int64_t sb2, const float scale, + const float qk_eps, const float qk_mul) { +#if defined(GDNC_ENABLED) + constexpr int SK = 136; + constexpr int NT = CL / 16; + constexpr int NTRI = NT * (NT + 1) / 2; + constexpr int MR = CL / 32; + constexpr int CPW = CL / 8; + constexpr int NQK = CL * 32 / 256; + constexpr int RMAX = 4; + __shared__ __align__(16) _Float16 Ks[CL * SK]; + __shared__ __align__(16) _Float16 Qs[CL * SK]; + __shared__ float KK[CL][CL + 1]; + __shared__ float QK[CL][CL + 1]; + __shared__ float A[CL][CL + 1]; + __shared__ float Gs[RMAX][CL], bs[RMAX][CL]; + const int hk = blockIdx.x, ch = blockIdx.y, tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int cl = lane & 15, rh = lane >> 4; + const int R = H / Hk; + const int t0 = ch * CL, Lc = min(CL, n_tokens - t0); + const int i4 = (tid & 31) * 4; // this thread's 4 channels of every row it handles (row = (idx >> 5)) + + { + float4 kv[NQK], qv[NQK]; + if (in.x) { + float wq[4][4], wk[4][4]; + gdnc_load_taps(in, in.qc + (int64_t) hk * in.sq1 + i4, wq); + gdnc_load_taps(in, in.kc + (int64_t) hk * in.sq1 + i4, wk); +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int r = (tid + 256 * j) >> 5; + const int64_t t = t0 + min(r, Lc - 1); + qv[j] = gdnc_conv4(in, wq, t, in.qc + (int64_t) hk * in.sq1 + i4); + kv[j] = gdnc_conv4(in, wk, t, in.kc + (int64_t) hk * in.sq1 + i4); + } + } else { +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int r = (tid + 256 * j) >> 5; + const int64_t o = (int64_t) (t0 + min(r, Lc - 1)) * in.sq2 + (int64_t) hk * in.sq1 + i4; + kv[j] = *(const float4 *) (in.k + o); + qv[j] = *(const float4 *) (in.q + o); + } + } + if (qk_eps >= 0.0f) { + // fused q/k L2 norm: a row's 128 values are the float4s of one wave (row = idx >> 5) +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const float sk = gdn_warp_reduce_sum<32>(kv[j].x * kv[j].x + kv[j].y * kv[j].y + kv[j].z * kv[j].z + kv[j].w * kv[j].w); + const float sq = gdn_warp_reduce_sum<32>(qv[j].x * qv[j].x + qv[j].y * qv[j].y + qv[j].z * qv[j].z + qv[j].w * qv[j].w); + const float fk = qk_mul * rsqrtf(sk + qk_eps), fq = qk_mul * rsqrtf(sq + qk_eps); + kv[j].x *= fk; kv[j].y *= fk; kv[j].z *= fk; kv[j].w *= fk; + qv[j].x *= fq; qv[j].y *= fq; qv[j].z *= fq; qv[j].w *= fq; + } + } +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int r = (tid + 256 * j) >> 5; + const bool ok = r < Lc; + const float4 z4 = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + const uint2 kh = gdnc_pack4h(ok ? kv[j] : z4), qh = gdnc_pack4h(ok ? qv[j] : z4); + *(uint2 *) (Ks + r * SK + i4) = kh; + *(uint2 *) (Qs + r * SK + i4) = qh; + if (ok) { + const int64_t o = ((int64_t) hk * n_tokens + t0 + r) * 128 + i4; + *(uint2 *) (Kh + o) = kh; + *(uint2 *) (Qh + o) = qh; + } + } + // v of the R value heads of this k-head -> Vh (fp16) + for (int jh = 0; jh < R; ++jh) { + const int h = hk + jh * Hk; + float wv[4][4]; + if (in.x) gdnc_load_taps(in, in.vc + (int64_t) h * in.sv1 + i4, wv); +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int r = (tid + 256 * j) >> 5; + if (r >= Lc) continue; + const int64_t t = t0 + r; + const float4 vv = in.x ? gdnc_conv4(in, wv, t, in.vc + (int64_t) h * in.sv1 + i4) + : *(const float4 *) (in.v + t * in.sv2 + (int64_t) h * in.sv1 + i4); + *(uint2 *) (Vh + ((int64_t) h * n_tokens + t) * 128 + i4) = gdnc_pack4h(vv); + } + } + } + for (int idx = tid; idx < R * CL; idx += 256) { + const int j = idx / CL, r = idx % CL, h = hk + j * Hk; + const int64_t o = (int64_t) (t0 + r) * sb2 + (int64_t) h * sb1; + Gs[j][r] = r < Lc ? g[o] : 0.0f; + bs[j][r] = r < Lc ? beta[o] : 0.0f; + } + __syncthreads(); + if (tid < R) { + float gv[CL]; +#pragma unroll + for (int r = 0; r < CL; ++r) gv[r] = Gs[tid][r]; +#pragma unroll + for (int r = 1; r < CL; ++r) gv[r] += gv[r - 1]; +#pragma unroll + for (int r = 0; r < CL; ++r) Gs[tid][r] = gv[r]; + } + // raw lower-triangular tiles of K K^T and Q K^T (shared by all value heads of this k-head) + for (int job = wave; job < 2 * NTRI; job += 8) { + const bool isq = job >= NTRI; + int jj = isq ? job - NTRI : job, rt = 0; + while (jj >= rt + 1) { jj -= rt + 1; ++rt; } + const int st = jj; + const _Float16 * As = isq ? Qs : Ks; + gdnc_v8f acc = {}; +#pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + acc = GDNC_WMMA(gdnc_frag(As, SK, rt * 16, kk, lane), gdnc_frag(Ks, SK, st * 16, kk, lane), acc); + } + float (*D)[CL + 1] = isq ? QK : KK; +#pragma unroll + for (int e = 0; e < 8; ++e) D[rt * 16 + 2 * e + rh][st * 16 + cl] = acc[e]; + } + __syncthreads(); + + for (int j = 0; j < R; ++j) { + const int h = hk + j * Hk; + _Float16 * Pp = Pb + ((int64_t) h * nch + ch) * CL * CL; + for (int idx = tid; idx < CL * CL; idx += 256) { + const int r = idx / CL, s = idx % CL; + const float d = s <= r ? expf(Gs[j][r] - Gs[j][s]) : 0.0f; + A[r][s] = s < r ? bs[j][r] * KK[r][s] * d : 0.0f; + Pp[idx] = (_Float16) (s <= r ? scale * QK[r][s] * d : 0.0f); + } + if (tid < CL) { + Gb[((int64_t) h * nch + ch) * CL + tid] = make_float2(expf(Gs[j][tid]), expf(Gs[j][CL - 1] - Gs[j][tid])); + } + __syncthreads(); + // T = (I + A)^-1 column by column, barrier-free: lane owns rows lane + 32 m, the wave CPW columns + { + const int j0 = wave * CPW; + float x[MR][CPW]; +#pragma unroll + for (int m = 0; m < MR; ++m) +#pragma unroll + for (int c = 0; c < CPW; ++c) x[m][c] = (lane + 32 * m == j0 + c) ? 1.0f : 0.0f; +#pragma unroll + for (int qq = 0; qq < CL - 1; ++qq) { + if (qq < j0) continue; + float a[MR]; +#pragma unroll + for (int m = 0; m < MR; ++m) a[m] = A[lane + 32 * m][qq]; +#pragma unroll + for (int c = 0; c < CPW; ++c) { + const float xq = __int_as_float(__builtin_amdgcn_readlane(__float_as_int(x[qq / 32][c]), qq % 32)); +#pragma unroll + for (int m = 0; m < MR; ++m) x[m][c] = fmaf(-a[m], xq, x[m][c]); + } + } + _Float16 * Tp = Tb + ((int64_t) h * nch + ch) * CL * CL; +#pragma unroll + for (int m = 0; m < MR; ++m) { + const int r = lane + 32 * m; + _Float16 tmp[CPW]; +#pragma unroll + for (int c = 0; c < CPW; ++c) tmp[c] = (_Float16) (x[m][c] * bs[j][j0 + c]); + __builtin_memcpy(Tp + r * CL + j0, tmp, sizeof(tmp)); + } + } + __syncthreads(); + } +#else + NO_DEVICE_CODE; +#endif +} + +// K2: one block (8 waves) per value head, wave w owns state columns [16 w, 16 w + 16). The chunk's q/k/T'/P are +// staged once per block in LDS (K^T transposed on the way in); the next chunk is prefetched into registers. +template +__global__ void __launch_bounds__(256, GDNC_SCAN_WAVES) +gdn_chunk_scan_kernel(const _Float16 * __restrict__ Qh, const _Float16 * __restrict__ Kh, const _Float16 * __restrict__ Vh, + const float * __restrict__ curr_state, float * __restrict__ dst, float * __restrict__ state, + const _Float16 * __restrict__ Tb, const _Float16 * __restrict__ Pb, const float2 * __restrict__ Gb, + const int H, const int Hk, const int n_tokens, const int nch, const float scale, + const float * __restrict__ on_w, const float * __restrict__ on_z, float * __restrict__ on_out, uint16_t * __restrict__ on_out16, + const float on_eps) { +#if defined(GDNC_ENABLED) + constexpr int NT = CL / 16; + constexpr int SQ = 136; + constexpr int ST = CL + 8; + constexpr int NQK = CL * 32 / 256; // 4-half groups of q (and of k) per thread per chunk + constexpr int NTP = 2 * CL * CL / 8 / 256; // uint4 of T'/P per thread per chunk + static_assert(NTP * 256 == 2 * CL * CL / 8, "T'/P staging must divide evenly"); + __shared__ __align__(16) _Float16 Qs[CL * SQ]; + __shared__ __align__(16) _Float16 Ks[CL * SQ]; + __shared__ __align__(16) _Float16 KTs[128 * ST]; + __shared__ __align__(16) _Float16 Ts[CL * ST]; + __shared__ __align__(16) _Float16 Ps[CL * ST]; + __shared__ float2 Gs[CL]; + __shared__ float red[CL][9]; // fused gated norm: per-row partial sums of squares of the 8 waves + const int h = blockIdx.x, tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int hk = h % Hk; + const int cl = lane & 15, rh = lane >> 4, col = wave * 16 + cl; + + const float * s_in = curr_state + (int64_t) h * 128 * 128; + gdnc_v8f S[8]; +#pragma unroll + for (int it = 0; it < 8; ++it) +#pragma unroll + for (int e = 0; e < 8; ++e) S[it][e] = s_in[col * 128 + it * 16 + 2 * e + rh]; + + const _Float16 * kh = Kh + (int64_t) hk * n_tokens * 128; // prep's fp16 rows: normalized q / k, v + const _Float16 * qh = Qh + (int64_t) hk * n_tokens * 128; + const _Float16 * vh = Vh + (int64_t) h * n_tokens * 128 + col; + + uint2 pk[NQK], pq[NQK]; + uint4 ptp[NTP]; + float2 pg = make_float2(0.0f, 0.0f); + float pv[NT][8]; + float pz[NT][8]; // fused output norm: the gate z of this lane's rows / column, prefetched with v + const float * zh = on_z ? on_z + (int64_t) h * 128 + col : nullptr; + auto prefetch = [&](const int ch) { + const int t0 = ch * CL; +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int idx = tid + 256 * j, r = idx >> 5, i = (idx & 31) * 4; + const int64_t o = (int64_t) min(t0 + r, n_tokens - 1) * 128 + i; + pk[j] = *(const uint2 *) (kh + o); + pq[j] = *(const uint2 *) (qh + o); + } + const _Float16 * Tp = Tb + ((int64_t) h * nch + ch) * CL * CL; + const _Float16 * Pp = Pb + ((int64_t) h * nch + ch) * CL * CL; +#pragma unroll + for (int j = 0; j < NTP; ++j) { + // 2 CL^2 / 8 is a multiple of 256 for CL = 32 / 64: no guard (a divergent guard forces vmcnt(0) at the join) + const int idx = tid + 256 * j; + const int w = idx % (CL * CL / 8); + ptp[j] = *(const uint4 *) ((idx < CL * CL / 8 ? Tp : Pp) + w * 8); + } + pg = Gb[((int64_t) h * nch + ch) * CL + (tid % CL)]; + }; + // v / z of a chunk: issued right after the previous chunk consumed its own (one register copy each) + auto prefetch_v = [&](const int ch) { + const int t0 = ch * CL; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) +#pragma unroll + for (int e = 0; e < 8; ++e) pv[rt][e] = (float) vh[(int64_t) min(t0 + rt * 16 + 2 * e + rh, n_tokens - 1) * 128]; + }; + auto prefetch_z = [&](const int ch) { + const int t0 = ch * CL; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) +#pragma unroll + for (int e = 0; e < 8; ++e) pz[rt][e] = zh[(int64_t) min(t0 + rt * 16 + 2 * e + rh, n_tokens - 1) * H * 128]; + }; + prefetch(0); + prefetch_v(0); + if (zh) prefetch_z(0); + + for (int ch = 0; ch < nch; ++ch) { + const int t0 = ch * CL; + __syncthreads(); +#pragma unroll + for (int j = 0; j < NQK; ++j) { + const int idx = tid + 256 * j, r = idx >> 5, i = (idx & 31) * 4; + *(uint2 *) (Ks + r * SQ + i) = pk[j]; + *(uint2 *) (Qs + r * SQ + i) = pq[j]; + _Float16 kk4[4]; + __builtin_memcpy(kk4, &pk[j], 8); + KTs[(i + 0) * ST + r] = kk4[0]; KTs[(i + 1) * ST + r] = kk4[1]; + KTs[(i + 2) * ST + r] = kk4[2]; KTs[(i + 3) * ST + r] = kk4[3]; + } +#pragma unroll + for (int j = 0; j < NTP; ++j) { + const int idx = tid + 256 * j; + const int w = idx % (CL * CL / 8), r = w / (CL / 8), c8 = (w % (CL / 8)) * 8; + *(uint4 *) ((idx < CL * CL / 8 ? Ts : Ps) + r * ST + c8) = ptp[j]; + } + if (tid < CL) Gs[tid] = pg; + __syncthreads(); + // unconditional (clamped): a branch here makes the compiler drain the prefetch before the compute + prefetch(min(ch + 1, nch - 1)); + + // K S0, Q S0 + gdnc_v8f ks[NT], qs[NT]; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) { ks[rt] = gdnc_v8f{}; qs[rt] = gdnc_v8f{}; } +#pragma unroll + for (int it = 0; it < 8; ++it) { + const gdnc_v16h sb = gdnc_acc_to_frag(S[it], lane); +#pragma unroll + for (int rt = 0; rt < NT; ++rt) { + ks[rt] = GDNC_WMMA(gdnc_frag(Ks, SQ, rt * 16, it * 16, lane), sb, ks[rt]); + qs[rt] = GDNC_WMMA(gdnc_frag(Qs, SQ, rt * 16, it * 16, lane), sb, qs[rt]); + } + } + // U = T' (V - diag(e^G) K S0) + gdnc_v16h YB[NT]; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) { + gdnc_v8f y; +#pragma unroll + for (int e = 0; e < 8; ++e) y[e] = fmaf(-Gs[rt * 16 + 2 * e + rh].x, ks[rt][e], pv[rt][e]); + YB[rt] = gdnc_acc_to_frag(y, lane); + } + prefetch_v(min(ch + 1, nch - 1)); + gdnc_v8f U[NT]; + gdnc_v16h UB[NT]; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) { + gdnc_v8f u = {}; +#pragma unroll + for (int st = 0; st <= rt; ++st) u = GDNC_WMMA(gdnc_frag(Ts, ST, rt * 16, st * 16, lane), YB[st], u); + U[rt] = u; + UB[rt] = gdnc_acc_to_frag(u, lane); + } + // S1 = e^{G_L} S0 + K^T diag(e^{G_L - G}) U + gdnc_v16h UB2[NT]; +#pragma unroll + for (int st = 0; st < NT; ++st) { + gdnc_v8f u; +#pragma unroll + for (int e = 0; e < 8; ++e) u[e] = U[st][e] * Gs[st * 16 + 2 * e + rh].y; + UB2[st] = gdnc_acc_to_frag(u, lane); + } + const float eGL = Gs[CL - 1].x; +#pragma unroll + for (int it = 0; it < 8; ++it) { + gdnc_v8f acc; +#pragma unroll + for (int e = 0; e < 8; ++e) acc[e] = S[it][e] * eGL; +#pragma unroll + for (int st = 0; st < NT; ++st) acc = GDNC_WMMA(gdnc_frag(KTs, ST, it * 16, st * 16, lane), UB2[st], acc); + S[it] = acc; + } + // O = scale diag(e^G) Q S0 + P U + gdnc_v8f O[NT]; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) { + gdnc_v8f o; +#pragma unroll + for (int e = 0; e < 8; ++e) o[e] = qs[rt][e] * (Gs[rt * 16 + 2 * e + rh].x * scale); +#pragma unroll + for (int st = 0; st <= rt; ++st) o = GDNC_WMMA(gdnc_frag(Ps, ST, rt * 16, st * 16, lane), UB[st], o); + O[rt] = o; + } + if (on_z == nullptr) { +#pragma unroll + for (int rt = 0; rt < NT; ++rt) +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int t = t0 + rt * 16 + 2 * e + rh; + if (t < n_tokens) dst[((int64_t) t * H + h) * 128 + col] = O[rt][e]; + } + } else { + // fused gated RMS norm (rms_rows_f32): v = rsqrt(mean(o^2) + eps) * o * w * sigmoid(z) +#pragma unroll + for (int rt = 0; rt < NT; ++rt) +#pragma unroll + for (int e = 0; e < 8; ++e) { + float p = O[rt][e] * O[rt][e]; + p += __shfl_xor(p, 1, 32); p += __shfl_xor(p, 2, 32); p += __shfl_xor(p, 4, 32); p += __shfl_xor(p, 8, 32); + if (cl == 0) red[rt * 16 + 2 * e + rh][wave] = p; + } + __syncthreads(); + const float wc = on_w[col]; +#pragma unroll + for (int rt = 0; rt < NT; ++rt) +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int r = rt * 16 + 2 * e + rh, t = t0 + r; + float ss = 0.0f; +#pragma unroll + for (int w8 = 0; w8 < 8; ++w8) ss += red[r][w8]; + if (t < n_tokens) { + const int64_t oi = ((int64_t) t * H + h) * 128 + col; + const float sc = rsqrtf(ss / 128.0f + on_eps); + const float v = (sc * O[rt][e] * wc) * (1.0f / (1.0f + expf(-pz[rt][e]))); + if (on_out) on_out[oi] = v; + if (on_out16) on_out16[oi] = ggml_cuda_f32_to_bf16_rne(v); + } + } + prefetch_z(min(ch + 1, nch - 1)); + } + } + + float * s_out = state + (int64_t) h * 128 * 128; +#pragma unroll + for (int it = 0; it < 8; ++it) +#pragma unroll + for (int e = 0; e < 8; ++e) s_out[col * 128 + it * 16 + 2 * e + rh] = S[it][e]; +#else + NO_DEVICE_CODE; +#endif +} + +static bool gdn_chunk_enabled() { + static const int v = getenv("GGML_GDN_CHUNK") ? atoi(getenv("GGML_GDN_CHUNK")) : 1; + return v != 0; +} + +static int gdn_chunk_min_tokens() { + static const int v = getenv("GGML_GDN_CHUNK_MIN") ? atoi(getenv("GGML_GDN_CHUNK_MIN")) : 256; + return v; +} + +// the dispatch condition of the chunked prefill path (single source of truth for the op and the q/k-norm fusion) +static bool gdn_chunk_eligible_raw(const int cc, const bool kda, const bool keep_rs, const int64_t S_v, const int64_t H, + const int64_t Hk, const int64_t n_tokens, const int64_t n_seqs, const int64_t sq1, const int64_t sq2, + const void * q_d, const void * k_d) { +#if defined(GGML_USE_HIP) + const bool aligned = sq1 % 4 == 0 && sq2 % 4 == 0 && ((uintptr_t) q_d) % 16 == 0 && ((uintptr_t) k_d) % 16 == 0; + return !kda && !keep_rs && GGML_CUDA_CC_IS_RDNA3_5(cc) && S_v == 128 && n_seqs == 1 && n_tokens >= gdn_chunk_min_tokens() && + aligned && Hk > 0 && H % Hk == 0 && H / Hk <= 4 && gdn_chunk_enabled(); +#else + GGML_UNUSED(cc); GGML_UNUSED(kda); GGML_UNUSED(keep_rs); GGML_UNUSED(S_v); GGML_UNUSED(H); GGML_UNUSED(Hk); + GGML_UNUSED(n_tokens); GGML_UNUSED(n_seqs); GGML_UNUSED(sq1); GGML_UNUSED(sq2); GGML_UNUSED(q_d); GGML_UNUSED(k_d); + return false; +#endif +} + +bool ggml_cuda_gdn_chunk_eligible(const ggml_tensor * gdn, const ggml_tensor * q, const ggml_tensor * k) { + const ggml_tensor * v = gdn->src[2]; + if (gdn->op != GGML_OP_GATED_DELTA_NET || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F32 || + !ggml_are_same_shape(q, k) || !ggml_are_same_stride(q, k) || !ggml_is_contiguous_rows(q) || q->ne[1] != gdn->src[0]->ne[1]) { + return false; + } + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool kda = gdn->src[3]->ne[0] == v->ne[0]; + const bool keep_rs = ggml_get_op_params_i32(gdn, 0) > 1; + return gdn_chunk_eligible_raw(cc, kda, keep_rs, v->ne[0], v->ne[1], q->ne[1], v->ne[2], v->ne[3], + (int64_t) (q->nb[1] / sizeof(float)), (int64_t) (q->nb[2] / sizeof(float)), q->data, k->data); +} + +struct gdn_qk_norm_ovr { + const ggml_tensor * q; + const ggml_tensor * k; + float eps; + float mul; +}; +static std::mutex g_gdn_qk_mtx; +static std::unordered_map g_gdn_qk; + +struct gdn_conv_ovr { + const float * x; + const float * st; + const float * cw; + int64_t C, qc, kc, vc; +}; +static std::unordered_map g_gdn_conv; + +void ggml_cuda_gdn_set_conv(const ggml_tensor * gdn, const float * x, const float * st, const float * cw, + int64_t C, int64_t qc, int64_t kc, int64_t vc) { + std::lock_guard lk(g_gdn_qk_mtx); + g_gdn_conv[gdn] = { x, st, cw, C, qc, kc, vc }; +} + +struct gdn_out_norm_ovr { + const float * w; + const float * z; + float * out; + uint16_t * out16; + float eps; +}; +static std::unordered_map g_gdn_on; + +void ggml_cuda_gdn_set_out_norm(const ggml_tensor * gdn, const float * w, const float * z, float * out, uint16_t * out16, float eps) { + std::lock_guard lk(g_gdn_qk_mtx); + g_gdn_on[gdn] = { w, z, out, out16, eps }; +} + +bool ggml_cuda_gdn_chunk_eligible_node(const ggml_tensor * gdn) { + const ggml_tensor * q = gdn->src[0], * k = gdn->src[1]; + { + std::lock_guard lk(g_gdn_qk_mtx); + auto it = g_gdn_qk.find(gdn); + if (it != g_gdn_qk.end()) { q = it->second.q; k = it->second.k; } + } + return ggml_cuda_gdn_chunk_eligible(gdn, q, k); +} + +void ggml_cuda_gdn_set_qk_norm(const ggml_tensor * gdn, const ggml_tensor * q_raw, const ggml_tensor * k_raw, float eps, float mul) { + std::lock_guard lk(g_gdn_qk_mtx); + g_gdn_qk[gdn] = { q_raw, k_raw, eps, mul }; +} + template static void launch_gated_delta_net( const float * q_d, const float * k_d, const float * v_d, @@ -178,14 +1112,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: @@ -222,8 +1196,36 @@ static void launch_gated_delta_net( static void ggml_cuda_op_gated_delta_net_impl( ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_cuda_gated_delta_net_fused_cache * cache) { - ggml_tensor * src_q = dst->src[0]; - ggml_tensor * src_k = dst->src[1]; + // q/k-norm fusion: the normalized q/k were never computed; read the raw views and normalize in the chunk kernels + gdn_qk_norm_ovr qk_ovr = { nullptr, nullptr, -1.0f, 1.0f }; + { + std::lock_guard lk(g_gdn_qk_mtx); + auto it = g_gdn_qk.find(dst); + if (it != g_gdn_qk.end()) { + qk_ovr = it->second; + g_gdn_qk.erase(it); + } + } + gdn_out_norm_ovr on_ovr = { nullptr, nullptr, nullptr, nullptr, 0.0f }; + { + std::lock_guard lk(g_gdn_qk_mtx); + auto it = g_gdn_on.find(dst); + if (it != g_gdn_on.end()) { + on_ovr = it->second; + g_gdn_on.erase(it); + } + } + gdn_conv_ovr cv_ovr = { nullptr, nullptr, nullptr, 0, 0, 0, 0 }; + { + std::lock_guard lk(g_gdn_qk_mtx); + auto it = g_gdn_conv.find(dst); + if (it != g_gdn_conv.end()) { + cv_ovr = it->second; + g_gdn_conv.erase(it); + } + } + ggml_tensor * src_q = qk_ovr.q ? (ggml_tensor *) qk_ovr.q : dst->src[0]; + ggml_tensor * src_k = qk_ovr.k ? (ggml_tensor *) qk_ovr.k : dst->src[1]; ggml_tensor * src_v = dst->src[2]; ggml_tensor * src_g = dst->src[3]; ggml_tensor * src_beta = dst->src[4]; @@ -294,6 +1296,40 @@ static void ggml_cuda_op_gated_delta_net_impl( state_slot_stride = cache->slot_stride; } +#if defined(GGML_USE_HIP) + { + const int dev = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[dev].cc; + const float qk_eps = qk_ovr.q ? qk_ovr.eps : -1.0f; + const float qk_mul = qk_ovr.q ? qk_ovr.mul : 1.0f; + const bool eligible = gdn_chunk_eligible_raw(cc, kda, keep_rs, S_v, H, neqk1, n_tokens, n_seqs, sq1, sq2, q_d, k_d); + if ((qk_ovr.q || on_ovr.z || cv_ovr.x) && !eligible) { + GGML_ABORT("gated_delta_net: conv / q/k / output norms were fused for %s but the chunked path is not taken", dst->name); + } + if (eligible) { + gdnc_in in = { q_d, k_d, v_d, sq1, sq2, sv1, sv2, cv_ovr.x, cv_ovr.st, cv_ovr.cw, cv_ovr.C, cv_ovr.qc, cv_ovr.kc, cv_ovr.vc }; + auto run = [&](auto CLc) { + constexpr int CL = decltype(CLc)::value; + const int nch = (int) ((n_tokens + CL - 1) / CL); + ggml_cuda_pool_alloc<_Float16> Tb(ctx.pool(), (size_t) H * nch * CL * CL); + ggml_cuda_pool_alloc<_Float16> Pb(ctx.pool(), (size_t) H * nch * CL * CL); + ggml_cuda_pool_alloc Gb(ctx.pool(), (size_t) H * nch * CL); + ggml_cuda_pool_alloc<_Float16> Qh(ctx.pool(), (size_t) neqk1 * n_tokens * 128); + ggml_cuda_pool_alloc<_Float16> Kh(ctx.pool(), (size_t) neqk1 * n_tokens * 128); + ggml_cuda_pool_alloc<_Float16> Vh(ctx.pool(), (size_t) H * n_tokens * 128); + gdn_chunk_prep_kernel<<>>(in, g_d, b_d, + Tb.get(), Pb.get(), Gb.get(), Qh.get(), Kh.get(), Vh.get(), (int) H, (int) neqk1, (int) n_tokens, nch, + sb1, sb2, scale, qk_eps, qk_mul); + gdn_chunk_scan_kernel<<>>(Qh.get(), Kh.get(), Vh.get(), s_d, dst_d, state_d, + Tb.get(), Pb.get(), Gb.get(), (int) H, (int) neqk1, (int) n_tokens, nch, scale, + on_ovr.w, on_ovr.z, on_ovr.out, on_ovr.out16, on_ovr.eps); + CUDA_CHECK(cudaGetLastError()); + }; + run(std::integral_constant{}); + return; + } + } +#endif if (kda) { if (keep_rs) { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, @@ -325,3 +1361,313 @@ 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. When args.qk_rms_scale is set, the q/k norms +// replay rms_norm_f32<256, false> + scale_f32 (the qwen4exp build_gdn_l2_norm form) with the same thread +// layout, reduction tree and single-rounding products. + +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 qk_red[8]; // q/k rms norm: the 4+4 per-warp partial sums of the standalone 256-thread block + + 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. q/k norm: l2_norm_dual_f32_s128 for q (warp 0) and k (warp 1), or, when the graph uses the + // qwen4exp build_gdn_l2_norm representation RMS_NORM(x, eps) -> SCALE(mul, add), the exact + // replay of the dispatched rms_norm_f32<256, false> + scale_f32 pair + if (!args.qk_rms_scale) { + 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]; + } + } + } else { + // rms_norm_f32<256, false>: 128 elements over 128 lanes (4 warps, one element per lane; warps 4..7 + // of the standalone 256-thread block hold exact zeros that drop out of the tree), the warp + // butterfly (xor 16/8/4/2/1), the cross-warp tree ((v0 + v2) + (v1 + v3)), mean = sum / 128, + // scale = rsqrtf(mean + eps), dst = scale * x; then scale_f32: dst = mul * dst + add. The standalone + // scale_f32 compiles its mul+add to a single FMA (v_fma_f32 in the RDNA3.5 build), so the same + // contraction is spelled explicitly here. + const int iqk = warp < 4 ? 0 : 1; + if (warp < 8) { + const float * x = iqk == 0 ? q_c : k_c; + const float xi = x[(warp & 3) * warp_size + lane]; + float part = fmaf(xi, xi, 0.0f); + part = warp_reduce_sum(part); + if (lane == 0) { + qk_red[warp] = part; + } + } + __syncthreads(); + const float total = iqk == 0 ? (qk_red[0] + qk_red[2]) + (qk_red[1] + qk_red[3]) + : (qk_red[4] + qk_red[6]) + (qk_red[5] + qk_red[7]); + const float scale = rsqrtf(total / (float) args.S + args.qk_rms_eps[iqk]); + if (warp < 8) { + const float * x = iqk == 0 ? q_c : k_c; + float * y = iqk == 0 ? q_n : k_n; + const int i = (warp & 3) * warp_size + lane; + const float rms = scale * x[i]; + y[i] = fmaf(args.qk_scale_mul[iqk], rms, args.qk_scale_add[iqk]); + } + } + + // 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, const float * z) { + 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); + const float normed = scale * xi * norm_w[tid]; + if (z != nullptr) { + // the following sigmoid(z) * normed of unary_gated_op_kernel, same expression + out[(int64_t) h * S + tid] = (1.0f / (1.0f + expf(-z[(int64_t) h * S + tid]))) * normed; + } else { + out[(int64_t) h * S + tid] = normed; + } +} + +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()); + if (args.n_tokens == 1 && args.K == 1) { + ggml_cuda_kernel_launch(gdn_decode_fused_cuda<128, 4>, launch_params, args); + return; + } + + GGML_ASSERT(args.n_tokens >= 1 && args.n_tokens <= args.K && args.K <= 8); + const int64_t conv_elems = (args.d_conv - 1) * args.S * (2 * args.H_k + args.H_v); + for (int64_t slot = args.n_tokens; slot < args.K; ++slot) { + CUDA_CHECK(cudaMemcpyAsync(args.conv_state_out + slot * args.conv_slot_stride, args.conv_state_in, + conv_elems * sizeof(float), cudaMemcpyDeviceToDevice, ctx.stream())); + } + for (int64_t t = 0; t < args.n_tokens; ++t) { + const int64_t slot = args.n_tokens - 1 - t; + ggml_cuda_gdn_decode_args step = args; + step.qkv += t * args.qkv_token_stride; + step.alpha += t * args.H_v; + step.beta += t * args.H_v; + step.conv_state_out += slot * args.conv_slot_stride; + step.state_out += slot * args.state_slot_stride; + step.attn_out += t * args.S * args.H_v; + if (t > 0) { + step.conv_state_in = args.conv_state_out + (slot + 1) * args.conv_slot_stride; + step.state_cache = args.state_out + (slot + 1) * args.state_slot_stride; + step.state_row_stride = 0; + } + ggml_cuda_kernel_launch(gdn_decode_fused_cuda<128, 4>, launch_params, step); + } +} + +void ggml_cuda_op_gdn_decode_fused(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args_in) { + ggml_cuda_op_gdn_decode_fused_gated(ctx, args_in, nullptr, nullptr); +} + +void ggml_cuda_op_gdn_decode_fused_gated(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args_in, + const float * z, float * gated_out) { + 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 * args_in.n_tokens); + 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()); + // Read all token inputs before the norm output can reuse their storage. + for (int64_t t = 0; t < args.n_tokens; ++t) { + const int64_t offset = t * args.S * args.H_v; + ggml_cuda_kernel_launch(gdn_decode_norm_cuda<128>, norm_params, attn.get() + offset, args.norm_w, + (z ? gated_out : args.out) + offset, args.eps_rms, z ? z + offset : nullptr); + } +} diff --git a/ggml/src/ggml-cuda/gated_delta_net.cuh b/ggml/src/ggml-cuda/gated_delta_net.cuh index f9bf43706789..77c34d2226fe 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cuh +++ b/ggml/src/ggml-cuda/gated_delta_net.cuh @@ -9,6 +9,69 @@ struct ggml_cuda_gated_delta_net_fused_cache { void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// Chunked prefill (RDNA3.5): the q/k L2 norms that precede GATED_DELTA_NET (RMS_NORM(x, eps/n) -> SCALE(1/sqrt(n)), +// i.e. x * mul * rsqrt(sum(x^2) + eps) per 128-row) are folded into the chunk kernels. ggml_cuda_gdn_chunk_eligible +// is the exact dispatch condition of the chunked path for this GDN node with these q/k tensors; set_qk_norm registers +// the raw q/k views for the next evaluation of that node (consumed by it). +bool ggml_cuda_gdn_chunk_eligible(const ggml_tensor * gdn, const ggml_tensor * q, const ggml_tensor * k); +void ggml_cuda_gdn_set_qk_norm(const ggml_tensor * gdn, const ggml_tensor * q_raw, const ggml_tensor * k_raw, float eps, float mul); +// chunked prefill: the gated RMS norm of the output (RMS_NORM -> MUL(w) -> MUL(sigmoid(z))) written by the scan kernel +// straight into out (+ optional BF16 copy out16); z / out laid out like the GDN output rows [S_v, H, T] +void ggml_cuda_gdn_set_out_norm(const ggml_tensor * gdn, const float * w, const float * z, float * out, uint16_t * out16, float eps); +// chunked prefill: the causal conv + SiLU producing q/k/v is computed by the chunk kernels from the raw conv input +// x [T, C] (row stride C floats), history st [C, 3] and taps cw [C, 4]; qc/kc/vc = channel of element 0 of each view +void ggml_cuda_gdn_set_conv(const ggml_tensor * gdn, const float * x, const float * st, const float * cw, + int64_t C, int64_t qc, int64_t kc, int64_t vc); +// ggml_cuda_gdn_chunk_eligible with the q/k views registered by set_qk_norm (if any) +bool ggml_cuda_gdn_chunk_eligible_node(const ggml_tensor * gdn); + // 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; + // q/k normalization: L2_NORM (eps_l2) by default, or, when qk_rms_scale is set, the qwen4exp + // build_gdn_l2_norm pair RMS_NORM(x, eps) -> SCALE(mul, add) replayed exactly (see the decode kernel) + bool qk_rms_scale = false; + float qk_rms_eps[2] = { 0.0f, 0.0f }; + float qk_scale_mul[2] = { 0.0f, 0.0f }; + float qk_scale_add[2] = { 0.0f, 0.0f }; + // 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; + int64_t n_tokens = 1, K = 1; + int64_t qkv_token_stride = 0; + int64_t conv_slot_stride = 0; + int64_t state_slot_stride = 0; +}; + +void ggml_cuda_op_gdn_decode_fused(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args); +// same, with the following sigmoid(z) * norm applied by the norm kernel, which then writes gated_out instead of args.out +void ggml_cuda_op_gdn_decode_fused_gated(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args, + const float * z, float * gated_out); +// 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..d80ed22a172a 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; @@ -404,6 +470,14 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_kq<32, dst_t, dequantize_mxfp4>(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + get_rows_cuda_kq<32, dst_t, dequantize_rocmfp4>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + get_rows_cuda_kq<32, dst_t, dequantize_rocmfp4_fast>(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; default: GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type)); break; @@ -488,3 +562,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 b50d05d5f76f..6b904b816466 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -61,6 +61,14 @@ #include "ggml-cuda/wkv.cuh" #include "ggml-cuda/gla.cuh" #include "ggml-cuda/gated_delta_net.cuh" +#include +#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" @@ -412,6 +420,8 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { static const int MAX_BUFFERS = 256; + // far below any KV-sized scratch, so small buffers are never rounded + static const size_t POOL_POW2_MIN_SIZE = 4ull*1024*1024; int device; struct ggml_cuda_buffer { @@ -451,6 +461,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { #endif size_t best_diff = 1ull << 36; int ibest = -1; + size_t largest_unusable = 0; for (int i = 0; i < MAX_BUFFERS; ++i) { ggml_cuda_buffer& b = buffer_pool[i]; if (b.ptr != nullptr) { @@ -458,7 +469,9 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { ++nnz; if (b.size > max_size) max_size = b.size; #endif - if (b.size >= size) { + if (b.size < size) { + if (b.size > largest_unusable) largest_unusable = b.size; + } else { size_t diff = b.size - size; if (diff < best_diff) { best_diff = diff; @@ -485,6 +498,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { void * ptr; size_t look_ahead_size = (size_t) (1.05 * size); look_ahead_size = 256 * ((look_ahead_size + 255)/256); + const size_t exact_size = look_ahead_size; + // Nothing cached can serve this request, and a cached block smaller than it + // never will, so a caller that grows its request step by step strands one block + // per step. A largest cached block just below the request is that signature: + // round up to a power of two so the series reuses one block per octave instead. + if (size > POOL_POW2_MIN_SIZE && largest_unusable >= size - size/4) { + size_t pow2 = POOL_POW2_MIN_SIZE; + while (pow2 < size && pow2 <= SIZE_MAX/2) { + pow2 *= 2; + } + if (pow2 >= size) { + look_ahead_size = pow2; + } + } ggml_cuda_set_device(device); cudaError_t err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); if (err == cudaErrorMemoryAllocation) { @@ -494,6 +521,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0); CUDA_CHECK(cudaDeviceSynchronize()); clear_pool(); + look_ahead_size = exact_size; // retry without the rounding err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); if (err == cudaSuccess) { GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device); @@ -510,6 +538,25 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { } void free(void * ptr, size_t size) override { + // GGML_CUDA_POOL_FREE_MIN_MB: blocks at or above this size return to the + // driver immediately instead of staying cached in the pool. Bounds the + // peak+caches ratchet of large transient scratch (e.g. QSA prefill/decode + // intermediates that scale with the KV length) at a small cudaFree cost + // per released block; small scratch keeps the hot cache path. + // Note: the legacy pool has no window to shrink by itself, so once a block + // is cached it is reused or replaced only by larger allocations until the + // process ends. Only the buffer stores freed here participate after the + // QSA/attention scratch is done each step. + // Applies to the legacy (non-VMM) pool; the VMM pool keeps its own + // management (GGML_HIP_NO_VMM=OFF GPU builds). + static const int free_min_mb = getenv("GGML_CUDA_POOL_FREE_MIN_MB") + ? std::atoi(getenv("GGML_CUDA_POOL_FREE_MIN_MB")) : 0; + if (free_min_mb > 0 && size >= (size_t) free_min_mb * 1024 * 1024) { + ggml_cuda_set_device(device); + CUDA_CHECK(cudaFree(ptr)); + pool_size -= size; + return; + } for (int i = 0; i < MAX_BUFFERS; ++i) { ggml_cuda_buffer& b = buffer_pool[i]; if (b.ptr == nullptr) { @@ -523,6 +570,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { CUDA_CHECK(cudaFree(ptr)); pool_size -= size; } + }; // pool with virtual memory @@ -697,6 +745,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)); } @@ -762,7 +827,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; @@ -1818,6 +1885,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_TENSOR_BINARY_OP_LOCALS + const int32_t hint = ggml_get_op_params_i32(dst, 1); if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } @@ -1835,12 +1903,24 @@ 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); return; } + // On gfx1151 a small output-row count that is not a multiple of the WMMA rows-per-block cannot use mul_mat_f, and the matvec + // gate above caps AMD at 3 columns, so a wider batch falls through to rocBLAS, which splits K to fill the grid. Measured with + // a 4-token MTP verification batch: the qwen4exp [10240, 4] BF16 hyper-connection inject matrices cost 230 us per call through + // rocBLAS (0.35 GB/s) against ~5 us for the matvec kernel. Drop this once rocBLAS stops losing that much on small-M GEMM. + if (ne11 > 3 && ne11 <= 8 && GGML_CUDA_CC_IS_AMD(cc) && !ggml_is_quantized(src0->type) && + src0->ne[1] <= 64 && + src0->ne[1] % (GGML_CUDA_CC_IS_CDNA(cc) ? MMF_ROWS_PER_BLOCK_CDNA : MMF_ROWS_PER_BLOCK) != 0 && + ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, 3)) { + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + } // A transposed vector can still use MMVQ (i.e. ne01 == 1) if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 && src0->type == GGML_TYPE_F32 @@ -1863,6 +1943,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; @@ -1931,6 +2015,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; @@ -2059,6 +2147,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); @@ -2586,9 +2684,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; @@ -2599,6 +2740,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; } @@ -2611,16 +2762,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; @@ -2970,7 +3112,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); @@ -2988,7 +3131,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]; } @@ -3029,6 +3172,387 @@ 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 one sequence with S == 128 and at most 8 tokens and rollback slots. +// 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 = 4; + const int64_t C = concat->ne[1]; + const int64_t n_tokens = concat->src[1]->ne[0]; + if (n_tokens < 1 || n_tokens > 8 || concat->ne[0] != d_conv - 1 + n_tokens || + 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]; // [n_tokens, 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] != n_tokens || qkv_t->ne[1] != C || qkv_t->ne[2] != 1 || qkv_t->ne[3] != 1 || + qkv_t->nb[1] % sizeof(float) != 0 || qkv_t->nb[0] % 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; + const ggml_tensor * conv_cpy[8]; + int n_conv_cpy = 0; + // qwen4exp (build_gdn_l2_norm) spells the q/k L2 norm as RMS_NORM(x, eps) -> SCALE(mul, add): accept the + // pair in place of each single L2_NORM at positions 4/5 without changing the 14-entry op list. The extra + // SCALE nodes are absorbed into positions 4/5 and only widen the skipped range, so idx[] keeps pointing + // at the logical ops (in particular idx[10] is still the gated delta net). + int qk_mode = -1; // -1: undecided, 0: L2_NORM, 1: RMS_NORM + SCALE + const ggml_tensor * qk_scale[2] = { nullptr, nullptr }; + for (int j = node_idx + 1; j < cgraph->n_nodes && found < n_ops; ) { + const ggml_tensor * t = cgraph->nodes[j]; + if (ggml_cuda_is_view_or_noop(t)) { + ++j; + continue; + } + if (t->type != GGML_TYPE_F32) { + return 0; + } + if (found == 1 && t->op == GGML_OP_CPY) { + if (n_conv_cpy == 8) { + return 0; + } + conv_cpy[n_conv_cpy++] = t; + last = j++; + continue; + } + if (found == 4 || found == 5) { + const int iqk = found - 4; + if (t->op == GGML_OP_L2_NORM) { + if (qk_mode == 1) { + return 0; + } + qk_mode = 0; + } else if (t->op == GGML_OP_RMS_NORM && t->src[1] == nullptr) { + if (qk_mode == 0) { + return 0; + } + qk_mode = 1; + int j_scale = j + 1; + while (j_scale < cgraph->n_nodes && ggml_cuda_is_view_or_noop(cgraph->nodes[j_scale])) { + ++j_scale; + } + if (j_scale >= cgraph->n_nodes) { + return 0; + } + const ggml_tensor * s = cgraph->nodes[j_scale]; + if (s->op != GGML_OP_SCALE || s->type != GGML_TYPE_F32 || s->src[0] != t || + s->src[1] != nullptr || !ggml_is_contiguous(s) || !ggml_are_same_shape(s, t)) { + return 0; + } + qk_scale[iqk] = s; + n[found] = t; + idx[found] = j; + last = j_scale; + ++found; + j = j_scale + 1; + continue; + } else { + return 0; + } + } else if (t->op != expected[found]) { + return 0; + } + if (found == 0) { + conv_cpy[n_conv_cpy++] = t; + } + n[found] = t; + idx[found] = j; + last = j; + ++found; + ++j; + } + 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; + } + } + for (int k = 0; k < 2; ++k) { + if (qk_scale[k] && (qk_scale[k]->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return 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); + }; + + const int64_t K = ggml_get_op_params_i32(gdn, 0); + if (K < n_tokens || K > 8 || n_conv_cpy != K || + ((n_tokens > 1 || K > 1) && getenv("GGML_CUDA_DISABLE_GDN_ROLLBACK") != nullptr)) { + return 0; + } + + const ggml_tensor * conv_dst = conv_cpy[0]->src[1]; + for (int k = 1; k < n_conv_cpy; ++k) { + if (conv_cpy[k]->src[1]->view_offs < conv_dst->view_offs) { + conv_dst = conv_cpy[k]->src[1]; + } + } + int64_t conv_slot_stride = 0; + if (K > 1) { + const ggml_tensor * cache = conv_dst->view_src; + if (!cache || cache->ne[0] != (d_conv - 1) * C || cache->ne[1] < K || cache->ne[1] % K != 0 || !ggml_is_contiguous(cache)) { + return 0; + } + conv_slot_stride = cache->nb[1] / sizeof(float) * (cache->ne[1] / K); + } + unsigned int conv_slots = 0; + for (int k = 0; k < n_conv_cpy; ++k) { + const ggml_tensor * src = conv_cpy[k]->src[0]; + const ggml_tensor * dst = conv_cpy[k]->src[1]; + const size_t delta = dst->view_offs - conv_dst->view_offs; + const int64_t slot = K == 1 ? 0 : delta / (conv_slot_stride * sizeof(float)); + if (slot >= K || delta != slot * conv_slot_stride * sizeof(float) || (conv_slots & (1u << slot)) || + (K > 1 && dst->view_src != conv_dst->view_src) || (conv_cpy[k]->flags & GGML_TENSOR_FLAG_OUTPUT) || + !is_view_of(src, concat, std::max(0, n_tokens - slot) * sizeof(float)) || + src->ne[0] != d_conv - 1 || src->ne[1] != C || src->ne[2] != 1 || src->ne[3] != 1 || + src->nb[0] != sizeof(float) || 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; + } + conv_slots |= 1u << slot; + } + + // 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 * n_tokens || + 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] != n_tokens || v_view->ne[3] != 1 || + 2 * S * H_k + S * H_v != C) { + return 0; + } + + // q/k/v are views into the silu output: [S, H, n_tokens, 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] == n_tokens && + t->ne[3] == 1 && t->nb[0] == sizeof(float) && t->nb[1] == S * sizeof(float) && + (n_tokens == 1 || t->nb[2] == C * 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 * n_tokens || add->src[0]->type != GGML_TYPE_F32 || add->src[1]->type != GGML_TYPE_F32 || + ggml_nelements(add->src[0]) != H_v * n_tokens || 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 * n_tokens || + ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID || sigmoid->src[0]->type != GGML_TYPE_F32 || + ggml_nelements(sigmoid->src[0]) != H_v * n_tokens || !ggml_is_contiguous(sigmoid->src[0]) || ggml_nelements(sigmoid) != H_v * n_tokens) { + return 0; + } + + // the gated delta net consumes the L2_NORM outputs, or the SCALE outputs of the RMS_NORM + SCALE pair + const ggml_tensor * qk_out[2] = { qk_mode == 1 ? qk_scale[0] : l2_q, qk_mode == 1 ? qk_scale[1] : l2_k }; + + // gated_delta_net(q, k, v, gate, beta, state) with the state gathered from the cache by get_rows + if (gdn->src[0] != qk_out[0] || gdn->src[1] != qk_out[1] || + !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 * n_tokens || ggml_nelements(gdn->src[4]) != H_v * n_tokens || + !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] || + (K > 1 && state_cpy.slot_stride < S * S * H_v)) { + 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] != n_tokens || x->ne[3] != 1 || + x->nb[0] != sizeof(float) || x->nb[1] != S * sizeof(float) || + (n_tokens > 1 && x->nb[2] != S * H_v * 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. + for (int k = 0; k < n_conv_cpy; ++k) { + const ggml_tensor * out = conv_cpy[k]->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 *) conv_dst->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.qk_rms_scale = qk_mode == 1; + if (qk_mode == 1) { + args.qk_rms_eps[0] = ggml_get_op_params_f32(l2_q, 0); + args.qk_rms_eps[1] = ggml_get_op_params_f32(l2_k, 0); + args.qk_scale_mul[0] = ggml_get_op_params_f32(qk_scale[0], 0); + args.qk_scale_mul[1] = ggml_get_op_params_f32(qk_scale[1], 0); + args.qk_scale_add[0] = ggml_get_op_params_f32(qk_scale[0], 1); + args.qk_scale_add[1] = ggml_get_op_params_f32(qk_scale[1], 1); + } + 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; + args.n_tokens = n_tokens; + args.K = K; + args.qkv_token_stride = qkv_t->nb[0] / sizeof(float); + args.conv_slot_stride = conv_slot_stride; + args.state_slot_stride = state_cpy.slot_stride; + + // 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 (std::find(conv_cpy, conv_cpy + n_conv_cpy, cgraph->nodes[j]) != conv_cpy + n_conv_cpy || + 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; @@ -3439,97 +3963,1326 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } // try and fuse nodes and return the number of nodes to skip -static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { +// 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; + }; - static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); - if (disable_fusion) { + 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; } - ggml_tensor * node = cgraph->nodes[i]; + 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; + } - 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)) { - ggml_cuda_op_moe_weighted_reduction( - *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); - return match.node_count - 1; + 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; + a.vis_tails = nullptr; a.vis_starts = nullptr; + for (auto & l : a.vis_leafs) { l = nullptr; } + if (getenv("GGML_CUDA_DISABLE_IDX_VIS") != nullptr || ns != 1) return count; + + auto vis_reject = [&]() { + for (auto & l : a.vis_leafs) { l = nullptr; } + return count; + }; + // qwen4exp_apply_compact_visibility with one sequence, emitted right after the head sum: + // starts = CPY(VIEW(limits) I32 [nb]) F32; tails = RESHAPE(CPY(VIEW(limits) I32 [nt]) F32, 1, nt) + // ADD(sum, LOG(STEP(SUB(REPEAT(tails, sum), starts)))) + // Every node between the head sum and the final ADD must belong to that chain, so skipping them is safe. + const int last = k - 1; + int fin = -1; + for (int j = last + 1; j < g->n_nodes && j <= last + 14; ++j) { + const ggml_tensor * n = g->nodes[j]; + if (n->op == GGML_OP_ADD && n->src[0] == prev) { fin = j; break; } + } + if (fin < 0) return vis_reject(); + const ggml_tensor * add = g->nodes[fin]; + const ggml_tensor * lg = add->src[1]; + auto is_unary = [](const ggml_tensor * t, ggml_unary_op op) { + return t && t->op == GGML_OP_UNARY && ggml_get_unary_op(t) == op && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t); + }; + if (add->type != GGML_TYPE_F32 || !ggml_is_contiguous(add) || !ggml_are_same_shape(add, prev)) return vis_reject(); + if (!lg || lg->op != GGML_OP_LOG || lg->type != GGML_TYPE_F32 || !ggml_is_contiguous(lg) || !ggml_are_same_shape(lg, prev)) return vis_reject(); + const ggml_tensor * st = lg->src[0]; + if (!is_unary(st, GGML_UNARY_OP_STEP) || !ggml_are_same_shape(st, prev)) return vis_reject(); + const ggml_tensor * sb = st->src[0]; + if (!sb || sb->op != GGML_OP_SUB || sb->type != GGML_TYPE_F32 || !ggml_is_contiguous(sb) || !ggml_are_same_shape(sb, prev)) return vis_reject(); + const ggml_tensor * rp = sb->src[0]; + const ggml_tensor * c1 = sb->src[1]; + if (!rp || rp->op != GGML_OP_REPEAT || rp->type != GGML_TYPE_F32 || !ggml_is_contiguous(rp) || !ggml_are_same_shape(rp, prev)) return vis_reject(); + auto is_i32_cast = [](const ggml_tensor * c, int64_t n) { + return c && c->op == GGML_OP_CPY && c->type == GGML_TYPE_F32 && ggml_is_contiguous(c) && ggml_nelements(c) == n && + c->src[0] && c->src[0]->type == GGML_TYPE_I32 && ggml_is_contiguous(c->src[0]) && + ggml_nelements(c->src[0]) == n; + }; + // starts: [nb] broadcast over the queries + if (!is_i32_cast(c1, nb) || c1->ne[0] != nb) return vis_reject(); + const ggml_tensor * tr = rp->src[0]; + const ggml_tensor * c2 = tr; + if (tr && tr->op == GGML_OP_RESHAPE) { c2 = tr->src[0]; } + if (!tr || tr->ne[0] != 1 || tr->ne[1] != nt || ggml_nelements(tr) != nt || !is_i32_cast(c2, nt)) return vis_reject(); + + // the chain, and nothing else, sits between the head sum and the final ADD + const ggml_tensor * chain[10] = { c1->src[0], c1, c2->src[0], c2, tr, rp, sb, st, lg, add }; + // The limits views and their I32->F32 casts read a graph input, which ggml_can_fuse_subgraph_ext rejects as an + // external view source (and the cast self-reference would skew its use count), so they are checked here: each + // must be read only by its successor in the chain. The fused kernel reads the limits input directly. The tails + // RESHAPE is a view of its cast and is checked the same way. When the scheduler copies the limits views to the + // device as split inputs, the casts read those copies (leafs) and the original VIEW nodes stay in the range + // unused; a view computes nothing, so skipping one is harmless whoever reads it. + const ggml_tensor * leafs[5] = { c1->src[0], c1, c2->src[0], c2, tr }; + const ggml_tensor * users[5] = { c1, sb, c2, tr, rp }; + int indices2[32]; ggml_op ops2[32]; + int count2 = 0; + int n_manual = 0; + bool seen[5] = {}; + for (int j = i; j <= fin; ++j) { + const ggml_tensor * n = g->nodes[j]; + if (j > last) { + int leaf = -1; + for (int l = 0; l < 5; ++l) { if (leafs[l] == n) leaf = l; } + bool member = leaf >= 0; + for (const ggml_tensor * c : chain) { member = member || c == n; } + if (!member) { + if (n->op != GGML_OP_VIEW && n->op != GGML_OP_RESHAPE && n->op != GGML_OP_PERMUTE && + n->op != GGML_OP_TRANSPOSE && n->op != GGML_OP_NONE) return vis_reject(); + if (n_manual >= 16) return vis_reject(); + a.vis_leafs[n_manual++] = n; + continue; + } + if (leaf >= 0) { + if (n->flags & GGML_TENSOR_FLAG_OUTPUT) return vis_reject(); + if (ggml_node_get_use_count(g, j) != 1 + (n->src[1] == n ? 1 : 0)) return vis_reject(); + int refs = 0; + for (int s2 = 0; s2 < GGML_MAX_SRC; ++s2) { if (users[leaf]->src[s2] == n) ++refs; } + if (refs != 1) return vis_reject(); + if (n_manual >= 16) return vis_reject(); + a.vis_leafs[n_manual++] = n; + seen[leaf] = true; + continue; } } + if (count2 >= 32) return vis_reject(); + indices2[count2] = j; ops2[count2] = n->op; ++count2; } - - // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache - if (node->op == GGML_OP_GATED_DELTA_NET) { - 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) { -#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); -#endif - ggml_cuda_op_gated_delta_net_fused_cache(*cuda_ctx, node, fused_state_cpy); - return nodes_to_skip; + // the casts and the tails RESHAPE must be nodes of the skipped range; the limits views may instead be leafs + // (device copies of split inputs), which are ready before the range starts + if (!seen[1] || !seen[3] || !seen[4] || tr == c2) return vis_reject(); + for (int l : { 0, 2 }) { + if (!seen[l]) { + for (int j = i; j < g->n_nodes; ++j) { if (g->nodes[j] == leafs[l]) return vis_reject(); } + } + } + if (!ggml_can_fuse_subgraph_ext(g, indices2, count2, ops2, &fin, 1)) { + for (auto & l : a.vis_leafs) { l = nullptr; } + return vis_reject(); + } + count2 = fin - i + 1; // nodes skipped: every node from the RELU to the final ADD + // the fused kernel reads the scores and the limits while it writes the result + for (const ggml_tensor * in : { src, (const ggml_tensor *) c1->src[0], (const ggml_tensor *) c2->src[0] }) { + if (in->data && add->data) { + const uintptr_t av = (uintptr_t) in->data, bv = (uintptr_t) add->data; + const bool overlap = av <= bv ? bv - av < ggml_nbytes(in) : av - bv < ggml_nbytes(add); + if (overlap) { return vis_reject(); } } } + a.dst = (ggml_tensor *) add; + a.vis_starts = c1->src[0]; + a.vis_tails = c2->src[0]; + return count2; +} - //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) { - ggml_cuda_topk_moe_args args; - const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); - std::vector ops; - ops.reserve(13); // max ops; avoids gcc -Wstringop-overflow false positive +// MUL_MAT(keys F32, q) -> RESHAPE [M, 4, nq] -> RELU/head sum (+ visibility): the scorer writes the summed scores +static int ggml_cuda_match_idx_gemm(ggml_backend_cuda_context & ctx, const ggml_cgraph * g, int i, ggml_cuda_idx_relu_sum_args & a) { + if (getenv("GGML_CUDA_DISABLE_IDX_GEMM") != nullptr || i + 3 >= g->n_nodes) return 0; + const ggml_tensor * mm = g->nodes[i]; + if (mm->op != GGML_OP_MUL_MAT || mm->src[2] != nullptr || mm->type != GGML_TYPE_F32 || !ggml_is_contiguous(mm)) return 0; + const ggml_tensor * rs = g->nodes[i + 1]; + if (rs->op != GGML_OP_RESHAPE || rs->src[0] != mm || rs->ne[0] != mm->ne[0] || rs->ne[1] != 4 || rs->ne[3] != 1 || + rs->ne[1] * rs->ne[2] != mm->ne[1] * mm->ne[2] * mm->ne[3]) return 0; + const ggml_tensor * relu = g->nodes[i + 2]; + if (relu->op != GGML_OP_UNARY || ggml_get_unary_op(relu) != GGML_UNARY_OP_RELU || relu->src[0] != rs) return 0; + const int c = ggml_cuda_match_idx_relu_sum(g, i + 2, a); + if (c <= 0 || a.score != rs || a.heads != 4) return 0; + if (!ggml_cuda_mmb_idx_score_supported(ctx, mm->src[0], mm->src[1], mm, a.heads)) return 0; + const int count = c + 2; + const int output = i + count - 1; + if (g->nodes[output] != a.dst) return 0; + int indices[32]; ggml_op ops[32]; + int n = 0; + for (int j = i; j <= output; ++j) { + const ggml_tensor * node = g->nodes[j]; + bool checked = false; + for (const ggml_tensor * l : a.vis_leafs) { checked = checked || (l && l == node); } + if (checked) continue; + if (n >= 32) return 0; + indices[n] = j; ops[n] = node->op; ++n; + } + if (!ggml_can_fuse_subgraph_ext(g, indices, n, ops, &output, 1)) return 0; + for (const ggml_tensor * in : { (const ggml_tensor *) mm->src[0], (const ggml_tensor *) mm->src[1], a.vis_tails, a.vis_starts }) { + if (in && in->data && a.dst->data) { + const uintptr_t av = (uintptr_t) in->data, bv = (uintptr_t) a.dst->data; + const bool overlap = av <= bv ? bv - av < ggml_nbytes(in) : av - bv < ggml_nbytes(a.dst); + if (overlap) { return 0; } + } + } + return count; +} - if (can_fuse) { - const ggml_tensor * logits = node->src[0]; - ggml_tensor * weights = nullptr; - ggml_tensor * ids = nullptr; - const ggml_tensor * bias = nullptr; - const ggml_tensor * clamp = nullptr; - const ggml_tensor * scale = nullptr; +// 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; +} - if (!args.delayed_softmax) { - int out_nodes[2]; // nodes which can't be elided +// node i is an HC gate GEMM that the fused gate mix can take: returns the mix node count after it, else 0 +// (callers also check the executor's reader rule, ggml_node_has_n_uses(cgraph, i, 1)) +static int ggml_cuda_hc_gate_mix_route(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph, const int i, ggml_cuda_hc_mix_args & ma) { + const ggml_tensor * node = cgraph->nodes[i]; + if (node->op != GGML_OP_MUL_MAT || i + 1 >= cgraph->n_nodes || !ggml_cuda_mmb_gatemix() || !GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc) || + !ggml_cuda_mmb_supported_mm(ctx, node->src[0], node->src[1], node)) { + return 0; + } + const int count = ggml_cuda_hc_mix_closed(cgraph, i + 1, ma); + return count > 0 && ma.gate == node && ggml_cuda_hc_gate_mix_supported(ctx, node->src[0], node->src[1], ma.xn, ma.dst, ma.hc) ? count : 0; +} - if (args.sigmoid) { - ops.insert(ops.end(), { GGML_OP_UNARY }); - } else if (args.sqrt_softplus) { - ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT }); - } else { - ops.insert(ops.end(), { GGML_OP_SOFT_MAX }); +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; + } } - const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation - if (args.prob_bias) { - bias = cgraph->nodes[i_probs + 2]->src[1]; - ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, - GGML_OP_GET_ROWS }); - out_nodes[0] = i_probs + 4; - } else { - ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); - out_nodes[0] = i_probs + 3; + 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; + } } - ids = cgraph->nodes[out_nodes[0]]; - if (args.norm) { - ops.insert(ops.end(), - { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }); - clamp = cgraph->nodes[i + ops.size() - 3]; + 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++; } - if (args.scale) { - ops.insert(ops.end(), { GGML_OP_SCALE }); - scale = cgraph->nodes[i + ops.size() - 1]; + 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; } - - weights = cgraph->nodes[i + ops.size() - 1]; - out_nodes[1] = i + ops.size() - 1; - - if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && - ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + 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; +} + +// qwen4exp chunked Gated DeltaNet prefill: the q/k L2 norms RMS_NORM(x, eps/n) -> SCALE(s) (x2) at node i whose only +// reader is a chunk-eligible GATED_DELTA_NET. Returns the index of the last matched node (-1: no match) and the GDN index, +// the raw q/k views, eps (n * eps/n) and the factor s*sqrt(n) of x * s*sqrt(n) * rsqrt(sum(x^2) + eps). +static int ggml_cuda_gdn_qk_norm_match(const ggml_cgraph * cgraph, const int i, int & ig, const ggml_tensor *& q_raw, + const ggml_tensor *& k_raw, float & eps, float & mul) { + const ggml_tensor * node = cgraph->nodes[i]; + if (node->op != GGML_OP_RMS_NORM) { + return -1; + } + auto next_compute = [&](int k) { + while (k < cgraph->n_nodes && ggml_cuda_is_view_or_noop(cgraph->nodes[k])) { + k++; + } + return k; + }; + const int iq_s = next_compute(i + 1); + const int ik_n = iq_s < cgraph->n_nodes ? next_compute(iq_s + 1) : cgraph->n_nodes; + const int ik_s = ik_n < cgraph->n_nodes ? next_compute(ik_n + 1) : cgraph->n_nodes; + if (ik_s >= cgraph->n_nodes) { + return -1; + } + const ggml_tensor * qs = cgraph->nodes[iq_s]; + const ggml_tensor * kn = cgraph->nodes[ik_n]; + const ggml_tensor * ks = cgraph->nodes[ik_s]; + q_raw = node->src[0]; + k_raw = kn->src[0]; + bool ok = qs->op == GGML_OP_SCALE && qs->src[0] == node && kn->op == GGML_OP_RMS_NORM && ks->op == GGML_OP_SCALE && + ks->src[0] == kn && q_raw && k_raw && node->ne[0] == 128 && q_raw->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && qs->type == GGML_TYPE_F32 && ks->type == GGML_TYPE_F32 && + ggml_are_same_shape(q_raw, k_raw) && + ggml_get_op_params_f32(node, 0) == ggml_get_op_params_f32(kn, 0) && + ggml_get_op_params_f32(qs, 0) == ggml_get_op_params_f32(ks, 0) && + ggml_get_op_params_f32(qs, 1) == 0.0f && ggml_get_op_params_f32(ks, 1) == 0.0f; + ig = -1; + for (int k = ik_s + 1; ok && k < cgraph->n_nodes && k <= ik_s + 24; ++k) { + const ggml_tensor * g = cgraph->nodes[k]; + if (g->op == GGML_OP_GATED_DELTA_NET && g->src[0] == qs && g->src[1] == ks) { + ig = k; + break; + } + } + ok = ok && ig > 0 && + ggml_node_get_use_count(cgraph, i) == 1 && ggml_node_get_use_count(cgraph, iq_s) == 1 && + ggml_node_get_use_count(cgraph, ik_n) == 1 && ggml_node_get_use_count(cgraph, ik_s) == 1 && + ((node->flags | qs->flags | kn->flags | ks->flags) & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_gdn_chunk_eligible(cgraph->nodes[ig], q_raw, k_raw); + if (!ok) { + return -1; + } + const float n = (float) node->ne[0]; + eps = ggml_get_op_params_f32(node, 0) * n; + mul = ggml_get_op_params_f32(qs, 0) * sqrtf(n); + return ik_s; +} + +// shared-expert merge after the MoE weighted reduction `red`: SIGMOID(gate [1, T]) -> MUL(shexp, .) -> ADD(red, .). +// One predicate for the fusion in ggml_cuda_try_fuse and the reorder in graph_optimize, so nodes are only moved +// when the fusion can fire (use counts are checked by each caller). +static bool ggml_cuda_moe_sgma_enabled() { + static const bool on = getenv("GGML_CUDA_DISABLE_MOE_SGMA") == nullptr && getenv("GGML_CUDA_DISABLE_SGMA") == nullptr && + !(getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"))); + return on; +} + +static bool ggml_cuda_moe_sgma_tail(const ggml_cuda_moe_weighted_reduction_match & m, const ggml_tensor * sig, + const ggml_tensor * mul, const ggml_tensor * add) { + const ggml_tensor * red = m.dst; + return sig->op == GGML_OP_UNARY && ggml_get_unary_op(sig) == GGML_UNARY_OP_SIGMOID && sig->ne[0] == 1 && + mul->op == GGML_OP_MUL && mul->src[1] == sig && add->op == GGML_OP_ADD && + ((add->src[0] == mul && add->src[1] == red) || (add->src[1] == mul && add->src[0] == red)) && + ggml_are_same_shape(mul->src[0], add) && ggml_are_same_shape(red, add) && ggml_nrows(add) == ggml_nelements(sig) && + (sig->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && (mul->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + (red->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_moe_weighted_reduction_sgma_supported(m.experts, m.expert_scale, m.weights, sig->src[0], mul->src[0], add); +} + +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")); + if (disable_fusion) { + 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]; + + // chains already computed by an earlier fused kernel (the chunked GDN writes its gated output norm itself) + static thread_local std::unordered_map s_precomputed; + if (!s_precomputed.empty()) { + auto it = s_precomputed.find(node); + if (it != s_precomputed.end()) { + const int count = it->second; + s_precomputed.erase(it); + return count > 0 ? count : GGML_CUDA_FUSED_SELF; + } + } + + if (node->op == GGML_OP_MUL_MAT && 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_gemm(*cuda_ctx, cgraph, i, args); + if (count > 0) { + ggml_cuda_mmb_idx_score(*cuda_ctx, node->src[0], node->src[1], args.dst, + args.vis_tails ? (const int32_t *) args.vis_tails->data : nullptr, + args.vis_starts ? (const int32_t *) args.vis_starts->data : nullptr); + return count - 1; + } + } + + // prefill shared-expert merge: SIGMOID(gate [1, T]) -> MUL(shexp [n, T], gate) -> ADD(moe_out, .) in one kernel + if (node->op == GGML_OP_UNARY && ggml_get_unary_op(node) == GGML_UNARY_OP_SIGMOID && i + 2 < cgraph->n_nodes && + node->type == GGML_TYPE_F32 && node->ne[0] == 1 && ggml_nrows(node) >= 2 && getenv("GGML_CUDA_DISABLE_SGMA") == nullptr) { + ggml_tensor * mul = cgraph->nodes[i + 1]; + ggml_tensor * add = cgraph->nodes[i + 2]; + const ggml_tensor * gate = node->src[0]; + if (mul->op == GGML_OP_MUL && add->op == GGML_OP_ADD && mul->src[1] == node && (add->src[0] == mul || add->src[1] == mul)) { + const ggml_tensor * src = mul->src[0]; + const ggml_tensor * other = add->src[0] == mul ? add->src[1] : add->src[0]; + auto aliases_ok = [&](const ggml_tensor * t) { // disjoint from out, or the same elements in place + const uintptr_t a0 = (uintptr_t) t->data, a1 = a0 + ggml_nbytes(t), b0 = (uintptr_t) add->data, b1 = b0 + ggml_nbytes(add); + return !(a0 < b1 && b0 < a1) || (t->data == add->data && ggml_are_same_layout(t, add)); + }; + if (gate->type == GGML_TYPE_F32 && ggml_is_contiguous(gate) && src->type == GGML_TYPE_F32 && other->type == GGML_TYPE_F32 && + add->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_is_contiguous(src) && ggml_is_contiguous(other) && + ggml_is_contiguous(add) && ggml_are_same_shape(src, add) && ggml_are_same_shape(other, add) && + ggml_are_same_shape(mul, add) && ggml_nrows(add) == ggml_nelements(node) && other != mul && + !ggml_cuda_mmb_is_bf16_only(*cuda_ctx, add) && aliases_ok(src) && aliases_ok(other) && aliases_ok(gate) && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL, GGML_OP_ADD }, { i + 2 })) { + ggml_cuda_op_sigmoid_gate_mul_add(*cuda_ctx, gate, src, other, add); + return 2; + } + } + } + + // qwen4exp GDN: ssm_beta / ssm_alpha, two narrow F32-weight GEMMs on the same activations, in one pass over them. + // The second node's output is written early, so it must not overlap anything the nodes in between touch. + if (node->op == GGML_OP_MUL_MAT && node->src[0]->type == GGML_TYPE_F32 && node->src[0]->ne[1] <= 64 && node->src[1]->type == GGML_TYPE_F32 && + node->src[1]->ne[1] * node->src[1]->ne[2] >= 512 && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && + getenv("GGML_CUDA_DISABLE_F32_DUAL") == nullptr && ggml_cuda_mmb_supported_mm(*cuda_ctx, node->src[0], node->src[1], node)) { + for (int j = i + 1; j < cgraph->n_nodes && j <= i + 8; ++j) { + ggml_tensor * n2 = cgraph->nodes[j]; + if (!(n2->op == GGML_OP_MUL_MAT && n2->src[1] == node->src[1] && n2->src[0]->type == GGML_TYPE_F32 && + n2->src[0]->ne[0] == node->src[0]->ne[0] && n2->src[0]->ne[1] <= 64 && n2->src[2] == nullptr && + ggml_cuda_mmb_supported_mm(*cuda_ctx, n2->src[0], n2->src[1], n2))) { + continue; + } + auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) { + if (!a || !b || !a->data || !b->data) return false; + const uintptr_t a0 = (uintptr_t) a->data, a1 = a0 + ggml_nbytes(a), b0 = (uintptr_t) b->data, b1 = b0 + ggml_nbytes(b); + return a0 < b1 && b0 < a1; + }; + bool ok = !overlaps(n2, node) && !overlaps(n2, node->src[1]); + for (int k = i + 1; ok && k < j; ++k) { + const ggml_tensor * t = cgraph->nodes[k]; + ok = !overlaps(n2, t); + for (int sidx = 0; ok && sidx < GGML_MAX_SRC && t->src[sidx]; ++sidx) ok = !overlaps(n2, t->src[sidx]); + } + if (ok && ggml_cuda_mmb_f32_dual(*cuda_ctx, node->src[0], n2->src[0], node->src[1], node, n2)) { + s_precomputed[n2] = 0; + return GGML_CUDA_FUSED_SELF; + } + break; + } + } + + // chunked GDN prefill: RMS_NORM(view of the GDN output) -> MUL(w) -> MUL(sigmoid(z)) is applied by the scan kernel + // as its epilogue (rms_rows_f32 math), so the pre-norm output never round-trips through memory + if (node->op == GGML_OP_GATED_DELTA_NET && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && + getenv("GGML_CUDA_DISABLE_GDN_OUTNORM") == nullptr) { + const ggml_tensor * v = node->src[2]; + int ix = -1; + for (int j = i + 1; j < cgraph->n_nodes && j <= i + 16; ++j) { + ggml_tensor * r = cgraph->nodes[j]; + if (r->op == GGML_OP_VIEW && r->view_src == node && r->view_offs == 0) { + ix = j; + } + if (r->op != GGML_OP_RMS_NORM) { + continue; + } + const ggml_tensor * x = r->src[0]; + if (!(x && x->view_src == node && x->view_offs == 0)) { + break; + } + ggml_cuda_norm_gated_match m; + const int sk = ggml_cuda_norm_gated_match_at(cgraph, j, m); + const int64_t S = v->ne[0], H = v->ne[1], T = v->ne[2]; + const bool ok = sk > 0 && !m.staged && m.pre < 0 && m.x == x && m.z && S == 128 && v->ne[3] == 1 && + x->ne[0] == S && x->ne[1] == H && x->ne[2] == T && x->nb[0] == 4 && x->nb[1] == (size_t) S * 4 && + x->nb[2] == (size_t) S * H * 4 && ggml_nelements(m.dst) == S * H * T && ggml_is_contiguous(m.dst) && + ggml_nelements(m.z) == S * H * T && ggml_is_contiguous(m.z) && ggml_nelements(m.w) == S && + (ix < 0 || ggml_node_get_use_count(cgraph, ix) == 1) && + (m.dst->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_gdn_chunk_eligible_node(node); + // the norm output is written during the scan, earlier than in the unfused schedule: it must not overlap the GDN + // tensor (output rows + state tail) and may alias z only element for element + auto overlap = [](const ggml_tensor * a, const ggml_tensor * b) { + const uintptr_t a0 = (uintptr_t) a->data, a1 = a0 + ggml_nbytes(a), b0 = (uintptr_t) b->data, b1 = b0 + ggml_nbytes(b); + return a->data && b->data && a0 < b1 && b0 < a1; + }; + const bool alias_ok = ok && !overlap(m.dst, node) && + (!overlap(m.dst, m.z) || (m.dst->data == m.z->data && ggml_nbytes(m.dst) == ggml_nbytes(m.z))); + if (ok && alias_ok) { + uint16_t * dst16 = ggml_cuda_bf16_for_mmb_reader(*cuda_ctx, cgraph, j + sk, m.dst); + // every reader takes the BF16 copy: skip the F32 store + float * dst32 = (dst16 && ggml_cuda_mmb_is_bf16_only(*cuda_ctx, m.dst)) ? nullptr : (float *) m.dst->data; + ggml_cuda_gdn_set_out_norm(node, (const float *) m.w->data, (const float *) m.z->data, dst32, dst16, m.eps); + s_precomputed[r] = sk; + } + break; + } + } + + // routed gate/up + SwiGLU on the BF16 WMMA path from 512 rows (32 with the small-batch opt-in): 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)) { + // chunked GDN prefill: if the SiLU output is read only through the q/k/v views of one chunk-eligible + // GATED_DELTA_NET (q/k through the fused L2 norms), the chunk kernels compute the conv themselves + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && getenv("GGML_CUDA_DISABLE_GDN_CONV") == nullptr && + getenv("GGML_CUDA_DISABLE_GDN_QKNORM") == nullptr && i + 1 < cgraph->n_nodes && cgraph->nodes[i + 1] == gm.out) { + const ggml_tensor * out = gm.out; + int i_rms = -1; + for (int k = i + 2; k < cgraph->n_nodes && k <= i + 24; ++k) { + const ggml_tensor * t = cgraph->nodes[k]; + if (t->op == GGML_OP_RMS_NORM && t->src[0] && t->src[0]->view_src == out) { i_rms = k; break; } + } + int ig = -1; const ggml_tensor * q_raw = nullptr; const ggml_tensor * k_raw = nullptr; float eps = 0.0f, mul = 1.0f; + const int last = i_rms > 0 ? ggml_cuda_gdn_qk_norm_match(cgraph, i_rms, ig, q_raw, k_raw, eps, mul) : -1; + const ggml_tensor * gdn = last > 0 ? cgraph->nodes[ig] : nullptr; + const ggml_tensor * v_raw = gdn ? gdn->src[2] : nullptr; + bool ok = gdn && v_raw && q_raw->view_src == out && k_raw->view_src == out && v_raw->view_src == out && + (out->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && ggml_node_get_use_count(cgraph, i + 1) == 3 && + out->ne[0] == gm.C && ggml_is_contiguous(out) && gm.x->ne[0] == gm.C && + q_raw->nb[2] == (size_t) gm.C * 4 && v_raw->nb[2] == (size_t) gm.C * 4 && v_raw->nb[1] == 128 * 4 && + q_raw->view_offs % 16 == 0 && k_raw->view_offs % 16 == 0 && v_raw->view_offs % 16 == 0 && gm.C % 4 == 0; + // the three views are the only readers of the SiLU output, and each view has exactly one reader + for (int k = i + 2; ok && k < cgraph->n_nodes; ++k) { + const ggml_tensor * t = cgraph->nodes[k]; + for (int sidx = 0; ok && sidx < GGML_MAX_SRC && t->src[sidx]; ++sidx) { + const ggml_tensor * sv = t->src[sidx]; + if (sv == out) { + ok = t == q_raw || t == k_raw || t == v_raw; + } else if (sv->view_src == out) { + ok = (sv == q_raw && t == cgraph->nodes[i_rms]) || (sv == k_raw && t->op == GGML_OP_RMS_NORM && t->src[0] == k_raw) || + (sv == v_raw && t == gdn); + } + } + } + if (ok) { + ggml_cuda_gdn_set_qk_norm(gdn, q_raw, k_raw, eps, mul); + ggml_cuda_gdn_set_conv(gdn, (const float *) gm.x->data, (const float *) gm.state->data, (const float *) gm.w->data, + gm.C, (int64_t) (q_raw->view_offs / 4), (int64_t) (k_raw->view_offs / 4), (int64_t) (v_raw->view_offs / 4)); + s_precomputed[cgraph->nodes[i_rms]] = last - i_rms; + return 1; + } + } + 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; + } + } + + // qwen4exp chunked Gated DeltaNet prefill: q/k L2 norms folded into the chunk kernels (ggml_cuda_gdn_qk_norm_match) + if (node->op == GGML_OP_RMS_NORM && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && + getenv("GGML_CUDA_DISABLE_GDN_QKNORM") == nullptr) { + int ig; const ggml_tensor * q_raw; const ggml_tensor * k_raw; float eps, mul; + const int last = ggml_cuda_gdn_qk_norm_match(cgraph, i, ig, q_raw, k_raw, eps, mul); + if (last > 0) { + ggml_cuda_gdn_set_qk_norm(cgraph->nodes[ig], q_raw, k_raw, eps, mul); + return last - 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; + // The folded gate path (rms_norm<128> + sigmoid + mul in the ssm_out matvec prologue) is only + // validated for the L2_NORM spelling. For the qwen4exp RMS_NORM+SCALE representation keep the + // plain fused kernel, which applies the gated norm itself: the chain fusion is the win and this + // avoids feeding the gate matvec tensors from a chain shape it was not exercised with. + if (getenv("GGML_CUDA_DISABLE_GDN_GATE") == nullptr && !args.qk_rms_scale && args.n_tokens == 1 && + 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; + } + // otherwise the norm kernel of the fused step also applies the sigmoid(z) gate (one launch less) + if (getenv("GGML_CUDA_DISABLE_GDN_NORMGATE") == nullptr && sig && mul && + sig->op == GGML_OP_UNARY && ggml_get_unary_op(sig) == GGML_UNARY_OP_SIGMOID && mul->op == GGML_OP_MUL && + ((mul->src[0] == sig && mul->src[1] == norm_out) || (mul->src[1] == sig && mul->src[0] == norm_out)) && + sig->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_is_contiguous(mul) && + sig->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(sig->src[0]) && + ggml_nelements(sig->src[0]) == ggml_nelements(norm_out) && + ggml_are_same_shape(sig, norm_out) && ggml_are_same_shape(mul, norm_out) && + ggml_nelements(norm_out) == args.S * args.H_v * args.n_tokens && args.S == 128 && + ggml_node_get_use_count(cgraph, norm_idx) == 1 && ggml_node_get_use_count(cgraph, sig_idx) == 1 && + (norm_out->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && (sig->flags & GGML_TENSOR_FLAG_OUTPUT) == 0) { + // z and the gated output are read/written element for element by the same thread: in-place is fine + ggml_cuda_op_gdn_decode_fused_gated(*cuda_ctx, args, (const float *) sig->src[0]->data, (float *) mul->data); + return mul_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)) { + 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; + } + } + } + } + // prefill shared-expert merge right after: SIGMOID(gate [1, T]) -> MUL(shexp, .) -> ADD(reduction, .), the reduction + // and the product used only there -> one kernel (bit-identical to reduction + sigmoid_gate_mul_add_f32) + if (ggml_cuda_moe_sgma_enabled() && !merge && i + count + 2 < cgraph->n_nodes && ggml_node_has_n_uses(cgraph, i + count - 1, 1)) { + ggml_tensor * sig = cgraph->nodes[i + count], * mul = cgraph->nodes[i + count + 1], * add = cgraph->nodes[i + count + 2]; + if (ggml_cuda_moe_sgma_tail(match, sig, mul, add) && + ggml_node_has_n_uses(cgraph, i + count, 1) && ggml_node_has_n_uses(cgraph, i + count + 1, 1)) { + const int out_idx = i + count + 2; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, count + 3, &out_idx, 1) && + ggml_cuda_op_moe_weighted_reduction_sgma(*cuda_ctx, match.experts, match.expert_scale, match.weights, + sig->src[0], mul->src[0], add)) { + return count + 2; + } + } + } + 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, merge); + return count - 1; + } + } + } + + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache + if (node->op == GGML_OP_GATED_DELTA_NET) { + 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); +#endif + ggml_cuda_op_gated_delta_net_fused_cache(*cuda_ctx, node, fused_state_cpy); + return nodes_to_skip; + } + } + + 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) { + ggml_cuda_topk_moe_args args; + const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); + std::vector ops; + ops.reserve(13); // max ops; avoids gcc -Wstringop-overflow false positive + + if (can_fuse) { + const ggml_tensor * logits = node->src[0]; + ggml_tensor * weights = nullptr; + ggml_tensor * ids = nullptr; + const ggml_tensor * bias = nullptr; + const ggml_tensor * clamp = nullptr; + const ggml_tensor * scale = nullptr; + + if (!args.delayed_softmax) { + int out_nodes[2]; // nodes which can't be elided + + if (args.sigmoid) { + ops.insert(ops.end(), { GGML_OP_UNARY }); + } else if (args.sqrt_softplus) { + ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT }); + } else { + ops.insert(ops.end(), { GGML_OP_SOFT_MAX }); + } + const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation + + if (args.prob_bias) { + bias = cgraph->nodes[i_probs + 2]->src[1]; + ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, + GGML_OP_GET_ROWS }); + out_nodes[0] = i_probs + 4; + } else { + ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); + out_nodes[0] = i_probs + 3; + } + ids = cgraph->nodes[out_nodes[0]]; + + if (args.norm) { + ops.insert(ops.end(), + { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }); + clamp = cgraph->nodes[i + ops.size() - 3]; + } + if (args.scale) { + ops.insert(ops.end(), { GGML_OP_SCALE }); + scale = cgraph->nodes[i + ops.size() - 1]; + } + + weights = cgraph->nodes[i + ops.size() - 1]; + out_nodes[1] = i + ops.size() - 1; + + if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && + ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && + ggml_cuda_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) { @@ -3543,8 +5296,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; } } @@ -3607,6 +5360,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; @@ -3643,6 +5581,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; @@ -3961,6 +5932,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]; @@ -4158,6 +6136,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; @@ -4178,6 +6179,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_node_has_n_uses(cgraph, i, 1)) { + ggml_cuda_hc_mix_args ma; + const int count = ggml_cuda_hc_gate_mix_route(*cuda_ctx, cgraph, i, ma); + if (count > 0 && ggml_cuda_hc_gate_mix(*cuda_ctx, node->src[0], node->src[1], 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); + // The fused reduction replays the unfused mul/add roundings exactly (hc_mix_reduce_f32 uses + // explicitly rounded mul/add in the same stream order), so it is bit-identical at decode too and + // removes ~6 tiny kernels per mix per layer (sigmoid/mul/3x add/scale). HC_NO_MIX_FUSE=1 disables. + if (n_ops > 0 && args.dst->ne[1] == 1 && getenv("HC_NO_MIX_FUSE")) { + n_ops = 0; + } + 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 })) { @@ -4195,6 +6402,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; } @@ -4340,6 +6576,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; @@ -4420,6 +6659,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) { @@ -4433,8 +6677,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); @@ -4489,6 +6761,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; } @@ -4517,13 +6790,351 @@ 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; } + } + // MoE prefill with a sigmoid-gated shared expert (qwen4exp, and any model with the same pattern): the weighted + // reduction is followed by the shared-expert GEMMs and only then by SIGMOID(gate) -> MUL(shexp, .) -> ADD(moe_out, .). + // Move the reduction nodes down to just before the SIGMOID so the fusion computes reduction + merge in one kernel. + // Only done when that fusion can fire (same predicate) and no node in between reads a reduction node or writes + // (in place / through a view) memory the reduction reads. Runs before allocation, so buffer lifetimes follow the new + // order. GGML_CUDA_DISABLE_MOE_SGMA=1 or GGML_CUDA_DISABLE_FUSION=1 keeps the original order. + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && ggml_cuda_moe_sgma_enabled()) { + auto root = [](const ggml_tensor * t) { return t->view_src ? t->view_src : t; }; + // readers of t in this graph (the use-count table is not guaranteed on a split's graph view) + auto n_uses = [&](const ggml_tensor * t) { + int c = 0; + for (int n = 0; n < cgraph->n_nodes; ++n) { + for (int s = 0; s < GGML_MAX_SRC && cgraph->nodes[n]->src[s]; ++s) { + c += cgraph->nodes[n]->src[s] == t; + } + } + return c; + }; + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_MUL) continue; + ggml_cuda_moe_weighted_reduction_match mm; + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, mm) || ggml_nrows(mm.dst) < 512) continue; + const int cnt = mm.node_count; + int k = -1; + for (int n = i + cnt; n + 2 < cgraph->n_nodes && n <= i + cnt + 24; ++n) { + const ggml_tensor * sig = cgraph->nodes[n], * mul = cgraph->nodes[n + 1], * add = cgraph->nodes[n + 2]; + if (ggml_cuda_moe_sgma_tail(mm, sig, mul, add)) { k = n; break; } + } + if (k <= i + cnt) continue; // not found, or already adjacent + if (n_uses(mm.dst) != 1 || n_uses(cgraph->nodes[k]) != 1 || n_uses(cgraph->nodes[k + 1]) != 1) continue; + bool indep = true; + for (int n = i + cnt; n < k && indep; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + for (int r = i; r < i + cnt && indep; ++r) { + const ggml_tensor * red = cgraph->nodes[r]; + for (int s = 0; s < GGML_MAX_SRC && t->src[s] && indep; ++s) { // t reads a reduction node + indep = t->src[s] != red && root(t->src[s]) != red; + } + for (int s = 0; s < GGML_MAX_SRC && red->src[s] && indep; ++s) { // t writes what the reduction reads + indep = !(t->view_src && root(red->src[s]) == t->view_src) && red->src[s] != t; + } + } + } + if (!indep) continue; + std::rotate(cgraph->nodes + i, cgraph->nodes + i + cnt, cgraph->nodes + k); // [between][reduction] SIGMOID MUL ADD + i = k + 2; + } + } + { // 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 by a stream mix that the fused + // gate mix takes; any other mix falls back to hc_mix_reduce, which reads the F32 xn + auto n_readers = [&](int from, const ggml_tensor * x) { int c = 0; for (int n = from; n < cgraph->n_nodes; ++n) c += reads(cgraph->nodes[n], x); return c; }; + 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 >= 2) { + const ggml_tensor * g = cgraph->nodes[n - 2]; + ggml_cuda_hc_mix_args ma, mg; + 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) && + ma.gate == g && ggml_cuda_hc_gate_mix_route(*cuda_ctx, cgraph, n - 2, mg) > 0 && + // the executor's reader rule: one reader in this split and none in another, not an output + n_readers(n - 1, g) == 1 && ggml_cuda_marks_readers_local(cgraph, g) && (g->flags & GGML_TENSOR_FLAG_OUTPUT) == 0) 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) { // gated norm of a GDN output read only by MMB GEMMs (ssm_out): the + ggml_cuda_norm_gated_match m; // chunked scan epilogue then stores only the BF16 copy + if (ggml_cuda_norm_gated_match_at(cgraph, i, m) <= 0 || !m.dst || !m.x) continue; + const ggml_tensor * xs = m.x->view_src; + if (!xs || xs->op != GGML_OP_GATED_DELTA_NET) continue; + const ggml_tensor * d = m.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); + } + // HC gate GEMM read only by the stream mix right after it. If the fused gate mix refuses, hc_mix_reduce reads the F32 + // gate unless it finds BF16 copies of both inputs. So drop the F32 output only when both copies are expected: this + // GEMM's gate copy, and the xn copy from the last combine-norm. A layer-0 mix has no combine-norm and keeps the F32 gate. + const ggml_tensor * cn_xn = nullptr; + 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) cn_xn = ca.out_xn; + 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 (n == i + 1 && cn_xn && 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 && (ma.xn == cn_xn || ma.xn->view_src == cn_xn)) 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 || t->src[0]->type == GGML_TYPE_Q8_0) && 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")); auto add_alloc_deps = [&](size_t start, size_t last_node) { - + // last_node is the inclusive index of the final group node. for (size_t i = start; i < last_node; ++i) { params->add_alloc_dep(params->user_data, cgraph->nodes[i], cgraph->nodes[last_node]); @@ -4538,6 +7149,144 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph if (!disable_fusion) { // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute + 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_MUL_MAT) continue; + ggml_cuda_idx_relu_sum_args args; + if (ggml_cuda_match_idx_gemm(*cuda_ctx, cgraph, i, args) > 0) { + const ggml_tensor * mm = cgraph->nodes[i]; + for (const ggml_tensor * in : { (const ggml_tensor *) mm->src[0], (const ggml_tensor *) mm->src[1], args.vis_tails, args.vis_starts }) { + if (!in) { continue; } + auto * root = const_cast(in->view_src ? in->view_src : in); + params->add_alloc_dep(params->user_data, root, args.dst); + } + } + } + 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) { + for (const ggml_tensor * in : { args.score, args.vis_tails, args.vis_starts }) { + if (!in) { continue; } + auto * root = const_cast(in->view_src ? in->view_src : in); + 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; + + // The chunked GDN can elide the conv/SiLU and read x, history and taps at the GDN node instead. + // Keeping them only through the SiLU allows the scheduler to reuse x (and corrupt q/k/v) before + // gdn_chunk_prep_kernel reads it. Restrict the longer lifetime to the fused prefill pattern. + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc) && + getenv("GGML_CUDA_DISABLE_GDN_CONV") == nullptr && + getenv("GGML_CUDA_DISABLE_GDN_QKNORM") == nullptr) { + int i_rms = -1; + for (int j = gm.conv_idx + 2; j < cgraph->n_nodes && j <= gm.conv_idx + 24; ++j) { + const ggml_tensor * t = cgraph->nodes[j]; + if (t->op == GGML_OP_RMS_NORM && t->src[0] && t->src[0]->view_src == gm.out) { + i_rms = j; + break; + } + } + int ig = -1; + const ggml_tensor * q_raw = nullptr, * k_raw = nullptr; + float eps = 0.0f, mul = 1.0f; + if (i_rms >= 0 && ggml_cuda_gdn_qk_norm_match(cgraph, i_rms, ig, q_raw, k_raw, eps, mul) > 0 && + ig >= 0 && cgraph->nodes[ig]->src[2] && cgraph->nodes[ig]->src[2]->view_src == gm.out) { + for (const ggml_tensor * input : {gm.x, gm.state, gm.w}) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, cgraph->nodes[ig]); + } + } + } + } + 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) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { @@ -4600,7 +7349,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, 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)) { - add_alloc_deps(i, i + ops.size()); + add_alloc_deps(i, i + ops.size() - 1); i += ops.size() - 1; } } @@ -4618,6 +7367,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; @@ -5242,6 +7994,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5295,6 +8049,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return true; case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: // 32-value sub-blocks, the row size does not guarantee // the QK_K super-blocks the get_rows kernel iterates on return op->src[0]->ne[0] % QK_K == 0; @@ -5580,9 +8336,21 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g 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_PRE: + // qwen4exp gated form on RDNA3.5: the fused CUDA kernels + // (dsv4_hc_pre_f32<>, dsv4-hc.cu) run on this backend. Supporting + // the gated form keeps the whole graph on GPU; refusing it used to + // force a CPU split of the graph and a multi-GiB host compute + // buffer plus a server-side decode stall (measured 0.3-1.6 tok/s + // at high context on Strix Halo with parallel 4). The hc_mix / + // hc_combine fusions this avoids are tile-level op fusions that + // need the unfused graph; the pp2048 delta measured earlier + // (Qwen3.8-Flash-Next Q2_0 ~8%) is smaller than the CPU split cost. + // Setting LLAMA_FUSED_DSV4_HC_PRE=0 keeps the older behaviour. 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_POST: + // identity-comb (qwen4exp) form on RDNA3.5: see GGML_OP_DSV4_HC_PRE + // above, same reasoning. 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] == nullptr || op->src[3]->type == GGML_TYPE_F32) && op->type == GGML_TYPE_F32; @@ -5761,8 +8529,26 @@ 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; +} + +// let the mmb path of this backend context take GEMMs from 32 rows instead of 512 (see mmb_min_t); call before the first graph +static void ggml_backend_cuda_set_mmb_small_batch(ggml_backend_t backend, bool enable) { + GGML_ASSERT(ggml_backend_is_cuda(backend)); + ((ggml_backend_cuda_context *) backend->context)->mmb_small_batch = 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_cuda_set_mmb_small_batch") == 0) { + return (void *)ggml_backend_cuda_set_mmb_small_batch; + } 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..8eef67d14b4e --- /dev/null +++ b/ggml/src/ggml-cuda/idx-relu-sum.cu @@ -0,0 +1,36 @@ +#include "idx-relu-sum.cuh" + +// Preserve the head addition order of the RELU/CONT/ADD graph. +// With tails/starts it also applies the compact QSA visibility exactly as CPY(I32->F32), REPEAT, SUB, STEP, LOG +// and ADD compute it, so the result is bit-identical to the unfused graph. +static __global__ void idx_relu_sum_f32(const float * __restrict__ src, float * __restrict__ dst, + const int n_blocks, const int heads, + const int32_t * __restrict__ tails, const int32_t * __restrict__ starts) { + 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); + } + if (tails) { + const float d = (float) tails[tt] - (float) starts[b]; + const float step = d > 0.0f; + acc = acc + logf(step); + } + 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); + const int32_t * tails = args.vis_tails ? (const int32_t *) args.vis_tails->data : nullptr; + const int32_t * starts = args.vis_starts ? (const int32_t *) args.vis_starts->data : nullptr; + idx_relu_sum_f32<<>>((const float *) s->data, (float *) args.dst->data, n_blocks, args.heads, + tails, starts); + 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..b796dbf254f5 --- /dev/null +++ b/ggml/src/ggml-cuda/idx-relu-sum.cuh @@ -0,0 +1,14 @@ +#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; + // optional fused compact visibility (single sequence): dst += log(step(float(tails[t]) - float(starts[b]))) + const ggml_tensor * vis_tails = nullptr; // I32 [n_tps] contiguous view of the QSA limits + const ggml_tensor * vis_starts = nullptr; // I32 [n_blocks] contiguous view of the QSA limits + // nodes of the visibility chain that the matcher checked itself (views/casts of the limits input); a caller that + // widens the fusion must leave them out of ggml_can_fuse_subgraph_ext + const ggml_tensor * vis_leafs[16] = {}; +}; +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..b6b004391e13 --- /dev/null +++ b/ggml/src/ggml-cuda/mmb-quant.cuh @@ -0,0 +1,549 @@ +#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}; } +}; + +// 8 bf16 weights from fp32 values, one 16-byte LDS store (RNE as mmb_f2bf) +__device__ __forceinline__ void mmb_store8(uint16_t * dst, const float * v) { + uint4 o; o.x = mmb_pack2(v[0], v[1]); o.y = mmb_pack2(v[2], v[3]); o.z = mmb_pack2(v[4], v[5]); o.w = mmb_pack2(v[6], v[7]); + *(uint4 *) dst = o; +} + +// 8 weights of an IQ2/IQ3_XXS grid group: fp32 d * grid * sign as in dequantize_*, RNE to bf16 +__device__ __forceinline__ void mmb_store_grid8(uint16_t * dst, const uint64_t grid, const float d, const uint32_t signs) { + float v[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) v[j] = d * (float)((grid >> (8 * j)) & 0xff) * (signs & (1u << j) ? -1.f : 1.f); + mmb_store8(dst, v); +} + +// 8 weights of a grid group as bf16: RNE each fp32 value, pack two per v_perm, then apply the signs as an XOR of +// the bf16 sign bits (RNE is sign-symmetric: bf16(-v) == bf16(v) ^ 0x8000) - bit-identical to d * grid * sign. +__device__ __forceinline__ uint4 mmb_grid8_bf16(const uint32_t lo, const uint32_t hi, const float d, const uint32_t sg) { + auto rb = [](const float x) { const uint32_t u = __float_as_uint(x); return u + 0x7fffu + ((u >> 16) & 1u); }; + auto pk = [&](const uint32_t g, const int j) { + return __builtin_amdgcn_perm(rb(d * (float)((g >> (8 * j + 8)) & 0xff)), rb(d * (float)((g >> (8 * j)) & 0xff)), 0x07060302u); }; + uint4 o; + o.x = pk(lo, 0) ^ ((sg & 1u) << 15 | (sg & 2u) << 30); + o.y = pk(lo, 2) ^ ((sg & 4u) << 13 | (sg & 8u) << 28); + o.z = pk(hi, 0) ^ ((sg & 16u) << 11 | (sg & 32u) << 26); + o.w = pk(hi, 2) ^ ((sg & 64u) << 9 | (sg & 128u) << 24); + return o; +} + +// ksigns_iq2xs[i] without the table: 7 sign bits plus an even-parity bit 7 +__device__ __forceinline__ uint32_t mmb_ksigns(const uint32_t i) { return i | ((__popc(i) & 1u) << 7); } + +// Routed GLU grid-based decoders: fetch() loads one lane's block fields (8 weights) in the WMMA loop, decode() turns +// them plus the LDS grid into 8 bf16. Same fp32 math as dequantize_* -> bit-identical; grid moves to LDS. +template struct mmb_lb { static constexpr bool ok = false; using grid_t = uint32_t; static constexpr int N = 1; }; + +template <> struct mmb_lb<32 + GGML_TYPE_IQ3_S> { + static constexpr bool ok = true; using grid_t = uint32_t; static constexpr int N = 512; + static __device__ __forceinline__ grid_t entry(const int i) { return iq3s_grid[i]; } + static __device__ __forceinline__ void fetch(const uint8_t * row, const int ks, const int sub, const int il, uint32_t & w0, uint32_t & w1) { + const block_iq3_s * x = (const block_iq3_s *) row + (ks * 64) / QK_K; + const int ib = ((ks * 64) % QK_K) / 32 + sub; + w0 = *(const uint16_t *)(x->qs + 8 * ib + 2 * il) | ((uint32_t) x->qh[ib] << 16) | ((uint32_t) x->signs[4 * ib + il] << 24); + w1 = (uint32_t) *(const uint16_t *) &x->d | ((uint32_t) ((x->scales[ib / 2] >> 4 * (ib % 2)) & 0xf) << 16); + } + static __device__ __forceinline__ uint4 decode(const grid_t * g, const uint32_t w0, const uint32_t w1, const int il) { + const uint32_t qh = (w0 >> 16) & 0xff; + return mmb_grid8_bf16(g[(w0 & 0xff) | ((qh << (8 - 2 * il)) & 256)], g[((w0 >> 8) & 0xff) | ((qh << (7 - 2 * il)) & 256)], + mmb_h2f((uint16_t) w1) * (1 + 2 * (int)(w1 >> 16)), w0 >> 24); + } +}; +template <> struct mmb_lb<32 + GGML_TYPE_IQ2_XXS> { + static constexpr bool ok = true; using grid_t = uint64_t; static constexpr int N = 256; + static __device__ __forceinline__ grid_t entry(const int i) { return iq2xxs_grid[i]; } + static __device__ __forceinline__ void fetch(const uint8_t * row, const int ks, const int sub, const int il, uint32_t & w0, uint32_t & w1) { + const block_iq2_xxs * x = (const block_iq2_xxs *) row + (ks * 64) / QK_K; + const uint16_t * q2 = x->qs + 4 * (((ks * 64) % QK_K) / 32 + sub); + w0 = q2[2] | ((uint32_t) q2[3] << 16); + w1 = (uint32_t) *(const uint16_t *) &x->d | ((uint32_t) ((const uint8_t *) q2)[il] << 16); + } + static __device__ __forceinline__ uint4 decode(const grid_t * g, const uint32_t w0, const uint32_t w1, const int il) { + const uint64_t grid = g[w1 >> 16]; + return mmb_grid8_bf16((uint32_t) grid, (uint32_t) (grid >> 32), mmb_h2f((uint16_t) w1) * (0.5f + (w0 >> 28)) * 0.25f, mmb_ksigns((w0 >> 7 * il) & 127)); + } +}; +template <> struct mmb_lb<32 + GGML_TYPE_IQ2_XS> { + static constexpr bool ok = true; using grid_t = uint64_t; static constexpr int N = 512; + static __device__ __forceinline__ grid_t entry(const int i) { return iq2xs_grid[i]; } + static __device__ __forceinline__ void fetch(const uint8_t * row, const int ks, const int sub, const int il, uint32_t & w0, uint32_t & w1) { + const block_iq2_xs * x = (const block_iq2_xs *) row + (ks * 64) / QK_K; + const int ib = ((ks * 64) % QK_K) / 32 + sub; + w0 = x->qs[4 * ib + il] | ((uint32_t) ((x->scales[ib] >> 4 * (il / 2)) & 0xf) << 16); + w1 = *(const uint16_t *) &x->d; + } + static __device__ __forceinline__ uint4 decode(const grid_t * g, const uint32_t w0, const uint32_t w1, const int il) { + GGML_UNUSED(il); + const uint64_t grid = g[w0 & 511]; + return mmb_grid8_bf16((uint32_t) grid, (uint32_t) (grid >> 32), mmb_h2f((uint16_t) w1) * (0.5f + (w0 >> 16)) * 0.25f, mmb_ksigns((w0 >> 9) & 127)); + } +}; +// IQ2_S: every grid byte is 8, 25 or 43, so the LDS table holds 2-bit codes (2 KB instead of 8 KB) +template <> struct mmb_lb<32 + GGML_TYPE_IQ2_S> { + static constexpr bool ok = true; using grid_t = uint16_t; static constexpr int N = 1024; + static __device__ __forceinline__ grid_t entry(const int i) { + const uint64_t g = iq2s_grid[i]; uint32_t c = 0; +#pragma unroll + for (int j = 0; j < 8; ++j) { const uint32_t b = (g >> (8 * j)) & 0xff; c |= ((uint32_t) (b > 8) + (uint32_t) (b > 25)) << (2 * j); } + return (grid_t) c; + } + static __device__ __forceinline__ void fetch(const uint8_t * row, const int ks, const int sub, const int il, uint32_t & w0, uint32_t & w1) { + const block_iq2_s * x = (const block_iq2_s *) row + (ks * 64) / QK_K; + const int ib = ((ks * 64) % QK_K) / 32 + sub; + w0 = (x->qs[4 * ib + il] | ((x->qh[ib] << (8 - 2 * il)) & 0x300)) | ((uint32_t) x->qs[QK_K / 8 + 4 * ib + il] << 16) + | ((uint32_t) ((x->scales[ib] >> 4 * (il / 2)) & 0xf) << 24); + w1 = *(const uint16_t *) &x->d; + } + static __device__ __forceinline__ uint32_t sel4(const uint32_t c8) { // 4 two-bit codes -> one per byte + const uint32_t t = (c8 | (c8 << 12)) & 0x000f000fu; + return (t | (t << 6)) & 0x03030303u; + } + static __device__ __forceinline__ uint4 decode(const grid_t * g, const uint32_t w0, const uint32_t w1, const int il) { + GGML_UNUSED(il); + const uint32_t c = g[w0 & 0x3ff]; + return mmb_grid8_bf16(__builtin_amdgcn_perm(0u, 0x002b1908u, sel4(c & 0xff)), __builtin_amdgcn_perm(0u, 0x002b1908u, sel4(c >> 8)), + mmb_h2f((uint16_t) w1) * (0.5f + (w0 >> 24)) * 0.25f, (w0 >> 16) & 0xff); + } +}; +template <> struct mmb_lb<32 + GGML_TYPE_IQ3_XXS> { + static constexpr bool ok = true; using grid_t = uint32_t; static constexpr int N = 256; + static __device__ __forceinline__ grid_t entry(const int i) { return iq3xxs_grid[i]; } + static __device__ __forceinline__ void fetch(const uint8_t * row, const int ks, const int sub, const int il, uint32_t & w0, uint32_t & w1) { + const block_iq3_xxs * x = (const block_iq3_xxs *) row + (ks * 64) / QK_K; + const int ib = ((ks * 64) % QK_K) / 32 + sub; + const uint16_t * gas = (const uint16_t *) (x->qs + QK_K / 4) + 2 * ib; + w0 = gas[0] | ((uint32_t) gas[1] << 16); + w1 = (uint32_t) *(const uint16_t *) &x->d | ((uint32_t) *(const uint16_t *) (x->qs + 8 * ib + 2 * il) << 16); + } + static __device__ __forceinline__ uint4 decode(const grid_t * g, const uint32_t w0, const uint32_t w1, const int il) { + return mmb_grid8_bf16(g[(w1 >> 16) & 0xff], g[w1 >> 24], mmb_h2f((uint16_t) w1) * (0.5f + (w0 >> 28)) * 0.5f, mmb_ksigns((w0 >> 7 * il) & 127)); + } +}; + +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) { + // QK = 64: the slice is one block, lane -> 8 consecutive weights (2 bytes of qs), (code - 1) * d as dequantize_q2_0 + static_assert(QK == 64, "Q2_0 slice decode assumes one block per 64-wide K slice"); + const block_q2_0 * x = (const block_q2_0 *) row + k0 / QK; + const float d = x->d; + const uint32_t q = x->qs[2 * lane] | (x->qs[2 * lane + 1] << 8); + float v[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) v[j] = ((int) ((q >> (2 * j)) & 3) - 1) * d; + mmb_store8(dst + 8 * lane, v); + } + 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) { + // QK_K = 256: a 64-wide slice is group il = (k0 % 256) / 64. lane < 4 takes low nibbles (is+0), lane >= 4 + // high nibbles (is+1), qh the 5th bit. Same fp32 math as dequantize_q5_K. + static_assert(QK == 256, "Q5_K slice decode assumes a 256-wide super-block"); + const block_q5_K * x = (const block_q5_K *) row + k0 / QK; + const int il = (k0 % QK) / 64, hi = lane >> 2, b0 = 8 * (hi ? lane - 4 : lane); + const uint8_t * ql = x->qs + 32 * il + b0; + const uint8_t * qh = x->qh + b0; + const float dall = __low2half(x->dm), dmin = __high2half(x->dm); + uint8_t sc, m; + get_scale_min_k4(2 * il + hi, x->scales, sc, m); + const float d = dall * sc, mn = dmin * m; + const uint32_t hm = 1u << (2 * il + hi); + float v[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int q = (hi ? (ql[j] >> 4) : (ql[j] & 0xF)) + ((qh[j] & hm) ? 16 : 0); + v[j] = d * (float) q - mn; + } + mmb_store8(dst + 8 * lane, v); + } + else if constexpr (TYPE == GGML_TYPE_Q6_K) { + // QK_K = 256: a 64-wide slice is half of one 128-value group (ip = (k0%256)/128, half = (k0%128)/64). half + // picks the nibble, the qh bit pair and the scale offset. Same fp32 math as dequantize_q6_K. + static_assert(QK == 256, "Q6_K slice decode assumes a 256-wide super-block"); + const block_q6_K * x = (const block_q6_K *) row + k0 / QK; + const int ip = (k0 % QK) / 128, half = (k0 % 128) / 64, L = lane & 7, Lh = (lane >> 2) ? L - 4 : L; + const uint8_t * ql = x->ql + 64 * ip + 8 * Lh + ((lane >> 2) ? 32 : 0); + const uint8_t * qh = x->qh + 32 * ip + 8 * Lh; + const int sh = 2 * (lane >> 2) + 4 * half; + const float d = x->d; + const int8_t * sc = x->scales + 8 * ip + Lh / 2 + (half ? ((lane >> 2) ? 6 : 4) : ((lane >> 2) ? 2 : 0)); + float v[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const int q = (half ? (ql[j] >> 4) : (ql[j] & 0xF)) | ((((qh[j] >> sh) & 3) << 4)); + v[j] = d * (float) sc[0] * (float) ((int8_t) q - 32); + } + mmb_store8(dst + 8 * lane, v); + } + 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) { + // a lane decodes one 8-weight group (il = lane & 3) of 32-block ib0 + (lane >> 2); same fp32 math as dequantize_iq2_xxs + static_assert(QK == 256, "IQ2_XXS slice decode assumes a 256-wide super-block"); + const block_iq2_xxs * x = (const block_iq2_xxs *) row + k0 / QK; + const int sub = lane >> 2, il = lane & 3, ib = (k0 % QK) / 32 + sub; + const uint16_t * q2 = x->qs + 4 * ib; + const uint32_t aux32 = q2[2] | (q2[3] << 16); + const float d = (float) x->d * (0.5f + (aux32 >> 28)) * 0.25f; + mmb_store_grid8(dst + 32 * sub + 8 * il, iq2xxs_grid[((const uint8_t *) q2)[il]], d, ksigns_iq2xs[(aux32 >> 7 * il) & 127]); + } + else if constexpr (TYPE == GGML_TYPE_IQ2_XS) { + static_assert(QK == 256, "IQ2_XS slice decode assumes a 256-wide super-block"); + const block_iq2_xs * x = (const block_iq2_xs *) row + k0 / QK; + const int sub = lane >> 2, il = lane & 3, ib = (k0 % QK) / 32 + sub; + const uint16_t q2 = x->qs[4 * ib + il]; + const float d = (float) x->d * (0.5f + ((x->scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; + mmb_store_grid8(dst + 32 * sub + 8 * il, iq2xs_grid[q2 & 511], d, ksigns_iq2xs[q2 >> 9]); + } + else if constexpr (TYPE == GGML_TYPE_IQ2_S) { + static_assert(QK == 256, "IQ2_S slice decode assumes a 256-wide super-block"); + const block_iq2_s * x = (const block_iq2_s *) row + k0 / QK; + const int sub = lane >> 2, il = lane & 3, ib = (k0 % QK) / 32 + sub; + const float d = (float) x->d * (0.5f + ((x->scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; + mmb_store_grid8(dst + 32 * sub + 8 * il, iq2s_grid[x->qs[4 * ib + il] | ((x->qh[ib] << (8 - 2 * il)) & 0x300)], d, + x->qs[QK_K / 8 + 4 * ib + il]); + } + else if constexpr (TYPE == GGML_TYPE_IQ3_XXS) { + static_assert(QK == 256, "IQ3_XXS slice decode assumes a 256-wide super-block"); + const block_iq3_xxs * x = (const block_iq3_xxs *) row + k0 / QK; + const int sub = lane >> 2, il = lane & 3, ib = (k0 % QK) / 32 + sub; + const uint8_t * q3 = x->qs + 8 * ib; + const uint16_t * gas = (const uint16_t *) (x->qs + QK_K / 4) + 2 * ib; + const uint32_t aux32 = gas[0] | (gas[1] << 16); + const float d = (float) x->d * (0.5f + (aux32 >> 28)) * 0.5f; + const uint64_t grid = (uint64_t) iq3xxs_grid[q3[2 * il + 0]] | ((uint64_t) iq3xxs_grid[q3[2 * il + 1]] << 32); + mmb_store_grid8(dst + 32 * sub + 8 * il, grid, d, ksigns_iq2xs[(aux32 >> 7 * il) & 127]); + } + else if constexpr (TYPE == GGML_TYPE_IQ3_S) { + // lane -> (il = lane>>1, ib = ib0 + (lane&1)): 8 consecutive weights, one 16-byte LDS store. + // Same arithmetic as dequantize_iq3_s (d * grid * sign in fp32, then RNE bf16) -> bit-identical. + const block_iq3_s * x = (const block_iq3_s *) row + k0 / QK; + const int ib0 = (k0 % QK) / 32, sub = lane & 1, il = lane >> 1, ib = ib0 + sub; + const uint16_t q2 = *(const uint16_t *)(x->qs + 8 * ib + 2 * il); + const int qh = x->qh[ib]; + const uint32_t g1 = iq3s_grid[(q2 & 0xff) | ((qh << (8 - 2 * il)) & 256)]; + const uint32_t g2 = iq3s_grid[(q2 >> 8) | ((qh << (7 - 2 * il)) & 256)]; + const float d = (float) x->d * (1 + 2 * ((x->scales[ib / 2] >> 4 * (ib % 2)) & 0xf)); + const int signs = x->signs[4 * ib + il]; + float v[8]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + v[j + 0] = d * (float)((g1 >> (8 * j)) & 0xff) * (signs & (1 << (j + 0)) ? -1.f : 1.f); + v[j + 4] = d * (float)((g2 >> (8 * j)) & 0xff) * (signs & (1 << (j + 4)) ? -1.f : 1.f); + } + uint4 o; o.x = mmb_pack2(v[0], v[1]); o.y = mmb_pack2(v[2], v[3]); o.z = mmb_pack2(v[4], v[5]); o.w = mmb_pack2(v[6], v[7]); + *(uint4 *)(dst + 32 * sub + 8 * il) = o; + } + else if constexpr (TYPE == GGML_TYPE_IQ4_XS) { + // decode only the two 32-blocks of this 64-wide K slice (tid = il*8 + ib): 1 call/lane instead of 4 + const mmb_quant_slice out{dst, k0 % QK}; + const int ib0 = (k0 % QK) / 32; + dequantize_iq4_xs(row, k0 / QK, out, (lane >> 1) * 8 + ib0 + (lane & 1)); + } + 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; + } + } +} + +// every type mmb_decode_slice() handles +static bool mmb_quant_type_any(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; + } +} + +// dense GEMMs: the two types measured faster than MMQ on every arch (#123) +static bool mmb_quant_type_dense(ggml_type type) { + return type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_NL; +} + +// routed experts: the types measured faster than MMQ at 512-4096 rows on gfx1151 (Q5_K / Q6_K decode the whole super-block and still won) +static bool mmb_quant_type_routed(ggml_type type) { + switch (type) { + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_MXFP4: + case GGML_TYPE_Q2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + return mmb_quant_type_any(type); + 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"); + } +} + +template +__device__ __forceinline__ void mmb_dq_q4k_slice(uint4 q0, uint4 q1, uint4 meta, int slice, uint32_t * out, int half = 0) { + 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 < (HALF ? 4 : 8); ++j) { + const int wj = HALF ? j + 4 * half : j; + const uint32_t q = words[j]; + out[2*wj] = mmb_pack2(d0 * (q & 15) - m0, d0 * ((q >> 8) & 15) - m0); + out[2*wj + 1] = mmb_pack2(d0 * ((q >> 16) & 15) - m0, d0 * ((q >> 24) & 15) - m0); + out[16 + 2*wj] = mmb_pack2(d1 * ((q >> 4) & 15) - m1, d1 * ((q >> 12) & 15) - m1); + out[16 + 2*wj + 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..7dd9f62b4c48 --- /dev/null +++ b/ggml/src/ggml-cuda/mmb.cu @@ -0,0 +1,1501 @@ +#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); } +// same RNE as mmb_f2bf on both halves, the two high halves joined by one v_perm_b32 (no mov_b16 + and_or). Unguarded +// like the other __builtin_amdgcn_perm uses in this file: device code here is AMD-only. +__device__ __forceinline__ uint32_t mmb_pack2(float a, float b) { + uint32_t ua = __float_as_uint(a), ub = __float_as_uint(b); + ua += 0x7fffu + ((ua >> 16) & 1u); ub += 0x7fffu + ((ub >> 16) & 1u); + return __builtin_amdgcn_perm(ub, ua, 0x07060302u); +} +__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); } + + // Q2_0 (routed down): a lane prefetches its 2 qs bytes and d in the WMMA loop, store_lds decodes its 8 weights + constexpr bool Q20 = WTYPE == 32 + GGML_TYPE_Q2_0; + constexpr int Q2R = Q20 ? (BM * 8) / MMB_NT : 1; + uint32_t q20[Q2R]; + + int weight_ks = 0; + auto load_regs = [&](const int ks) { + weight_ks = ks; + if constexpr (Q20) { +#pragma unroll + for (int r = 0; r < Q2R; ++r) { + const int row = (tid >> 3) + r * (MMB_NT / 8); + const block_q2_0 * x = (const block_q2_0 *)(Wbase + (size_t)row * wrow_bytes) + ks; + q20[r] = row < a_rows ? *(const uint16_t *)(x->qs + 2 * (tid & 7)) | ((uint32_t)*(const uint16_t *)&x->d << 16) : 0u; + } + } +#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 (Q20) { +#pragma unroll + for (int r = 0; r < Q2R; ++r) { + const int row = (tid >> 3) + r * (MMB_NT / 8); + const float d = mmb_h2f((uint16_t)(q20[r] >> 16)); + float v[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) v[j] = ((int) ((q20[r] >> (2 * j)) & 3) - 1) * d; + mmb_store8(As + row * MMB_LDS_STRIDE + 8 * (tid & 7), v); + } + } else 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]; + uint4 aw[WTYPE == 2 ? 8 : 1]; + 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); } + if constexpr (WTYPE == 2) if (tid < BM) { const uint4 * p = (const uint4 *)(arow + (size_t)ks * MMB_BK * 2); +#pragma unroll + for (int q = 0; q < 8; ++q) aw[q] = p[q]; } +#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 if constexpr (WTYPE == 2) { if (tid < BM) { +#pragma unroll + for (int q = 0; q < 8; ++q) *(uint4 *)(As + tid * MMB_LDS_STRIDE + q * 8) = aw[q]; } } + 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); } + // grid-based types (mmb_lb): prefetch each lane's block fields into registers during the WMMA loop, grid in + // LDS, decode in store_lds. Same fp32 math as dequantize_* -> bit-identical. + using LBT = mmb_lb; + constexpr bool LB = LBT::ok; + constexpr int LBR = LB ? (BM * 8) / MMB_NT : 1; + uint32_t lbg[LBR][2], lbu[LBR][2]; + __shared__ typename LBT::grid_t lbgrid[LBT::N]; + if constexpr (LB) { for (int i = tid; i < LBT::N; i += MMB_NT) lbgrid[i] = LBT::entry(i); } + int weight_ks = 0; + auto load_regs = [&](const int ks) { + weight_ks = ks; + if constexpr (LB) { + const int l8 = tid & 7, sub = l8 & 1, il = l8 >> 1; +#pragma unroll + for (int r = 0; r < LBR; ++r) { + const int row = (tid >> 3) + r * (MMB_NT / 8); + if (row < a_rows) { + LBT::fetch(Wg + (size_t)row * wrow_bytes, ks, sub, il, lbg[r][0], lbg[r][1]); + LBT::fetch(Wu + (size_t)row * wrow_bytes, ks, sub, il, lbu[r][0], lbu[r][1]); + } else { lbg[r][0] = lbg[r][1] = lbu[r][0] = lbu[r][1] = 0u; } // d = 0 -> zero rows (never stored) + } + } +#pragma unroll + for (int i = 0; i < A_ITEMS; ++i) { + const int row = Q4K_HALF ? (tid + i * MMB_NT) / 2 : 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 + (Q4K_HALF ? (tid & 1) * 16 : 0); + g0[i] = *(const uint4 *)(pg + offset); u0[i] = *(const uint4 *)(pu + offset); + if constexpr (!Q4K_HALF) { + g1[i] = *(const uint4 *)(pg + offset + 16); 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 (LB) { + // branchless (padding rows decode to zero); same fp32 arithmetic and the same RNE + sign-XOR packing + // as dequantize_*, so the BF16 weights are bit-identical to the generic decoder path + const int l8 = tid & 7, sub = l8 & 1, il = l8 >> 1; +#pragma unroll + for (int r = 0; r < LBR; ++r) { + const int row = (tid >> 3) + r * (MMB_NT / 8); + *(uint4 *)(Ag + row * MMB_LDS_STRIDE + 32 * sub + 8 * il) = LBT::decode(lbgrid, lbg[r][0], lbg[r][1], il); + *(uint4 *)(Au + row * MMB_LDS_STRIDE + 32 * sub + 8 * il) = LBT::decode(lbgrid, lbu[r][0], lbu[r][1], il); + } + } else 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 = Q4K_HALF ? (tid + i * MMB_NT) / 2 : 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], Q4K_HALF ? g0[i] : g1[i], gm[i], weight_ks & 3, (uint32_t *)(Ag + row * MMB_LDS_STRIDE), tid & 1); + mmb_dq_q4k_slice(u0[i], Q4K_HALF ? u0[i] : u1[i], um[i], weight_ks & 3, (uint32_t *)(Au + row * MMB_LDS_STRIDE), tid & 1); + } + } + } +#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); if constexpr (LB) __syncthreads(); 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))); +} +// EPI == 1 (qwen4exp QSA indexer scorer): the T columns are (query, head) pairs with 4 heads, and D is [M, T/4] with +// D[q*M + m] = ((relu(h0) + relu(h1)) + relu(h2)) + relu(h3) (+ log(step(float(tails[q]) - float(starts[m]))) when +// tails != nullptr): the exact expression the separate GEMM, idx_relu_sum and visibility chain evaluate. +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, + const float * __restrict__ W2 = nullptr, float * __restrict__ D2 = nullptr, const int M1 = 0, + const int32_t * __restrict__ tails = nullptr, const int32_t * __restrict__ starts = nullptr) { + // W2 != nullptr: two GEMMs on the same X, rows [0, M1) from W -> D (ld M1) and rows [M1, M) from W2 -> D2 (ld M - M1) +#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; + // one register stage: the next K slice is fetched while the current one is multiplied (unguarded loads with + // clamped rows, zeroed at the LDS store: a divergent guard would make the compiler drain them at the join) + constexpr int A_IT = (A_CH + MMB_NT - 1) / MMB_NT, B_IT = (B_CH + MMB_NT - 1) / MMB_NT; + static_assert(A_CH % MMB_NT == 0 && B_CH % MMB_NT == 0, "f32split staging assumes whole passes"); + float4 ra[A_IT], rb[B_IT]; + auto gload = [&](const int k0) { +#pragma unroll + for (int i = 0; i < A_IT; ++i) { const int idx = tid + i * MMB_NT, row = idx >> 3, c4 = (idx & 7) * 4; + const int gr = min(m0 + row, M - 1); + const float * wr = (W2 && gr >= M1) ? W2 + (size_t) (gr - M1) * K : W + (size_t) gr * K; + ra[i] = *(const float4 *)(wr + k0 + c4); } +#pragma unroll + for (int i = 0; i < B_IT; ++i) { const int idx = tid + i * MMB_NT, row = idx >> 3, c4 = (idx & 7) * 4; + rb[i] = *(const float4 *)(X + (size_t) min(t0 + row, T - 1) * K + k0 + c4); } + }; + auto lstore = [&]() { +#pragma unroll + for (int i = 0; i < A_IT; ++i) { const int idx = tid + i * MMB_NT, row = idx >> 3, c4 = (idx & 7) * 4; + const float4 v = m0 + row < M ? ra[i] : make_float4(0.f,0.f,0.f,0.f); + 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)); } +#pragma unroll + for (int i = 0; i < B_IT; ++i) { const int idx = tid + i * MMB_NT, row = idx >> 3, c4 = (idx & 7) * 4; + const float4 v = t0 + row < T ? rb[i] : make_float4(0.f,0.f,0.f,0.f); + 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)); } + }; + gload(0); + for (int k0 = 0; k0 < K; k0 += BKs) { + lstore(); + __syncthreads(); + gload(min(k0 + BKs, K - BKs)); + 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; + if constexpr (EPI == 1) { + // lane holds t = base + 4p + cn (e = 2p) and base + 4p + 2 + cn (e = 2p + 1): heads cn and 2 + cn of query + // base/4 + p; lane ^ 16 holds the other two heads of the same (m, query). The shuffles run on every lane. + const int Q = T >> 2; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int m = m0 + wm * WTM + i * 16 + cm; + const float sm = (tails && m < M) ? (float) starts[m] : 0.0f; +#pragma unroll + for (int j = 0; j < TN; ++j) { + float o[8]; +#pragma unroll + for (int e = 0; e < 8; ++e) o[e] = __shfl_xor(acc[i][j][e], 16, 32); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const float h0 = cn == 0 ? acc[i][j][2 * p] : o[2 * p]; + const float h1 = cn == 0 ? o[2 * p] : acc[i][j][2 * p]; + const float h2 = cn == 0 ? acc[i][j][2 * p + 1] : o[2 * p + 1]; + const float h3 = cn == 0 ? o[2 * p + 1] : acc[i][j][2 * p + 1]; + float s = fmaxf(h0, 0.0f); + s = s + fmaxf(h1, 0.0f); + s = s + fmaxf(h2, 0.0f); + s = s + fmaxf(h3, 0.0f); + const int q = ((t0 + wn * WTN + j * 16) >> 2) + p; + if ((p >> 1) == cn && m < M && q < Q) { + if (tails) { + const float d = (float) tails[q] - sm; + const float step = d > 0.0f; + s = s + logf(step); + } + D[(size_t) q * M + m] = s; + } + } + } + } + return; + } +#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) { + if (W2) { if (m < M1) D[(size_t) t * M1 + m] = acc[i][j][e]; else D2[(size_t) t * (M - M1) + (m - M1)] = acc[i][j][e]; } + else D[(size_t) t * M + m] = acc[i][j][e]; } } } +#endif +} + + +// qwen4exp QSA indexer scorer with the keys stationary. mmb_f32split_kernel<..., EPI=1> re-reads its key tile for every +// 128-column tile of the 2048 (query, head) columns, which spills the ~32 MB of pooled F32 keys per layer (250k depth) +// out of the Infinity Cache. Here each workgroup rounds its BM key rows to F16 into LDS once and walks all T columns, +// staging each 128-column chunk of the queries (pre-rounded to F16 once per call) in LDS. Every output tile issues the +// same WMMAs on the same F16 fragments in the same K order as the split kernel (TWO, !XSPLIT: hi parts only), and the +// epilogue computes the same values, so the result matches it bit for bit. +__global__ void mmb_f32_to_f16_kernel(const float * __restrict__ x, uint16_t * __restrict__ y, const int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = __builtin_bit_cast(uint16_t, (_Float16) x[i]); +} + +template +__global__ void __launch_bounds__(MMB_NT, 1) +mmb_idx_score_kernel(const float * __restrict__ W, const uint16_t * __restrict__ Xh, float * __restrict__ D, + const int M, const int K, const int T, const int32_t * __restrict__ tails, const int32_t * __restrict__ starts) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; +#else + constexpr int KMAX = 128, LS = KMAX + 8; + constexpr int WAVES_N = MMB_NT / 32 / WAVES_M, WTM = BM / WAVES_M, WTN = BN / WAVES_N, TM = WTM / 16, TN = WTN / 16; + static_assert(WAVES_M * WAVES_N * 32 == MMB_NT && TM >= 1 && TN >= 1 && WTM % 16 == 0 && WTN % 16 == 0, "tile must use every wave"); + __shared__ __align__(16) uint16_t Ah[BM * LS]; + __shared__ __align__(16) uint16_t Bh[BN * LS]; + 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; + const int K4 = K / 4, K8 = K / 8; + + // keys -> F16 (the hi part of mmb_split2), rows past M are zero + for (int idx = tid; idx < BM * K4; idx += MMB_NT) { + const int row = idx / K4, c4 = (idx % K4) * 4; + const float4 v = m0 + row < M ? *(const float4 *)(W + (size_t) (m0 + row) * K + c4) : make_float4(0.f, 0.f, 0.f, 0.f); + const uint32_t h0 = __builtin_bit_cast(uint16_t, (_Float16) v.x), h1 = __builtin_bit_cast(uint16_t, (_Float16) v.y); + const uint32_t h2 = __builtin_bit_cast(uint16_t, (_Float16) v.z), h3 = __builtin_bit_cast(uint16_t, (_Float16) v.w); + *(uint2 *)(Ah + row * LS + c4) = make_uint2(h0 | (h1 << 16), h2 | (h3 << 16)); + } + + // query chunks: BN rows x K halfs, 8 halfs per load; columns past T are zero + constexpr int B_IT = BN * KMAX / 8 / MMB_NT; + uint4 rb[B_IT]; + const int b_loads = BN * K8; + auto gload = [&](const int t0) { +#pragma unroll + for (int it = 0; it < B_IT; ++it) { + const int idx = tid + it * MMB_NT; + const int row = idx / K8, c8 = (idx % K8) * 8; + rb[it] = (idx < b_loads && t0 + row < T) ? *(const uint4 *)(Xh + (size_t) (t0 + row) * K + c8) : make_uint4(0, 0, 0, 0); + } + }; + auto lstore = [&]() { +#pragma unroll + for (int it = 0; it < B_IT; ++it) { + const int idx = tid + it * MMB_NT; + if (idx < b_loads) { const int row = idx / K8, c8 = (idx % K8) * 8; *(uint4 *)(Bh + row * LS + c8) = rb[it]; } + } + }; + + const int r = lane & 15, cm = lane & 15, cn = lane >> 4; + const int Q = T >> 2; + float sm[TM]; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int m = m0 + wm * WTM + i * 16 + cm; sm[i] = (tails && m < M) ? (float) starts[m] : 0.0f; } + + gload(0); + for (int t0 = 0; t0 < T; t0 += BN) { + __syncthreads(); // the previous chunk is consumed (and the keys are stored on the first pass) + lstore(); + __syncthreads(); + if (t0 + BN < T) gload(t0 + 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; + for (int k0 = 0; k0 < K; k0 += 16) { + v16s ah[TM], bh[TN]; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int off = (wm * WTM + i * 16 + r) * LS + k0; + ah[i] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Ah + off), *(const uint4 *)(Ah + off + 8)}); } +#pragma unroll + for (int j = 0; j < TN; ++j) { const int off = (wn * WTN + j * 16 + r) * LS + k0; + bh[j] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Bh + off), *(const uint4 *)(Bh + 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]); + } + // epilogue of mmb_f32split_kernel<..., EPI=1> +#pragma unroll + for (int i = 0; i < TM; ++i) { const int m = m0 + wm * WTM + i * 16 + cm; +#pragma unroll + for (int j = 0; j < TN; ++j) { + float o[8]; +#pragma unroll + for (int e = 0; e < 8; ++e) o[e] = __shfl_xor(acc[i][j][e], 16, 32); +#pragma unroll + for (int p = 0; p < 4; ++p) { + const float h0 = cn == 0 ? acc[i][j][2 * p] : o[2 * p]; + const float h1 = cn == 0 ? o[2 * p] : acc[i][j][2 * p]; + const float h2 = cn == 0 ? acc[i][j][2 * p + 1] : o[2 * p + 1]; + const float h3 = cn == 0 ? o[2 * p + 1] : acc[i][j][2 * p + 1]; + float s = fmaxf(h0, 0.0f); + s = s + fmaxf(h1, 0.0f); + s = s + fmaxf(h2, 0.0f); + s = s + fmaxf(h3, 0.0f); + const int q = ((t0 + wn * WTN + j * 16) >> 2) + p; + if ((p >> 1) == cn && m < M && q < Q) { + if (tails) { + // log(step(d)) is exactly 0 or -inf, and s >= +0, so this is s + logf(step) + const float d = (float) tails[q] - sm[i]; + s = d > 0.0f ? s + 0.0f : -INFINITY; + } + D[(size_t) q * M + m] = s; + } + } + } + } + } +#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_q8_0_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 * 34; + const float d = mmb_h2f((uint16_t)(p[0] | (p[1] << 8))); + uint32_t * o = (uint32_t *) (out + b * 32); +#pragma unroll + for (int j = 0; j < 16; ++j) o[j] = mmb_pack2(d * (float)(int8_t) p[2 + 2 * j], d * (float)(int8_t) p[3 + 2 * j]); +} + +// HC Q8_0 matrices (hc_*_down [10240 -> 320], hc_*_up [320 -> 10240]) get a BF16 shadow: their GEMMs re-decode the +// whole matrix per token tile. MMB_SHADOW_Q8=0 disables. +static bool mmb_is_hc_q8(const ggml_tensor * w) { + static const bool on = getenv("MMB_SHADOW_Q8") ? atoi(getenv("MMB_SHADOW_Q8")) != 0 : true; + // HC down [10240 -> 320] reads its weights once per token tile: Q8_0 (68 B / 64 weights) beats the BF16 shadow + // (128 B) there, so only the HC up / gate weights get a shadow unless MMB_HCD_SHADOW=1 + static const bool down = getenv("MMB_HCD_SHADOW") ? atoi(getenv("MMB_HCD_SHADOW")) != 0 : false; + if (!on || !w || w->type != GGML_TYPE_Q8_0 || w->op != GGML_OP_NONE || !w->data || !w->buffer || w->ne[2] != 1 || w->ne[3] != 1 || !ggml_is_contiguous(w)) return false; + return (down && w->ne[0] >= 4096 && w->ne[1] <= 384) || (w->ne[0] <= 384 && w->ne[1] >= 4096); +} + +__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). +// llama now opts in every backend context it creates, limited to the weight types in mmb_quant_type_dense() / mmb_quant_type_routed(); only the 32-row small-batch gate (mmb_min_t) depends on the arch. +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); +} +// Smallest GEMM row count for the MMB consumers: 32 in a context with the small-batch opt-in (qwen4exp), else 512. The QSA indexer score keeps 512, see below. +// gfx1151 / ROCm 10, Qwen3.8-Flash-Next UD-Q4_K_XL and UD-IQ4_XS prefill: MMB is about 5% slower than MMQ at 16 tokens, 11-14% faster at 32. +// Retest when the MMQ or MMB tiles or the mmb_quant_type_*() gates change. +int mmb_min_t(const ggml_backend_cuda_context & ctx) { return ctx.mmb_small_batch ? 32 : 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() { static const bool v = getenv("MMB_GLU") ? atoi(getenv("MMB_GLU")) != 0 : true; return v; } + +} // 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(ctx)) 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_dense(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]; + // graph-computed F32 src0 (not a weight), such as the QSA indexer score with heads x tokens GEMM rows: keep the 512-row gate. + // At 8-31 tokens and 32K depth, MMB changed the scores with no measured speed gain. + const bool graph_src0 = f32w && src0->op != GGML_OP_NONE; + if (T < (graph_src0 ? 512 : mmb_min_t(ctx)) || T > INT32_MAX / 4) return false; + return ggml_nrows(dst) == T; +} + +// true when act is (or derives from) a gate/up view of one fused gate_up MUL_MAT_ID output (ne[0] == 2*n_ff). Split- +// tensor models build gate and up as their own mul_mat_id, so they are not matched and keep the GLU fusion. +static bool mmb_act_from_fused_gateup(const ggml_tensor * act, const int n_ff, const int depth = 0) { + if (!act || depth > 3) return false; + if (act->op == GGML_OP_VIEW) { + const ggml_tensor * s = act->view_src; + return s && s->op == GGML_OP_MUL_MAT_ID && s->ne[0] == 2 * n_ff; + } + if (act->op == GGML_OP_MUL_MAT_ID) return false; // split path: gate / up are their own projections + for (int i = 0; i < 4; ++i) { + if (!act->src[i]) break; + if (mmb_act_from_fused_gateup(act->src[i], n_ff, depth + 1)) return true; + } + return false; +} + +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; + // Fused gate_up is split into gate/up views, so the SwiGLU fusion never fires and MMB pays BF16 precision with no + // fused-kernel payoff. gemma-4-26B-A4B / gfx1151 / ROCm 10.0: KL vs master 0.19-0.31, top-1 90-93% (split-tensor + // models 0.002-0.065, 97-99%); pp4096 +0.9%. Keep on MMQ; Q8_0 / IQ4_NL already eligible, so exempt. + const char * nm = ggml_get_name(dst); + if (nm && strncmp(nm, "ffn_moe_gate_up", 15) == 0 && + src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_IQ4_NL) return false; + // down projection of a fused gate_up model: its activation is the SwiGLU of two gate_up views, so the whole + // expert FFN stays on MMQ (same Q8_0 / IQ4_NL exemption). + if (mmb_act_from_fused_gateup(src1, (int) src0->ne[0]) && + src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_IQ4_NL) return false; + if (!mmb_quant_type_routed(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(ctx) || n_used > 64 || (T * n_used) >> 16 >= 1024) return false; // tile index must fit in 16 bits per expert + return true; +} + +// geometry ggml_cuda_mul_mat_mmb picks for an F32 weight with M > 64; the fused scorer only takes 0 and 5 +static int mmb_f32_tile_for(int M) { + static const int FT_env = getenv("MMB_F32_TILE") ? atoi(getenv("MMB_F32_TILE")) : -1; + return FT_env >= 0 ? FT_env : (M <= 1024 ? 5 : 0); +} + +bool ggml_cuda_mmb_idx_score_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * mm, int heads) { + if (heads != 4 || src0->type != GGML_TYPE_F32 || !ggml_cuda_mmb_supported_mm(ctx, src0, src1, mm)) return false; + const int64_t M = src0->ne[1], T = src1->ne[1] * src1->ne[2] * src1->ne[3]; + if (M <= 64 || T % 4 != 0 || src1->ne[2] != 1 || src1->ne[3] != 1) return false; + const int FT = mmb_f32_tile_for((int) M); + return FT == 0 || FT == 5; +} + +void ggml_cuda_mmb_idx_score(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, + const int32_t * tails, const int32_t * starts) { + 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]); + const float * W = (const float *) src0->data, * X = (const float *) src1->data; float * D = (float *) dst->data; + // keys-stationary scorer for the wide case (MMB_IDX_SCORE_V2=0 keeps the split kernel; same result) + static const bool v2 = getenv("MMB_IDX_SCORE_V2") == nullptr || atoi(getenv("MMB_IDX_SCORE_V2")) != 0; + if (v2 && mmb_f32_tile_for(M) == 0 && K <= 128 && K % 16 == 0) { + ggml_cuda_pool_alloc xh(ctx.pool(), (size_t) T * K); + const int64_t n = (int64_t) T * K; + mmb_f32_to_f16_kernel<<<(unsigned) ((n + 255) / 256), 256, 0, stream>>>(X, xh.get(), n); + // 64 key rows x 64 query columns per step (2 x 4 waves): fastest of the shapes tried on gfx1151 at 250k depth + mmb_idx_score_kernel<64, 64, 2><<<(M + 63) / 64, MMB_NT, 0, stream>>>(W, xh.get(), D, M, K, T, tails, starts); + CUDA_CHECK(cudaGetLastError()); + return; + } + if (mmb_f32_tile_for(M) == 5) { + dim3 g((M + 31) / 32, (T + 63) / 64); + mmb_f32split_kernel<32, 64, 16, 16, true, false, 1><<>>(W, X, D, M, K, T, nullptr, nullptr, 0, tails, starts); + } else { + dim3 g((M + 127) / 128, (T + 127) / 128); + mmb_f32split_kernel<128, 128, 32, 64, true, false, 1><<>>(W, X, D, M, K, T, nullptr, nullptr, 0, tails, starts); + } + CUDA_CHECK(cudaGetLastError()); +} + +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; + // MMB_F32_TILE narrows the M>64 F32 GEMM (e.g. the M=512 MoE router) for more blocks per ubatch chunk. + // 0 keeps master's 128x128 tile. Valid geometries satisfy (8/WAVES_M)*WTN == BN (the kernel's static_assert). + // Default: 5 (32x64) for M <= 1024 (the M=512 router: 87 -> 77 ms per pp4096 on gfx1151, about -0.4% GPU time), + // 128x128 above, where the narrow tile would reload the activations ~3x more per output. Every geometry keeps the + // per-output K order, so the output is the same. + static const int FT_env = getenv("MMB_F32_TILE") ? atoi(getenv("MMB_F32_TILE")) : -1; + const int FT = FT_env >= 0 ? FT_env : (M <= 1024 ? 5 : 0); + // MMB_F32_SMALL: M <= 64 (ssm_alpha / ssm_beta [2560 -> 48]) launches only T/128 blocks at 64x128 (32 at 4096 tokens + // on 40 CUs); narrower token tiles give 2-4x the blocks with the same per-output K order (bit-identical) + static const int FS = getenv("MMB_F32_SMALL") ? atoi(getenv("MMB_F32_SMALL")) : 2; + if (M <= 64 && FS == 2) { + dim3 grid((M + 63) / 64, (T + 31) / 32); + mmb_f32split_kernel<64, 32, 16, 16, true, false><<>>(W, X, D, M, K, T); + } else if (M <= 64 && FS == 1) { + dim3 grid((M + 63) / 64, (T + 63) / 64); + mmb_f32split_kernel<64, 64, 16, 32, true, false><<>>(W, X, D, M, K, T); + } else 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 if (FT == 1) { dim3 g((M + 63) / 64, (T + 127) / 128); mmb_f32split_kernel< 64, 128, 32, 32, true, false><<>>(W, X, D, M, K, T); } + else if (FT == 4) { dim3 g((M + 31) / 32, (T + 127) / 128); mmb_f32split_kernel< 32, 128, 32, 16, true, false><<>>(W, X, D, M, K, T); } + else if (FT == 5) { dim3 g((M + 31) / 32, (T + 63) / 64); mmb_f32split_kernel< 32, 64, 16, 16, true, false><<>>(W, X, D, M, K, T); } + else if (FT == 6) { dim3 g((M + 63) / 64, (T + 63) / 64); mmb_f32split_kernel< 64, 64, 32, 16, true, false><<>>(W, X, D, M, K, T); } + else { dim3 g((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; + } + // Q8_0 HC down [10240 -> 320]: at 128x128 tiles only 3x(T/128) blocks with a K=10240 loop (occupancy-bound). + // Narrower tiles keep the per-output K accumulation order (bit-identical) but launch 3-10x more blocks. + // MMB_HCD_TILE: 0 = generic path, 1 <64,128,32,32>, 2 <32,128,32,16>, 3 <64,64,32,16>, 4 <32,64,16,16>, 5 <64,32,16,16> + const uint16_t * hc_shadow = mmb_is_hc_q8(src0) ? mmb_shadow_lookup(ctx, src0) : nullptr; + static const int HCD = getenv("MMB_HCD_TILE") ? atoi(getenv("MMB_HCD_TILE")) : 3; + if (HCD && src0->type == GGML_TYPE_Q8_0 && M <= 384 && K >= 4096 && T >= 512) { + const bool hbf = ggml_cuda_mmb_blk16() && ggml_cuda_mmb_is_bf16_only(ctx, dst) && (M & 7) == 0; + uint16_t * Dh2 = hbf ? (uint16_t *) dst->data : nullptr; const bool sf = !hbf; + auto go = [&](auto tag, const uint8_t * WP) { + constexpr int WT = decltype(tag)::value; + if (HCD == 1) { dim3 g((M + 63) / 64, (T + 127) / 128); mmb_dense_kernel<64, 128, 32, 32, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + else if (HCD == 2) { dim3 g((M + 31) / 32, (T + 127) / 128); mmb_dense_kernel<32, 128, 32, 16, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + else if (HCD == 3) { dim3 g((M + 63) / 64, (T + 63) / 64); mmb_dense_kernel<64, 64, 32, 16, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + else if (HCD == 4) { dim3 g((M + 31) / 32, (T + 63) / 64); mmb_dense_kernel<32, 64, 16, 16, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + else if (HCD == 6) { dim3 g(1, (T + 31) / 32); mmb_dense_kernel<384, 32, 96, 16, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + else { dim3 g((M + 63) / 64, (T + 31) / 32); mmb_dense_kernel<64, 32, 16, 16, WT><<>>(WP, xhp, D, Dh2, sf, M, K, T); } + }; + if (hc_shadow) go(std::integral_constant{}, (const uint8_t *) hc_shadow); + else go(std::integral_constant{}, W); + 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; + // MMB_DENSE_TILE narrows the shared-expert / dense GEMM tiles for more blocks per ubatch chunk. +#define MMB_SMALL_DENSE(WT, WP) do { \ + static const int DT_ = getenv("MMB_DENSE_TILE") ? atoi(getenv("MMB_DENSE_TILE")) : 0; \ + if (DT_ == 1) { dim3 gd((M + 63) / 64, (T + 127) / 128); mmb_dense_kernel<64, 128, 32, 32, WT><<>>(WP, xhp, D, Dh, store_f32, M, K, T); } \ + else if (DT_ == 2) { dim3 gd((M + 31) / 32, (T + 127) / 128); mmb_dense_kernel<32, 128, 32, 16, WT><<>>(WP, xhp, D, Dh, store_f32, M, K, T); } \ + else if (DT_ == 3) { dim3 gd((M + 63) / 64, (T + 63) / 64); mmb_dense_kernel<64, 64, 32, 16, WT><<>>(WP, xhp, D, Dh, store_f32, M, K, T); } \ + else { mmb_dense_kernel<128, 128, 32, 64, WT><<>>(WP, xhp, D, Dh, store_f32, M, K, T); } \ +} while (0) + + 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_SMALL_DENSE(2, (const uint8_t *) shadow); } + } 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_SMALL_DENSE(0, W); } + } 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_SMALL_DENSE(1, W); } + } else if (mmb_quant_type_dense(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_SMALL_DENSE(WT, W); } + }); + } else { + if (big) mmb_dense_kernel<128, 256, 64, 64, 2><<>>(W, xhp, D, Dh, store_f32, M, K, T); + else { MMB_SMALL_DENSE(2, W); } + } + CUDA_CHECK(cudaGetLastError()); +} + +// two F32-weight GEMMs [K -> M1] and [K -> M2] (M1 + M2 <= 128) on the same F32 activations, one pass over X +bool ggml_cuda_mmb_f32_dual(ggml_backend_cuda_context & ctx, const ggml_tensor * w1, const ggml_tensor * w2, const ggml_tensor * x, + ggml_tensor * d1, ggml_tensor * d2) { + const int K = (int) w1->ne[0], M1 = (int) w1->ne[1], M2 = (int) w2->ne[1]; + const int T = (int) (x->ne[1] * x->ne[2] * x->ne[3]); + if (w1->type != GGML_TYPE_F32 || w2->type != GGML_TYPE_F32 || x->type != GGML_TYPE_F32 || w2->ne[0] != K || x->ne[0] != K || + M1 + M2 > 128 || M1 % 16 != 0 || K % 32 != 0 || !ggml_is_contiguous(w1) || !ggml_is_contiguous(w2) || !ggml_is_contiguous(x) || + !ggml_is_contiguous(d1) || !ggml_is_contiguous(d2) || ggml_nelements(d1) != (int64_t) M1 * T || ggml_nelements(d2) != (int64_t) M2 * T) { + return false; + } + const int M = M1 + M2; + dim3 grid((M + 63) / 64, (T + 31) / 32); + mmb_f32split_kernel<64, 32, 16, 16, true, false><<>>((const float *) w1->data, (const float *) x->data, + (float *) d1->data, M, K, T, (const float *) w2->data, (float *) d2->data, M1); + CUDA_CHECK(cudaGetLastError()); + return true; +} + +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_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * w, const ggml_tensor * lo, const ggml_tensor * xn, const ggml_tensor * dst, const int hc) { + if (!mmb_gatemix_flag() || hc != 4 || !mmb_quant_type_dense(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; + return !(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(ctx)); +} +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 (!ggml_cuda_hc_gate_mix_supported(ctx, w, lo, xn, dst, hc)) return false; + const int K = (int) w->ne[0], E = (int) dst->ne[0]; const int T = (int) ggml_nrows(dst); + 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); + const uint16_t * hc_shadow = mmb_is_hc_q8(w) ? mmb_shadow_lookup(ctx, w) : nullptr; + if (hc_shadow) { + hc_gate_mix_kernel<4, 2><<>>((const uint8_t *) hc_shadow, lo16, xn16, (float *) dst->data, outh, store_f32, E, K, T, scale, bias); + } else + 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=32 (was 128): more MoE experts use the wide BN tile during prefill. + // Measured on qwen3.8-flash-next / gfx1151, -p 2048 -d 0,12000,32000,64000: + // +3.6% d0, +4.7% d12k, +4.0% d32k, +1.9% d64k vs THRESH=128; byte-identical + // output, test-backend-ops 29643/29643. Env MMB_THRESH overrides. + static const int THRESH = getenv("MMB_THRESH") ? atoi(getenv("MMB_THRESH")) : 32; + 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; + static const int RT = getenv("MMB_ROUTED_TILE") ? atoi(getenv("MMB_ROUTED_TILE")) : 0; + // small tile mapped WAVES_M=8 x WAVES_N=1 (<128,32,16,32>), not 4 x 2 (<128,32,32,16>): a <=16-token tile (the + // common small-expert case) leaves 4 waves WMMA-idle under the old shape. Bit-identical: same (m,n), same K order. + 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; + if (RT == 1) { + const dim3 gb((M + 63) / 64, nbig_max), gs((M + 63) / 64, nsmall_max); + mmb_routed_kernel<64, BN, 32, 32, WT><<>>(W, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + mmb_routed_kernel<64, BN_SMALL, 16, 16, WT><<>>(W, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_small.get(), M, K); + } else { + 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, 16, 32, 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_routed(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; + static const int BN_sel = getenv("MMB_BN") ? atoi(getenv("MMB_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=32 (was 128): more MoE experts use the wide BN tile during prefill. + // Measured on qwen3.8-flash-next / gfx1151, -p 2048 -d 0,12000,32000,64000: + // +3.6% d0, +4.7% d12k, +4.0% d32k, +1.9% d64k vs THRESH=128; byte-identical + // output, test-backend-ops 29643/29643. Env MMB_THRESH overrides. + static const int THRESH = getenv("MMB_THRESH") ? atoi(getenv("MMB_THRESH")) : 32; + const int BN_rt = BN_sel; + // actual tile counts for the chosen BN; pool sized for the smallest BN (largest count) so any BN_sel fits + const int nbig_max = n_rows / BN_rt + E + 1; + const int nbig_cap = n_rows / 64 + E + 1; + const int nsmall_max = E * ((THRESH + BN_SMALL - 1) / BN_SMALL) + 1; + ggml_cuda_pool_alloc desc_big(ctx.pool(), nbig_cap); + 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_cap, nsmall_max, BN_rt, 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]; + // MMB_BM_BIG / MMB_BM_SMALL widen the M tile (n_ff_exp rows per block) to halve the M-tile count and the + // B-tile reloads at small rows/expert. BM=128 needs a matching (WTM,WTN) to keep the 8 waves full. + static const int BM_BIG = getenv("MMB_BM_BIG") ? atoi(getenv("MMB_BM_BIG")) : 64; + static const int BM_SMALL = getenv("MMB_BM_SMALL") ? atoi(getenv("MMB_BM_SMALL")) : 64; + // gfx1151 / ROCm 10, Qwen3.8 Q4_K pp4096: routed GLU 510 -> 476 ms in a GPU trace. + // Retest if quant decode or lane layout changes; MMB_Q4K_HALF=0 restores one lane per row. + static const bool q4k_half = getenv("MMB_Q4K_HALF") ? atoi(getenv("MMB_Q4K_HALF")) != 0 : true; + mmb_dispatch_quant(gw->type, [&](auto tag) { + constexpr int WT = decltype(tag)::value; + if (BM_BIG == 128 && BN_sel != 256) { + const dim3 gbig((M + 127) / 128, nbig_max); + mmb_routed_glu_kernel<128, 128, 64, 32, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + } else { + auto launch_big = [&](auto bn_tag) { + constexpr int BN = decltype(bn_tag)::value; + const dim3 gbig((M + 63) / 64, nbig_max); + if constexpr (WT == 32 + GGML_TYPE_Q4_K && BN == 128) { + if (q4k_half) { + mmb_routed_glu_kernel<64, BN, 32, BN / 4, WT, true><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + } else { + mmb_routed_glu_kernel<64, BN, 32, BN / 4, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + } + } else { + mmb_routed_glu_kernel<64, BN, 32, BN / 4, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + } + }; + if (BN_sel == 256) launch_big(std::integral_constant{}); + else launch_big(std::integral_constant{}); + } + if (BM_SMALL == 128) { + const dim3 gsmall((M + 127) / 128, nsmall_max); + mmb_routed_glu_kernel<128, BN_SMALL, 32, 16, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_small.get(), M, K); + } else { + const dim3 gsmall((M + 63) / 64, nsmall_max); + if constexpr (WT == 32 + GGML_TYPE_Q4_K) { + if (q4k_half) { + mmb_routed_glu_kernel<64, BN_SMALL, 16, 16, WT, true><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_small.get(), M, K); + } else { + 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); + } + } else { + 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_hc_q8(w)) { + if (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()) 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_q8_0_bf16_kernel<<<(unsigned) ((n / 32 + 255) / 256), 256, 0, ctx.stream()>>>((const uint8_t *) w->data, buf, n / 32); + CUDA_CHECK(cudaGetLastError()); + mmb_state(ctx).shadow[w->data] = buf; mmb_state(ctx).shadow_bytes += bytes; + 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..b573712b653f --- /dev/null +++ b/ggml/src/ggml-cuda/mmb.cuh @@ -0,0 +1,38 @@ +#pragma once +#include "common.cuh" +// Quantized-weight BF16 WMMA GEMM on gfx1151, from 512 GEMM rows up, or 32 with the small-batch opt-in (see mmb_min_t). +// 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); +// qwen4exp QSA indexer scorer: F32 GEMM + relu + 4-head sum (+ compact visibility) in one kernel, dst [M, T/4] +bool ggml_cuda_mmb_idx_score_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * mm, int heads); +void ggml_cuda_mmb_idx_score(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, + const int32_t * tails, const int32_t * starts); +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(); +// two F32-weight GEMMs (M1 + M2 <= 128 rows) on the same F32 activations in one pass (qwen4exp ssm_beta + ssm_alpha) +bool ggml_cuda_mmb_f32_dual(ggml_backend_cuda_context & ctx, const ggml_tensor * w1, const ggml_tensor * w2, const ggml_tensor * x, ggml_tensor * d1, ggml_tensor * d2); +bool ggml_cuda_mmb_down16(); +bool ggml_cuda_mmb_res16(); +bool ggml_cuda_mmb_blk16(); +// shape and type checks of ggml_cuda_hc_gate_mix (which also needs the BF16 xn copy) +bool ggml_cuda_hc_gate_mix_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * w, const ggml_tensor * lo, const ggml_tensor * xn, const ggml_tensor * dst, int hc); +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-ampere.cuh b/ggml/src/ggml-cuda/mmq-config-ampere.cuh index 9f9fd197382f..fa62f34d6876 100644 --- a/ggml/src/ggml-cuda/mmq-config-ampere.cuh +++ b/ggml/src/ggml-cuda/mmq-config-ampere.cuh @@ -346,21 +346,37 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf // --------------------------------------------------------------------------------------------- CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, true); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_MXFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, true, false); CASE(GGML_TYPE_NVFP4, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); CASE(GGML_TYPE_NVFP4, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, true); diff --git a/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh b/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh index 180b2d9370d1..aaadfaa1f3bf 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); @@ -261,17 +261,29 @@ 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_Q4_0_ROCMFP4_FAST, 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_Q4_0_ROCMFP4_FAST, 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_Q4_0_ROCMFP4_FAST, 256, 2, 128, 64, 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_Q4_0_ROCMFP4_FAST, 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_Q4_0_ROCMFP4_FAST, 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_Q4_0_ROCMFP4_FAST, 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); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_NVFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); CASE(GGML_TYPE_NVFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); @@ -288,3 +300,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 e19f4f24d78b..7173ba98d8a3 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( @@ -1766,3 +1862,73 @@ template static __device__ __forceinline_ x_u32_scale[i*sram_stride] = get_int_b4(bxi->d, 0); } } + +// Q4_0_ROCMFP4_FAST: 32 values / block, one UE4M3 scale per 32 values - same tile structure as MXFP4. +// Differences from the mxfp4 loader above: block_rocmfp4_fast layout, kvalues_rocmfp4 codebook +// (Codebook10), and the UE4M3 scale decoder without the e8m0 0.5 factor. +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0_rocmfp4_fast( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + 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); + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0_ROCMFP4_FAST, I); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_ROCMFP4); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI_ROCMFP4; + const int kqsx = txi % QI_ROCMFP4; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbx; + + const int aux_q4 = get_int_b1(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_rocmfp4); + const int k0 = kbx * (2 * QI_ROCMFP4) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*sram_stride + k0 + 0] = v.x; + x_qs[i*sram_stride + k0 + QI_ROCMFP4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_ROCMFP4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_ROCMFP4; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*sram_stride + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#else + x_df[i*(MMQ_TILE_NE_K/QI_ROCMFP4) + i/QI_ROCMFP4 + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} 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 4ca6542d80bf..d9e99d7b3e05 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 3e1721bbe686..af5383e622b9 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -5,7 +5,27 @@ #include -static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream, const ggml_prec prec_src1) { +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, const ggml_prec prec_src1 = GGML_PREC_Q8) { switch (args.type_x) { case GGML_TYPE_Q1_0: mul_mat_q_case(ctx, args, stream); @@ -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: @@ -85,12 +113,150 @@ 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_Q4_0_ROCMFP4_FAST: + mul_mat_q_case(ctx, args, stream); + break; default: GGML_ABORT("fatal error"); break; } } +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); + } +} + // overrides the src1 precision requested by the graph, "auto" keeps the requested one static ggml_prec ggml_cuda_mmq_get_prec_env() { const char * env_c = getenv("GGML_CUDA_MMQ_PREC"); @@ -132,8 +298,9 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm return GGML_PREC_Q4; } -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) { +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. @@ -152,6 +319,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; @@ -204,6 +379,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); @@ -249,8 +430,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()); } @@ -281,6 +465,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); @@ -315,6 +505,17 @@ void ggml_cuda_mul_mat_q( ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); } +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; @@ -350,6 +551,15 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t case GGML_TYPE_NVFP4: mmq_supported = true; break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + // MMQ tile rows exist only for RDNA3.5 (gfx1151, validated) and the + // Ampere config table (validated on RTX 4090 / sm_89). Gate to exactly + // those: Blackwell/Rubin have no rows of their own (native FP4 MMQ is + // MXFP4/NVFP4-only there) and would otherwise fall through to the + // Ampere table unvalidated. Everywhere else falls back to dequant + BLAS. + mmq_supported = (ampere_mma_available(cc) && cc < GGML_CUDA_CC_BLACKWELL) || + GGML_CUDA_CC_IS_RDNA3_5(cc); + break; default: mmq_supported = false; break; @@ -420,6 +630,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 4b50d1dc333b..46658b243dd8 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -75,6 +75,8 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_NVFP4: return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_Q2_K: return MMQ_Q8_1_DS_LAYOUT_D2S6; case GGML_TYPE_Q3_K: @@ -406,6 +408,7 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K; case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K; @@ -489,18 +492,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; @@ -680,6 +689,12 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func ggml_cuda_mmq_load_tiles_nvfp4, ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a, ggml_cuda_mmq_write_back_dp4a); + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return ggml_cuda_mmq_util_funcs( + VDR_ROCMFP4_FAST_Q8_1_MMQ, + ggml_cuda_mmq_load_tiles_q4_0_rocmfp4_fast, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a, + ggml_cuda_mmq_write_back_dp4a); default: return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); } @@ -850,6 +865,12 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func ggml_cuda_mmq_load_tiles_nvfp4, ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma, ggml_cuda_mmq_write_back_mma); + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return ggml_cuda_mmq_util_funcs( + -1, + ggml_cuda_mmq_load_tiles_q4_0_rocmfp4_fast, + ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma, + ggml_cuda_mmq_write_back_mma); default: return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr); } @@ -877,6 +898,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, @@ -913,6 +950,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); { @@ -982,21 +1099,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, prec_src1)) { const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); @@ -1045,6 +1148,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)); @@ -1068,6 +1183,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, prec_src1); constexpr int blocks_per_iter = ITER_K / qk; @@ -1245,6 +1372,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, prec_src1)/2, 1) static __global__ void mul_mat_q_stream_k_fixup( @@ -1400,6 +1622,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(); @@ -1434,6 +1732,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) && + !config.stream_k && prec_src1 == GGML_PREC_Q8 && 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 (!config.stream_k) { mul_mat_q<<>> (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, @@ -1482,6 +1797,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(); @@ -1491,6 +1815,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, prec_src1); + 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, prec_src1); + 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, prec_src1); if (config.type == GGML_TYPE_COUNT) { @@ -1608,6 +1957,7 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); // ----------------------------------------- extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); extern DECL_MMQ_CASE_W4A4(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4); @@ -1616,4 +1966,9 @@ extern DECL_MMQ_CASE_W4A4(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 a85155360a7b..3d1c913826ba 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" @@ -48,6 +49,8 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4: return vec_dot_q4_0_rocmfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return vec_dot_q4_0_rocmfp4_fast_q8_1; case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; @@ -77,6 +80,8 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4: return VDR_ROCMFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return VDR_ROCMFP4_Q8_1_MMVQ; case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; @@ -98,6 +103,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 +112,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 +136,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)) { @@ -450,6 +461,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: @@ -576,7 +590,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: @@ -742,12 +1430,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); + } } } } @@ -844,39 +1556,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) { @@ -1177,6 +2642,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; @@ -1221,6 +2733,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, @@ -1322,6 +2863,18 @@ static void mul_mat_vec_q_switch_type( nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; case GGML_TYPE_NVFP4: mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, @@ -1418,6 +2971,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) { @@ -1497,6 +3329,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..95c75603daaf 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cu +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -1,4 +1,10 @@ #include "moe-weighted-reduction.cuh" +#include "mmb.cuh" +#include +#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 +30,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 +148,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 +163,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,9 +179,130 @@ 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, (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, stream); CUDA_CHECK(cudaGetLastError()); } + +// prefill: reduction + shared-expert merge in one pass, out = sum_k (e_k * scale_k) * w_k + shexp * sigmoid(gate[t]). +// The sum is the moe_weighted_reduction_*_v4 expression; the merge is sigmoid_gate_mul_add_f32's (product and sum rounded +// separately, sigmoid as op_sigmoid) -> bit-identical to the two kernels, without the F32 round trip of the reduction. +#if defined(GGML_USE_HIP) +static __device__ __forceinline__ float moe_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 moe_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 moe_mul_rn(const float a, const float b) { return __fmul_rn(a, b); } +static __device__ __forceinline__ float moe_add_rn(const float a, const float b) { return __fadd_rn(a, b); } +#endif +template +static __global__ void moe_weighted_reduction_sgma_v4(const void * __restrict__ experts_v, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + const float * __restrict__ gate, + const float * __restrict__ shexp, + 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; + auto ld = [&](const uint64_t row) -> float4 { + if constexpr (EIN16) { + const ushort4 h = *(const ushort4 *)((const uint16_t *) experts_v + row * n_embd + col4); + return make_float4(moe_bf2f(h.x), moe_bf2f(h.y), moe_bf2f(h.z), moe_bf2f(h.w)); + } else { + return *(const float4 *)((const float *) experts_v + row * n_embd + col4); + } + }; + 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 = ld(first_row); + 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 = ld(row); + sum.x += (e.x * scale) * w; sum.y += (e.y * scale) * w; sum.z += (e.z * scale) * w; sum.w += (e.w * scale) * w; + } + const float g = 1.0f / (1.0f + expf(-gate[token])); + const float4 a = *(const float4 *)(shexp + token * n_embd + col4); + const float4 o = make_float4(moe_add_rn(sum.x, moe_mul_rn(a.x, g)), moe_add_rn(sum.y, moe_mul_rn(a.y, g)), + moe_add_rn(sum.z, moe_mul_rn(a.z, g)), moe_add_rn(sum.w, moe_mul_rn(a.w, g))); + *(float4 *)(dst + token * n_embd + col4) = o; +} + +bool ggml_cuda_moe_weighted_reduction_sgma_supported(const ggml_tensor * experts, const ggml_tensor * expert_scale, + const ggml_tensor * weights, const ggml_tensor * gate, const ggml_tensor * shexp, const ggml_tensor * dst) { + const int64_t n_embd = experts->ne[0]; + const int64_t n_tokens = experts->ne[2] * experts->ne[3]; + return n_embd % 4 == 0 && experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && shexp->type == GGML_TYPE_F32 && + gate->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && + (expert_scale == nullptr || (expert_scale->type == GGML_TYPE_F32 && ggml_is_contiguous(expert_scale))) && + ggml_is_contiguous(experts) && ggml_is_contiguous(weights) && ggml_is_contiguous(shexp) && ggml_is_contiguous(gate) && + ggml_is_contiguous(dst) && ggml_nelements(gate) == n_tokens && ggml_nelements(shexp) == n_embd * n_tokens && + ggml_nelements(dst) == n_embd * n_tokens; +} + +bool ggml_cuda_op_moe_weighted_reduction_sgma(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, + const ggml_tensor * expert_scale, const ggml_tensor * weights, const ggml_tensor * gate, const ggml_tensor * shexp, + ggml_tensor * dst) { + const int64_t n_embd = experts->ne[0]; + const int64_t n_expert_used = experts->ne[1]; + const int64_t n_tokens = experts->ne[2] * experts->ne[3]; + if (!ggml_cuda_moe_weighted_reduction_sgma_supported(experts, expert_scale, weights, gate, shexp, dst) || + ((uintptr_t) shexp->data % 16) != 0 || ((uintptr_t) dst->data % 16) != 0 || ((uintptr_t) experts->data % 16) != 0) { + return false; + } + // a BF16-only sum stays on the unfused path, which writes it; this kernel only stores F32 + if (ggml_cuda_mmb_blk16() && ggml_cuda_mmb_is_bf16_only(ctx, dst)) { + return false; + } + // same condition as the unfused reduction's BF16-input path + const bool ein16 = ggml_cuda_mmb_down16() && ggml_cuda_mmb_is_bf16_only(ctx, experts); + // one block per token when its n_embd/4 lanes fit a block (640 for n_embd 2560: no idle lanes, where 256 left 128 of + // 768 idle), else 256. MOE_SGMA_THREADS overrides, rounded to whole warps within [32, 1024]. + static const int env_threads = getenv("MOE_SGMA_THREADS") ? atoi(getenv("MOE_SGMA_THREADS")) : 0; + const int64_t lanes = n_embd / 4; + int threads = (lanes <= 1024 && lanes % 32 == 0) ? (int) lanes : 256; + if (env_threads > 0) { + threads = std::min(1024, std::max(32, (env_threads + 31) / 32 * 32)); + } + const dim3 blocks(n_tokens, (n_embd / 4 + threads - 1) / threads, 1); + const float * sc = expert_scale ? (const float *) expert_scale->data : nullptr; + cudaStream_t stream = ctx.stream(); +#define MOE_SGMA_LAUNCH(A) moe_weighted_reduction_sgma_v4<<>>(experts->data, sc, \ + (const float *) weights->data, (const float *) gate->data, (const float *) shexp->data, (float *) dst->data, n_embd, (int) n_expert_used) + if (ein16) { MOE_SGMA_LAUNCH(true); } else { MOE_SGMA_LAUNCH(false); } +#undef MOE_SGMA_LAUNCH + CUDA_CHECK(cudaGetLastError()); + return true; +} diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh index b72f947ab398..cff331ba80e3 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh @@ -4,4 +4,14 @@ 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); + +// types/shapes/contiguity the fused kernel below needs (no data pointers, so graph_optimize can ask before allocation) +bool ggml_cuda_moe_weighted_reduction_sgma_supported(const ggml_tensor * experts, const ggml_tensor * expert_scale, + const ggml_tensor * weights, const ggml_tensor * gate, const ggml_tensor * shexp, const ggml_tensor * dst); + +// reduction + SIGMOID(gate [1, T]) -> MUL(shexp, .) -> ADD(reduction, .) in one kernel; false if the shapes/types don't fit +bool ggml_cuda_op_moe_weighted_reduction_sgma(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, + const ggml_tensor * expert_scale, const ggml_tensor * weights, const ggml_tensor * gate, const ggml_tensor * shexp, + ggml_tensor * dst); 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 5543307be697..13efa82d9725 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -281,6 +281,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) { @@ -366,7 +404,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), 1.0f); + } 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, @@ -729,3 +774,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 95618df5fd4d..3df46a140af6 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -18,3 +18,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..1215357aaf02 --- /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] > GGML_QSA_MAX_KEYS || 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..2bf6289eba91 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-prefill.cu @@ -0,0 +1,576 @@ +#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 + +// padding sentinel in the union's block-id list. The list is uint32 so the id space is not what limits the selected +// depth (uint16 did: 65535 real blocks * ratio 4 = 262140 keys, which collided with the 0xFFFF padding sentinel). +#define QSA3_BLK_SENT 0xFFFFFFFFu +__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, const int np2) { + 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; + if (np2 == 0) { + 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; + } + } else { + // bitonic sort of the row padded to np2 with the sentinel (sorts last): O(n log^2 n) instead of the O(n^2) + // rank sort; equal entries are equal values, so the sorted row is the same as the rank sort's + for (int j = ns + tid; j < np2; j += 256) { ent[j] = QSA3_SENT; } + __syncthreads(); + for (int k = 2; k <= np2; k <<= 1) { + for (int jj = k >> 1; jj > 0; jj >>= 1) { + for (int i = tid; i < np2; i += 256) { + const int ixj = i ^ jj; + if (ixj > i) { + const int a = ent[i], b = ent[ixj]; + const bool asc = (i & k) == 0; + if ((a > b) == asc) { ent[i] = b; ent[ixj] = a; } + } + } + __syncthreads(); + } + } + for (int j = tid; j < ns; j += 256) { dst[j] = ent[j]; } + } + } + 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, + uint32_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; } } + } + uint32_t * ob = ublk + (size_t) g * cap; uint16_t * 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] = (uint32_t) b; om[o] = (uint16_t) mk; ++o; + } + } + if (lane == QSA3_MERGE_LANES - 1) { + int t = total; + while (t & 3) { ob[t] = QSA3_BLK_SENT; om[t] = 0; ++t; } + ucount[g] = t; + } +} + +struct qsa3_layout { + size_t q1, q2, o1, o2, m1; + size_t knb1, vnb1, knb2, vnb2; // raw f16 cache strides (DIRECT reads only) + int nk, n_q, gqa, cap; + float scale; +}; + +template __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 uint32_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 = nullptr; + const uint16_t * pvg = nullptr; + const char * kbase = nullptr; + const char * vbase = nullptr; + if constexpr (DIRECT) { + kbase = reinterpret_cast(pk) + (size_t) kvh * s.knb2; + vbase = reinterpret_cast(pv) + (size_t) kvh * s.vnb2; + } else { + pkg = pk + (size_t) kvh * nblk * 1024; + pvg = pv + (size_t) kvh * nblk * 1024; + } + const uint32_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 * b, uint32_t & m01, uint32_t & m23) { + const uint4 bb = *reinterpret_cast(gblk + 4*c); + const uint2 mm = *reinterpret_cast(gmsk + 4*c); + b[0] = bb.x; b[1] = bb.y; b[2] = bb.z; b[3] = bb.w; m01 = mm.x; m23 = mm.y; + }; + // a padded slot reads block 0 but carries no mask bits, so its contribution is discarded downstream + auto blk_of = [&](const uint32_t * b, const int j) -> int { + return b[j] == QSA3_BLK_SENT ? 0 : (int) b[j]; + }; + auto load_k = [&](const uint32_t * b, 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(b, j); + if constexpr (DIRECT) { + // the packed 16-half chunk at (key, w, t) is dims (2w+t)*16 .. +15 of that key's row + const int key = 4*kb + (r & 3); + const char * p = kbase + (size_t)key * s.knb1 + (size_t)((2*w) * 16) * 2; +#pragma unroll + for (int t = 0; t < 2; ++t) { + const char * q = p + (size_t)t * 32; + kf[t] = __builtin_bit_cast(v16s, (uint4[2]){*reinterpret_cast(q), *reinterpret_cast(q + 16)}); + } + } else { + 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 b[4] = { 0, 0, 0, 0 }; + uint32_t m01 = 0, m23 = 0; + v16s kf[2]; + if (nchunks > 0) { load_desc(0, b, m01, m23); load_k(b, 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(b, 0), kb1 = blk_of(b, 1), kb2 = blk_of(b, 2), kb3 = blk_of(b, 3); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int d = 32*w + 16*t + r; + if constexpr (DIRECT) { + // the packed uint2 holds the four keys of block kb at dim d; gather them from the rows + const char * col = vbase + (size_t)d * 2; + auto kv4 = [&](const int kb) -> uint2 { + const char * b = col + (size_t)(4*kb) * s.vnb1; + const uint32_t a0 = *reinterpret_cast(b); + const uint32_t a1 = *reinterpret_cast(b + s.vnb1); + const uint32_t a2 = *reinterpret_cast(b + 2*s.vnb1); + const uint32_t a3 = *reinterpret_cast(b + 3*s.vnb1); + return make_uint2(a0 | (a1 << 16), a2 | (a3 << 16)); + }; + vf[t] = __builtin_bit_cast(v16s, (uint2[4]){ + mask_values(kv4(kb0), m01 & 0xffffu), mask_values(kv4(kb1), m01 >> 16), + mask_values(kv4(kb2), m23 & 0xffffu), mask_values(kv4(kb3), m23 >> 16)}); + } else { + 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 nb[4] = { 0, 0, 0, 0 }; + uint32_t nm01 = 0, nm23 = 0; + v16s kfn[2]; + if (c + 1 < nchunks) { load_desc(c + 1, nb, nm01, nm23); load_k(nb, 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(b, 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]; + b[0] = nb[0]; b[1] = nb[1]; b[2] = nb[2]; b[3] = nb[3]; 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]<=GGML_QSA_MAX_KEYS && 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]; + // QSA_DIRECT: -1/auto picks direct f16 reads when the attention is small against the cache + // (the pack is O(n_kv) per ubatch), 0 always packs, 1 always reads direct. Both paths are bit-identical. + static const int qsa_direct_cfg = getenv("QSA_DIRECT") ? atoi(getenv("QSA_DIRECT")) : -1; + const bool qsa_direct = qsa_direct_cfg == 1 || (qsa_direct_cfg < 0 && (size_t) n_q * 32 < (size_t) nk); + ggml_cuda_pool_alloc packed_k(ctx.pool(), qsa_direct ? 0 : ggml_nelements(k)); + ggml_cuda_pool_alloc packed_v(ctx.pool(), qsa_direct ? 0 : ggml_nelements(v)); + if (!qsa_direct) { + 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); + { + // bitonic row sort over the next power of two (GGML_QSA_RANKSORT=1: the O(n^2) rank sort) + static const bool ranksort = getenv("GGML_QSA_RANKSORT") && atoi(getenv("GGML_QSA_RANKSORT")) != 0; + int np2 = 1; while (np2 < ns) np2 <<= 1; + if (ranksort || (size_t) np2 * sizeof(int) > 60 * 1024) np2 = 0; + const ggml_cuda_kernel_launch_params launch(dim3(n_q), dim3(256), (size_t) (np2 ? np2 : 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, np2); + 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, + k->nb[1], v->nb[1], k->nb[2], v->nb[2], nk, n_q, 12, cap, scale}; + const ggml_cuda_kernel_launch_params launch(dim3(ngroups, k->ne[2]), dim3(256), 0, ctx.stream()); + if (qsa_direct) { + ggml_cuda_kernel_launch(qsa3_attn_kernel, launch, (const float *) q->data, (const uint16_t *) k->data, (const uint16_t *) v->data, + m ? (const uint16_t *) m->data : nullptr, (const uint32_t *) ublk.get(), (const uint16_t *) umask.get(), + (const int *) ucount.get(), (float *) dst->data, layout); + } else { + 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 uint32_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/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index 7be409bd36ab..b88c9294d6b7 100755 --- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -40,7 +40,7 @@ "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_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S", - "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4" + "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4", "GGML_TYPE_Q4_0_ROCMFP4_FAST" ] SOURCE_MMQ = """// This file has been autogenerated by generate_cu_files.py, do not edit manually. diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu new file mode 100644 index 000000000000..df63d09ae0c6 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index c7a0c831788d..4843f15af71f 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,202 @@ 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) +// One workgroup per row (RDNA3.5 prefill, e.g. the qwen4exp QSA block selection: 512 of ~62k blocks per query). +// The radix path above reads every row ~5.5 times (four 8-bit histogram passes, a gather and a tie gather); here a +// 12-bit LDS histogram locates the threshold bucket, a second pass writes the keys above it and compacts the bucket +// into LDS, and the remaining 20 bits are resolved on the compacted candidates. A bucket larger than the LDS buffer +// falls back to further histogram passes over the row. The selected set is the same as top_k_radix_cuda with +// stable_ties: every key above the k-th largest, then the lowest-column keys equal to it. Like that path, the order +// of the keys above the threshold is unspecified; the equal keys follow them in column order. +#define TOPK_WG_THREADS 1024 +#define TOPK_WG_CAP 4096 + +// bin b of hist (NBINS bins, descending search) with sum(hist[> b]) < rank <= sum(hist[>= b]); returns b and the +// rank left inside it. Every thread gets the result. +template +static __device__ __forceinline__ void top_k_wg_find(const int * hist, int rank, int * red, int & bin, int & rank_in) { + constexpr int PER = NBINS / TOPK_WG_THREADS; + const int tid = threadIdx.x; + // thread t owns bins [NBINS - (t+1)*PER, NBINS - t*PER): thread 0 the highest + int own = 0; +#pragma unroll + for (int i = 0; i < PER; ++i) { own += hist[NBINS - 1 - (tid * PER + i)]; } + red[tid] = own; + __syncthreads(); + // inclusive scan over threads (Hillis-Steele) + for (int off = 1; off < TOPK_WG_THREADS; off <<= 1) { + const int v = tid >= off ? red[tid - off] : 0; + __syncthreads(); + red[tid] += v; + __syncthreads(); + } + const int incl = red[tid], excl = incl - own; + __shared__ int s_bin, s_rank; + if (excl < rank && rank <= incl) { + int r = rank - excl, b = NBINS - 1 - tid * PER; + for (int i = 0; i < PER; ++i, --b) { + const int h = hist[b]; + if (r <= h) break; + r -= h; + } + s_bin = b; s_rank = r; + } + __syncthreads(); + bin = s_bin; rank_in = s_rank; + __syncthreads(); +} + +static __global__ void __launch_bounds__(TOPK_WG_THREADS) top_k_wg_kernel( + const float * __restrict__ src, int * __restrict__ dst, const int ncols, const int k) { + const int row = blockIdx.x, tid = threadIdx.x; + const float * x = src + (size_t) row * ncols; + int * out = dst + (size_t) row * k; + + __shared__ int hist[4096]; + __shared__ int red[TOPK_WG_THREADS]; + __shared__ uint32_t ckey[TOPK_WG_CAP]; + __shared__ int ccol[TOPK_WG_CAP]; + __shared__ int n_out, n_cand; + + for (int i = tid; i < 4096; i += TOPK_WG_THREADS) hist[i] = 0; + if (tid == 0) { n_out = 0; n_cand = 0; } + __syncthreads(); + + // pass 1: top 12 bits (four independent loads in flight per thread) + { + int c = tid; + for (; c + 3 * TOPK_WG_THREADS < ncols; c += 4 * TOPK_WG_THREADS) { + const float v0 = x[c], v1 = x[c + TOPK_WG_THREADS], v2 = x[c + 2 * TOPK_WG_THREADS], v3 = x[c + 3 * TOPK_WG_THREADS]; + atomicAdd(&hist[top_k_float_to_ordered(v0) >> 20], 1); + atomicAdd(&hist[top_k_float_to_ordered(v1) >> 20], 1); + atomicAdd(&hist[top_k_float_to_ordered(v2) >> 20], 1); + atomicAdd(&hist[top_k_float_to_ordered(v3) >> 20], 1); + } + for (; c < ncols; c += TOPK_WG_THREADS) { + atomicAdd(&hist[top_k_float_to_ordered(x[c]) >> 20], 1); + } + } + __syncthreads(); + int b12, rank; + top_k_wg_find<4096>(hist, k, red, b12, rank); + const int bucket = hist[b12]; + uint32_t prefix = (uint32_t) b12 << 20, pmask = 0xFFF00000u; + __syncthreads(); + + // pass 2: keys above the bucket go out, the bucket is compacted when it fits + const bool fits = bucket <= TOPK_WG_CAP; + auto place = [&](const int c, const float v) { + const uint32_t key = top_k_float_to_ordered(v); + const uint32_t top = key >> 20; + if (top > (uint32_t) b12) { + out[atomicAdd(&n_out, 1)] = c; + } else if (fits && top == (uint32_t) b12) { + const int p = atomicAdd(&n_cand, 1); + ckey[p] = key; ccol[p] = c; + } + }; + { + int c = tid; + for (; c + 3 * TOPK_WG_THREADS < ncols; c += 4 * TOPK_WG_THREADS) { + const float v0 = x[c], v1 = x[c + TOPK_WG_THREADS], v2 = x[c + 2 * TOPK_WG_THREADS], v3 = x[c + 3 * TOPK_WG_THREADS]; + place(c, v0); place(c + TOPK_WG_THREADS, v1); place(c + 2 * TOPK_WG_THREADS, v2); place(c + 3 * TOPK_WG_THREADS, v3); + } + for (; c < ncols; c += TOPK_WG_THREADS) place(c, x[c]); + } + __syncthreads(); + + // remaining 20 bits: 10 + 10, on the candidates or on the row + const int ncand = fits ? n_cand : 0; + for (int shift = 10; shift >= 0; shift -= 10) { + for (int i = tid; i < 1024; i += TOPK_WG_THREADS) hist[i] = 0; + __syncthreads(); + if (fits) { + for (int i = tid; i < ncand; i += TOPK_WG_THREADS) { + const uint32_t key = ckey[i]; + if ((key & pmask) == prefix) atomicAdd(&hist[(key >> shift) & 1023], 1); + } + } else { + for (int c = tid; c < ncols; c += TOPK_WG_THREADS) { + const uint32_t key = top_k_float_to_ordered(x[c]); + if ((key & pmask) == prefix) atomicAdd(&hist[(key >> shift) & 1023], 1); + } + } + __syncthreads(); + int b10, r2; + top_k_wg_find<1024>(hist, rank, red, b10, r2); + rank = r2; + prefix |= (uint32_t) b10 << shift; + pmask |= 1023u << shift; + __syncthreads(); + } + const uint32_t thr = prefix; // the k-th largest key; `rank` keys equal to it are taken + + // keys above the threshold inside the bucket + if (fits) { + for (int i = tid; i < ncand; i += TOPK_WG_THREADS) { + if (ckey[i] > thr) out[atomicAdd(&n_out, 1)] = ccol[i]; + } + } else { + for (int c = tid; c < ncols; c += TOPK_WG_THREADS) { + const uint32_t key = top_k_float_to_ordered(x[c]); + if ((key >> 20) == (uint32_t) b12 && key > thr) out[atomicAdd(&n_out, 1)] = c; + } + } + __syncthreads(); + const int base = n_out; // == k - rank + + // equal keys, lowest columns first + if (fits) { + // gather the equal columns, then pick the `rank` smallest by counting (each is unique) + __shared__ int n_eq; + if (tid == 0) n_eq = 0; + __syncthreads(); + // the histogram is free now and holds up to TOPK_WG_CAP (>= ncand) columns + for (int i = tid; i < ncand; i += TOPK_WG_THREADS) { + if (ckey[i] == thr) { const int p = atomicAdd(&n_eq, 1); hist[p] = ccol[i]; } + } + __syncthreads(); + const int ne = n_eq; + for (int i = tid; i < ne; i += TOPK_WG_THREADS) { + const int col = hist[i]; + int less = 0; + for (int j = 0; j < ne; ++j) less += hist[j] < col; + if (less < rank) out[base + less] = col; + } + } else { + const int lane = tid % warpSize, warp = tid / warpSize; + constexpr int NW = TOPK_WG_THREADS / 32; + int taken = 0; + for (int c0 = 0; c0 < ncols && taken < rank; c0 += TOPK_WG_THREADS) { + const int c = c0 + tid; + const bool eq = c < ncols && top_k_float_to_ordered(x[c]) == thr; + const unsigned long long m = __ballot(eq); + if (lane == 0) red[warp] = __popcll(m); + __syncthreads(); + int before = taken; + for (int w = 0; w < NW; ++w) { if (w < warp) before += red[w]; } + int total = 0; + for (int w = 0; w < NW; ++w) total += red[w]; + const int pos = before + __popcll(m & ((1ULL << lane) - 1)); + if (eq && pos < rank) out[base + pos] = c; + taken += total; + __syncthreads(); + } + } +} + +#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 +465,23 @@ 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) { + // GGML_CUDA_TOPK_WG=0 keeps the multi-pass radix kernels (same selected set) + static const bool wg = getenv("GGML_CUDA_TOPK_WG") == nullptr || atoi(getenv("GGML_CUDA_TOPK_WG")) != 0; + if (wg && k <= ncols && nrows <= INT_MAX) { + top_k_wg_kernel<<<(unsigned) nrows, TOPK_WG_THREADS, 0, stream>>>(src0_d, dst_d, (int) ncols, (int) k); + CUDA_CHECK(cudaGetLastError()); + return; + } + 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 +517,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 ee903bf21c93..d0687e177bb4 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, @@ -177,10 +177,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; } @@ -243,7 +244,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) { @@ -259,17 +260,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; @@ -339,8 +340,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..3cdbc5709419 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,87 @@ 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"); + } +} + +/* prefill shared-expert merge: out = other + src * sigmoid(gate[t]) with one gate value per row (token); the product + and the sum are rounded separately (no fma contraction), as the unfused MUL and ADD nodes -> bit-identical */ +#if defined(GGML_USE_HIP) +static __device__ __forceinline__ float sgma_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 sgma_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 sgma_mul_rn(const float a, const float b) { return __fmul_rn(a, b); } +static __device__ __forceinline__ float sgma_add_rn(const float a, const float b) { return __fadd_rn(a, b); } +#endif +static __global__ void sigmoid_gate_mul_add_f32(const float * __restrict__ src, const float * __restrict__ other, + const float * __restrict__ gate, float * __restrict__ out, const int64_t ncols, const int64_t n) { + const int64_t i = ((int64_t) blockIdx.x * blockDim.x + threadIdx.x) * 4; + if (i >= n) { + return; + } + if (i + 4 <= n && (ncols % 4) == 0) { + const float g = op_sigmoid(gate[i / ncols]); + const float4 a = *(const float4 *) (src + i), b = *(const float4 *) (other + i); + *(float4 *) (out + i) = make_float4(sgma_add_rn(b.x, sgma_mul_rn(a.x, g)), sgma_add_rn(b.y, sgma_mul_rn(a.y, g)), + sgma_add_rn(b.z, sgma_mul_rn(a.z, g)), sgma_add_rn(b.w, sgma_mul_rn(a.w, g))); + } else { + for (int64_t j = i; j < n && j < i + 4; ++j) { + out[j] = sgma_add_rn(other[j], sgma_mul_rn(src[j], op_sigmoid(gate[j / ncols]))); + } + } +} + +void ggml_cuda_op_sigmoid_gate_mul_add(ggml_backend_cuda_context & ctx, const ggml_tensor * gate, const ggml_tensor * src, + const ggml_tensor * other, ggml_tensor * out) { + const int64_t n = ggml_nelements(out), ncols = out->ne[0]; + sigmoid_gate_mul_add_f32<<<(unsigned) ((n / 4 + 255) / 256 + 1), 256, 0, ctx.stream()>>>((const float *) src->data, + (const float *) other->data, (const float *) gate->data, (float *) out->data, ncols, n); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 04f3af6443a7..07192f7975aa 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,10 @@ __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); + +// out = other + src * sigmoid(gate[row]) (gate: one value per row) +void ggml_cuda_op_sigmoid_gate_mul_add(ggml_backend_cuda_context & ctx, const ggml_tensor * gate, const ggml_tensor * src, + const ggml_tensor * other, ggml_tensor * out); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index f2a6f2009c9a..83606ab529e7 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -306,6 +306,7 @@ template static __device__ __forceinline__ float vec_dot_q8_0_16_q8_1_ #define VDR_MXFP4_Q8_1_MMVQ 2 #define VDR_MXFP4_Q8_1_MMQ 4 +#define VDR_ROCMFP4_FAST_Q8_1_MMQ 4 static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { @@ -328,6 +329,56 @@ static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( return d * sumi; } +#define VDR_ROCMFP4_Q8_1_MMVQ 2 + +// Dual UE4M3 half-block scales: low nibbles of qs[j] are weights j (scale e[0]), +// high nibbles are weights j + 16 (scale e[1]). Two int-dot accumulators, matching +// the q8_1 slot pairing used by vec_dot_mxfp4_q8_1 above. +static __device__ __forceinline__ float vec_dot_q4_0_rocmfp4_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4 * bq4 = (const block_rocmfp4 *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b1(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_rocmfp4); + + sumi0 = ggml_cuda_dp4a(v.x, q8[l + 0], sumi0); + sumi1 = ggml_cuda_dp4a(v.y, q8[l + 4], sumi1); + } + + const float dq8 = __low2float(bq8_1->ds); + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[0]) * dq8; + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[1]) * dq8; + return d0 * sumi0 + d1 * sumi1; +} + +static __device__ __forceinline__ float vec_dot_q4_0_rocmfp4_fast_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4_fast * bq4 = (const block_rocmfp4_fast *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b1(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_rocmfp4); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + const float d = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e) * __low2float(bq8_1->ds); + return d * sumi; +} + #define VDR_NVFP4_Q8_1_MMVQ 4 #define VDR_NVFP4_Q8_1_MMQ 8 @@ -1204,6 +1255,14 @@ static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( return d * sumi; } +#if defined(GGML_USE_HIP) +// strixllama: 0xFF in byte i when sign bit i is set, for the 4-weight halves of an IQ3_S sign byte +static __device__ const uint32_t strixllama_iq3s_sign_mask[16] = { + 0x00000000u, 0x000000FFu, 0x0000FF00u, 0x0000FFFFu, 0x00FF0000u, 0x00FF00FFu, 0x00FFFF00u, 0x00FFFFFFu, + 0xFF000000u, 0xFF0000FFu, 0xFF00FF00u, 0xFF00FFFFu, 0xFFFF0000u, 0xFFFF00FFu, 0xFFFFFF00u, 0xFFFFFFFFu, +}; +#endif + #define VDR_IQ3_S_Q8_1_MMVQ 2 #define VDR_IQ3_S_Q8_1_MMQ 2 @@ -1228,17 +1287,27 @@ static __device__ __forceinline__ float vec_dot_iq3_s_q8_1( iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); +#if defined(GGML_USE_HIP) + // strixllama: __vcmpne4 / __vsub4 are byte loops on HIP. The grid magnitudes are positive, so split + // each word by the sign mask and take two native dp4a instead: identical integer result. + const uint32_t m0 = strixllama_iq3s_sign_mask[signs_packed_8[l0/2] & 0x0F]; + const uint32_t m1 = strixllama_iq3s_sign_mask[signs_packed_8[l0/2] >> 4]; + sumi = ggml_cuda_dp4a((int) ((uint32_t) grid_pos.x & ~m0), u0, sumi); + sumi -= ggml_cuda_dp4a((int) ((uint32_t) grid_pos.x & m0), u0, 0); + sumi = ggml_cuda_dp4a((int) ((uint32_t) grid_pos.y & ~m1), u1, sumi); + sumi -= ggml_cuda_dp4a((int) ((uint32_t) grid_pos.y & m1), u1, 0); +#else 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); +#endif } sumi *= 1 + 2*((bq3->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F); @@ -1247,6 +1316,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 da7e3da0b9ca..dabda9cd2070 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-hip/CMakeLists.txt b/ggml/src/ggml-hip/CMakeLists.txt index a6a6b7271c54..938dbc6ad12f 100644 --- a/ggml/src/ggml-hip/CMakeLists.txt +++ b/ggml/src/ggml-hip/CMakeLists.txt @@ -39,6 +39,12 @@ else() if(GPU_TARGETS AND NOT CMAKE_HIP_ARCHITECTURES) set(CMAKE_HIP_ARCHITECTURES ${GPU_TARGETS}) endif() + + # This fork only supports Strix Halo (gfx1151), also covering Gorgon Halo. + # Without this, cmake/clang defaults to the full ROCm arch list. + if(NOT DEFINED CMAKE_HIP_ARCHITECTURES) + set(CMAKE_HIP_ARCHITECTURES "gfx1151") + endif() cmake_minimum_required(VERSION 3.21) enable_language(HIP) endif() 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 55db802c0aeb..b4897f778530 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 // @@ -5512,6 +5538,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-buffers.cpp b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp index 4d4c84951371..39c71bf52b30 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp @@ -621,7 +621,13 @@ void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, size_t sp // 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); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-common.h b/ggml/src/ggml-vulkan/ggml-vulkan-common.h index 125e7608a3dc..db00765dfb0f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h @@ -31,7 +31,7 @@ bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t low // shaders void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc); -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); +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, ggml_type k_type, ggml_type v_type, bool use_dynamic_kv = false, bool nan_safe_v = false); uint32_t get_subgroup_size(const std::string &pipeline_name, const vk_device_architecture &arch); void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested = nullptr); bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type k_type, ggml_type v_type); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 037bcee825c3..e05343d7b218 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h @@ -692,6 +692,56 @@ struct vk_op_lightning_indexer_push_constants { static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128); +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_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); + +// shared by flash_attn_sel_prep.comp, flash_attn_sel.comp and flash_attn_sel_cm.comp +struct vk_op_flash_attn_sel_push_constants { + uint32_t n_kv, n_sel, list_stride; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk1, nbk2, nbk3; + uint32_t nbv1, nbv2, nbv3; + uint32_t nbi1, nbi3; + uint32_t nb1, nb2, nb3; + float scale; + uint32_t n_split, chunk, n_head, n_tok, n_stream; +}; +static_assert(sizeof(vk_op_flash_attn_sel_push_constants) <= 128); + struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -705,6 +755,12 @@ struct vk_op_gated_delta_net_push_constants { uint32_t K; }; +// gdn_chunk_prep.comp / gdn_chunk_scan.comp: the sequential kernel's constants plus the chunk count +struct vk_op_gdn_chunk_push_constants { + vk_op_gated_delta_net_push_constants gdn; + uint32_t nch; +}; + struct vk_op_ssm_scan_push_constants { uint32_t nb02, nb03, nb12, nb13; uint32_t nb21, nb22, nb31; @@ -865,7 +921,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; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index ee50f51a3836..2d75dc5e1c2c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -76,6 +76,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #include #include +#include #include @@ -430,6 +431,14 @@ enum FaCodePath { FA_COOPMAT2, }; +// vk_fa_pipeline_state::flags bits beyond the ones get_fa_pipeline_state derives from its arguments (1..32); +// the shader side is flash_attn_base.glsl / flash_attn_cm1_mr.comp +static constexpr uint32_t FA_FLAG_V_TRANSPOSED = 64; // V is read from the per-head V^T scratch +static constexpr uint32_t FA_FLAG_MULTI_ROW = 128; // host only: build the pipeline from flash_attn_cm1_mr +static constexpr uint32_t FA_FLAG_MR_LAZY = 256; // multi-row: skip identity O rescales +static constexpr uint32_t FA_FLAG_MR_NOEB = 512; // multi-row: no loop-end barrier +static constexpr uint32_t FA_FLAG_MR_VPOST = 1024; // multi-row: V^T loads before the softmax barrier + struct vk_fa_pipeline_state { uint32_t HSK, HSV; uint32_t Br, Bc; @@ -641,6 +650,8 @@ enum rms_norm_mode { RMS_NORM_MUL_ROPE, RMS_NORM_MUL_ROPE_VIEW_SET_ROWS, RMS_NORM_VIEW_SET_ROWS, + RMS_NORM_MUL_MUL, // short rows only: normalised * gamma * gate + RMS_NORM_SCALE, // short rows only: scale(rms_norm(x)) RMS_NORM_COUNT, }; @@ -675,6 +686,22 @@ static constexpr std::array lightning_indexer_k_types = { class vk_memory_logger; +// 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 inline 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; @@ -732,6 +759,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; @@ -810,6 +838,7 @@ struct vk_device_struct { vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV + vk_pipeline pipeline_dequant_f16_transpose_vt; // f16 V -> V^T per head for the coopmat1 FA (V_TRANSPOSED) vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT]; @@ -844,6 +873,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]; @@ -874,6 +904,10 @@ struct vk_device_struct { vk_pipeline pipeline_rms_norm_mul_add_mul_partials_f32; vk_pipeline pipeline_rms_norm_set_rows_f32_f32; vk_pipeline pipeline_rms_norm_set_rows_f32_f16; + vk_pipeline pipeline_rms_norm_small_f32; // one subgroup per row, ne00 <= 256 + vk_pipeline pipeline_rms_norm_mul_small_f32; + vk_pipeline pipeline_rms_norm_mul_mul_small_f32; // + second multiplier (gated norm) + vk_pipeline pipeline_rms_norm_scale_small_f32; // + SCALE (the GDN q/k l2 norm) vk_pipeline pipeline_rms_norm_partials_f32; vk_pipeline pipeline_rms_norm_mul_partials_f32; vk_pipeline pipeline_rms_norm_mul_rope_f32_f32; @@ -984,11 +1018,32 @@ 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]; + vk_pipeline pipeline_gdn_chunk_prep, pipeline_gdn_chunk_scan; // chunked prefill form (GGML_VK_GDN_CHUNK) + uint32_t gdn_chunk_scan_nsplit = 1; // scan workgroups per value head + // 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_sel_f16[3]; // maskless selected-key FA, GQA 4, 8, 12 (scalar, portable) + vk_pipeline pipeline_flash_attn_sel_cm_f16[3]; // the same on RDNA3 + RADV (WMMA, gfx11 fragment layout) + vk_pipeline pipeline_flash_attn_sel_prep; // cleans the selection rows for both + 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_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; vk_pipeline pipeline_ssm_conv_silu_f32; vk_pipeline pipeline_ssm_conv_bias_silu_f32; + // CONCAT(state, x^T) + SSM_CONV + SILU without the transposed concat: the conv reads x and the state + // directly, and the concat node only writes the tail columns its other readers (state copies) use + vk_pipeline pipeline_ssm_conv_direct_silu_f32; + vk_pipeline pipeline_ssm_conv_tail_f32; vk_pipeline pipeline_opt_step_adamw_f32; vk_pipeline pipeline_opt_step_sgd_f32; std::map pipeline_conv2d_f32[CONV_SHAPE_COUNT]; @@ -1233,6 +1288,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); + } + private: std::map> timings; std::map> flops; @@ -1248,6 +1310,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 {}; @@ -1304,6 +1378,21 @@ struct ggml_backend_vk_context { bool fused_topk_qsa {}; bool fused_hc_post_gate {}; rms_norm_mode fused_rms_norm_mode {RMS_NORM_COUNT}; + // direct GDN conv: concats matched in this graph (decided at the concat, consumed at its SSM_CONV) + struct ssm_conv_direct_match { + const ggml_tensor * concat; + const ggml_tensor * state; // [ncs, C, S], contiguous + const ggml_tensor * x; // x^T's storage: [C, T, S], contiguous + const ggml_tensor * kern; // [ncs + 1, C] + ggml_tensor * out; // the SILU node + int conv_idx; + int64_t tail_from; // first concat column read by anything but the conv + }; + std::vector ssm_conv_direct_matches; + int fused_concat_tail {-1}; // current CONCAT writes only its tail (index into ssm_conv_direct_matches) + int fused_ssm_conv_direct {-1}; // current SSM_CONV (+SILU) runs direct (index into ssm_conv_direct_matches) + // set while the current node is a direct conv: srcs read outside node->src (x, state) for the sync tracking + const ggml_tensor * fused_extra_srcs[2] {}; // for GGML_VK_PERF_LOGGER std::unique_ptr perf_logger; @@ -1312,6 +1401,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 {}; }; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 38ec1f9e82ab..1b9682d0454b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -201,6 +201,20 @@ void ggml_vk_print_device_fault_info(const vk_device& device) { (unsigned long long)info.vendorFaultData); } } +// 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; +} + 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]; @@ -1289,8 +1303,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); @@ -1372,7 +1424,8 @@ vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, } 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, bool nan_safe_v) { 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); @@ -1380,7 +1433,8 @@ vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_ (use_mask ? 2 : 0) | (use_logit_softcap ? 4 : 0) | (old_amd_windows ? 8 : 0) | - (use_sparse ? 16 : 0); + (use_dynamic_kv ? 16 : 0) | + (nan_safe_v ? 32 : 0); const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size; @@ -1415,6 +1469,46 @@ static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& s }; } +// SHMEM_STRIDE_PAD for mul_mm.comp (constantID=12, 13 for MULMAT_QUANT), 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; @@ -1449,15 +1543,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]; @@ -1504,6 +1609,8 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std: case GGML_TYPE_Q8_0: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + dm case GGML_TYPE_IQ4_XS: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + d 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_IQ4_NL: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + d case GGML_TYPE_NVFP4: block_a_size = std430_size({{32, 4}, {fp2_size, fp2_align}}); break; // qs[8] + d_scales(vec2) 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) @@ -1660,6 +1767,16 @@ static bool ggml_vk_fa_type_needs_shmem(ggml_type type) { } } +// int8 MMQ (coopmat1) f32 output: wave32 RDNA3 lanes swap accumulator rows with one shuffle pair so each lane +// writes four adjacent rows as one vec4 store. GGML_VK_MMQ_PACKSTORE=0 keeps the per-element store. +static bool ggml_vk_mmq_packstore_enabled() { + static const bool on = [] { + const char * e = getenv("GGML_VK_MMQ_PACKSTORE"); + return e == nullptr || atoi(e) != 0; + }(); + return on; +} + static bool ggml_vk_fa_scalar_uses_mmq(const vk_device& device, ggml_type k_type, ggml_type v_type) { #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) return device->integer_dot_product && device->subgroup_clustered && @@ -1807,13 +1924,16 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { return cm1_sg * (bm / std::min(cm1_sg, bm)) * (bn / 32); }; - l_warptile_mmq_cm1_int = { cm1_bs(128, 128), 128, 128, 32, std::min(cm1_sg, 128u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; - m_warptile_mmq_cm1_int = { cm1_bs( 64, 64), 64, 64, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; - s_warptile_mmq_cm1_int = { cm1_bs( 32, 32), 32, 32, 32, std::min(cm1_sg, 32u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; + // spec constant 12 of mul_mmq_cm1.comp + const uint32_t packstore = ggml_vk_mmq_packstore_enabled() ? 1u : 0u; + + l_warptile_mmq_cm1_int = { cm1_bs(128, 128), 128, 128, 32, std::min(cm1_sg, 128u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; + m_warptile_mmq_cm1_int = { cm1_bs( 64, 64), 64, 64, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; + s_warptile_mmq_cm1_int = { cm1_bs( 32, 32), 32, 32, 32, std::min(cm1_sg, 32u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; - l_warptile_mmq_cm1_int_k = { cm1_bs( 64, 128), 64, 128, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; - m_warptile_mmq_cm1_int_k = { cm1_bs( 64, 64), 64, 64, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; - s_warptile_mmq_cm1_int_k = { cm1_bs( 32, 32), 32, 32, 32, std::min(cm1_sg, 32u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture }; + l_warptile_mmq_cm1_int_k = { cm1_bs( 64, 128), 64, 128, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; + m_warptile_mmq_cm1_int_k = { cm1_bs( 64, 64), 64, 64, 32, std::min(cm1_sg, 64u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; + s_warptile_mmq_cm1_int_k = { cm1_bs( 32, 32), 32, 32, 32, std::min(cm1_sg, 32u), 32, 2, itm, itn, itk, cm1_sg, (uint32_t)device->architecture, packstore }; l_mmq_cm1_wg_denoms_k = { l_warptile_mmq_cm1_int_k[1], l_warptile_mmq_cm1_int_k[2], 1 }; m_mmq_cm1_wg_denoms_k = { m_warptile_mmq_cm1_int_k[1], m_warptile_mmq_cm1_int_k[2], 1 }; @@ -2109,6 +2229,13 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; } else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; } name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1"; + if (fa.first.flags & FA_FLAG_MULTI_ROW) { + // multi-row prefill flash attention (GGML_VK_FA_MR): f32 accumulation, aligned only + GGML_ASSERT(f32acc && aligned); + spv_data = flash_attn_f32_f16_mr_cm1_data; + spv_size = flash_attn_f32_f16_mr_cm1_len; + name = "flash_attn_f32_f16_mr_cm1"; + } } ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, @@ -2156,10 +2283,12 @@ 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 + 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; }; @@ -2167,10 +2296,12 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { 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=13: SHMEM_STRIDE_PAD + spec.push_back(intel_slm ? 1u : 0u); // constantID=14: APPLY_SLM_A_RESHAPE } return spec; }; @@ -2189,7 +2320,12 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { 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) + X(GGML_TYPE_IQ4_NL, iq4_nl) \ + X(GGML_TYPE_Q4_0_ROCMFP4, rocmfp4) \ + X(GGML_TYPE_Q4_0_ROCMFP4_FAST, rocmfp4_fast) \ + X(GGML_TYPE_Q3_0_ROCMFPX, rocmfpx_fp3) \ + X(GGML_TYPE_Q6_0_ROCMFPX, rocmfpx_fp6) \ + X(GGML_TYPE_Q8_0_ROCMFPX, rocmfpx_fp8) #define FOR_EACH_LUT_FP4_TYPE(X) \ X(GGML_TYPE_MXFP4, mxfp4) \ X(GGML_TYPE_NVFP4, nvfp4) @@ -2198,6 +2334,8 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { FOR_EACH_LUT_FP4_TYPE(X) const int mul_mat_id_param_count = 5; + // int8 MMQ coopmat1 MUL_MAT_ID: a sixth binding carries the router weights of the fused MUL_MAT_ID+MUL epilogue + const int mul_mat_id_mmq_cm1_param_count = 6; using spec_fn_t = std::function(const std::vector&, bool)>; auto const &create_mm_pipelines = [&]( @@ -2553,19 +2691,19 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #undef X_CM1_ID if (device->coopmat_int_support && (rdna3 || rdna4)) { - cm1_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_0_q8_1", matmul_id_subgroup_q4_0_q8_1_cm1_len, matmul_id_subgroup_q4_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_1_q8_1", matmul_id_subgroup_q4_1_q8_1_cm1_len, matmul_id_subgroup_q4_1_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_0_q8_1", matmul_id_subgroup_q5_0_q8_1_cm1_len, matmul_id_subgroup_q5_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_1_q8_1", matmul_id_subgroup_q5_1_q8_1_cm1_len, matmul_id_subgroup_q5_1_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q8_0_q8_1", matmul_id_subgroup_q8_0_q8_1_cm1_len, matmul_id_subgroup_q8_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_IQ4_NL, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_iq4_nl_q8_1", matmul_id_subgroup_iq4_nl_q8_1_cm1_len, matmul_id_subgroup_iq4_nl_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_IQ4_XS, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_iq4_xs_q8_1", matmul_id_subgroup_iq4_xs_q8_1_cm1_len, matmul_id_subgroup_iq4_xs_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_MXFP4, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_mxfp4_q8_1", matmul_id_subgroup_mxfp4_q8_1_cm1_len, matmul_id_subgroup_mxfp4_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_q3_k_q8_1", matmul_id_subgroup_q3_k_q8_1_cm1_len, matmul_id_subgroup_q3_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_k_q8_1", matmul_id_subgroup_q4_k_q8_1_cm1_len, matmul_id_subgroup_q4_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_k_q8_1", matmul_id_subgroup_q5_k_q8_1_cm1_len, matmul_id_subgroup_q5_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_q6_k_q8_1", matmul_id_subgroup_q6_k_q8_1_cm1_len, matmul_id_subgroup_q6_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - if (!rdna4) { cm1_create_mmq({GGML_TYPE_NVFP4, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_nvfp4_q8_1", matmul_id_subgroup_nvfp4_q8_1_cm1_len, matmul_id_subgroup_nvfp4_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); } + cm1_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_0_q8_1", matmul_id_subgroup_q4_0_q8_1_cm1_len, matmul_id_subgroup_q4_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_1_q8_1", matmul_id_subgroup_q4_1_q8_1_cm1_len, matmul_id_subgroup_q4_1_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_0_q8_1", matmul_id_subgroup_q5_0_q8_1_cm1_len, matmul_id_subgroup_q5_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_1_q8_1", matmul_id_subgroup_q5_1_q8_1_cm1_len, matmul_id_subgroup_q5_1_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q8_0_q8_1", matmul_id_subgroup_q8_0_q8_1_cm1_len, matmul_id_subgroup_q8_0_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_IQ4_NL, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_iq4_nl_q8_1", matmul_id_subgroup_iq4_nl_q8_1_cm1_len, matmul_id_subgroup_iq4_nl_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_IQ4_XS, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_iq4_xs_q8_1", matmul_id_subgroup_iq4_xs_q8_1_cm1_len, matmul_id_subgroup_iq4_xs_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_MXFP4, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_mxfp4_q8_1", matmul_id_subgroup_mxfp4_q8_1_cm1_len, matmul_id_subgroup_mxfp4_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_q3_k_q8_1", matmul_id_subgroup_q3_k_q8_1_cm1_len, matmul_id_subgroup_q3_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q4_k_q8_1", matmul_id_subgroup_q4_k_q8_1_cm1_len, matmul_id_subgroup_q4_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int, "matmul_id_subgroup_q5_k_q8_1", matmul_id_subgroup_q5_k_q8_1_cm1_len, matmul_id_subgroup_q5_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + cm1_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_q6_k_q8_1", matmul_id_subgroup_q6_k_q8_1_cm1_len, matmul_id_subgroup_q6_k_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); + if (!rdna4) { cm1_create_mmq({GGML_TYPE_NVFP4, GGML_TYPE_Q8_1, true, false}, tc_mmq_cm1_int_k, "matmul_id_subgroup_nvfp4_q8_1", matmul_id_subgroup_nvfp4_q8_1_cm1_len, matmul_id_subgroup_nvfp4_q8_1_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_mmq_cm1_param_count); } } } else #endif // defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) @@ -2645,6 +2783,8 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { 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); sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_iq3_s_q8_1", matmul_iq3_s_q8_1_len, matmul_iq3_s_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); + sg_create_mmq({GGML_TYPE_Q4_0_ROCMFP4, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_rocmfp4_q8_1", matmul_rocmfp4_q8_1_len, matmul_rocmfp4_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); + sg_create_mmq({GGML_TYPE_Q4_0_ROCMFP4_FAST, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_rocmfp4_fast_q8_1", matmul_rocmfp4_fast_q8_1_len, matmul_rocmfp4_fast_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); } #endif @@ -2758,6 +2898,8 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { 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); sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_iq3_s_q8_1", matmul_iq3_s_q8_1_fp32_len, matmul_iq3_s_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); + sg_create_mmq({GGML_TYPE_Q4_0_ROCMFP4, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_rocmfp4_q8_1", matmul_rocmfp4_q8_1_fp32_len, matmul_rocmfp4_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); + sg_create_mmq({GGML_TYPE_Q4_0_ROCMFP4_FAST, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_rocmfp4_fast_q8_1", matmul_rocmfp4_fast_q8_1_fp32_len, matmul_rocmfp4_fast_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); } #endif @@ -2950,6 +3092,11 @@ 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); @@ -2979,6 +3126,11 @@ 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); @@ -2995,6 +3147,9 @@ 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); @@ -3036,6 +3191,11 @@ 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); @@ -3052,6 +3212,9 @@ 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); @@ -3086,15 +3249,22 @@ 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); + // one 64 kv x 64 hs tile per workgroup, dispatched as (HS / 64, KV / 64, heads) + ggml_vk_create_pipeline(device, device->pipeline_dequant_f16_transpose_vt, "dequant_f16_transpose_vt", dequant_f16_transpose_vt_len, dequant_f16_transpose_vt_data, "main", 2, 5 * sizeof(uint32_t), {64, 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); @@ -3104,6 +3274,12 @@ 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); @@ -3134,6 +3310,11 @@ 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); @@ -3164,6 +3345,11 @@ 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); @@ -3218,22 +3404,6 @@ 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 { @@ -3260,6 +3430,14 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_partials_f32, "rms_norm_mul_add_mul_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f32, "rms_norm_set_rows_f32_f32", rms_norm_set_rows_f32_f32_len, rms_norm_set_rows_f32_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f16, "rms_norm_set_rows_f32_f16", rms_norm_set_rows_f32_f16_len, rms_norm_set_rows_f32_f16_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); + if (device->subgroup_arithmetic && device->subgroup_size >= 32 && 256 % device->subgroup_size == 0) { + // rms_norm_small.comp: one subgroup per row; pin the subgroup size the host counts rows with + const uint32_t sg = device->subgroup_size_control ? device->subgroup_size : 0; + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_small_f32, "rms_norm_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true, false, sg); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_small_f32, "rms_norm_mul_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true, false, sg); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_mul_small_f32, "rms_norm_mul_mul_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true, false, sg); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_scale_small_f32, "rms_norm_scale_small_f32", rms_norm_small_f32_len, rms_norm_small_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0, 0, 1}, 1, true, false, sg); + } ggml_vk_create_pipeline(device, device->pipeline_rms_norm_partials_f32, "rms_norm_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_partials_f32, "rms_norm_mul_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); @@ -3302,6 +3480,11 @@ 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); \ @@ -3314,7 +3497,12 @@ 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) @@ -3331,6 +3519,11 @@ 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; @@ -3373,6 +3566,8 @@ 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); @@ -3722,6 +3917,131 @@ 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); + } + // chunked gated delta net (default on, GGML_VK_GDN_CHUNK=0 off): hand-built wave32 WMMA fragments in the RDNA3 layout + // (RDNA4 lays its fragments out differently) as RADV/ACO lowers them, so RDNA3 + RADV + wave32 only + if (device->vendor_id == VK_VENDOR_ID_AMD && device->architecture == vk_device_architecture::AMD_RDNA3 && + device->driver_id == vk::DriverId::eMesaRadv && + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size) { + ggml_vk_create_pipeline(device, device->pipeline_gdn_chunk_prep, "gdn_chunk_prep_f32", gdn_chunk_prep_f32_len, gdn_chunk_prep_f32_data, + "main", 5, sizeof(vk_op_gdn_chunk_push_constants), {1, 1, 1}, {}, 1, true, true, 32); + // one wave (16 state columns) per workgroup, operands straight from the scratch + device->gdn_chunk_scan_nsplit = 8; + ggml_vk_create_pipeline(device, device->pipeline_gdn_chunk_scan, "gdn_chunk_scan_f32", gdn_chunk_scan_f32_len, gdn_chunk_scan_f32_data, + "main", 4, sizeof(vk_op_gdn_chunk_push_constants), {1, 1, 1}, {32}, 1, true, true, 32); + } + 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); + // maskless selected-key attention (QSA block selection). The kernels hard-code 64-lane subgroups, so the + // device must run them at 64 (a required size it supports, or 64 by default without size control). + const bool sel_sg64 = device->subgroup_size_control + ? device->subgroup_min_size <= 64 && 64 <= device->subgroup_max_size + : device->subgroup_size == 64; + if (sel_sg64 && device->subgroup_shuffle && device->subgroup_arithmetic && device->fp16) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_sel_prep, + "flash_attn_sel_prep", flash_attn_sel_prep_len, flash_attn_sel_prep_data, "main", 2, + sizeof(vk_op_flash_attn_sel_push_constants), {1, 1, 1}, {}, 1, true); + static const char * sel_names[3] = { "flash_attn_sel_f16_gqa4", "flash_attn_sel_f16_gqa8", "flash_attn_sel_f16_gqa12" }; + for (uint32_t i = 0; i < 3; ++i) { + const uint32_t gqa = 4 * (i + 1); + // shared memory: Q (1 KiB per head), a K and a V block (16 KiB each), P and the block's keys: + // 46,720 bytes at GQA 12 (the selection row is read from the pre-pass output, not kept here) + const uint32_t sel_shmem = 4 * gqa * 256 + 2 * 2 * 256 * 32 + 4 * gqa * 32 + 4 * 32; + if (64 * gqa <= device->properties.limits.maxComputeWorkGroupInvocations && + sel_shmem <= device->properties.limits.maxComputeSharedMemorySize) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_sel_f16[i], + sel_names[i], flash_attn_sel_f16_len, flash_attn_sel_f16_data, "main", 5, + sizeof(vk_op_flash_attn_sel_push_constants), {1, 1, 1}, {64 * gqa, 64, gqa}, 1, true, true, 64); + } + } +#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + // WMMA kernel: builds its fragments with the gfx11 register layout of RADV's KHR cooperative matrix + // lowering, so RDNA3 + RADV only (GGML_VK_SEL_CM=0 forces the scalar kernel). 64 threads, 8.4 KB shared. + static const bool sel_cm_env = [] { const char * e = getenv("GGML_VK_SEL_CM"); return !(e && e[0] == '0'); }(); + if (sel_cm_env && device->subgroup_ballot && device->coopmat_support && device->coopmat_support_16x16x16_f32acc && + device->vendor_id == VK_VENDOR_ID_AMD && device->architecture == AMD_RDNA3 && + device->driver_id == vk::DriverId::eMesaRadv) { + static const char * sel_cm_names[3] = { "flash_attn_sel_cm_f16_gqa4", "flash_attn_sel_cm_f16_gqa8", "flash_attn_sel_cm_f16_gqa12" }; + for (uint32_t i = 0; i < 3; ++i) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_sel_cm_f16[i], + sel_cm_names[i], flash_attn_sel_cm_f16_len, flash_attn_sel_cm_f16_data, "main", 5, + sizeof(vk_op_flash_attn_sel_push_constants), {1, 1, 1}, {4 * (i + 1)}, 1, true, true, 64); + } + } +#endif + } + 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 + } + } + 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); @@ -3733,6 +4053,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_f32, "ssm_conv_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 4, sizeof(vk_op_ssm_conv_push_constants), {32, 16, 1}, {32, 16, 0, 0}, 1); ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_silu_f32, "ssm_conv_silu_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 4, sizeof(vk_op_ssm_conv_push_constants), {32, 16, 1}, {32, 16, 0, 1}, 1); ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_bias_silu_f32, "ssm_conv_bias_silu_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 4, sizeof(vk_op_ssm_conv_push_constants), {32, 16, 1}, {32, 16, 1, 1}, 1); + // direct conv: 64 channels x 4 token groups of 8 tokens per workgroup; tail: 256 (channel, column) pairs + ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_direct_silu_f32, "ssm_conv_direct_silu_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 4, sizeof(vk_op_ssm_conv_push_constants), {64, 32, 1}, {64, 4, 0, 1, 1, 8}, 1); + ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_tail_f32, "ssm_conv_tail_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 4, sizeof(vk_op_ssm_conv_push_constants), {256, 1, 1}, {256, 1, 0, 0, 2, 1}, 1); ggml_vk_create_pipeline(device, device->pipeline_opt_step_adamw_f32, "opt_step_adamw_f32", opt_step_adamw_f32_len, opt_step_adamw_f32_data, "main", 5, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1); @@ -4256,6 +4579,15 @@ 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; @@ -5424,6 +5756,11 @@ 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; @@ -5480,8 +5817,13 @@ vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type 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; @@ -5504,6 +5846,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: @@ -5545,8 +5890,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; @@ -5592,6 +5942,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: @@ -5633,8 +5986,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; @@ -5987,14 +6345,20 @@ 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 * scale = nullptr) { 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 << ")"); + // bit 0 of the last push constant: hoisted row ids; bit 1 (int8 MMQ cm1 only): scale rows by the router weights 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) | (scale ? 2u : 0u) }; + if (pipeline->parameter_count == 6) { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf, scale ? *scale : d }, pc, { m, nei1, n_as }); + } else { + GGML_ASSERT(scale == nullptr); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); + } } bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { @@ -6103,6 +6467,11 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t 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; @@ -6119,6 +6488,11 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t 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; @@ -6543,6 +6917,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) { @@ -6614,7 +6994,7 @@ 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, bool swap_inputs = false) { 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]; @@ -6632,7 +7012,7 @@ 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]; @@ -6653,7 +7033,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; @@ -6702,6 +7082,24 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& const uint64_t y_sz = quantize_y ? (ggml_vk_align_size(y_ne, 128) * ggml_type_size(GGML_TYPE_Q8_1) / ggml_blck_size(GGML_TYPE_Q8_1)) : (f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne); + // Swapped inputs with more than one output row: the mat-vec writes dst transposed + // ([ne1, ne0] row-major, the vector index is the slow dimension), so it goes to the + // split_k scratch buffer and a strided copy puts it in place. dst is ne0 x ne1 with + // ne0 <= mul_mat_vec_max_cols, so the copy is a few KB. + const bool transpose_d = swap_inputs && dst->ne[0] > 1; + ggml_tensor dst_t; + vk_pipeline cpy_d = nullptr; + if (transpose_d) { + GGML_ASSERT(dst->type == GGML_TYPE_F32 && ggml_is_contiguous(dst) && ctx->num_additional_fused_ops == 0); + dst_t = *dst; + dst_t.nb[0] = dst->ne[1] * sizeof(float); + dst_t.nb[1] = sizeof(float); + dst_t.nb[2] = ggml_nbytes(dst); + dst_t.nb[3] = ggml_nbytes(dst); + cpy_d = ggml_vk_get_cpy_pipeline(ctx, &dst_t, nullptr, GGML_TYPE_F32); + GGML_ASSERT(cpy_d != nullptr); + } + { if ( (qx_needs_dequant && x_sz > ctx->device->properties.limits.maxStorageBufferRange) || @@ -6716,6 +7114,10 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& ctx->prealloc_size_y = y_sz; ggml_vk_preallocate_buffers(ctx, subctx); } + if (transpose_d && ctx->prealloc_size_split_k < ggml_nbytes(dst)) { + ctx->prealloc_size_split_k = ggml_nbytes(dst); + ggml_vk_preallocate_buffers(ctx, subctx); + } // Request descriptor sets if (qx_needs_dequant) { @@ -6727,11 +7129,20 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& if (quantize_y) { ggml_pipeline_request_descriptor_sets(ctx, to_q8_1, 1); } + if (transpose_d) { + ggml_pipeline_request_descriptor_sets(ctx, cpy_d, 1); + } } 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) { @@ -6795,6 +7206,14 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); } + vk_subbuffer d_Out = d_D; + if (transpose_d) { + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + d_Out = { ctx->prealloc_split_k, 0, ggml_nbytes(dst) }; + } + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; uint32_t groups_x = ne01; @@ -6841,7 +7260,7 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_X, d_Y, - d_D, + d_Out, d_F0, d_F1, }, @@ -6849,6 +7268,12 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& base_work_group_y += groups_y; } + if (transpose_d) { + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_cpy_to_contiguous(ctx, subctx, cpy_d, &dst_t, d_Out, d_D); + ctx->prealloc_split_k_need_sync = true; + } + if (x_non_contig) { ctx->prealloc_x_need_sync = true; } @@ -6857,6 +7282,80 @@ 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, bool swap_inputs = false) { + ggml_tensor * dst = cgraph->nodes[node_idx]; + if (swap_inputs) { + ggml_vk_mul_mat_vec_q_f16_cols(ctx, subctx, cgraph, node_idx, 0, (uint32_t) dst->src[0]->ne[1], /*allow_quantize_y=*/ true, true); + return; + } + 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]; @@ -7260,19 +7759,26 @@ void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const st 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. + // 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. + // With one output row, B^T*A has the same flat output as A^T*B. With a few output rows + // (up to mul_mat_vec_max_cols, e.g. a [k, 4] f32 projection over a whole ubatch), the + // src0 rows become the mat-vec columns and the result is transposed into dst, instead + // of padding a tiny M to a full mul_mm tile. } 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 && + dst->ne[0] <= mul_mat_vec_max_cols && dst->ne[1] > mul_mat_vec_max_cols && + // more than one output row: transposed write, F32 dst only, and src1 (the swapped mat-vec matrix) + // limited to the f32/f16 that test-backend-ops can check against the CPU reference + (dst->ne[0] == 1 || (dst->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16))) && 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 { @@ -7280,7 +7786,8 @@ void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const st } } -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) { +// scale (optional): the router weights of a fused MUL_MAT_ID+MUL; dst is then the MUL's destination +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 * scale = 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]; @@ -7634,6 +8141,13 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); } + vk_subbuffer d_scale; + if (scale != nullptr) { + // ggml_vk_mmid_scale_epilogue_ok admitted this node only for the int8 MMQ coopmat1 pipelines + GGML_ASSERT(quantize_y && !qx_needs_dequant && pipeline->parameter_count == 6); + d_scale = ggml_vk_tensor_subbuffer(ctx, scale); + } + // compute ggml_vk_matmul_id( ctx, subctx, pipeline, @@ -7641,7 +8155,8 @@ 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, + scale != nullptr ? &d_scale : nullptr ); // NOLINT if (x_non_contig || qx_needs_dequant) { @@ -7905,7 +8420,13 @@ void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx, const VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); 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 if (ctx->num_additional_fused_ops == 1) { + // MUL_MAT_ID+MUL: the router weights scale each row in the GEMM epilogue, written to the MUL's destination + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + GGML_ASSERT(mul->op == GGML_OP_MUL && mul->src[0] == dst); + ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, mul, mul->src[1]); } else { + GGML_ASSERT(ctx->num_additional_fused_ops == 0); ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst); } } @@ -7984,8 +8505,8 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_ const uint32_t qstride = hsk_pad / 4 + 2; const uint32_t Qf = Br * qstride * f16vec4; - const uint32_t psh_stride = Br / 4 + 2; - const uint32_t Psh = Bc * psh_stride * 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; @@ -8001,7 +8522,10 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_ const uint32_t slope = Br * acctype; - const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope; + // dead_stamp[Bc] + dead_blk (hidden-key tracking; only NAN_SAFE_V pipelines use it, counted for all) + const uint32_t live = (Bc + 1) * sizeof(uint32_t); + + const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope + live; 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); @@ -8009,166 +8533,1253 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_ return supported; } -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) { - VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; - std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; - std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; - std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; - if (sinks) { - std::cerr << "), (" << sinks << ", name=" << sinks->name << ", type=" << sinks->type << ", ne0=" << sinks->ne[0] << ", ne1=" << sinks->ne[1] << ", ne2=" << sinks->ne[2] << ", ne3=" << sinks->ne[3] << ", nb0=" << sinks->nb[0] << ", nb1=" << sinks->nb[1] << ", nb2=" << sinks->nb[2] << ", nb3=" << sinks->nb[3]; +// 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; } - std::cerr << "))"); - - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - 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; +// 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; - const uint32_t HSK = nek0; - const uint32_t HSV = nev0; - uint32_t N = neq1; - const uint32_t KV = nek1; + if (device->uma) { + return true; + } - GGML_ASSERT(ne0 == HSV); - GGML_ASSERT(ne2 == N); + // 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; + } - // input tensor rows must be contiguous - GGML_ASSERT(nbq0 == ggml_type_size(q->type)); - GGML_ASSERT(nbk0 == ggml_type_size(k->type)); - GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + static const uint64_t reserve = [] { + const char * env = getenv("GGML_VK_FA_DEQUANT_RESERVE_MB"); + return (uint64_t)(env ? atoi(env) : 1024) * 1024 * 1024; + }(); - GGML_ASSERT(neq0 == HSK); + bool fits = false; - GGML_ASSERT(neq1 == N); + // 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; + } - GGML_ASSERT(nev1 == nek1); + 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)); + } - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + ctx->fa_dequant_gate_sz = scratch_sz; + ctx->fa_dequant_gate_fits = fits; + return fits; +} - assert(dst->type == GGML_TYPE_F32); - assert(q->type == GGML_TYPE_F32); - uint32_t gqa_ratio = 1; - uint32_t qk_ratio = neq2 / nek2; - uint32_t workgroups_x = (uint32_t)neq1; - uint32_t workgroups_y = (uint32_t)neq2; - uint32_t workgroups_z = (uint32_t)neq3; +// 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++); +} - const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; +// maskless selected-key attention (QSA block selection): the pipeline that takes this op, or nullptr +static vk_pipeline ggml_vk_flash_attn_sel_pipeline(const vk_device & device, const ggml_tensor * op) { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + const ggml_tensor * ids = op->src[5]; + if (ids == nullptr || op->src[3] != nullptr || op->src[4] != nullptr || ggml_get_op_params_i32(op, 4) != 0) { + return nullptr; + } + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f || op->type != GGML_TYPE_F32 || !ggml_is_contiguous(op) || + q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || ids->type != GGML_TYPE_I32 || + q->ne[0] != 256 || k->ne[0] != 256 || v->ne[0] != 256 || + k->ne[1] != v->ne[1] || k->ne[2] != v->ne[2] || k->ne[3] != q->ne[3] || v->ne[3] != q->ne[3] || + k->ne[2] < 1 || q->ne[2] % k->ne[2] != 0 || + ids->ne[0] < 1 || ids->ne[0] > 4096 || ids->ne[1] < q->ne[1] || ids->ne[2] != 1 || + (ids->ne[3] != 1 && ids->ne[3] != q->ne[3]) || + q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t) || + ids->nb[0] != sizeof(int32_t)) { + return nullptr; + } + for (int i = 1; i < 4; ++i) { + if (q->nb[i] % sizeof(float) || k->nb[i] % sizeof(ggml_fp16_t) || v->nb[i] % sizeof(ggml_fp16_t) || + ids->nb[i] % sizeof(int32_t)) { + return nullptr; + } + } + // the shader indexes with 32-bit element offsets + if (ggml_nbytes(k) / sizeof(ggml_fp16_t) > UINT32_MAX || ggml_nbytes(v) / sizeof(ggml_fp16_t) > UINT32_MAX || + ggml_nbytes(q) / sizeof(float) > UINT32_MAX || ggml_nbytes(op) / sizeof(float) > UINT32_MAX) { + return nullptr; + } + switch (q->ne[2] / k->ne[2]) { + case 4: return device->pipeline_flash_attn_sel_f16[0]; + case 8: return device->pipeline_flash_attn_sel_f16[1]; + case 12: return device->pipeline_flash_attn_sel_f16[2]; + default: return nullptr; + } +} - // dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride - auto is_dense_kv_cache = [](const ggml_tensor * t) { - return t->nb[0] == ggml_type_size(t->type) && - t->nb[2] == ggml_row_size(t->type, t->ne[0]) && - t->nb[1] == t->nb[2] * t->ne[2] && - (t->ne[3] == 1 || t->nb[3] == t->nb[1] * t->ne[1]); +static bool ggml_vk_flash_attn_sel(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, ggml_tensor * dst) { + vk_pipeline pipeline = ggml_vk_flash_attn_sel_pipeline(ctx->device, dst); + if (!pipeline) { + return false; + } + const ggml_tensor * ids = dst->src[5]; + float scale = 1.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + + const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); + const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); + const vk_subbuffer v_buf = ggml_vk_tensor_subbuffer(ctx, v); + const vk_subbuffer ids_buf = ggml_vk_tensor_subbuffer(ctx, ids); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + + // the WMMA kernel reads K and V rows in 16-byte pieces + const int gqa_idx = (int) (q->ne[2] / k->ne[2]) / 4 - 1; + vk_pipeline pipeline_cm = ctx->device->pipeline_flash_attn_sel_cm_f16[gqa_idx]; + bool use_cm = pipeline_cm != nullptr && k_buf.offset % 16 == 0 && v_buf.offset % 16 == 0; + for (int i = 1; i < 4; ++i) { + use_cm = use_cm && k->nb[i] % 16 == 0 && v->nb[i] % 16 == 0; + } + // one workgroup per (query, split, K/V head) on x + use_cm = use_cm && (uint64_t) q->ne[1] * k->ne[2] * CEIL_DIV(ids->ne[0], 32) <= ctx->device->properties.limits.maxComputeWorkGroupCount[0]; + if (use_cm) { + pipeline = pipeline_cm; + } + + const uint32_t n_sel = (uint32_t) ids->ne[0]; + const uint32_t N = (uint32_t) q->ne[1]; + const uint32_t NS = (uint32_t) q->ne[3]; + const uint32_t NH = (uint32_t) q->ne[2]; + const uint32_t D = 256; + // the cleaned rows: n_sel entries padded to whole 64-key blocks, then one count per row + const uint32_t list_stride = ROUNDUP_POW2(n_sel, 64u); + const size_t list_size = ((size_t) N * NS * list_stride + (size_t) N * NS) * sizeof(uint32_t); + + // few queries (decode) give too few workgroups to fill the GPU: split each row into whole key blocks. The + // WMMA kernel runs one subgroup per workgroup, the scalar one GQA subgroups. + const uint32_t n_wg = N * (uint32_t) k->ne[2] * NS; + const uint32_t min_wg = use_cm ? 128u : 64u; + uint32_t n_split = n_wg < min_wg ? std::min(CEIL_DIV(min_wg, n_wg), CEIL_DIV(n_sel, 64u)) : 1u; + const uint32_t chunk = ROUNDUP_POW2(CEIL_DIV(n_sel, std::max(n_split, 1u)), 32u); // whole key blocks + n_split = CEIL_DIV(n_sel, chunk); + const size_t split_size = n_split > 1 ? ((size_t) D * NH * N * NS * n_split + (size_t) NH * 2 * N * NS * n_split) * sizeof(float) : 0; + // the partials go after the rows, 256-byte aligned + const size_t split_off = ROUNDUP_POW2(list_size, 256); + + const vk_op_flash_attn_sel_push_constants pc = { + (uint32_t) k->ne[1], n_sel, list_stride, + (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[1] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (v->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[2] / sizeof(ggml_fp16_t)), (uint32_t) (v->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (ids->nb[1] / sizeof(int32_t)), ids->ne[3] == 1 ? 0u : (uint32_t) (ids->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, n_split, chunk, NH, N, NS, }; - const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; - const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; - const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 && - is_dense_kv_cache(k) && is_dense_kv_cache(v) && - (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - ctx->device->pipeline_dequant_transpose[k->type] != nullptr && - ctx->device->pipeline_dequant_transpose[v->type] != nullptr && - // coopmat2 path does not benefit from the f16 scratch - !ctx->device->coopmat2 && - // Intel Xe1 regresses, see PR 25494 - (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || - (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)); - const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type; - const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type; - // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga. - // For coopmat2 FA, we always use the small size (which is still pretty large for gqa). - vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc); - const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u); + vk_pipeline prep = ctx->device->pipeline_flash_attn_sel_prep; + ggml_pipeline_request_descriptor_sets(ctx, prep, 1); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + if (n_split > 1) { + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); + } + if (ctx->prealloc_size_split_k < split_off + split_size) { + ctx->prealloc_size_split_k = split_off + split_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_subbuffer list_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + ggml_vk_dispatch_pipeline(ctx, subctx, prep, { ids_buf, list_buf }, pc, { N, NS, 1 }); + ggml_vk_sync_buffers(ctx, subctx); - if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa && - qk_ratio * nek2 == neq2 && nek2 == nev2 && nem2 <= 1) { - // grouped query attention - make the N dimension equal to gqa_ratio, reduce - // workgroups proportionally in y dimension. The shader will detect gqa_ratio > 1 - // and change addressing calculations to index Q's dimension 2. - gqa_ratio = qk_ratio; - N = gqa_ratio; - workgroups_y /= gqa_ratio; + // the WMMA kernel takes the K/V heads on x (next to each other), the scalar one on y + GGML_ASSERT(!use_cm || (uint64_t) N * n_split * k->ne[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + const std::array grid = use_cm ? std::array{ N * n_split * (uint32_t) k->ne[2], 1, NS } + : std::array{ N * n_split, (uint32_t) k->ne[2], NS }; + if (n_split == 1) { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { q_buf, k_buf, v_buf, list_buf, dst_buf }, pc, grid); + ctx->prealloc_split_k_need_sync = true; + return true; } + const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, split_off); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { q_buf, k_buf, v_buf, list_buf, split_buf }, pc, grid); + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = { D, NH, N, N, NS, n_split, 0 }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + { split_buf, q_buf, dst_buf }, reduce_pc, { NH, D, N * NS }); + ctx->prealloc_split_k_need_sync = true; + return true; +} - tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); +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 = 1.0f; - float max_bias = 0.0f; + 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 (logit_softcap != 0) { - scale /= logit_softcap; + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; } - // 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)); - - // For F32, the shader treats it as a block of size 4 (for vec4 loads) - if (k->type == GGML_TYPE_F32) { - k_stride /= 4; + 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; } - if (v->type == GGML_TYPE_F32) { - v_stride /= 4; + 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; } - uint32_t nbk2_eff = (uint32_t)nbk2, nbk3_eff = (uint32_t)nbk3; - uint32_t nbv2_eff = (uint32_t)nbv2, nbv3_eff = (uint32_t)nbv3; - if (use_dequant_kv) { - k_stride = HSK; - v_stride = HSV; - nbk2_eff = (uint32_t)((uint64_t)HSK * KV * sizeof(ggml_fp16_t)); - nbk3_eff = (uint32_t)((uint64_t)HSK * KV * nek2 * sizeof(ggml_fp16_t)); - nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t)); - nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t)); - } - const uint32_t alignment = tuning_params.block_cols; - bool aligned = (KV % alignment) == 0 && - // the "aligned" shader variant will forcibly align strides, for performance - (q_stride & 7) == 0 && (k_stride & 7) == 0 && (v_stride & 7) == 0; + // ---- 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_check(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:; + + // 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; + } + + // 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 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); + + 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); + } + + 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); + + // 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, + }; + + 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)); + } + } + } + + 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; +} + +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) { + VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; + std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; + std::cerr << "), (" << v << ", name=" << v->name << ", type=" << v->type << ", ne0=" << v->ne[0] << ", ne1=" << v->ne[1] << ", ne2=" << v->ne[2] << ", ne3=" << v->ne[3] << ", nb0=" << v->nb[0] << ", nb1=" << v->nb[1] << ", nb2=" << v->nb[2] << ", nb3=" << v->nb[3]; + std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; + if (sinks) { + std::cerr << "), (" << sinks << ", name=" << sinks->name << ", type=" << sinks->type << ", ne0=" << sinks->ne[0] << ", ne1=" << sinks->ne[1] << ", ne2=" << sinks->ne[2] << ", ne3=" << sinks->ne[3] << ", nb0=" << sinks->nb[0] << ", nb1=" << sinks->nb[1] << ", nb2=" << sinks->nb[2] << ", nb3=" << sinks->nb[3]; + } + std::cerr << "))"); + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + 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; + uint32_t KV = nek1; + + GGML_ASSERT(ne0 == HSV); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == HSK); + + GGML_ASSERT(neq1 == N); + + GGML_ASSERT(nev1 == nek1); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + assert(dst->type == GGML_TYPE_F32); + assert(q->type == GGML_TYPE_F32); + if (dst->src[5] != nullptr && mask == nullptr) { + // supports_op only takes a maskless selected-key op when flash_attn_sel does + const bool done = ggml_vk_flash_attn_sel(ctx, subctx, q, k, v, dst); + GGML_ASSERT(done); + return; + } + 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; + uint32_t workgroups_y = (uint32_t)neq2; + uint32_t workgroups_z = (uint32_t)neq3; + + const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; + + // 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]) && + t->nb[1] == t->nb[2] * t->ne[2] && + (t->ne[3] == 1 || t->nb[3] == t->nb[1] * t->ne[1]); + }; + const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; + const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; + const 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) && + 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 reads its native types directly; non-native still needs the scratch + (kv_needs_dequant || !ctx->device->coopmat2) && + // Intel Xe1 regresses, see PR 25494 + (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). + vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc); + const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u); + + if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa && + qk_ratio * nek2 == neq2 && nek2 == nev2 && nem2 <= 1) { + // grouped query attention - make the N dimension equal to gqa_ratio, reduce + // workgroups proportionally in y dimension. The shader will detect gqa_ratio > 1 + // and change addressing calculations to index Q's dimension 2. + gqa_ratio = qk_ratio; + N = gqa_ratio; + workgroups_y /= gqa_ratio; + } + + tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); + + const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type)); + uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); + 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) { + k_stride /= 4; + } + if (v->type == GGML_TYPE_F32) { + v_stride /= 4; + } + + uint32_t nbk2_eff = (uint32_t)nbk2, nbk3_eff = (uint32_t)nbk3; + uint32_t nbv2_eff = (uint32_t)nbv2, nbv3_eff = (uint32_t)nbv3; + if (use_dequant_kv) { + k_stride = HSK; + v_stride = HSV; + nbk2_eff = (uint32_t)((uint64_t)HSK * KV * sizeof(ggml_fp16_t)); + nbk3_eff = (uint32_t)((uint64_t)HSK * KV * nek2 * sizeof(ggml_fp16_t)); + nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t)); + nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t)); + } + + const uint32_t alignment = tuning_params.block_cols; + bool aligned = (KV % alignment) == 0 && + // the "aligned" shader variant will forcibly align strides, for performance + (q_stride & 7) == 0 && (k_stride & 7) == 0 && (v_stride & 7) == 0; // Need to use the coopmat2 variant that clamps loads when HSK/HSV aren't sufficiently aligned. if (((HSK | HSV) % 16) != 0 && tuning_params.path == FA_COOPMAT2) { 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); + // Selected-key attention (src[5]) may name any cell, so masked keys must not leak NaN/inf from V; dense + // attention relies on the KV cache zeroing freed cells. GGML_VK_FA_NAN_SAFE=1 hardens every masked FA + // (e.g. a unified multi-sequence cache, whose masked cells belong to live sequences). + static const bool fa_nan_safe_all = [] { const char * e = getenv("GGML_VK_FA_NAN_SAFE"); return e && atoi(e) != 0; }(); + const bool nan_safe_v = mask != nullptr && (dst->src[5] != nullptr || fa_nan_safe_all); + + const bool radv_rdna3 = ctx->device->driver_id == vk::DriverId::eMesaRadv && + ctx->device->architecture == vk_device_architecture::AMD_RDNA3; + + // V^T for the coopmat1 prefill FA: transpose f16 V per head into a [HSV][KV + 128] scratch so a P x V + // B-operand fragment is one contiguous 32-byte run per lane instead of 16 kv-strided 16-bit loads. + // GGML_VK_FA_VT: unset = on for the cm1 kernel on RADV RDNA3, 1 = on everywhere, 0 = off (this also turns + // off the multi-row kernel, which reads V only from the V^T scratch). The cm1 kernel cannot use it under + // NAN_SAFE_V or shmem staging, which stage V rows from the row-major layout. + static const int fa_vt_env = [] { + const char * e = getenv("GGML_VK_FA_VT"); + return e ? atoi(e) : -1; + }(); + // dest row stride KV + 128: 2 * KV bytes is a multiple of 4 KB for every KV admitted here, which would put + // the 16 hd rows of a fragment on one memory channel; +256 B moves consecutive rows to the next channel + const uint32_t vt_kv_pad = KV + 128; + const uint64_t vt_size = (uint64_t)HSV * vt_kv_pad * nev2 * sizeof(ggml_fp16_t); + const uint64_t vt_offset = use_dequant_kv ? kv_f16_sz : 0; + const bool vt_ok = fa_vt_env != 0 && tuning_params.path == FA_COOPMAT1 && aligned && !tuning_params.shmem_staging && + v_type_eff == GGML_TYPE_F16 && !fa_compact.active && gqa_ratio == 1 && neq1 >= 64 && nev3 == 1 && + (KV % 64) == 0 && (HSV % 64) == 0 && KV < 65536 * 64 && + ctx->device->pipeline_dequant_f16_transpose_vt != nullptr && + (use_dequant_kv || ((nbv1 % 8) == 0 && (nbv2 % 8) == 0)) && + vt_size <= ctx->device->properties.limits.maxStorageBufferRange && + vt_offset + vt_size <= ctx->device->max_buffer_size; + + // Multi-row prefill FA (flash_attn_cm1_mr.comp): one workgroup of eight wave32 subgroups covers 64 rows = + // H heads x 64/H tokens that share a KV head. Q stays in registers and each 16-key K/V tile serves all 64 + // rows (cm1 pins Br at 16). K is read as per-head f16 rows (the contiguized scratch for a KV cache view), + // V from the V^T scratch. GGML_VK_FA_MR_H sets H (default 2 when it divides the GQA ratio, else 1). + // GGML_VK_FA_MR: unset = on for head size 256 on RADV RDNA3 (where it was measured: Qwen3.8-27B, FA per + // pp2048 ubatch 2067 -> 954 ms at d32768 on gfx1151), 1 = on for every eligible shape, 0 = off. + static const int fa_mr_env = [] { + const char * e = getenv("GGML_VK_FA_MR"); + return e ? atoi(e) : -1; + }(); + static const int fa_mr_h_env = [] { + const char * e = getenv("GGML_VK_FA_MR_H"); + return e ? atoi(e) : 0; + }(); + const bool fa_mr_on = fa_mr_env > 0 || (fa_mr_env < 0 && HSK == 256 && radv_rdna3); + const uint32_t qk_ratio_mr = nek2 > 0 ? (uint32_t)(neq2 / nek2) : 0; + const uint32_t mr_h = fa_mr_h_env > 0 ? (uint32_t)fa_mr_h_env : ((qk_ratio_mr % 2) == 0 ? 2u : 1u); + const bool use_mr = fa_mr_on && vt_ok && f32acc && + k_type_eff == GGML_TYPE_F16 && HSK == HSV && (HSK == 128 || HSK == 256) && + (mr_h == 1 || mr_h == 2 || mr_h == 4) && qk_ratio_mr > 0 && (qk_ratio_mr % mr_h) == 0 && + neq2 == (int64_t)qk_ratio_mr * nek2 && nek2 == nev2 && + (mask == nullptr || (mask->type == GGML_TYPE_F16 && nem2 <= 1 && nem3 <= 1 && + mask->nb[1] == (size_t)KV * sizeof(ggml_fp16_t))) && + (q_stride % 4) == 0 && (nbq2 % 16) == 0 && (nbq3 % 16) == 0 && + (k_stride % 8) == 0 && (nbk2_eff % 16) == 0 && (nbk3_eff % 16) == 0 && + ctx->device->subgroup_size_control && ctx->device->subgroup_min_size <= 32 && 32 <= ctx->device->subgroup_max_size && + ctx->device->subgroup_clustered && ctx->device->subgroup_vote; + const bool use_vt = vt_ok && (use_mr || (!nan_safe_v && (fa_vt_env > 0 || radv_rdna3))); + if (use_mr) { + tuning_params.block_rows = 64 / mr_h; + tuning_params.block_cols = 64; + tuning_params.row_split = mr_h; + tuning_params.subgroup_size = 32; + tuning_params.workgroup_size = 256; + tuning_params.shmem_staging = 0; + tuning_params.disable_subgroups = false; + tuning_params.limit_occupancy_shmem = 0; + workgroups_y = (uint32_t)neq2 / mr_h; + static const bool fa_mr_log = [] { + const char * e = getenv("GGML_VK_FA_MR_LOG"); + return e != nullptr && atoi(e) != 0; + }(); + if (fa_mr_log) { + GGML_LOG_INFO("ggml_vulkan: multi-row FA N=%u KV=%u hs=%u heads=%u H=%u gqa=%u mask=%d mask_opt=%d contig=%d\n", + N, KV, HSK, (uint32_t)neq2, mr_h, qk_ratio_mr, mask != nullptr, (int)use_mask_opt, (int)use_dequant_kv); + } + } + 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, nan_safe_v); + if (use_vt) { + fa_pipeline_state.flags |= FA_FLAG_V_TRANSPOSED; + } + if (use_mr) { + // schedule knobs, each default on (=0 disables); Qwen3.8-27B pp2048 at d32768, FA op time on gfx1151: + // GGML_VK_FA_MR_LAZY: skip the O rescale of row blocks whose running max did not move 1171 -> 1095 ms + // GGML_VK_FA_MR_NOEB: drop the loop-end barrier (no LDS hazard needs it) 1099 -> 1008 ms + // GGML_VK_FA_MR_VPOST: issue the V^T loads between the S store and the softmax barrier 1099 -> 969 ms + // all three: 954 ms + const auto fa_mr_knob = [](const char * name) { + const char * e = getenv(name); + return e == nullptr || atoi(e) != 0; + }; + static const bool fa_mr_lazy = fa_mr_knob("GGML_VK_FA_MR_LAZY"); + static const bool fa_mr_noeb = fa_mr_knob("GGML_VK_FA_MR_NOEB"); + static const bool fa_mr_vpost = fa_mr_knob("GGML_VK_FA_MR_VPOST"); + fa_pipeline_state.flags |= FA_FLAG_MULTI_ROW | + (fa_mr_lazy ? FA_FLAG_MR_LAZY : 0) | + (fa_mr_noeb ? FA_FLAG_MR_NOEB : 0) | + (fa_mr_vpost ? FA_FLAG_MR_VPOST : 0); + } vk_pipeline pipeline = nullptr; @@ -8215,18 +9826,8 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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); + if (use_mr) { + // multi-row FA: prefill only, enough workgroups without a split (and the kernel writes dst directly) } else if (gqa_ratio > 1 && workgroups_x <= Br) { split_k = shader_core_count * 2 / (workgroups_x * workgroups_y * workgroups_z); } else if (gqa_ratio <= 1) { @@ -8236,12 +9837,13 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const } } - 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); split_k = CEIL_DIV(KV, split_kv); - xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype; + // the Xe kernels read V row-major and use prealloc_x as their own scratch + xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype && !use_vt; if (xe_fa_opt) { const uint32_t split_p_size = 32; const size_t max_dim = (nek1 + split_p_size - 1) / split_p_size; @@ -8320,24 +9922,6 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const } } - // 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); @@ -8348,16 +9932,23 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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); const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp; const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp; - if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) { - ctx->prealloc_size_x = k_f16_sz + v_f16_sz; + // the V^T region after K and V is reserved here: growing prealloc_x after the dispatches below would + // replace the buffer they wrote + const uint64_t need = k_f16_sz + v_f16_sz + (use_vt ? vt_size : 0); + if (ctx->prealloc_size_x < need) { + ctx->prealloc_size_x = need; ggml_vk_preallocate_buffers(ctx, subctx); } vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; @@ -8378,6 +9969,39 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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)"); + } + + if (use_vt) { + // V -> V^T [nev2][HSV][KV + 128] in prealloc_x, after the K/V contiguize scratch when that ran + if (ctx->prealloc_size_x < vt_offset + vt_size) { + // only reachable without the contiguize pass, which reserves the V^T region itself + GGML_ASSERT(!use_dequant_kv); + ctx->prealloc_size_x = vt_offset + vt_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + vk_pipeline tr_vt = ctx->device->pipeline_dequant_f16_transpose_vt; + ggml_pipeline_request_descriptor_sets(ctx, tr_vt, 1); + if (ctx->prealloc_x_need_sync && !use_dequant_kv) { + ggml_vk_sync_buffers(ctx, subctx); + } + vk_subbuffer vt_dst = vk_subbuffer{ ctx->prealloc_x, vt_offset, vt_size }; + // source element strides: the contiguized scratch is [HSV][KV][nev2], the tensor has nbv1 / nbv2 + const uint32_t src_kv_stride = use_dequant_kv ? HSV : (uint32_t)(nbv1 / sizeof(ggml_fp16_t)); + const uint32_t src_head_stride = use_dequant_kv ? HSV * KV : (uint32_t)(nbv2 / sizeof(ggml_fp16_t)); + const std::vector tr_pc = { HSV, vt_kv_pad, src_kv_stride, src_head_stride, 0 }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_vt, { v_buf, vt_dst }, tr_pc, { HSV, KV / 64, (uint32_t)nev2 }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_x_need_sync = true; + v_buf = vt_dst; + v_stride = vt_kv_pad; + nbv2_eff = (uint32_t)((uint64_t)HSV * vt_kv_pad * sizeof(ggml_fp16_t)); + nbv3_eff = (uint32_t)vt_size; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_V_TRANSPOSE (sub-op)"); } uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2; @@ -8402,23 +10026,11 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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, @@ -8427,8 +10039,8 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const (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 }; @@ -8513,11 +10125,11 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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) }); @@ -8528,16 +10140,13 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const 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) { @@ -8572,6 +10181,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: @@ -8664,6 +10300,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; @@ -9101,6 +10742,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]; @@ -9589,6 +11276,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. @@ -10126,6 +11818,8 @@ void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& subctx 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); + 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]; @@ -10135,6 +11829,19 @@ void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx 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]; @@ -10227,11 +11934,80 @@ void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, K }; + // chunked prefill form (gdn_chunk_prep/scan.comp): S_v 128, scalar gate, final state only (K == 1), q/k/v + // rows with unit element stride, q and k with the same strides. f16 WMMA operands, f32 accumulation: not + // bit-identical to the sequential scan. Default on where the pipelines exist (RDNA3 RADV wave32); GGML_VK_GDN_CHUNK=0 + // opts out. Batches under 64 tokens keep the sequential kernel. + static const bool chunk_enabled = [] { const char * e = getenv("GGML_VK_GDN_CHUNK"); return !e || atoi(e) != 0; }(); + constexpr uint32_t chunk_min = 64; + const bool kda = dst->src[3]->ne[0] == (int64_t) S_v; + if (chunk_enabled && ctx->device->pipeline_gdn_chunk_scan && S_v == 128 && !kda && K == 1 && n_tokens >= chunk_min && + src_q->nb[0] == sizeof(float) && dst->src[1]->nb[0] == sizeof(float) && src_v->nb[0] == sizeof(float) && + dst->src[1]->nb[1] == src_q->nb[1] && dst->src[1]->nb[2] == src_q->nb[2] && dst->src[1]->nb[3] == src_q->nb[3] && + H % neq1 == 0) { + constexpr uint32_t CL = 32; + const uint32_t nch = (n_tokens + CL - 1) / CL; + const size_t qku = 3 * CL * 128 / 2; + const size_t tpu = CL * CL + 2 * CL; + const size_t scr_size = sizeof(uint32_t) * ((size_t) n_seqs * neq1 * nch * qku + (size_t) n_seqs * H * nch * tpu); + if (ctx->prealloc_size_split_k < scr_size) { + ctx->prealloc_size_split_k = scr_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + vk_subbuffer scr_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + const vk_op_gdn_chunk_push_constants cpc = { pc, nch }; + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_gdn_chunk_prep, 1); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_gdn_chunk_scan, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_gdn_chunk_prep, + { src_buf[0], src_buf[1], src_buf[3], src_buf[4], scr_buf }, cpc, { neq1, nch, n_seqs }); + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_gdn_chunk_scan, + { src_buf[2], src_buf[5], dst_buf, scr_buf }, cpc, { H * ctx->device->gdn_chunk_scan_nsplit, n_seqs, 1 }); + ctx->prealloc_split_k_need_sync = true; + return; + } + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {src_buf[0], src_buf[1], src_buf[2], src_buf[3], src_buf[4], src_buf[5], dst_buf}, 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]}); +} + 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]; @@ -10289,7 +12065,64 @@ void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_te pc, elements); } +// direct conv + SiLU (see ggml_vk_ssm_conv_direct_check): x [C, T, S] and the state [3, C, S] in place of the concat +static void ggml_vk_ssm_conv_direct(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_backend_vk_context::ssm_conv_direct_match & m) { + const uint32_t C = (uint32_t) m.kern->ne[1]; + const uint32_t T = (uint32_t) m.out->ne[1]; + const uint32_t S = (uint32_t) m.out->ne[2]; + + vk_pipeline pipeline = ctx->device->pipeline_ssm_conv_direct_silu_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_ssm_conv_push_constants pc = { + (uint32_t) (C * sizeof(float)), (uint32_t) ((size_t) C * T * sizeof(float)), + (uint32_t) m.kern->nb[1], + (uint32_t) m.out->nb[0], (uint32_t) m.out->nb[1], (uint32_t) m.out->nb[2], + (uint32_t) m.kern->ne[0], + (uint32_t) m.state->ne[0], + C, T, S, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { ggml_vk_tensor_subbuffer(ctx, m.x), ggml_vk_tensor_subbuffer(ctx, m.kern), + ggml_vk_tensor_subbuffer(ctx, m.state), ggml_vk_tensor_subbuffer(ctx, m.out) }, + pc, { C, T, S }); +} + +// the concat columns [tail_from, T + 3) that the conv-state copies read; the conv itself runs direct +static void ggml_vk_concat_tail(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_backend_vk_context::ssm_conv_direct_match & m) { + const ggml_tensor * cc = m.concat; + const uint32_t C = (uint32_t) cc->ne[1]; + const uint32_t T = (uint32_t) cc->ne[0] - 3; + const uint32_t S = (uint32_t) cc->ne[2]; + const uint32_t ncols = (uint32_t) (cc->ne[0] - m.tail_from); + if (ncols == 0) { + return; + } + + vk_pipeline pipeline = ctx->device->pipeline_ssm_conv_tail_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_ssm_conv_push_constants pc = { + (uint32_t) (C * sizeof(float)), (uint32_t) ((size_t) C * T * sizeof(float)), + 0, + (uint32_t) cc->nb[0], (uint32_t) cc->nb[1], (uint32_t) cc->nb[2], + (uint32_t) m.tail_from, + 3, + C, T, S, + }; + // binding 1 is unused in this mode + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { ggml_vk_tensor_subbuffer(ctx, m.x), ggml_vk_tensor_subbuffer(ctx, m.state), + ggml_vk_tensor_subbuffer(ctx, m.state), ggml_vk_tensor_subbuffer(ctx, cc) }, + pc, { ncols * C, S, 1 }); +} + void ggml_vk_ssm_conv(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { + if (ctx->fused_ssm_conv_direct >= 0) { + GGML_ASSERT(ctx->num_additional_fused_ops == 1); + ggml_vk_ssm_conv_direct(ctx, subctx, ctx->ssm_conv_direct_matches[ctx->fused_ssm_conv_direct]); + return; + } ggml_tensor * conv = cgraph->nodes[node_idx]; const ggml_tensor * src0 = conv->src[0]; const ggml_tensor * src1 = conv->src[1]; @@ -10710,6 +12543,84 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * return rope; } +// Short-row RMS norm (rms_norm_small.comp, one subgroup per row): plain, +MUL, and the RMS_NORM+MUL+MUL / +// RMS_NORM+SCALE fusions that only this kernel has. GGML_VK_NORM_SMALL=0 disables all of it. +static bool ggml_vk_norm_small_enabled() { + static const bool enabled = [] { + const char * e = getenv("GGML_VK_NORM_SMALL"); + return e == nullptr || atoi(e) != 0; + }(); + return enabled; +} + +// RMS_NORM+MUL+MUL and RMS_NORM+SCALE fusions on the short-row kernel. GGML_VK_FUSE_NORM_EXTRA=0 disables both. +static bool ggml_vk_fuse_norm_extra_enabled() { + static const bool enabled = [] { + const char * e = getenv("GGML_VK_FUSE_NORM_EXTRA"); + return e == nullptr || atoi(e) != 0; + }(); + return ggml_vk_norm_small_enabled() && enabled; +} + +static bool ggml_vk_rms_norm_small_ok(const ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * dst) { + return ggml_vk_norm_small_enabled() && ctx->device->pipeline_rms_norm_small_f32 != nullptr && + src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && src0->ne[0] <= 256 && + ggml_is_contiguous_rows(src0) && ggml_is_contiguous(dst) && + !ctx->do_add_rms_partials; +} + +static void ggml_vk_rms_norm_small(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float eps) { + ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * src0 = rms->src[0]; + const ggml_tensor * src1 = src0; + const ggml_tensor * src2 = src0; // second multiplier (RMS_NORM+MUL+MUL), else unused + ggml_tensor * dst = rms; + float post_scale = 0.0f; + vk_pipeline pipeline = ctx->device->pipeline_rms_norm_small_f32; + + switch (ctx->fused_rms_norm_mode) { + case RMS_NORM_MUL: + case RMS_NORM_MUL_MUL: { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + src1 = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + dst = mul; + pipeline = ctx->device->pipeline_rms_norm_mul_small_f32; + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_MUL) { + dst = cgraph->nodes[node_idx + 2]; + src2 = dst->src[0] == mul ? dst->src[1] : dst->src[0]; + pipeline = ctx->device->pipeline_rms_norm_mul_mul_small_f32; + } + } break; + case RMS_NORM_SCALE: + dst = cgraph->nodes[node_idx + 1]; + post_scale = ggml_get_op_params_f32(dst, 0); + pipeline = ctx->device->pipeline_rms_norm_scale_small_f32; + break; + default: + GGML_ASSERT(ctx->fused_rms_norm_mode == RMS_NORM_COUNT); + break; + } + GGML_ASSERT(pipeline != nullptr); + + vk_op_binary_push_constants pc = ggml_vk_rms_norm_push_constants(src0, src1, dst, eps, 0); + pc.param2 = post_scale; + init_pushconst_tensor_offsets(ctx, pc, src0, src1, nullptr, nullptr, dst); + + const uint32_t rows_per_wg = 256 / ctx->device->subgroup_size; + const uint64_t nrows = (uint64_t)src0->ne[1] * src0->ne[2] * src0->ne[3]; + // the shader strides over rows, so capping the grid only costs parallelism + const uint32_t wgs = (uint32_t)std::min(CEIL_DIV(nrows, rows_per_wg), 65535); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, src1, true), + ggml_vk_tensor_subbuffer(ctx, dst, true), + ggml_vk_tensor_subbuffer(ctx, src2, true), + }, pc, { wgs, 1, 1 }); +} + static void ggml_vk_rms_norm_finish(ggml_backend_vk_context * ctx, const ggml_tensor * src0) { if (ctx->do_add_rms_partials_offset_calculation) { ctx->prealloc_size_add_rms_partials_offset += ggml_vk_rms_partials_size(ctx, src0); @@ -10722,6 +12633,14 @@ void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const s ggml_tensor * rms = cgraph->nodes[node_idx]; const ggml_tensor * src0 = rms->src[0]; + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_MUL || ctx->fused_rms_norm_mode == RMS_NORM_SCALE || + ((ctx->fused_rms_norm_mode == RMS_NORM_COUNT || ctx->fused_rms_norm_mode == RMS_NORM_MUL) && + ggml_vk_rms_norm_small_ok(ctx, src0, cgraph->nodes[node_idx + ctx->num_additional_fused_ops]))) { + ggml_vk_rms_norm_small(ctx, subctx, cgraph, node_idx, op_params[0]); + ggml_vk_rms_norm_finish(ctx, src0); + return; + } + if (ctx->fused_rms_norm_mode == RMS_NORM_VIEW_SET_ROWS) { GGML_ASSERT(ctx->num_additional_fused_ops == 2); ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; @@ -11224,7 +13143,7 @@ void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const gg std::array elements; elements[0] = ncolsp2; - elements[1] = std::min((uint32_t)ggml_nrows(src0), ctx->device->properties.limits.maxComputeWorkGroupCount[1]); + elements[1] = std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]); elements[2] = 1; // First dispatch initializes tmp_idx and does the first N passes where @@ -11268,6 +13187,9 @@ void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, const gg } } +// 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). 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); @@ -12239,6 +14161,12 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in } } } + // tensors a fused kernel reads that are not srcs of its nodes (the direct conv reads x and the state) + for (const ggml_tensor * extra : ctx->fused_extra_srcs) { + if (extra && !need_sync && overlaps_unsynced(extra, ctx->unsynced_nodes_written)) { + need_sync = true; + } + } if (need_sync) { if (vk_enable_sync_logger) { @@ -12268,6 +14196,11 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in ctx->unsynced_nodes_read.push_back(cur_node->src[j]); } } + for (const ggml_tensor * extra : ctx->fused_extra_srcs) { + if (extra) { + ctx->unsynced_nodes_read.push_back(extra); + } + } } if (vk_enable_sync_logger) { for (int i = 0; i < ctx->num_additional_fused_ops + 1; ++i) { @@ -12344,6 +14277,10 @@ bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, in break; case GGML_OP_CONCAT: + if (ctx->fused_concat_tail >= 0) { + ggml_vk_concat_tail(ctx, compute_ctx, ctx->ssm_conv_direct_matches[ctx->fused_concat_tail]); + break; + } ggml_vk_concat(ctx, compute_ctx, src0, src1, node); break; @@ -12826,8 +14763,11 @@ 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; @@ -13458,6 +15398,59 @@ static bool ggml_vk_can_fuse_hc_post_gate(const struct ggml_cgraph * cgraph, int ggml_are_same_shape(scale_in->src[0], scale_out); } +// MUL_MAT_ID+MUL(router weights) on the GEMM path: the int8 MMQ coopmat1 epilogue multiplies each output row by +// its (expert slot, token) weight before the store, so the f32 expert outputs are written once and never re-read. +// GGML_VK_MMID_EPILOGUE=0 disables (mat-vec MUL_MAT_ID+MUL fusion is unaffected). +static bool ggml_vk_mmid_epilogue_enabled() { + static const bool on = [] { + const char * e = getenv("GGML_VK_MMID_EPILOGUE"); + return e == nullptr || atoi(e) != 0; + }(); + return on; +} + +// true when ggml_vk_mul_mat_id_q_f16 will run this MUL_MAT_ID on an int8 MMQ coopmat1 pipeline (the only MUL_MAT_ID +// GEMM with the router-weight binding) and the MUL is a per-row scale it can apply +static bool ggml_vk_mmid_scale_epilogue_ok(const ggml_backend_vk_context * cctx, const ggml_tensor * mmid, const ggml_tensor * mul) { + if (!ggml_vk_mmid_epilogue_enabled()) { + return false; + } + ggml_backend_vk_context * ctx = const_cast(cctx); + const ggml_tensor * src0 = mmid->src[0]; + const ggml_tensor * src1 = mmid->src[1]; + const ggml_tensor * scale = mul->src[1]; + if (mmid->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32 || scale->type != GGML_TYPE_F32 || + !ggml_is_quantized(src0->type) || src1->type != GGML_TYPE_F32) { + return false; + } + // same conditions as quantize_y / !qx_needs_dequant in ggml_vk_mul_mat_id_q_f16 + if (!ctx->device->integer_dot_product || ctx->device->coopmat2 || + !ggml_vk_dim01_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_vk_dim01_contiguous(src1) || + (src1->ne[0] * src1->ne[1]) % 4 != 0) { + return false; + } + if (ctx->device->coopmat_support && ctx->device->vendor_id == VK_VENDOR_ID_INTEL) { + return false; + } + const std::vector * mmp = ggml_vk_get_mul_mat_mat_pipeline_map(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)mmid->op_params[0], true); + if (mmp == nullptr) { + return false; + } + for (const auto & pair : *mmp) { + if (!pair.unaligned || pair.unaligned->parameter_count != 6) { + return false; + } + } + // weights [1, n_expert_used, n_tokens], contiguous, one per output row; no broadcast of the expert outputs + if (!ggml_are_same_shape(mul, mmid) || !ggml_is_contiguous(mul) || !ggml_is_contiguous(mmid) || + !ggml_is_contiguous(scale) || get_misalign_bytes(ctx, scale) != 0 || + scale->ne[0] != 1 || scale->ne[1] != mul->ne[1] || scale->ne[2] != mul->ne[2] || scale->ne[3] != 1 || + mul->ne[3] != 1) { + return false; + } + return true; +} + bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL) { return ggml_vk_can_fuse_unary_mul_pair(cgraph, node_idx); @@ -13467,6 +15460,40 @@ bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgr return false; } + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_SCALE) { + // RMS_NORM+SCALE (short rows only): build_gdn_l2_norm is scale(rms_norm(x, eps/n), 1/sqrt(n)). + // The kernel applies the factor as it stores, so the norm's output is never written or read back. + const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor *scale = cgraph->nodes[node_idx + 1]; + return ggml_vk_fuse_norm_extra_enabled() && ctx->device->pipeline_rms_norm_scale_small_f32 != nullptr && + scale->src[0] == rms_norm && ggml_get_op_params_f32(scale, 1) == 0.0f && + scale->type == GGML_TYPE_F32 && ggml_are_same_shape(scale, rms_norm) && + ggml_vk_rms_norm_small_ok(ctx, rms_norm->src[0], scale); + } + + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL && ops.begin()[2] == GGML_OP_MUL) { + // RMS_NORM+MUL+MUL (short rows only): the gated norm, normalised * gamma * gate. The gate is + // indexed like the destination, so it must be contiguous, dst-shaped and aligned. + if (!ggml_vk_fuse_norm_extra_enabled() || ctx->device->pipeline_rms_norm_mul_mul_small_f32 == nullptr || + !ggml_vk_can_fuse(ctx, cgraph, node_idx, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { + return false; + } + const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; + const ggml_tensor *mul2 = cgraph->nodes[node_idx + 2]; + if (mul2->src[0] != mul && mul2->src[1] != mul) { + return false; + } + const ggml_tensor *gate = mul2->src[0] == mul ? mul2->src[1] : mul2->src[0]; + if (gate == mul || gate->type != GGML_TYPE_F32 || mul2->type != GGML_TYPE_F32 || + !ggml_is_contiguous(gate) || !ggml_is_contiguous(mul) || !ggml_is_contiguous(mul2) || + !ggml_are_same_shape(gate, mul2) || !ggml_are_same_shape(mul, mul2) || !ggml_are_same_shape(rms_norm, mul) || + get_misalign_bytes(ctx, gate) != 0) { + return false; + } + return ggml_vk_rms_norm_small_ok(ctx, rms_norm->src[0], mul2); + } + if ((ops.size() == 2 || ops.size() == 3 || ops.size() == 4) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { // additional constraints specific to this fusion @@ -13576,9 +15603,9 @@ bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct ggml_cgr if (mmid != mul->src[0]) { return false; } - // mat-vec only + // mat-vec, or the int8 MMQ GEMM epilogue (MUL_MAT_ID+MUL only) if (!ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { - return false; + return mmid->op == GGML_OP_MUL_MAT_ID && ggml_vk_mmid_scale_epilogue_ok(ctx, mmid, mul); } // shaders assume the types match if (mmid->type != scale->type) { @@ -14048,6 +16075,217 @@ bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b, bool return false; } +// Direct GDN conv: CONCAT(state [3, C, S], transpose(x [C, T, S]), dim 0) -> SSM_CONV(width 4) -> SILU. +// The concat transposes x into an [T+3, C, S] tensor (84 MB at T = 2048 for Flash-Next) that the conv reads once. +// Instead the conv reads x in its channel-contiguous layout and the state directly, and the concat node writes +// only the tail columns that its other readers (the recurrent conv-state copies, which read the last 3 columns +// once per rollback slot) need. The graph is unchanged: other backends and the GGML_VK_SSMCONV_DIRECT=0 arm run +// the same nodes. The decision is made at the concat and consumed at its SSM_CONV, so both see one answer. +static bool ggml_vk_ssm_conv_direct_enabled() { + static const bool enabled = [] { + const char * e = getenv("GGML_VK_SSMCONV_DIRECT"); + return e == nullptr || atoi(e) != 0; + }(); + return enabled; +} + +// decode and short verify ubatches keep the concat: nothing to save there, and the reader scan stays off that path +static constexpr int64_t ssm_conv_direct_min_tokens = 8; + +static bool ggml_vk_ssm_conv_direct_reject(const char * why, int node) { + static const bool dbg = getenv("GGML_VK_FUSION_DEBUG") != nullptr; + if (dbg) { + fprintf(stderr, "ggml_vulkan: direct ssm conv not matched at node %d: %s\n", node, why); + } + return false; +} + +static bool ggml_vk_node_reads(const ggml_tensor * t, const ggml_tensor * root) { + for (int s = 0; s < GGML_MAX_SRC; ++s) { + if (t->src[s] && (t->src[s] == root || t->src[s]->view_src == root)) { + return true; + } + } + return false; +} + +// a view of the concat covering H whole columns of every row: returns the first column, or -1 +static int64_t ggml_vk_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; +} + +// the conv-state copies: a tail view, a CONT of one, or a CPY of either +static bool ggml_vk_conv_tail_reader(const ggml_tensor * t, const ggml_tensor * cc, int64_t H, int64_t & tail_from) { + const ggml_tensor * v = nullptr; + if (t->op == GGML_OP_VIEW) { + v = t; + } else if (t->op == GGML_OP_CONT || t->op == GGML_OP_CPY) { + v = t->src[0]; + if (v && t->op == GGML_OP_CPY && v->op == GGML_OP_CONT) { + v = v->src[0]; + } + } + const int64_t col = v ? ggml_vk_conv_tail_view_col(v, cc, H) : -1; + if (col < 0) { + return false; + } + tail_from = std::min(tail_from, col); + return true; +} + +// the whole concat, directly or through a view/reshape covering all of it +static bool ggml_vk_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)); +} + +// Structural match at a CONCAT node. With runtime = true it also checks the allocated addresses: the SiLU output +// must not overlap anything the direct conv reads, and nothing between the concat and the SiLU may write x or +// the state (graph_optimize asks the allocator to keep both alive, this is the check that it did). +static bool ggml_vk_ssm_conv_direct_check(const ggml_backend_vk_context * ctx, const ggml_cgraph * cgraph, int i, + ggml_backend_vk_context::ssm_conv_direct_match & m, bool runtime) { + 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)) { + return false; + } + 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 || tr->view_offs != 0) { + return ggml_vk_ssm_conv_direct_reject("concat srcs are not (state, transpose(x))", i); + } + const int64_t T = tr->ne[0]; + const int64_t C = tr->ne[1]; + const int64_t S = tr->ne[2]; + const ggml_tensor * x = tr->view_src; + // x^T must be a plain transpose of the whole contiguous x: token stride C, channel stride 1, sequence stride C*T + if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || ggml_nelements(x) != C * T * S || tr->ne[3] != 1 || + tr->nb[0] != (size_t) C * sizeof(float) || tr->nb[1] != sizeof(float) || tr->nb[2] != (size_t) C * T * sizeof(float)) { + return ggml_vk_ssm_conv_direct_reject("x layout", i); + } + if (T < ssm_conv_direct_min_tokens) { + return false; + } + if (st->ne[0] != 3 || st->ne[1] != C || st->ne[2] != S || st->ne[3] != 1 || !ggml_is_contiguous(st)) { + return ggml_vk_ssm_conv_direct_reject("state layout", i); + } + if (cc->ne[0] != T + 3 || cc->ne[1] != C || cc->ne[2] != S || cc->ne[3] != 1 || !ggml_is_contiguous(cc)) { + return ggml_vk_ssm_conv_direct_reject("concat layout", i); + } + + // one SSM_CONV reads the whole concat; everything else must read only tail columns + 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_vk_node_reads(t, cc)) { + continue; + } + if (t->op == GGML_OP_SSM_CONV && ggml_vk_conv_whole(t->src[0], cc)) { + if (conv >= 0) { + return ggml_vk_ssm_conv_direct_reject("two convs read the concat", n); + } + conv = n; + continue; + } + if ((t->op == GGML_OP_RESHAPE || t->op == GGML_OP_VIEW) && ggml_vk_conv_whole(t, cc)) { + continue; // shape plumbing into the conv + } + if (!ggml_vk_conv_tail_reader(t, cc, 3, tail_from)) { + return ggml_vk_ssm_conv_direct_reject("the concat has a reader that is neither the conv nor a tail copy", n); + } + } + if (conv < 0 || conv + 1 >= cgraph->n_nodes) { + return ggml_vk_ssm_conv_direct_reject("no conv in this graph", i); + } + const ggml_tensor * cv = cgraph->nodes[conv]; + const ggml_tensor * w = cv->src[1]; + if (cv->type != GGML_TYPE_F32 || !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] != S || cv->ne[3] != 1 || (cv->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return ggml_vk_ssm_conv_direct_reject("conv shape", conv); + } + // the conv output is never written, so the SiLU right after it must be its only reader + 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_are_same_shape(su, cv) || (cv->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || + (su->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return ggml_vk_ssm_conv_direct_reject("no SILU right after the conv", conv); + } + for (int n = conv + 2; n < cgraph->n_nodes; ++n) { + if (ggml_vk_node_reads(cgraph->nodes[n], cv)) { + return ggml_vk_ssm_conv_direct_reject("the conv output has another reader", n); + } + } + + if (runtime) { + for (const ggml_tensor * t : { (const ggml_tensor *) x, st, w, (const ggml_tensor *) cc, (const ggml_tensor *) su }) { + if (!t->buffer || get_misalign_bytes(ctx, t) != 0) { + return ggml_vk_ssm_conv_direct_reject("unallocated or misaligned tensor", i); + } + } + for (const ggml_tensor * t : { (const ggml_tensor *) x, st, w, (const ggml_tensor *) cc }) { + if (ggml_vk_tensors_overlap(t, su, false)) { + return ggml_vk_ssm_conv_direct_reject("the SiLU output overlaps an input", conv + 1); + } + } + // x and the state are read at the conv, later than the concat that names them: nothing in between may + // overwrite them + for (int n = i + 1; n <= conv; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (t->buffer == nullptr || t->op == GGML_OP_NONE || ggml_op_is_empty(t->op) || ggml_is_empty(t)) { + continue; + } + if (ggml_vk_tensors_overlap(t, x, false) || ggml_vk_tensors_overlap(t, st, false)) { + return ggml_vk_ssm_conv_direct_reject("x or the state is overwritten before the conv", n); + } + } + } + + m.concat = cc; + m.state = st; + m.x = x; + m.kern = w; + m.out = su; + m.conv_idx = conv; + m.tail_from = std::max(0, tail_from); + return true; +} + +static bool ggml_vk_ssm_conv_direct_at_concat(ggml_backend_vk_context * ctx, const ggml_cgraph * cgraph, int i) { + if (!ggml_vk_ssm_conv_direct_enabled() || cgraph->nodes[i]->op != GGML_OP_CONCAT) { + return false; + } + ggml_backend_vk_context::ssm_conv_direct_match m; + if (!ggml_vk_ssm_conv_direct_check(ctx, cgraph, i, m, true)) { + return false; + } + ctx->ssm_conv_direct_matches.push_back(m); + ctx->fused_concat_tail = (int) ctx->ssm_conv_direct_matches.size() - 1; + return true; +} + +static bool ggml_vk_ssm_conv_direct_at_conv(ggml_backend_vk_context * ctx, const ggml_cgraph * cgraph, int i) { + if (cgraph->nodes[i]->op != GGML_OP_SSM_CONV) { + return false; + } + for (size_t k = 0; k < ctx->ssm_conv_direct_matches.size(); ++k) { + const auto & m = ctx->ssm_conv_direct_matches[k]; + if (m.conv_idx == i && ggml_vk_conv_whole(cgraph->nodes[i]->src[0], m.concat)) { + ctx->fused_ssm_conv_direct = (int) k; + ctx->fused_extra_srcs[0] = m.x; + ctx->fused_extra_srcs[1] = m.state; + return true; + } + } + return false; +} + bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { const ggml_tensor *rms = cgraph->nodes[node_idx + 0]; @@ -14164,6 +16402,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->prealloc_size_add_rms_partials_offset = 0; ctx->do_add_rms_partials = false; ctx->do_add_rms_partials_offset_calculation = false; + ctx->ssm_conv_direct_matches.clear(); int last_node = cgraph->n_nodes - 1; @@ -14196,14 +16435,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; @@ -14230,6 +16477,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; @@ -14267,6 +16515,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; } @@ -14281,9 +16530,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; @@ -14294,6 +16545,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 @@ -14307,6 +16559,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_qsa = false; ctx->fused_hc_post_gate = false; ctx->fused_rms_norm_mode = RMS_NORM_COUNT; + ctx->fused_concat_tail = -1; + ctx->fused_ssm_conv_direct = -1; + ctx->fused_extra_srcs[0] = ctx->fused_extra_srcs[1] = nullptr; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -14386,6 +16641,17 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_rms_norm_mode = RMS_NORM_VIEW_SET_ROWS; fusion_string = "RMS_NORM_VIEW_SET_ROWS"; std::fill_n(op_srcs_fused_elementwise, 3, false); + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_MUL })) { + ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_MUL; + fusion_string = "RMS_NORM_MUL_MUL"; + // one subgroup reads its whole row before writing it; the gate shares the destination index + std::fill_n(op_srcs_fused_elementwise, 3, true); + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_SCALE })) { + ctx->num_additional_fused_ops = 1; + ctx->fused_rms_norm_mode = RMS_NORM_SCALE; + fusion_string = "RMS_NORM_SCALE"; + std::fill_n(op_srcs_fused_elementwise, 2, true); } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; ctx->fused_rms_norm_mode = RMS_NORM_MUL; @@ -14404,6 +16670,14 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } op_srcs_fused_elementwise[0] = true; op_srcs_fused_elementwise[1] = true; + } else if (ggml_vk_ssm_conv_direct_at_conv(ctx, cgraph, i)) { + // decided at the concat, which then wrote only its tail: this must not fall back + ctx->num_additional_fused_ops = 1; + fusion_string = "CONCAT_SSM_CONV_SILU"; + op_srcs_fused_elementwise[0] = false; + op_srcs_fused_elementwise[1] = true; + } else if (ggml_vk_ssm_conv_direct_at_concat(ctx, cgraph, i)) { + fusion_string = "CONCAT_TAIL"; } else if (ggml_vk_can_fuse_ssm_conv(ctx, cgraph, i, 2)) { ctx->num_additional_fused_ops = 2; fusion_string = "SSM_CONV_BIAS_SILU"; @@ -14542,6 +16816,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } } if (need_disable) { + // the direct conv's concat already skipped writing the full tensor; its match checks every overlap + GGML_ASSERT(ctx->fused_ssm_conv_direct < 0); ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 1; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; @@ -14557,6 +16833,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); @@ -14568,6 +16845,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 { @@ -14609,14 +16887,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 @@ -14810,6 +17096,91 @@ void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, continue; } + // Gated norm (RMS_NORM_MUL_MUL): rms_norm -> mul(gamma) -> mul(gate). The gate chain (e.g. a + // reshape + sigmoid) sits between the two MULs in the graph, so schedule it first and keep the + // three nodes adjacent. Returns the new order (gate chain, norm, MUL, MUL), or empty. + auto const &gated_norm_order = [&](int r) -> std::vector { + std::vector order; + if (!ggml_vk_fuse_norm_extra_enabled() || graph->nodes[r]->op != GGML_OP_RMS_NORM) { + return order; + } + const int lim = std::min(r + 20, graph->n_nodes); + int j = r + 1; + while (j < lim && (used[j] || is_empty(graph->nodes[j]))) { + j++; + } + if (j >= lim || graph->nodes[j]->op != GGML_OP_MUL || + (graph->nodes[j]->src[0] != graph->nodes[r] && graph->nodes[j]->src[1] != graph->nodes[r]) || + !empty_or_scheduled_between(r + 1, j)) { + return order; + } + int k = -1; + for (int m = j + 1; m < lim; ++m) { + if (!used[m] && graph->nodes[m]->op == GGML_OP_MUL && + (graph->nodes[m]->src[0] == graph->nodes[j] || graph->nodes[m]->src[1] == graph->nodes[j])) { + k = m; + break; + } + } + if (k < 0) { + return order; + } + const ggml_tensor * gate = graph->nodes[k]->src[0] == graph->nodes[j] ? graph->nodes[k]->src[1] : graph->nodes[k]->src[0]; + if (gate == graph->nodes[j]) { + return order; + } + // nodes in (j, k) the gate needs; none may depend on unused nodes up to the gamma MUL + std::vector need(k - j, false); + for (int m = k - 1; m > j; --m) { + if (used[m]) { + continue; + } + if (graph->nodes[m] == gate) { + need[m - j] = true; + } + if (!need[m - j]) { + // left behind: the gated MUL must not depend on it + if (is_src_of(graph->nodes[k], graph->nodes[m])) { + return order; + } + continue; + } + for (int c = first_unused; c < m; ++c) { + if (used[c] || !is_src_of(graph->nodes[m], graph->nodes[c])) { + continue; + } + if (c <= j) { + return order; + } + need[c - j] = true; + } + } + for (int m = j + 1; m < k; ++m) { + if (!used[m] && need[m - j]) { + order.push_back(m); + } + } + order.push_back(r); + order.push_back(j); + order.push_back(k); + return order; + }; + + { + const std::vector order = gated_norm_order(first_unused); + if (!order.empty()) { + for (int m : order) { + new_order.push_back(graph->nodes[m]); + used_node_set.insert(graph->nodes[m]); + used[m] = true; + } + while (first_unused < graph->n_nodes && used[first_unused]) { + first_unused++; + } + continue; + } + } + // First, grab the next unused node. current_set.push_back(first_unused); @@ -14849,7 +17220,8 @@ void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, match_pattern(rms_norm_mul_rope_view_set_rows_pattern, j) || match_pattern(rms_norm_view_set_rows_pattern, j) || match_pattern(rope_view_set_rows_pattern, j) || - match_pattern(hc_post_gate_pattern, j)) { + match_pattern(hc_post_gate_pattern, j) || + !gated_norm_order(j).empty()) { continue; } bool ok = true; @@ -14863,7 +17235,10 @@ void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_MUL_MAT_ID && graph->nodes[j]->op == GGML_OP_MUL) && !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_ADD && graph->nodes[j]->op == GGML_OP_ADD) && !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_SSM_CONV && graph->nodes[j]->op == GGML_OP_ADD) && - !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_SSM_CONV && graph->nodes[j]->op == GGML_OP_UNARY)) { + !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_SSM_CONV && graph->nodes[j]->op == GGML_OP_UNARY) && + // RMS_NORM -> SCALE (RMS_NORM_SCALE, the GDN q/k l2 norm): keep the scale behind its norm + !(j == c+1 && c == current_set.back() && graph->nodes[c]->op == GGML_OP_RMS_NORM && graph->nodes[j]->op == GGML_OP_SCALE && + graph->nodes[j]->src[0] == graph->nodes[c] && ggml_vk_fuse_norm_extra_enabled())) { ok = false; break; } @@ -15034,6 +17409,20 @@ void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, for (int i = 0; i < graph->n_nodes; ++i) { graph->nodes[i] = new_order[i]; } + + // direct GDN conv: x and the state are read at the conv, after the concat that names them, and while the + // SiLU output is written; keep them allocated through the SiLU (the runtime match re-checks the addresses) + if (ggml_vk_ssm_conv_direct_enabled() && !ctx->device->disable_fusion && params && params->add_alloc_dep) { + for (int i = 0; i < graph->n_nodes; ++i) { + ggml_backend_vk_context::ssm_conv_direct_match m; + if (graph->nodes[i]->op == GGML_OP_CONCAT && ggml_vk_ssm_conv_direct_check(ctx, graph, i, m, false)) { + for (const ggml_tensor * input : { m.x, m.state }) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, m.out); + } + } + } + } } static void ggml_backend_vk_event_record(ggml_backend_t backend, ggml_backend_event_t event) { @@ -15403,8 +17792,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; @@ -15436,6 +17830,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_FLASH_ATTN_EXT: { bool coopmat2 = device->coopmat2; + // maskless selected-key attention (src[5] without a mask): only flash_attn_sel reads the list, the dense kernel would attend to all cells + if (op->src[5] != nullptr && op->src[3] == nullptr) { + return ggml_vk_flash_attn_sel_pipeline(device, op) != nullptr; + } uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; if ((HSK % 8) != 0 || (HSV % 8) != 0) { @@ -15453,25 +17851,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; } @@ -15510,8 +17917,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: @@ -15538,6 +17950,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; @@ -15563,6 +17980,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; @@ -15580,6 +18002,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; @@ -16414,6 +18841,20 @@ std::string vk_perf_logger::get_node_fusion_name(const ggml_tensor * node, const 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]; 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_f16_transpose_vt.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose_vt.comp new file mode 100644 index 000000000000..3bafa4146712 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose_vt.comp @@ -0,0 +1,49 @@ +#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[];}; + +// f16 V -> V^T per head for the coopmat1 flash attention (V_TRANSPOSED): source V[nh][kv][hs] with element +// strides p.stride_a (kv) and p.stride_b (head), hs contiguous; dest [nh][hs][kv_pad] with p.K = kv_pad, the +// dest row stride, and p.M = HS. The pad keeps the row stride off a multiple of 4 KB, which would put every +// hs row of a coopmat fragment on the same memory channel. +// One workgroup moves one 64 kv x 64 hs tile through LDS, so reads and writes are both coalesced: +// gl_WorkGroupID = (hs tile, kv tile, head). Host gate: KV % 64 == 0, HS % 64 == 0. +shared uint16_t tile[64][65]; // [kv][hs], padded to spread the transposed reads over the banks + +void main() { + const uint HS = p.M, KVs = p.K; + const uint ths = gl_WorkGroupID.x; + const uint tkv = gl_WorkGroupID.y; + const uint nh = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationID.x; + // read: 64 kv rows x 16 f16vec4 of hs, 16 threads per row, 4 passes + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = i * 256 + tid; + const uint r = idx / 16; + const uint c = idx % 16; + const uint src = (nh * p.stride_b + (tkv * 64 + r) * p.stride_a + ths * 64) / 4 + c; + const f16vec4 v = data_a[src]; + tile[r][4 * c + 0] = float16BitsToUint16(v.x); + tile[r][4 * c + 1] = float16BitsToUint16(v.y); + tile[r][4 * c + 2] = float16BitsToUint16(v.z); + tile[r][4 * c + 3] = float16BitsToUint16(v.w); + } + barrier(); + // write: 64 hs rows x 16 f16vec4 of kv + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = i * 256 + tid; + const uint h = idx / 16; + const uint c = idx % 16; + const f16vec4 v = f16vec4(uint16BitsToFloat16(tile[4 * c + 0][h]), uint16BitsToFloat16(tile[4 * c + 1][h]), + uint16BitsToFloat16(tile[4 * c + 2][h]), uint16BitsToFloat16(tile[4 * c + 3][h])); + const uint dst = (nh * HS * KVs + (ths * 64 + h) * KVs + tkv * 64) / 4 + c; + data_b[dst] = v; + } +} 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..61bb4aaf9922 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl @@ -1360,6 +1360,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 +1638,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/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..f992eb031708 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); } @@ -458,7 +455,8 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { FLOAT_TYPE mvf = masksh[(c * cols_per_iter + col_tid) * masksh_stride + tile_row(r)]; - Sf[r][c] += slope[r]*mvf; + // a masked key scores -inf outright, so non-finite K data cannot turn it into NaN + Sf[r][c] = float(mvf) <= FA_MASK_DEAD ? ACC_TYPE(uintBitsToFloat(0xFF800000)) : Sf[r][c] + slope[r]*mvf; } } } @@ -496,15 +494,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,16 +512,18 @@ 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; } FLOAT_TYPE Pf[rows_per_thread]; + bool dead[rows_per_thread]; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Pf[r] = FLOAT_TYPE(exp(float(Sf[r][c]) - Mf[r])); Lf[r] += Pf[r]; + // masked keys add nothing, even where their V is not finite (0 * NaN) + dead[r] = NAN_SAFE_V && MASK_ENABLE && mask_opt_bits != MASK_OPT_ALL_ZERO && + float(masksh[(c * cols_per_iter + col_tid) * masksh_stride + tile_row(r)]) <= FA_MASK_DEAD; } [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { @@ -532,15 +531,17 @@ 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); + if (!dead[r]) { + Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf); + } } } } @@ -666,10 +667,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 +681,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 +695,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 +703,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 +749,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..ac0125e8a040 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,18 @@ 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; +// Masked keys contribute nothing even where their V rows hold NaN/inf (P is 0, but 0 * NaN is not): set for +// selected-key (sparse) attention, whose selections can name any cell. Costs occupancy in cm1 (V staging). +const bool NAN_SAFE_V = (Flags & 32) != 0; +// V is supplied transposed per head ([HSV][v_stride] with KV contiguous, from dequant_f16_transpose_vt), so a +// coopmat1 P x V B-operand fragment is one contiguous 32-byte run per lane instead of 16 kv-strided halves. +// coopmat1 only, f16 V, aligned, no shmem staging and no NAN_SAFE_V (both stage V rows from the row-major layout). +const bool V_TRANSPOSED = (Flags & 64) != 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; @@ -83,10 +93,11 @@ 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 +// a mask value at or below this (f16 -inf, or the lowest finite f16 some mask writers use for it) hides the key +#define FA_MASK_DEAD (-65504.0) #define MASK_OPT_ALL_ZERO 2 #define BINDING_IDX_K 0 @@ -109,8 +120,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 +139,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 +152,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 +202,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 +223,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..20522d53398a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -13,6 +13,7 @@ #extension GL_KHR_shader_subgroup_basic : enable #extension GL_KHR_shader_subgroup_arithmetic : enable #extension GL_KHR_shader_subgroup_vote : enable +#extension GL_KHR_shader_subgroup_shuffle : enable #extension GL_KHR_memory_scope_semantics : enable #extension GL_KHR_cooperative_matrix : enable @@ -44,8 +45,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; @@ -62,6 +66,11 @@ shared O_TYPEV4 pvsh[MatBc * osh_stride]; shared ACC_TYPE slope[Br]; +// Keys of block j hidden from every row of the tile: dead_stamp[c] == j + 1, and dead_blk == j + 1 when +// block j has any. Stamps, so nothing is reset between blocks; only hidden keys write. +shared uint dead_stamp[Bc]; +shared uint dead_blk; + void main() { #ifdef NEEDS_INIT_IQ_SHMEM if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { @@ -100,6 +109,12 @@ void main() { Qf[r * qstride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); } } + if (NAN_SAFE_V && tid < Bc) { + dead_stamp[tid] = 0; + } + if (NAN_SAFE_V && tid == 0) { + dead_blk = 0; + } barrier(); O_TYPEV4 Of[rows_per_thread][d_per_thread]; @@ -139,9 +154,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; } @@ -152,6 +167,7 @@ void main() { [[dont_unroll]] for (uint32_t j = start_j; j < end_j; ++j) { + bool dead_rows = false; [[unroll]] for (uint32_t idx = 0; idx < mask_cache.length(); ++idx) { mask_cache[idx] = f16vec4(0); @@ -169,23 +185,17 @@ 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); + bool part_dead = false; 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 +223,27 @@ 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); + // rows past the last query (a partial tile) see nothing; their zero mask fill must not + // count as visibility + vec4 mv = vec4(m); + if (!(gqa_ratio > 1)) { + [[unroll]] for (uint32_t k = 0; k < 4; ++k) { + if (i * Br + r * 4 + k >= N) { + mv[k] = FA_MASK_DEAD; + } + } + } + part_dead = NAN_SAFE_V && all(lessThanEqual(mv, vec4(FA_MASK_DEAD))); + } + } + // the Br / 4 == 4 consecutive lanes holding key c agree it is hidden from all Br rows + if (NAN_SAFE_V) { + uint col_dead = part_dead ? 1u : 0u; + col_dead &= subgroupShuffleXor(col_dead, 1); + col_dead &= subgroupShuffleXor(col_dead, 2); + if (col_dead != 0 && r == 0) { + dead_stamp[c] = j + 1; + atomicMax(dead_blk, j + 1); } } } @@ -231,6 +260,7 @@ void main() { if (max_mask <= NEG_FLT_MAX_OVER_2) { continue; } + dead_rows = NAN_SAFE_V && dead_blk == j + 1; } } @@ -240,19 +270,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 +305,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 +313,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]); } } @@ -348,7 +374,9 @@ void main() { // Mask nem1 bounds check is handled when loading masks ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; + // a masked key scores -inf outright, so non-finite K data cannot turn it into NaN + sfsh[c * sfshstride + r] = mix(sfsh[c * sfshstride + r] + slopes * masks, + ACC_TYPEV4(uintBitsToFloat(0xFF800000)), lessThanEqual(vec4(masks), vec4(FA_MASK_DEAD))); } } } @@ -395,15 +423,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 +446,19 @@ 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)) { + // keys hidden from every row stay zero: P is 0 there, but 0 * NaN is not + if ((!KV_bounds_check || j * Bc + c < KV) && (HSV == HSV_pad || d < HSV / 4) && + !(dead_rows && dead_stamp[c] == j + 1)) { #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 +470,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 +495,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 || dead_rows; 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)) && !(dead_rows && dead_stamp[row] == j + 1)) { #if !defined(BFLOAT16) if (USE_DECODE_V) { kvsh[row * vsh_stride + col] = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); @@ -492,14 +530,18 @@ 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 && !dead_rows) { // 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; - coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); + if (V_TRANSPOSED) { + // B[k = kv][n = hd] is column-major in the V^T scratch: each lane reads 16 contiguous kv + const uint vt_tile_offset = v_offset / 4 + hsv_offset * (v_stride / 4) + v_tile_row / 4; + coopMatLoad(QMat, data_vv4, vt_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutColumnMajor); + } else { + const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; + coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); + } } else { const uint v_tile_offset = bc_chunk * MatBr * v_cols + gl_SubgroupID * (MatBc / 4); coopMatLoad(QMat, kvsh, v_tile_offset, vsh_stride, gl_CooperativeMatrixLayoutRowMajor); @@ -509,7 +551,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 +591,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 +607,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 +620,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 +630,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 +679,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_cm1_mr.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1_mr.comp new file mode 100644 index 000000000000..6c74113d4958 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1_mr.comp @@ -0,0 +1,397 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require + +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_shader_subgroup_clustered : enable +#extension GL_KHR_shader_subgroup_vote : enable +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_cooperative_matrix : enable + +// Multi-row coopmat1 prefill flash attention (GGML_VK_FA_MR), after the gufo engine's attention_wmma kernel. +// +// The cm1 kernel pins Br at 16 query rows and splits Bc over its subgroups, so every 16 rows stream +// the whole K/V range and each subgroup needs every Q fragment (re-read from LDS per KV block). +// Here one workgroup of eight wave32 subgroups owns 64 rows = H heads x (64 / H) tokens (the heads +// share one KV head under GQA), and the work is split the other way round: +// S = Q K^T per 16-key tile: subgroup s owns row block s % 4 and half s / 4 of the head dim, +// so its Q fragments (HSK / 32 of them) stay in registers for the whole KV loop. The two halves' +// partial S tiles meet in LDS. +// softmax: four threads per row, running max / sum in registers, P (f16) and the per-row rescale +// factor to LDS. +// O += P V: subgroup s owns head-dim tiles s, s + 8, ... for all four row blocks, so each V^T +// fragment is loaded once per subgroup and reused four times. +// K fragments are read directly from per-head [KV][HSK] f16 rows (the contiguized scratch for a KV cache view) +// and V fragments from the V^T scratch, so neither goes through LDS. +// The accumulator rescale multiplies by a row-broadcast matrix loaded from LDS with a zero column +// stride, so no accumulator element-to-row layout is assumed. +// +// Spec constants reuse the cm1 slots: Br = tokens per workgroup (64 / H), Bc = 64 (the mask_opt +// block), row_split = H. Host gates: f16 K (contiguous rows) + V^T, HSK == HSV in {128, 256}, +// KV % 64 == 0, f32 accumulation, no split_k, mask shared across heads. + +#include "types.glsl" +#include "flash_attn_base.glsl" + +const uint32_t NSG = 8; // wave32 subgroups +const uint32_t ROWS = 64; // rows per workgroup +const uint32_t RB = 4; // 16-row blocks +const uint32_t KT = 16; // keys per inner tile +const uint32_t H = row_split; // heads per workgroup +const uint32_t TOKS = Br; // tokens per workgroup (ROWS / H) +const uint32_t QF = HSK / 32; // Q fragments per subgroup (half the head dim) +const uint32_t OT = HSV / (16 * NSG); // O dim tiles per subgroup + +layout (binding = 0) readonly buffer Q {float data_q[];}; +layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];}; +layout (binding = 1) readonly buffer K {float16_t data_k[];}; +layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];}; +layout (binding = 2) readonly buffer V {float16_t data_v[];}; +layout (binding = 2) readonly buffer VV4 {f16vec4 data_vv4[];}; +layout (binding = 3) readonly buffer M {float16_t data_m[];}; + +// one 16 x 16 f16 Q tile per subgroup at the start, then the P tile [ROWS][KT] (f16vec4 units) +const uint32_t pstride = KT / 4 + 2; +const uint32_t stage_len = (NSG * 64 > ROWS * pstride) ? NSG * 64 : ROWS * pstride; +shared f16vec4 stage_sh[stage_len]; + +// partial S tiles of the two head-dim halves: [2][ROWS][sstride] f32; reused for the partial-block output +const uint32_t sstride = KT + 4; +const uint32_t s_len = (2 * ROWS * sstride > NSG * 256) ? 2 * ROWS * sstride : NSG * 256; +shared float s_sh[s_len]; + +// per-row rescale factor of the current tile (then the final 1/L factor) +shared float scale_sh[ROWS]; +// MR_LAZY: per softmax subgroup, whether any of its eight rows moved its running max this tile +shared uint32_t resc_sh[NSG]; + +// bit k of vis_sh[tile parity]: key k of the 16-key tile is unmasked in at least one row that reaches the output. +// RDNA3 WMMA rounds a dot product toward -inf when any product term is -0.0, and a key masked in every row has +// P = +0 there, so a negative V in that cell (a discarded prompt branch, say) moved the output by an ulp. Tiles +// with such a key restage their V^T tiles through vstage_sh with those keys zeroed (one 16 x 16 tile per subgroup). +shared uint32_t vis_sh[2]; +shared f16vec4 vstage_sh[NSG * 16 * (KT / 4)]; + +// skip the O rescale of a row block whose rows all kept their running max (eM == 1 exactly, so the +// skipped multiply is an identity and the result is bit-identical); at depth the max rarely moves +const bool MR_LAZY = (Flags & 256) != 0; +// drop the loop-end barrier (hazard argument at the loop end) +const bool MR_NOEB = (Flags & 512) != 0; +// issue the tile's V^T loads right after the S store, before the softmax barrier +const bool MR_VPOST = (Flags & 1024) != 0; + + +// 2-bit mask_opt state of 64-key block j for this workgroup's tokens +uint32_t block_bits(const uint32_t j, const uint32_t mo_offset) { + return (data_mask_opt[mo_offset + j / 16] >> ((j % 16) * 2)) & 0x3; +} + +// first block after j (j = ~0: the first block) that the mask does not skip entirely; end_j if none +uint32_t next_block(uint32_t j, const uint32_t end_j, const uint32_t mo_offset) { + j = j + 1; + if (MASK_ENABLE && USE_MASK_OPT) { + while (j < end_j && block_bits(j, mo_offset) == MASK_OPT_ALL_NEG_INF) { + j++; + } + } + return j; +} + +void main() { + const uint32_t tid = gl_LocalInvocationIndex; + const uint32_t sg = gl_SubgroupID; + + const uint32_t tok0 = gl_WorkGroupID.x * TOKS; + const uint32_t head0 = gl_WorkGroupID.y * H; + const uint32_t iq3 = gl_WorkGroupID.z; + const uint32_t N = p.N; + const uint32_t KV = p.KV; + + const uint32_t rk2 = p.neq2 / p.nek2, rk3 = p.neq3 / p.nek3; + const uint32_t rv2 = p.neq2 / p.nev2, rv3 = p.neq3 / p.nev3; + const uint32_t k_offset = ((head0 / rk2) * p.nb12 + (iq3 / rk3) * p.nb13) / 2; // f16 elements + const uint32_t v_offset = ((head0 / rv2) * p.nb22 + (iq3 / rv3) * p.nb23) / 2; + const uint32_t k_stride = p.nb11; // elements (HSK for the contiguized scratch) + const uint32_t v_stride = p.nb21; // V^T row stride (KV + pad) + + // row r -> (row block r / 16, head (r / 16) % H, token tok0 + ((r / 16) / H) * 16 + r % 16) +#define ROW_HEAD(r) (head0 + ((r) / 16) % H) +#define ROW_TOK(r) (tok0 + (((r) / 16) / H) * 16 + (r) % 16) + + const uint32_t s_rb = sg % RB; // S phase: row block + const uint32_t s_half = sg / RB; // S phase: head-dim half + + // ---- Q -> registers ---- + // Each subgroup converts its 16 x 16 Q tiles (zero rows past N) into a private LDS slice and loads the + // fragment back with one unconditional coopMatLoad: a fragment converted in registers stays sixteen 16-bit + // values that ACO repacks before every WMMA (130 v_mov_b16 per tile in the loop), and a conditional load + // leaves a phi that it copies at the loop header. + coopmat QA[QF]; + { + const uint32_t q_tok0 = ROW_TOK(s_rb * 16); + const uint32_t q_off = (ROW_HEAD(s_rb * 16) * p.nb02 + iq3 * p.nb03) / 4; + const uint32_t lane = gl_SubgroupInvocationID; + [[unroll]] for (uint32_t f = 0; f < QF; ++f) { + // lane owns (row lane / 2, dims (lane % 2) * 8 .. + 7) of the tile: two vec4 reads, two f16vec4 writes + const uint32_t r = lane / 2; + const uint32_t c4 = (lane % 2) * 2; + const uint32_t d0 = (s_half * QF + f) * 16; + vec4 q0 = vec4(0.0), q1 = vec4(0.0); + if (q_tok0 + r < N) { + const uint32_t o = (q_off + (q_tok0 + r) * p.nb01 + d0) / 4 + c4; + q0 = data_qv4[o] * p.scale; + q1 = data_qv4[o + 1] * p.scale; + } + stage_sh[sg * 64 + r * 4 + c4] = f16vec4(q0); + stage_sh[sg * 64 + r * 4 + c4 + 1] = f16vec4(q1); + subgroupMemoryBarrierShared(); + subgroupBarrier(); + coopMatLoad(QA[f], stage_sh, sg * 64, 4, gl_CooperativeMatrixLayoutRowMajor); + subgroupBarrier(); + } + } + barrier(); + + // ---- softmax thread state: four threads per row ---- + const uint32_t sm_row = tid / 4; + const uint32_t sm_seg = tid % 4; + const uint32_t sm_tok = ROW_TOK(sm_row); + const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF); + float Mf = NEG_FLT_MAX_OVER_2; + float Lf = 0.0; + + float slope = 1.0; + if (p.max_bias > 0.0f) { + const uint32_t h = ROW_HEAD(sm_row); + const uint32_t n_head_log2 = p.mask_n_head_log2 & N_LOG2_MASK; + const float base = h < n_head_log2 ? p.m0 : p.m1; + const int exph = int(h < n_head_log2 ? h + 1 : 2 * (h - n_head_log2) + 1); + slope = pow(base, float(exph)); + } + + // mask: shared by all heads (host gate nem2 == nem3 == 1); rows past nem1 read as zero + const uint32_t m_row_off = sm_tok * KV; + const bool m_row_ok = sm_tok < p.nem1; + + const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc); + const uint32_t mo_offset = mo_stride * gl_WorkGroupID.x; + + coopmat OAcc[RB][OT]; + [[unroll]] for (uint32_t rb = 0; rb < RB; ++rb) { + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + OAcc[rb][t] = coopmat(0.0); + } + } + + const uint32_t end_j = KV / Bc; + + uint32_t j = next_block(0xFFFFFFFFu, end_j, mo_offset); + uint32_t key0 = j * Bc; + bool read_mask = false; + + if (tid < 2) { + vis_sh[tid] = 0u; + } + barrier(); + // vis_sh[tp & 1] is written in this tile's softmax and read in its PV; vis_sh[(tp + 1) & 1] was last read in the + // previous tile's PV, which every subgroup finished before this tile's S barrier, so the softmax may clear it + uint32_t tp = 0; + + [[dont_unroll]] + while (j < end_j) { + const uint32_t kt = (key0 / KT) % (Bc / KT); + if (kt == 0) { + read_mask = MASK_ENABLE && (!USE_MASK_OPT || block_bits(j, mo_offset) != MASK_OPT_ALL_ZERO); + } + // the next tile: the rest of this block, then the next block the mask does not skip + const uint32_t nj = (kt + 1 < Bc / KT) ? j : next_block(j, end_j, mo_offset); + const uint32_t nkey0 = (kt + 1 < Bc / KT) ? key0 + KT : nj * Bc; + + coopmat VB[OT]; + { + // ---- S = Q K^T (this subgroup: row block s_rb, head-dim half s_half) ---- + coopmat SAcc[2]; + SAcc[0] = coopmat(0.0); + SAcc[1] = coopmat(0.0); + [[unroll]] for (uint32_t f = 0; f < QF; ++f) { + coopmat KB; + // B[k = dim][n = key] at key * k_stride + dim: column-major, one 32-byte run per lane + const uint32_t coord = (k_offset + key0 * k_stride + (s_half * QF + f) * 16) / 4; + coopMatLoad(KB, data_kv4, coord, k_stride / 4, gl_CooperativeMatrixLayoutColumnMajor); + SAcc[f % 2] = coopMatMulAdd(QA[f], KB, SAcc[f % 2]); + } + SAcc[0] = SAcc[0] + SAcc[1]; + coopMatStore(SAcc[0], s_sh, (s_half * ROWS + s_rb * 16) * sstride, sstride, gl_CooperativeMatrixLayoutRowMajor); + if (MR_VPOST) { + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + const uint32_t coord = (v_offset + ((t * NSG + sg) * 16) * v_stride + key0) / 4; + coopMatLoad(VB[t], data_vv4, coord, v_stride / 4, gl_CooperativeMatrixLayoutColumnMajor); + } + } + barrier(); + + // ---- online softmax ---- + { + vec4 s = vec4(s_sh[sm_row * sstride + sm_seg * 4 + 0], s_sh[sm_row * sstride + sm_seg * 4 + 1], + s_sh[sm_row * sstride + sm_seg * 4 + 2], s_sh[sm_row * sstride + sm_seg * 4 + 3]) + + vec4(s_sh[(ROWS + sm_row) * sstride + sm_seg * 4 + 0], s_sh[(ROWS + sm_row) * sstride + sm_seg * 4 + 1], + s_sh[(ROWS + sm_row) * sstride + sm_seg * 4 + 2], s_sh[(ROWS + sm_row) * sstride + sm_seg * 4 + 3]); + if (LOGIT_SOFTCAP) { + s = p.logit_softcap * tanh(s); + } + if (read_mask && m_row_ok) { + const uint32_t mo = m_row_off + key0 + sm_seg * 4; + const vec4 m = vec4(data_m[mo], data_m[mo + 1], data_m[mo + 2], data_m[mo + 3]); + // a masked key scores -inf outright, so non-finite K data cannot turn it into NaN + s = mix(s + slope * m, vec4(uintBitsToFloat(0xFF800000)), lessThanEqual(m, vec4(FA_MASK_DEAD))); + } + if (read_mask) { + // rows past N never reach the output; rows past nem1 read no mask, so every key is visible there + uint32_t bits = 0u; + if (sm_tok < N) { + [[unroll]] for (uint32_t k = 0; k < 4; ++k) { + const bool vis = !m_row_ok || float(data_m[m_row_off + key0 + sm_seg * 4 + k]) > NEG_FLT_MAX_OVER_2; + bits |= vis ? 1u << (sm_seg * 4 + k) : 0u; + } + } + bits = subgroupOr(bits); + if (subgroupElect()) { + atomicOr(vis_sh[tp & 1], bits); + } + } + if (tid == 0) { + vis_sh[(tp + 1) & 1] = 0u; + } + float rowmax = max(max(s.x, s.y), max(s.z, s.w)); + rowmax = subgroupClusteredMax(rowmax, 4) + FATTN_KQ_MAX_OFFSET; + const float Mold = Mf; + Mf = max(rowmax, Mold); + const float eM = exp(Mold - Mf); + const vec4 P = exp(s - Mf); + Lf = eM * Lf + (P.x + P.y + P.z + P.w); + stage_sh[sm_row * pstride + sm_seg] = f16vec4(P); + if (sm_seg == 0) { + scale_sh[sm_row] = eM; + } + if (MR_LAZY) { + const bool moved = subgroupAny(eM != 1.0); + if (subgroupElect()) { + resc_sh[sg] = moved ? 1u : 0u; + } + } + } + barrier(); + + // ---- O = diag(eM) O + P V ---- + if (!MR_VPOST) { + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + // V^T [HSV][v_stride]: B[k = key][n = dim] at dim * v_stride + key, column-major + const uint32_t coord = (v_offset + ((t * NSG + sg) * 16) * v_stride + key0) / 4; + coopMatLoad(VB[t], data_vv4, coord, v_stride / 4, gl_CooperativeMatrixLayoutColumnMajor); + } + } + if (read_mask && vis_sh[tp & 1] != 0xFFFFu) { + const uint32_t vis = vis_sh[tp & 1]; + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + for (uint32_t e = gl_SubgroupInvocationID; e < 16 * (KT / 4); e += gl_SubgroupSize) { + const uint32_t hd = e / (KT / 4); + const uint32_t kq = e % (KT / 4); + f16vec4 vt = data_vv4[(v_offset + ((t * NSG + sg) * 16 + hd) * v_stride + key0) / 4 + kq]; + [[unroll]] for (uint32_t k = 0; k < 4; ++k) { + if (((vis >> (kq * 4 + k)) & 1u) == 0u) { + vt[k] = float16_t(0.0); + } + } + vstage_sh[sg * 16 * (KT / 4) + e] = vt; + } + subgroupMemoryBarrierShared(); + subgroupBarrier(); + coopMatLoad(VB[t], vstage_sh, sg * 16 * (KT / 4), KT / 4, gl_CooperativeMatrixLayoutColumnMajor); + subgroupBarrier(); + } + } + [[unroll]] for (uint32_t rb = 0; rb < RB; ++rb) { + coopmat PA; + coopMatLoad(PA, stage_sh, rb * 16 * pstride, pstride, gl_CooperativeMatrixLayoutRowMajor); + // softmax rows of row block rb live in subgroups 2 rb and 2 rb + 1 (four threads per row) + if (!MR_LAZY || (resc_sh[2 * rb] | resc_sh[2 * rb + 1]) != 0) { + // element (r, c) = scale_sh[rb * 16 + r]: column-major with a zero column stride + coopmat SC; + coopMatLoad(SC, scale_sh, rb * 16, 0, gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + OAcc[rb][t] = OAcc[rb][t] * SC; + } + } + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + OAcc[rb][t] = coopMatMulAdd(PA, VB[t], OAcc[rb][t]); + } + } + // the next tile's S store and softmax overwrite s_sh / stage_sh / scale_sh. MR_NOEB: neither needs + // this barrier. A subgroup at the next S store has passed this tile's softmax barrier, so every + // softmax read of s_sh is done; the next softmax writes P / scale_sh / resc_sh only after the next + // S barrier, which every subgroup reaches after finishing this PV. + if (!MR_NOEB) { + barrier(); + } + } + j = nj; + key0 = nkey0; + ++tp; + } + if (MR_NOEB) { + barrier(); // the last PV may still read scale_sh + } + + // ---- final per-row factor: sinks, 1 / L ---- + { + float L = subgroupClusteredAdd(Lf, 4); + float ms = 1.0; + if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) { + const float sink = data_s[ROW_HEAD(sm_row)]; + float vs = 1.0; + if (sink > Mf) { + ms = exp(Mf - sink); + } else { + vs = exp(sink - Mf); + } + L = L * ms + vs; + } + if (sm_seg == 0) { + scale_sh[sm_row] = (L == 0.0) ? 0.0 : ms / L; + } + } + barrier(); + + const uint32_t o_base = iq3 * p.ne2 * p.ne1 * HSV; + [[unroll]] for (uint32_t rb = 0; rb < RB; ++rb) { + coopmat SC; + coopMatLoad(SC, scale_sh, rb * 16, 0, gl_CooperativeMatrixLayoutColumnMajor); + const uint32_t tokb = ROW_TOK(rb * 16); + const uint32_t h = ROW_HEAD(rb * 16); + [[unroll]] for (uint32_t t = 0; t < OT; ++t) { + OAcc[rb][t] = OAcc[rb][t] * SC; + const uint32_t d0 = (t * NSG + sg) * 16; + if (tokb + 16 <= N) { + // dst [HSV][ne1 heads][ne2 tokens]: row = token (stride ne1 * HSV), col = dim + coopMatStore(OAcc[rb][t], data_o, o_base + (tokb * p.ne1 + h) * HSV + d0, p.ne1 * HSV, gl_CooperativeMatrixLayoutRowMajor); + } else if (tokb < N) { + // partial row block: stage through this subgroup's private LDS slice and store the valid rows + coopMatStore(OAcc[rb][t], s_sh, sg * 256, 16, gl_CooperativeMatrixLayoutRowMajor); + subgroupMemoryBarrierShared(); + subgroupBarrier(); + [[unroll]] for (uint32_t e = gl_SubgroupInvocationID; e < 256; e += 32) { + const uint32_t r = e / 16, c = e % 16; + if (tokb + r < N) { + data_o[o_base + ((tokb + r) * p.ne1 + h) * HSV + d0 + c] = D_TYPE(s_sh[sg * 256 + e]); + } + } + subgroupBarrier(); + } + } + } +} 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_sel.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp new file mode 100644 index 000000000000..d7af1a482009 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel.comp @@ -0,0 +1,176 @@ +#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 +#extension GL_KHR_shader_subgroup_shuffle : require + +// Maskless selected-key flash attention (QSA complete-block selection), portable scalar kernel: each query's row names +// the cells it may attend to. flash_attn_sel_prep.comp has already cleaned every row (invalid entries and repeated +// cells dropped, survivors first in their original order, count at the end of the buffer). One workgroup per (query, +// K/V head, stream), one subgroup per query head of that K/V head, so each K/V row is read once per query. +// Keys go in blocks of BK: two lanes compute one key's dot product (half the dims each), so the softmax needs two +// subgroup reductions per block instead of one per key. With n_split > 1 (few queries) each workgroup takes one +// chunk of the cleaned row and writes a partial result in the layout of flash_attn_split_k_reduce.comp. +// flash_attn_sel_cm.comp is the fast path on RDNA3; this kernel runs everywhere else. + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 768; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(constant_id = 2) const uint GQA = 12; +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 VBuf { float16_t data_v[]; }; +layout(binding = 3) readonly buffer ListBuf { uint data_list[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_sel; + uint list_stride; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk2; + uint nbk3; + uint nbv1; + uint nbv2; + uint nbv3; + uint nbi1; + uint nbi3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint n_split; + uint chunk; + uint n_head; + uint n_tok; + uint n_stream; +} p; + +// The pipeline is created with subgroup size 64 and head size 256 only (see ggml_vk_flash_attn_sel_pipeline). +const uint HEAD_SIZE = 256; +const uint LANES = 64; +const uint PER_LANE = HEAD_SIZE / LANES; +const uint BK = 32; +const uint HALF = HEAD_SIZE / 2; +const uint NONE = 0xffffffffu; +const float FLT_MAX = 3.402823466e38; + +shared uint block_key[BK]; +shared float q_sh[GQA * HEAD_SIZE]; +shared float16_t kt_sh[HEAD_SIZE * BK]; // K block, transposed: [dim][key], so the per-lane key reads hit different banks +shared float16_t v_sh[BK * HEAD_SIZE]; // V block: [key][dim] +shared float p_sh[GQA * BK]; + +void main() { + const float NEG_INF = uintBitsToFloat(0xff800000); + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint sg = gl_SubgroupID; + const uint split = gl_WorkGroupID.x % p.n_split; + const uint query = gl_WorkGroupID.x / p.n_split; + const uint kv_head = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + const uint head = kv_head * GQA + sg; + + // No early return: every invocation must reach every barrier below. + + const uint row = stream * p.n_tok + query; + const uint list_base = row * p.list_stride; + const uint count = data_list[p.n_tok * p.n_stream * p.list_stride + row]; + const uint q_base = stream * p.nbq3 + head * p.nbq2 + query * p.nbq1; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + q_sh[sg * HEAD_SIZE + lane + i * LANES] = data_q[q_base + lane + i * LANES]; + } + barrier(); + + float acc[PER_LANE]; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] = 0.0; + } + float row_max = NEG_INF; + float row_sum = 0.0; + + const uint k_base = stream * p.nbk3 + kv_head * p.nbk2; + const uint v_base = stream * p.nbv3 + kv_head * p.nbv2; + const uint kcol = lane % BK; // the key this lane scores + const uint khalf = lane / BK; // and which half of its dims + + const uint kb_begin = split * p.chunk; + const uint kb_end = min(count, kb_begin + p.chunk); + for (uint kb = kb_begin; kb < kb_end; kb += BK) { + if (tid < BK) { + const uint s = kb + tid; + block_key[tid] = s < kb_end ? data_list[list_base + s] : NONE; + } + barrier(); + + for (uint idx = tid; idx < BK * HEAD_SIZE; idx += WORKGROUP_SIZE) { + const uint col = idx / HEAD_SIZE; + const uint dim = idx % HEAD_SIZE; + const uint key = block_key[col]; + kt_sh[dim * BK + col] = key != NONE ? data_k[k_base + key * p.nbk1 + dim] : float16_t(0.0); + v_sh[idx] = key != NONE ? data_v[v_base + key * p.nbv1 + dim] : float16_t(0.0); + } + barrier(); + + float s = 0.0; + const uint qoff = sg * HEAD_SIZE + khalf * HALF; + const uint koff = khalf * HALF * BK + kcol; + for (uint d = 0; d < HALF; ++d) { + s += q_sh[qoff + d] * float(kt_sh[koff + d * BK]); + } + s += subgroupShuffleXor(s, BK); + const bool valid = block_key[kcol] != NONE; + s = valid ? s * p.scale : NEG_INF; + + const float new_max = max(row_max, subgroupMax(s)); + const float old_scale = new_max == row_max ? 1.0 : exp(row_max - new_max); + const float weight = valid ? exp(s - new_max) : 0.0; + row_sum = row_sum * old_scale + subgroupAdd(khalf == 0 ? weight : 0.0); + row_max = new_max; + if (khalf == 0) { + p_sh[sg * BK + kcol] = weight; + } + subgroupMemoryBarrierShared(); + subgroupBarrier(); + + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] *= old_scale; + } + for (uint j = 0; j < BK; ++j) { + const float w = p_sh[sg * BK + j]; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + acc[i] += w * float(v_sh[j * HEAD_SIZE + lane + i * LANES]); + } + } + barrier(); + } + + if (p.n_split == 1) { + const uint dst_base = stream * p.nb3 + query * 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 < PER_LANE; ++i) { + data_dst[dst_base + lane + i * LANES] = acc[i] * inv_sum; + } + } else { + // partial: unnormalized O, then L and M per head (an empty split writes M = -FLT_MAX, not -inf, so the + // reduce never computes exp(-inf - -inf)) + const uint slot = split + p.n_split * (query + p.n_tok * stream); + const uint o_base = HEAD_SIZE * p.n_head * slot + HEAD_SIZE * head; + [[unroll]] for (uint i = 0; i < PER_LANE; ++i) { + data_dst[o_base + lane + i * LANES] = acc[i]; + } + if (lane == 0) { + const uint lm = HEAD_SIZE * p.n_head * p.n_tok * p.n_stream * p.n_split + p.n_head * 2 * slot + head; + data_dst[lm] = row_sum; + data_dst[lm + p.n_head] = max(row_max, -FLT_MAX); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_cm.comp new file mode 100644 index 000000000000..25e59779dfa8 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_cm.comp @@ -0,0 +1,269 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_ballot : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_cooperative_matrix : require + +// Maskless selected-key flash attention (QSA complete-block selection), RDNA3 coopmat kernel. Same op and same +// cleaned selection rows as flash_attn_sel.comp (see flash_attn_sel_prep.comp); WMMA instead of scalar FMAs. +// +// One 64-lane subgroup per (query, K/V head, stream, split): the GQA query heads of that K/V head share the query's +// selection, so they are the 16 columns of every tile (12 used at GQA 12). Keys go in steps of 16: +// S^T (16 keys x 16 heads) = K (16 keys x 16 dims) * Q^T (16 dims x 16 heads), 16 dim chunks +// O^T (16 dims x 16 heads) += V^T (16 dims x 16 keys) * P^T (16 keys x 16 heads), 16 dim tiles +// so every accumulator holds one head per lane column, and the online softmax state (max, sum, rescale) of a head +// lives in the lanes of that column: no shared memory and no barrier in the key loop. +// +// The K and V fragments are built in registers straight from the gathered cache rows. That relies on the gfx11 WMMA +// fragment layout as RADV lowers KHR cooperative matrices (A: lane L holds row L % 16, element i = column i; B: lane +// L holds column L % 16, element i = row i; f32 accumulator: lane L holds column L % 16, rows L / 16 + 4 i), which +// the spec does not guarantee, so the host creates this pipeline only for RADV on RDNA3 and uses +// flash_attn_sel.comp elsewhere. +// K: lane L reads its key's 16 contiguous dims of chunk c with two 16-byte loads. +// V: one 16-byte load gives 8 contiguous dims of one key, i.e. element e of 8 tiles. Tile t, row r is therefore +// dim (t / 8) * 128 + 8 r + t % 8; the output store undoes this permutation. +// The rows are cleaned and padded by the pre-pass: a column at or past the row's count reads a live row and gets +// score -inf (P exactly 0), so no load needs a bounds test. + +layout(constant_id = 0) const uint GQA = 12; +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { uvec4 data_k4[]; }; +layout(binding = 2) readonly buffer VBuf { uvec4 data_v4[]; }; +layout(binding = 3) readonly buffer ListBuf { uint data_list[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_sel; + uint list_stride; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk2; + uint nbk3; + uint nbv1; + uint nbv2; + uint nbv3; + uint nbi1; + uint nbi3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint n_split; + uint chunk; + uint n_head; + uint n_tok; + uint n_stream; +} p; + +const uint D = 256; // host gate: head size 256 +const uint QS = D / 8 + 1; // Q row stride in uvec4, padded against bank conflicts +const uint STEP = 16; +const float FLT_MAX = 3.402823466e38; + +shared uvec4 q_sh[16 * QS]; // Q * scale in f16, [head][dim], heads >= GQA zero + +#define ACC coopmat +#define MAT_A coopmat +#define MAT_B coopmat + +// 16 contiguous f16 (two uvec4) into elements 0..15 +#define FILL16(M, a, b) { \ + f16vec2 t_; \ + t_ = unpackFloat2x16(a.x); M[ 0] = t_.x; M[ 1] = t_.y; \ + t_ = unpackFloat2x16(a.y); M[ 2] = t_.x; M[ 3] = t_.y; \ + t_ = unpackFloat2x16(a.z); M[ 4] = t_.x; M[ 5] = t_.y; \ + t_ = unpackFloat2x16(a.w); M[ 6] = t_.x; M[ 7] = t_.y; \ + t_ = unpackFloat2x16(b.x); M[ 8] = t_.x; M[ 9] = t_.y; \ + t_ = unpackFloat2x16(b.y); M[10] = t_.x; M[11] = t_.y; \ + t_ = unpackFloat2x16(b.z); M[12] = t_.x; M[13] = t_.y; \ + t_ = unpackFloat2x16(b.w); M[14] = t_.x; M[15] = t_.y; } + +void main() { + const float NEG_INF = uintBitsToFloat(0xff800000); + const uint lane = gl_SubgroupInvocationID; + const uint col = lane & 15; // fragment lane: K row (key), Q^T / P^T / accumulator column (head), V^T row (dim) + const uint grp = lane >> 4; // accumulator rows grp + 4 i + // the K/V heads of a query are neighbours in the grid, so they run together: they read the same cells, which + // sit next to each other in an interleaved cache + const uint n_kv_head = p.n_head / GQA; + const uint kv_head = gl_WorkGroupID.x % n_kv_head; + const uint split = (gl_WorkGroupID.x / n_kv_head) % p.n_split; + const uint query = gl_WorkGroupID.x / (n_kv_head * p.n_split); + const uint stream = gl_WorkGroupID.z; + + // Q * scale -> f16 in shared memory + [[unroll]] for (uint i = 0; i < 16 * (D / 8) / 64; ++i) { + const uint idx = lane + 64 * i; + const uint h = idx / (D / 8); + const uint c8 = idx % (D / 8); + uvec4 w = uvec4(0); + if (h < GQA) { + const uint qb = stream * p.nbq3 + (kv_head * GQA + h) * p.nbq2 + query * p.nbq1 + c8 * 8; + w.x = packFloat2x16(f16vec2(data_q[qb + 0] * p.scale, data_q[qb + 1] * p.scale)); + w.y = packFloat2x16(f16vec2(data_q[qb + 2] * p.scale, data_q[qb + 3] * p.scale)); + w.z = packFloat2x16(f16vec2(data_q[qb + 4] * p.scale, data_q[qb + 5] * p.scale)); + w.w = packFloat2x16(f16vec2(data_q[qb + 6] * p.scale, data_q[qb + 7] * p.scale)); + } + q_sh[h * QS + c8] = w; + } + barrier(); + + const uint row = stream * p.n_tok + query; + const uint list_base = row * p.list_stride; + const uint count = data_list[p.n_tok * p.n_stream * p.list_stride + row]; + const uint kb_begin = split * p.chunk; + const uint kb_end = min(count, kb_begin + p.chunk); + + // host gate: K/V offsets and strides are multiples of 8 halves (16 bytes) + const uint k_base4 = (stream * p.nbk3 + kv_head * p.nbk2) / 8; + const uint v_base4 = (stream * p.nbv3 + kv_head * p.nbv2) / 8; + const uint nbk1_4 = p.nbk1 / 8; + const uint nbv1_4 = p.nbv1 / 8; + + ACC O[16]; + [[unroll]] for (uint t = 0; t < 16; ++t) { + O[t] = ACC(0.0); + } + float m = NEG_INF; // running max and sum of head col (identical in the 4 lanes of a column) + float l = 0.0; + + uint next = kb_begin < kb_end ? data_list[list_base + kb_begin + col] : 0; + for (uint kb = kb_begin; kb < kb_end; kb += STEP) { + const uint cur = next; // the cache row of key kb + col + if (kb + STEP < kb_end) { + next = data_list[list_base + kb + STEP + col]; + } + + ACC S = ACC(0.0); + const uint kr = k_base4 + cur * nbk1_4; + [[unroll]] for (uint c = 0; c < D / 16; ++c) { + const uvec4 ka = data_k4[kr + 2 * c]; + const uvec4 kb2 = data_k4[kr + 2 * c + 1]; + const uvec4 qa = q_sh[col * QS + 2 * c]; + const uvec4 qb = q_sh[col * QS + 2 * c + 1]; + MAT_A KA; + MAT_B QB; + FILL16(KA, ka, kb2); + FILL16(QB, qa, qb); + S = coopMatMulAdd(KA, QB, S); + } + + // online softmax for head col over keys kb + grp + 4 i + float s[4]; + float mx = NEG_INF; + [[unroll]] for (uint i = 0; i < 4; ++i) { + s[i] = kb + grp + 4 * i < kb_end ? S[i] : NEG_INF; + mx = max(mx, s[i]); + } + mx = max(mx, subgroupShuffleXor(mx, 16)); + mx = max(mx, subgroupShuffleXor(mx, 32)); + const float m_new = max(m, mx); // finite: key kb is always live + const float corr = exp(m - m_new); + float pv[4]; + float ps = 0.0; + [[unroll]] for (uint i = 0; i < 4; ++i) { + pv[i] = exp(s[i] - m_new); + ps += pv[i]; + } + ps += subgroupShuffleXor(ps, 16); + ps += subgroupShuffleXor(ps, 32); + l = l * corr + ps; + m = m_new; + + // P^T as a B fragment: element e = key e of the step. Key e = g + 4 i sits in pv[i] of the lanes of group g, + // and a register of the fragment packs keys (2k, 2k+1) = groups (2(k & 1), 2(k & 1) + 1), i = k / 2. + // Pair the groups with a xor-16 exchange, then trade halves with a xor-32 exchange. + const uint own01 = packFloat2x16(f16vec2(pv[0], pv[1])); + const uint own23 = packFloat2x16(f16vec2(pv[2], pv[3])); + const uint par01 = subgroupShuffleXor(own01, 16); + const uint par23 = subgroupShuffleXor(own23, 16); + const bool even = (grp & 1) == 0; + const f16vec2 o01 = unpackFloat2x16(own01), o23 = unpackFloat2x16(own23); + const f16vec2 q01 = unpackFloat2x16(par01), q23 = unpackFloat2x16(par23); + // w[i] = keys (2(grp/2) + 4 i, +1): register k = grp / 2 + 2 i of the fragment + uint w[4]; + w[0] = even ? packFloat2x16(f16vec2(o01.x, q01.x)) : packFloat2x16(f16vec2(q01.x, o01.x)); + w[1] = even ? packFloat2x16(f16vec2(o01.y, q01.y)) : packFloat2x16(f16vec2(q01.y, o01.y)); + w[2] = even ? packFloat2x16(f16vec2(o23.x, q23.x)) : packFloat2x16(f16vec2(q23.x, o23.x)); + w[3] = even ? packFloat2x16(f16vec2(o23.y, q23.y)) : packFloat2x16(f16vec2(q23.y, o23.y)); + uint x[4]; + [[unroll]] for (uint i = 0; i < 4; ++i) { + x[i] = subgroupShuffleXor(w[i], 32); + } + const uint base = grp >> 1; + MAT_B PB; + [[unroll]] for (uint k = 0; k < 8; ++k) { + const f16vec2 t = unpackFloat2x16((k & 1) == base ? w[k >> 1] : x[k >> 1]); + PB[2 * k] = t.x; + PB[2 * k + 1] = t.y; + } + + // once the running maxima settle, most steps leave every head's max unchanged: skip the rescale (uniform) + if (subgroupBallot(corr != 1.0) != uvec4(0)) { + [[unroll]] for (uint t = 0; t < 16; ++t) { + [[unroll]] for (uint i = 0; i < O[t].length(); ++i) { + O[t][i] *= corr; + } + } + } + + // V^T fragments, 8 tiles per half of the head + [[unroll]] for (uint hh = 0; hh < 2; ++hh) { + MAT_A VA[8]; + [[unroll]] for (uint e = 0; e < STEP; ++e) { + const uint vr = subgroupBroadcast(cur, e); // lane e holds key e's row + const uvec4 v = data_v4[v_base4 + vr * nbv1_4 + hh * 16 + col]; + f16vec2 t; + t = unpackFloat2x16(v.x); VA[0][e] = t.x; VA[1][e] = t.y; + t = unpackFloat2x16(v.y); VA[2][e] = t.x; VA[3][e] = t.y; + t = unpackFloat2x16(v.z); VA[4][e] = t.x; VA[5][e] = t.y; + t = unpackFloat2x16(v.w); VA[6][e] = t.x; VA[7][e] = t.y; + } + [[unroll]] for (uint j = 0; j < 8; ++j) { + O[hh * 8 + j] = coopMatMulAdd(VA[j], PB, O[hh * 8 + j]); + } + } + } + + const uint head = kv_head * GQA + col; + if (p.n_split == 1) { + if (col < GQA) { + const uint dst_base = stream * p.nb3 + query * p.nb2 + head * p.nb1; + const float inv_l = l == 0.0 ? 0.0 : 1.0 / l; + [[unroll]] for (uint t = 0; t < 16; ++t) { + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint dim = (t / 8) * 128 + 8 * (grp + 4 * i) + t % 8; + data_dst[dst_base + dim] = O[t][i] * inv_l; + } + } + } + } else { + // partial: unnormalized O, then L and M per head (an empty split writes M = -FLT_MAX, not -inf, so the + // reduce never computes exp(-inf - -inf)) + if (col < GQA) { + const uint slot = split + p.n_split * (query + p.n_tok * stream); + const uint o_base = D * p.n_head * slot + D * head; + [[unroll]] for (uint t = 0; t < 16; ++t) { + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint dim = (t / 8) * 128 + 8 * (grp + 4 * i) + t % 8; + data_dst[o_base + dim] = O[t][i]; + } + } + if (grp == 0) { + const uint lm = D * p.n_head * p.n_tok * p.n_stream * p.n_split + p.n_head * 2 * slot + head; + data_dst[lm] = l; + data_dst[lm + p.n_head] = max(m, -FLT_MAX); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_prep.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_prep.comp new file mode 100644 index 000000000000..3f7a637b9fca --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sel_prep.comp @@ -0,0 +1,153 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +// Maskless selected-key attention, pre-pass: one workgroup per (query, stream) row of the selection list. It drops +// invalid entries (< 0 or >= n_kv) and repeated cells (the first occurrence is kept, as the CPU path counts a cell +// once), and writes the surviving cells to the front of a row of list_stride entries, in their original order. +// The rest of the row is padded with the first surviving cell, so the attention kernels can read whole key blocks +// without bounds tests (a padded column is masked by its position, never by its contents). The surviving count +// goes to the end of the buffer, one per row. Each row is cleaned once here instead of once per K/V head and split +// in the attention kernel. +// +// Duplicates are found with an open-addressing hash table of list positions in shared memory: a position whose +// cell is already in the table lowers the slot to the smaller position (atomicMin), so the first occurrence wins +// whatever the thread order. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer IdsBuf { int data_ids[]; }; +layout(binding = 1) writeonly buffer ListBuf { uint data_list[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_sel; + uint list_stride; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk2; + uint nbk3; + uint nbv1; + uint nbv2; + uint nbv3; + uint nbi1; + uint nbi3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint n_split; + uint chunk; + uint n_head; + uint n_tok; + uint n_stream; +} p; + +const uint WG = 256; +const uint MAX_SEL = 4096; // host gate: n_sel <= 4096 +const uint PER_THREAD = MAX_SEL / WG; +const uint TABLE = 4096; // >= n_sel, so every distinct cell finds a slot +const uint NONE = 0xffffffffu; + +shared uint key_sh[MAX_SEL]; // the row's cells, NONE for an invalid entry +shared uint table_sh[TABLE]; // list positions, NONE = empty slot +shared uint scan_sh[WG]; +shared uint first_sh; // the lowest valid position + +uint slot_hash(uint key) { + return (key * 2654435761u) >> 20; // 12 bits: TABLE == 4096 +} + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint query = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.y; + const uint row = stream * p.n_tok + query; + const uint ids_base = stream * p.nbi3 + query * p.nbi1; + + // thread tid owns the contiguous positions [tid * PER_THREAD, tid * PER_THREAD + PER_THREAD), so the scan below + // keeps the original order + const uint first = tid * PER_THREAD; + [[unroll]] for (uint i = 0; i < PER_THREAD; ++i) { + const uint s = first + i; + uint key = NONE; + if (s < p.n_sel) { + const int id = data_ids[ids_base + s]; + if (id >= 0 && uint(id) < p.n_kv) { + key = uint(id); + } + } + key_sh[s] = key; + } + for (uint i = tid; i < TABLE; i += WG) { + table_sh[i] = NONE; + } + if (tid == 0) { + first_sh = NONE; + } + barrier(); + + uint slots[PER_THREAD]; + [[unroll]] for (uint i = 0; i < PER_THREAD; ++i) { + const uint s = first + i; + const uint key = key_sh[s]; + uint slot = NONE; + if (key != NONE) { + uint h = slot_hash(key); + // terminates: at most n_sel <= TABLE distinct cells + for (;;) { + const uint c = atomicCompSwap(table_sh[h], NONE, s); + if (c == NONE) { + break; + } + if (key_sh[c] == key) { + atomicMin(table_sh[h], s); + break; + } + h = (h + 1) & (TABLE - 1); + } + slot = h; + atomicMin(first_sh, s); + } + slots[i] = slot; + } + barrier(); + + uint kept = 0; + [[unroll]] for (uint i = 0; i < PER_THREAD; ++i) { + if (slots[i] != NONE && table_sh[slots[i]] == first + i) { + kept |= 1u << i; + } + } + // inclusive scan of the per-thread counts (Hillis-Steele) + scan_sh[tid] = bitCount(kept); + barrier(); + for (uint off = 1; off < WG; off <<= 1) { + const uint v = tid >= off ? scan_sh[tid - off] : 0; + barrier(); + scan_sh[tid] += v; + barrier(); + } + const uint count = scan_sh[WG - 1]; + uint pos = scan_sh[tid] - bitCount(kept); + + const uint list_base = row * p.list_stride; + [[unroll]] for (uint i = 0; i < PER_THREAD; ++i) { + if ((kept & (1u << i)) != 0) { + data_list[list_base + pos] = key_sh[first + i]; + pos++; + } + } + // pad with the row's first valid entry, which is always a first occurrence (so it survives) + if (count < p.list_stride) { + const uint pad = count > 0 ? key_sh[first_sh] : 0; + for (uint s = count + tid; s < p.list_stride; s += WG) { + data_list[list_base + s] = pad; + } + } + if (tid == 0) { + data_list[p.n_tok * p.n_stream * p.list_stride + row] = count; + } +} 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/gdn_chunk_prep.comp b/ggml/src/ggml-vulkan/vulkan-shaders/gdn_chunk_prep.comp new file mode 100644 index 000000000000..cec6a3062cb5 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/gdn_chunk_prep.comp @@ -0,0 +1,210 @@ +#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_int16 : 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_arithmetic : require +#extension GL_KHR_shader_subgroup_ballot : require + +// Chunked gated delta net, part 1 of 2 (part 2: gdn_chunk_scan.comp). Fully parallel: one workgroup per +// (chunk of CL tokens, q/k head, sequence), serving every value head that reads that q/k head. +// With G the in-chunk inclusive cumsum of the gate g and D[r][s] = exp(G_r - G_s): +// A[r][s] = beta_r (k_r . k_s) D[r][s] (s < r), T = (I + A)^-1 diag(beta) +// P[r][s] = scale (q_r . k_s) D[r][s] (s <= r) +// Scratch written for the scan (f16 unless noted), per (seq, chunk, q/k head): Qh [CL][128], Kh [CL][128], +// KhT [128][CL]; per (seq, chunk, value head): T [CL][CL], P [CL][CL], gam = exp(G) [CL] f32, +// dec = exp(G_last - G) [CL] f32. Chunk-major, so concurrent workgroups touch adjacent memory (head-major +// scratch put every workgroup's rows one head stride apart: DRAM channel camping, 3.6x slower). +// The f16 matrices are stored as WMMA A fragments (FRAG_H): 16x16 tiles row-block-major, each tile 16 rows of +// 16 contiguous halves, so a scan lane reads its fragment as two adjacent 16-byte units. +// Tokens past n_tokens are padded with q = k = 0, g = 0, beta = 0: their T rows/columns and updates vanish. +// Wave32 required (hand-built RDNA3 WMMA fragments, see gdn_chunk_scan.comp). +// Idea from halo-box/strix-llama.cpp#91 (Gaetan Puleo). + +#define S_V 128 +#define CL 32 +#define NVMAX 4 +#define QS 68 // LDS row stride of Q/K in uints (64 + 4 pad) +#define AS 36 // LDS row stride of A in floats (16-byte aligned rows) +#define QKU (3 * CL * S_V / 2) // scratch uints per (q/k head, chunk) +#define TPU (CL * CL + 2 * CL) // scratch uints per (value head, chunk) + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint H; + uint n_tokens; + uint n_seqs; + uint s_off; + uint sq1, sq2, sq3; + uint sv1, sv2, sv3; + uint sb1, sb2, sb3; + uint neq1, rq3; + float scale; + uint K; + uint nch; +}; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float data_k[]; }; +layout(binding = 2) readonly buffer GBuf { float data_g[]; }; +layout(binding = 3) readonly buffer BetaBuf { float data_beta[]; }; +layout(binding = 4) writeonly buffer ScrBuf { uint data_scr[]; }; +layout(binding = 4) writeonly buffer ScrBuf16 { float16_t data_scr16[]; }; + +// half index of element (row, col) of a matrix with NCB 16-column blocks, in fragment-tile order +#define FRAG_H(row, col, NCB) (((((row) / 16) * (NCB) + (col) / 16) * 16 + (row) % 16) * 16 + (col) % 16) + +// Q/K rows as packed f16 pairs (K rows 0..CL-1, Q rows CL..2CL-1); reused as A[NVMAX][CL][AS] f32 afterwards +shared uint U[NVMAX * CL * AS > 2 * CL * QS ? NVMAX * CL * AS : 2 * CL * QS]; +shared float KK[CL * (CL + 1)]; +shared float QK[CL * (CL + 1)]; +shared float Gs[NVMAX][CL]; +shared float Bs[NVMAX][CL]; + +#define ACC coopmat +#define MATA coopmat +#define MATB coopmat + +void main() { + const uint hk = gl_WorkGroupID.x; // heads fastest: concurrent workgroups stream adjacent rows + const uint c = gl_WorkGroupID.y; + const uint seq = gl_WorkGroupID.z; // value sequence (q/k sequence seq / rq3) + const uint tid = gl_LocalInvocationIndex; + const uint w = gl_SubgroupID; + const uint lane = gl_SubgroupInvocationID; + const uint t0 = c * CL; + const uint nv = H / neq1; // value heads per q/k head (hk, hk + neq1, ...) + const uint iq3 = seq / rq3; + + const uint qk_scr = ((seq * nch + c) * neq1 + hk) * QKU; + const uint tp_region = n_seqs * nch * neq1 * QKU; + + // Q/K rows -> f16 pairs in LDS and in the scratch (Qh at +0, Kh at +CL*64) + for (uint p = tid; p < CL * S_V / 2; p += 256) { + const uint t = p / (S_V / 2); + const uint kp = p % (S_V / 2); + vec2 kv = vec2(0.0); + vec2 qv = vec2(0.0); + if (t0 + t < n_tokens) { + const uint off = iq3 * sq3 + (t0 + t) * sq2 + hk * sq1 + 2 * kp; + kv = vec2(data_k[off], data_k[off + 1]); + qv = vec2(data_q[off], data_q[off + 1]); + } + const uint ku = packFloat2x16(f16vec2(kv)); + const uint qu = packFloat2x16(f16vec2(qv)); + U[t * QS + kp] = ku; + U[(CL + t) * QS + kp] = qu; + data_scr[qk_scr + FRAG_H(t, 2 * kp, S_V / 16) / 2] = qu; + data_scr[qk_scr + CL * S_V / 2 + FRAG_H(t, 2 * kp, S_V / 16) / 2] = ku; + } + barrier(); + + // KhT [128][CL] at +CL*128 (in uints): thread (k, t pair) + for (uint p = tid; p < S_V * CL / 2; p += 256) { + const uint k = p / (CL / 2); + const uint tp = p % (CL / 2); + const uint a = U[(2 * tp) * QS + k / 2]; + const uint b = U[(2 * tp + 1) * QS + k / 2]; + const uint lo = (k & 1) != 0 ? (a >> 16) : (a & 0xFFFFu); + const uint hi = (k & 1) != 0 ? (b >> 16) : (b & 0xFFFFu); + data_scr[qk_scr + CL * S_V + FRAG_H(k, 2 * tp, CL / 16) / 2] = lo | (hi << 16); + } + + // KK^T and QK^T (CL x CL, f32 accumulation): wave w -> product w / 4 (0 KK, 1 QK), tile (w / 2 % 2, w % 2) + { + const uint lr = lane % 16; + const uint prod = w / 4; + const uint rt = (w / 2) % 2; + const uint st = w % 2; + const uint arow = (prod * CL + rt * 16 + lr) * QS; + const uint brow = (st * 16 + lr) * QS; + ACC acc = ACC(0.0); + [[unroll]] for (uint kt = 0; kt < S_V / 16; kt++) { + MATA a; + MATB b; + [[unroll]] for (uint e = 0; e < 8; e++) { + const f16vec2 av = unpackFloat2x16(U[arow + kt * 8 + e]); + const f16vec2 bv = unpackFloat2x16(U[brow + kt * 8 + e]); + a[2 * e] = av.x; a[2 * e + 1] = av.y; + b[2 * e] = bv.x; b[2 * e + 1] = bv.y; + } + acc = coopMatMulAdd(a, b, acc); + } + if (prod == 0) { + coopMatStore(acc, KK, rt * 16 * (CL + 1) + st * 16, CL + 1, gl_CooperativeMatrixLayoutRowMajor); + } else { + coopMatStore(acc, QK, rt * 16 * (CL + 1) + st * 16, CL + 1, gl_CooperativeMatrixLayoutRowMajor); + } + } + + for (uint j0 = 0; j0 < nv; j0 += NVMAX) { + barrier(); + // wave j: gate cumsum and beta of value head j0 + j, lane = token + if (w < NVMAX && j0 + w < nv) { + const uint h = hk + (j0 + w) * neq1; + float g = 0.0; + float b = 0.0; + if (lane < CL && t0 + lane < n_tokens) { + const uint gb_off = seq * sb3 + (t0 + lane) * sb2 + h * sb1; + g = data_g[gb_off]; + b = data_beta[gb_off]; + } + const float G = subgroupInclusiveAdd(g); + const float GL = subgroupBroadcast(G, CL - 1); + if (lane < CL) { + Gs[w][lane] = G; + Bs[w][lane] = b; + const uint tp = tp_region + ((seq * nch + c) * H + h) * TPU; + data_scr[tp + CL * CL + lane] = floatBitsToUint(exp(G)); + data_scr[tp + CL * CL + CL + lane] = floatBitsToUint(exp(GL - G)); + } + } + barrier(); + // A (f32, LDS) and P (f16, scratch); thread (j, r, s) + for (uint i = tid; i < NVMAX * CL * CL; i += 256) { + const uint j = i / (CL * CL); + const uint r = (i / CL) % CL; + const uint s = i % CL; + if (j0 + j < nv) { + const uint h = hk + (j0 + j) * neq1; + const float e = s <= r ? exp(Gs[j][r] - Gs[j][s]) : 0.0; + U[(j * CL + r) * AS + s] = floatBitsToUint(s < r ? Bs[j][r] * KK[r * (CL + 1) + s] * e : 0.0); + const uint tp = tp_region + ((seq * nch + c) * H + h) * TPU; + data_scr16[2 * (tp + CL * CL / 2) + FRAG_H(r, s, CL / 16)] = float16_t(scale * QK[r * (CL + 1) + s] * e); + } + } + barrier(); + // T = (I + A)^-1 diag(beta) by forward substitution: thread (j, col) keeps its column in registers + if (tid < NVMAX * CL) { + const uint j = tid / CL; + const uint col = tid % CL; + if (j0 + j < nv) { + const uint h = hk + (j0 + j) * neq1; + const uint tp = tp_region + ((seq * nch + c) * H + h) * TPU; + const float b = Bs[j][col]; +#if CL != 32 +#error TROW list assumes CL 32 +#endif + float x[CL]; + // fully unrolled (a rolled row loop leaves one LDS round trip per FMA): row R of the column +#define TROW(R) { \ + float v = (R) == col ? 1.0 : 0.0; \ + [[unroll]] for (uint ss = 0; ss < (R); ss++) { \ + v -= uintBitsToFloat(U[(j * CL + (R)) * AS + ss]) * x[ss]; \ + } \ + x[R] = v; \ + data_scr16[2 * tp + FRAG_H(R, col, CL / 16)] = float16_t(v * b); } + TROW(0) TROW(1) TROW(2) TROW(3) TROW(4) TROW(5) TROW(6) TROW(7) + TROW(8) TROW(9) TROW(10) TROW(11) TROW(12) TROW(13) TROW(14) TROW(15) + TROW(16) TROW(17) TROW(18) TROW(19) TROW(20) TROW(21) TROW(22) TROW(23) + TROW(24) TROW(25) TROW(26) TROW(27) TROW(28) TROW(29) TROW(30) TROW(31) +#undef TROW + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/gdn_chunk_scan.comp b/ggml/src/ggml-vulkan/vulkan-shaders/gdn_chunk_scan.comp new file mode 100644 index 000000000000..5e54a07fb3b3 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/gdn_chunk_scan.comp @@ -0,0 +1,236 @@ +#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 + +// Chunked gated delta net, part 2 of 2: sequential over chunks of CL tokens, after gdn_chunk_prep.comp. +// The v columns of the state S (k x v) evolve independently: every wave32 owns 16 columns of one value head, +// NW waves per workgroup, 8 / NW workgroups per head. Per chunk: +// R = V - diag(gam) K S, Delta = T R +// O = scale diag(gam) Q S + P Delta +// S' = gam_last S + K^T diag(dec) Delta +// S stays in f32 WMMA accumulators; it is fed back as an f16 B operand by swapping rows between the two +// half-waves (RDNA3 layout: an f32 accumulator lane L holds column L%16, rows 2i + L/16; a B operand lane L +// holds column L%16, rows 0..15). The A operands (Q, K, K^T, T, P) are stored by the prep in fragment order, +// so a lane's fragment is two adjacent 16-byte units, read straight from the scratch (one wave per workgroup: +// 384 waves at the Flash-Next shape). Staging K, K^T, T, P and the gate factors in LDS per workgroup measured +// equal in-model, 6% slower op-level and 1.5x slower at 4 heads. Operands f16, accumulation f32: not +// bit-identical to the sequential kernel. Idea from halo-box/strix-llama.cpp#91 (Gaetan Puleo). + +layout(constant_id = 0) const uint WGT = 128; // NW waves of 32 +const uint NW = WGT / 32; + +#define S_V 128 +#define CL 32 +#define NT (CL / 16) +#define KT (S_V / 16) +#define QKU (3 * CL * S_V / 2) +#define TPU (CL * CL + 2 * CL) + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform Parameters { + uint H; + uint n_tokens; + uint n_seqs; + uint s_off; + uint sq1, sq2, sq3; + uint sv1, sv2, sv3; + uint sb1, sb2, sb3; + uint neq1, rq3; + float scale; + uint K; + uint nch; +}; + +layout(binding = 0) readonly buffer VBuf { float data_v[]; }; +layout(binding = 1) readonly buffer StateBuf { float data_state[]; }; +layout(binding = 2) buffer DstBuf { float data_dst[]; }; +layout(binding = 3) readonly buffer ScrBuf { uvec4 scr4[]; }; + +// one chunk in 16-byte units, fragment order (tile (rb, cb) of 32 units, lane row r at +2r): +// q/k block: Qh [0, 512), Kh [512, 1024), KhT [1024, 1536); T/P block: T [0, 128), P [128, 256), gam/dec [256, 272) +#define G_K 512 +#define G_KT 1024 +#define G_P 128 +#define G_G 256 +#define QFRAG(j, kt) (((j) * KT + (kt)) * 32) // Qh/Kh tile (row block j, k block kt) +#define KTFRAG(kt, j) (((kt) * NT + (j)) * 32) // KhT tile (k block kt, t block j) +#define TFRAG(j, i) (((j) * NT + (i)) * 32) // T/P tile + +#define A_K(u) scr4[qk4 + G_K + (u)] +#define A_KT(u) scr4[qk4 + G_KT + (u)] +#define A_T(u) scr4[tp4 + (u)] +#define A_P(u) scr4[tp4 + G_P + (u)] +#define gd(i) uintBitsToFloat(scr4[tp4 + G_G + (i) / 4][(i) % 4]) +#define A_Q(u) scr4[qk4 + (u)] + +#define ACC coopmat +#define MATA coopmat +#define MATB coopmat + +// accumulator element i of this lane: row AROW(i) of the tile, column lr +#define AROW(i) (2 * (i) + hi) + +// A fragment from two 16-byte units (16 halves of this lane's row) +#define SET_A(a, x0_, x1_) { \ + f16vec2 v_; \ + v_ = unpackFloat2x16(x0_.x); a[0] = v_.x; a[1] = v_.y; \ + v_ = unpackFloat2x16(x0_.y); a[2] = v_.x; a[3] = v_.y; \ + v_ = unpackFloat2x16(x0_.z); a[4] = v_.x; a[5] = v_.y; \ + v_ = unpackFloat2x16(x0_.w); a[6] = v_.x; a[7] = v_.y; \ + v_ = unpackFloat2x16(x1_.x); a[8] = v_.x; a[9] = v_.y; \ + v_ = unpackFloat2x16(x1_.y); a[10] = v_.x; a[11] = v_.y; \ + v_ = unpackFloat2x16(x1_.z); a[12] = v_.x; a[13] = v_.y; \ + v_ = unpackFloat2x16(x1_.w); a[14] = v_.x; a[15] = v_.y; } +#define LOAD_A(a, SRC, frag) { const uvec4 y0_ = SRC((frag) + la); const uvec4 y1_ = SRC((frag) + la + 1); SET_A(a, y0_, y1_) } + +// f32 accumulator -> f16 B operand of the same 16x16 tile: rows 4m+hi, 4m+2+hi are packed into one dword, +// swapped with the other half-wave, and interleaved into rows 4m..4m+3. (A hi/lo f16 split of S, R and +// Delta changed neither the test NMSE nor the Flash-Next KLD: the error is set by the f16 Q/K/T/P.) +#define TO_B(b, acc) { \ + [[unroll]] for (uint m_ = 0; m_ < 4; m_++) { \ + const uint own_ = packFloat2x16(f16vec2(acc[2 * m_], acc[2 * m_ + 1])); \ + const uint oth_ = subgroupShuffleXor(own_, 16); \ + const f16vec2 ev_ = unpackFloat2x16(hi == 0 ? own_ : oth_); \ + const f16vec2 od_ = unpackFloat2x16(hi == 0 ? oth_ : own_); \ + b[4 * m_] = ev_.x; \ + b[4 * m_ + 1] = od_.x; \ + b[4 * m_ + 2] = ev_.y; \ + b[4 * m_ + 3] = od_.y; \ + } } + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint lr = lane % 16; + const uint hi = lane / 16; + const uint la = 2 * lr; // this lane's row within a fragment tile + const uint tid = gl_LocalInvocationIndex; + const uint nsplit = 8 / NW; + const uint h = gl_WorkGroupID.x / nsplit; + const uint seq = gl_WorkGroupID.y; + const uint iq1 = h % neq1; + const uint v0 = ((gl_WorkGroupID.x % nsplit) * NW + gl_SubgroupID) * 16; + + const uint state_base = (seq * H + h) * S_V * S_V + (v0 + lr) * S_V; + // token-row addressing: lane part (column, half-wave row parity) + a uniform row offset + const uint v_lane = seq * sv3 + h * sv1 + v0 + lr + hi * sv2; + const uint o_lane = (seq * n_tokens * H + h) * S_V + v0 + lr + hi * (S_V * H); + // scratch is chunk-major (all heads of a chunk adjacent), see gdn_chunk_prep.comp + const uint qk_base4 = (seq * nch * neq1 + iq1) * (QKU / 4); + const uint tp_base4 = (n_seqs * nch * neq1 * QKU + (seq * nch * H + h) * TPU) / 4; + + ACC S[KT]; + [[unroll]] for (uint kt = 0; kt < KT; kt++) { + [[unroll]] for (uint i = 0; i < 8; i++) { + S[kt][i] = data_state[state_base + kt * 16 + AROW(i)]; + } + } + + + for (uint c = 0; c < nch; c++) { + const uint t0 = c * CL; + const uint qk4 = qk_base4 + c * neq1 * (QKU / 4); + const uint tp4 = tp_base4 + c * H * (TPU / 4); + // V rows of this chunk, issued before phase 1 (rows past n_tokens clamp to the last row and are masked) + float vv[NT][8]; + [[unroll]] for (uint j = 0; j < NT; j++) { + [[unroll]] for (uint i = 0; i < 8; i++) { + const uint t = t0 + 16 * j + 2 * i; // uniform; the lane adds hi + vv[j][i] = data_v[v_lane + min(t, n_tokens - 1 - hi) * sv2]; + } + } + const float gam_last = gd(CL - 1); + + // phase 1, all reads of the chunk-start state S: KS_j = K_j S, QS_j = Q_j S; S *= gam_last + // separate declarations: older glslang (shaderc 2023.8) drops the matrix use from the second declarator + ACC KS[NT]; + ACC QS[NT]; + [[unroll]] for (uint j = 0; j < NT; j++) { + KS[j] = ACC(0.0); + QS[j] = ACC(0.0); + } + [[unroll]] for (uint kt = 0; kt < KT; kt++) { + MATB sb; + TO_B(sb, S[kt]); + S[kt] = S[kt] * gam_last; + [[unroll]] for (uint j = 0; j < NT; j++) { + MATA a; + LOAD_A(a, A_K, QFRAG(j, kt)); + KS[j] = coopMatMulAdd(a, sb, KS[j]); + LOAD_A(a, A_Q, QFRAG(j, kt)); + QS[j] = coopMatMulAdd(a, sb, QS[j]); + } + } + + // R_j = V_j - diag(gam) KS_j + MATB RB[NT]; + [[unroll]] for (uint j = 0; j < NT; j++) { + ACC r; + [[unroll]] for (uint i = 0; i < 8; i++) { + const uint t = 16 * j + AROW(i); + const float v = t0 + t < n_tokens ? vv[j][i] : 0.0; + r[i] = v - gd(t) * KS[j][i]; + } + TO_B(RB[j], r); + } + + MATB DB[NT]; + [[unroll]] for (uint j = 0; j < NT; j++) { + // Delta_j = sum_{i <= j} T_ji R_i + ACC d = ACC(0.0); + [[unroll]] for (uint i = 0; i <= j; i++) { + MATA a; + LOAD_A(a, A_T, TFRAG(j, i)); + d = coopMatMulAdd(a, RB[i], d); + } + TO_B(DB[j], d); + + // S += K_j^T diag(dec_j) Delta_j + ACC ds; + [[unroll]] for (uint i = 0; i < 8; i++) { + ds[i] = d[i] * gd(CL + 16 * j + AROW(i)); + } + MATB dsb; + TO_B(dsb, ds); + [[unroll]] for (uint kt = 0; kt < KT; kt++) { + MATA a; + LOAD_A(a, A_KT, KTFRAG(kt, j)); + S[kt] = coopMatMulAdd(a, dsb, S[kt]); + } + + // O_j = scale diag(gam) QS_j + sum_{i <= j} P_ji Delta_i + ACC o = QS[j]; + [[unroll]] for (uint i = 0; i < 8; i++) { + o[i] *= scale * gd(16 * j + AROW(i)); + } + [[unroll]] for (uint i = 0; i <= j; i++) { + MATA a; + LOAD_A(a, A_P, TFRAG(j, i)); + o = coopMatMulAdd(a, DB[i], o); + } + if (t0 + CL <= n_tokens) { + [[unroll]] for (uint i = 0; i < 8; i++) { + data_dst[o_lane + (t0 + 16 * j + 2 * i) * (S_V * H)] = o[i]; + } + } else { + [[unroll]] for (uint i = 0; i < 8; i++) { + if (t0 + 16 * j + AROW(i) < n_tokens) { + data_dst[o_lane + (t0 + 16 * j + 2 * i) * (S_V * H)] = o[i]; + } + } + } + } + } + + [[unroll]] for (uint kt = 0; kt < KT; kt++) { + [[unroll]] for (uint i = 0; i < 8; i++) { + data_dst[s_off + state_base + kt * 16 + AROW(i)] = S[kt][i]; + } + } +} 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 3383aa88ef91..ee32ede00b67 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp @@ -14,7 +14,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) || defined(DATA_A_IQ4_XS) #define K_PER_ITER 32 @@ -38,9 +43,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 03c348bf9e3c..a76f4a783570 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_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 588fb73546a3..69b059d9f6e5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -17,6 +17,54 @@ 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 @@ -323,6 +371,63 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin 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; + + 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; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp index 7cab9a119541..bf33deb33d54 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp @@ -8,6 +8,7 @@ #extension GL_KHR_shader_subgroup_basic : require #extension GL_KHR_cooperative_matrix : require #extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_shader_subgroup_shuffle : enable #if defined(MUL_MAT_ID_USE_SUBGROUPS) #extension GL_KHR_shader_subgroup_ballot : enable @@ -37,9 +38,16 @@ layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32 layout (binding = 1) readonly buffer B {block_q8_1_x4_packed128 data_b[];}; layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; +#ifndef D16 +// f32 output: four adjacent rows of a column go out as one 128-bit store +layout (binding = 2) writeonly buffer D4 {vec4 data_d4[];}; +#endif + #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// MUL_MAT_ID + MUL(router weights) epilogue: one f32 weight per (expert slot, token), [1, nei0, nei1] contiguous +layout (binding = 5) readonly buffer Scale {float data_scale[];}; #endif layout (push_constant) uniform parameter @@ -61,7 +69,7 @@ layout (push_constant) uniform parameter uint nbi1; uint ne11; uint n_experts; - uint hoist_row_ids; + uint hoist_row_ids; // bit 0: hoisted row ids, bit 1: scale each output row by data_scale (MMID_SCALE_ROWS) #else uint base_work_group_z; uint num_batches; @@ -85,6 +93,7 @@ layout (constant_id = 9) const uint TK = 16; layout (constant_id = 10) const uint WARP = 32; layout (constant_id = 11) const uint DEVICE_ARCH = 0; // vk_device_architecture (ggml-vulkan.cpp) #define VK_ARCH_AMD_RDNA4 5u +layout (constant_id = 12) const uint PACKSTORE = 1; // 0: per-element f32 stores (GGML_VK_MMQ_PACKSTORE=0) #define BK 32 #ifdef MUL_MAT_ID @@ -126,6 +135,25 @@ const uint CM_ELEMS = (TM * TN) / WARP; #define ACC_BIAS_F 12582912.0f const bool USE_MAGIC_BIAS = WARP != 32; +// RDNA3/3.5 wave32 accumulator: lane l holds column l % 16, rows (l / 16) + 2e. Two shuffles with the lane 16 +// away give each lane four adjacent rows of its column: for element group e..e+3 the low half stores rows +// 2e..2e+3 and the high half rows 2e+4..2e+7, so every lane writes one vec4 per four elements instead of four +// scattered dwords. All lanes must reach the shuffles, so the caller only branches on uniform values. +#ifndef D16 +const uint PACK_F32 = PACKSTORE * uint(DEVICE_ARCH != VK_ARCH_AMD_RDNA4) * uint(WARP == 32) * uint(TM == 16) * uint(TN == 16); + +void store_f32x4(const uint idx, const uint row_g, const vec4 v) { + if (row_g + 3 < p.M) { + data_d4[idx / 4] = v; + } else { + // ragged tail: never write past the row end + if (row_g < p.M) { data_d[idx] = v.x; } + if (row_g + 1 < p.M) { data_d[idx + 1] = v.y; } + if (row_g + 2 < p.M) { data_d[idx + 2] = v.z; } + } +} +#endif + // Accumulator row for element e: RDNA4 blocked, RDNA3/3.5 interleaved. uint cm_elem_row(uint e) { const uint row_half = gl_SubgroupInvocationID / TN; @@ -142,6 +170,8 @@ ACC_TYPE cm1_accumulate(ACC_TYPE prev, int acc_e, float scale_a, float nbias_a, } #ifdef MUL_MAT_ID +#define MMID_HOIST_ROW_IDS 1u +#define MMID_SCALE_ROWS 2u #define NUM_WARPS (BLOCK_SIZE / WARP) #include "mul_mm_id_funcs.glsl" #endif @@ -231,7 +261,7 @@ void main() { const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK; #ifdef MUL_MAT_ID - if (p.hoist_row_ids != 0) { + if ((p.hoist_row_ids & MMID_HOIST_ROW_IDS) != 0) { load_row_ids_hoisted(expert_idx, ic); } else { #ifdef MUL_MAT_ID_USE_SUBGROUPS @@ -475,6 +505,36 @@ void main() { const uint dc = ic * BN + b_col0; #ifdef MUL_MAT_ID + const bool scale_rows = (p.hoist_row_ids & MMID_SCALE_ROWS) != 0; +#ifndef D16 + if (PACK_F32 != 0) { + if ((p.stride_d % 4) == 0 && (p.batch_stride_d % 4) == 0) { + const bool lo_half = gl_SubgroupInvocationID < 16; + [[unroll]] for (uint r = 0; r < cms_per_row; r++) { + [[unroll]] for (uint c = 0; c < cms_per_col; c++) { + const uint tile_idx = r * cms_per_col + c; + const uint col_i = dc + c * TN + elem_col0; + [[unroll]] for (uint e = 0; e < CM_ELEMS; e += 4) { + const float a0 = float(sums[tile_idx * CM_ELEMS + e]); + const float a1 = float(sums[tile_idx * CM_ELEMS + e + 1]); + const float a2 = float(sums[tile_idx * CM_ELEMS + e + 2]); + const float a3 = float(sums[tile_idx * CM_ELEMS + e + 3]); + const float r0 = subgroupShuffleXor(lo_half ? a2 : a0, 16); + const float r1 = subgroupShuffleXor(lo_half ? a3 : a1, 16); + const vec4 v = lo_half ? vec4(a0, r0, a1, r1) : vec4(r0, a2, r1, a3); + const uint row_g = dr + r * TM + 2u * e + (lo_half ? 0u : 4u); + if (col_i < _ne1) { + const u16vec2 row_idx = row_ids[col_i - ic * BN]; + const vec4 vs = scale_rows ? v * data_scale[uint(row_idx.y) * p.nei0 + uint(row_idx.x)] : v; + store_f32x4(row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + row_g, row_g, vs); + } + } + } + } + return; + } + } +#endif [[unroll]] for (uint r = 0; r < cms_per_row; r++) { [[unroll]] for (uint c = 0; c < cms_per_col; c++) { const uint tile_idx = r * cms_per_col + c; @@ -487,13 +547,45 @@ void main() { const u16vec2 row_idx = row_ids[col_i - ic * BN]; const uint store_offset = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + row_g; - data_d[store_offset] = D_TYPE(sums[tile_idx * CM_ELEMS + e]); + float val = float(sums[tile_idx * CM_ELEMS + e]); + if (scale_rows) { + val *= data_scale[uint(row_idx.y) * p.nei0 + uint(row_idx.x)]; + } + data_d[store_offset] = D_TYPE(val); } } } #else const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches; +#ifndef D16 + if (PACK_F32 != 0) { + if ((p.stride_d % 4) == 0 && (offsets % 4) == 0) { + const bool lo_half = gl_SubgroupInvocationID < 16; + [[unroll]] for (uint r = 0; r < cms_per_row; r++) { + [[unroll]] for (uint c = 0; c < cms_per_col; c++) { + const uint tile_idx = r * cms_per_col + c; + const uint col_g = dc + c * TN + elem_col0; + [[unroll]] for (uint e = 0; e < CM_ELEMS; e += 4) { + const float a0 = float(sums[tile_idx * CM_ELEMS + e]); + const float a1 = float(sums[tile_idx * CM_ELEMS + e + 1]); + const float a2 = float(sums[tile_idx * CM_ELEMS + e + 2]); + const float a3 = float(sums[tile_idx * CM_ELEMS + e + 3]); + const float r0 = subgroupShuffleXor(lo_half ? a2 : a0, 16); + const float r1 = subgroupShuffleXor(lo_half ? a3 : a1, 16); + const vec4 v = lo_half ? vec4(a0, r0, a1, r1) : vec4(r0, a2, r1, a3); + const uint row_g = dr + r * TM + 2u * e + (lo_half ? 0u : 4u); + if (col_g < p.N) { + store_f32x4(offsets + col_g * p.stride_d + row_g, row_g, v); + } + } + } + } + return; + } + } +#endif + [[unroll]] for (uint r = 0; r < cms_per_row; r++) { [[unroll]] for (uint c = 0; c < cms_per_col; c++) { const uint tile_idx = r * cms_per_col + c; 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 136d5b74532e..aa5a67c5966d 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 56784cff38a2..4ae282c97ff7 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 @@ -65,6 +65,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_IQ3_S) #define QUANT_R_MMQ 2 struct block_a_cache { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_small.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_small.comp new file mode 100644 index 000000000000..590febc3c43e --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_small.comp @@ -0,0 +1,85 @@ +#version 450 + +// RMS norm for short rows (ne00 <= 256): one subgroup per row instead of one 512-thread workgroup +// per row. rms_norm.comp leaves 3/4 of its lanes idle on a 128-wide row and pays a shared-memory +// reduction with a barrier per step; the GDN head norms ((128,48,n) gated, (128,16,n) q/k l2) are +// that shape. Same push constants and do_multiply semantics as rms_norm.comp (broadcast weight +// through src1_idx / fastmod), contiguous destination. Rows are the flattened (i01, i02, i03) and a +// subgroup strides over them, so any workgroup count and subgroup size covers every row. + +#include "generic_binary_head.glsl" +#include "types.glsl" + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_shader_subgroup_basic : enable + +#define BLOCK_SIZE 256 + +layout (constant_id = 1) const bool do_multiply = false; +// RMS_NORM + MUL + MUL (the gated norm, normalised * gamma * gate): the second multiplier is +// contiguous and the same shape as the destination, so it shares the destination index +layout (constant_id = 2) const bool do_multiply2 = false; +// RMS_NORM + SCALE (the GDN q/k l2 norm): multiply by p.param2 on the store. The normalised value +// is rounded to f32 first, the same two roundings as the separate SCALE op. +layout (constant_id = 3) const bool do_post_scale = false; + +layout (binding = 3) readonly buffer C {float data_c[];}; + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint ncols = p.ne00; + const uint nrows = p.ne01 * p.ne02 * p.ne03; + const uint lane = gl_SubgroupInvocationID; + const uint sgsize = gl_SubgroupSize; + const uint stride = gl_NumWorkGroups.x * gl_NumSubgroups; + + for (uint r = gl_WorkGroupID.x * gl_NumSubgroups + gl_SubgroupID; r < nrows; r += stride) { + const uint row = r % p.ne01; + const uint channel = (r / p.ne01) % p.ne02; + const uint samp = r / (p.ne01 * p.ne02); + + const uint a_offset = samp*p.nb03 + channel*p.nb02 + row*p.nb01 + get_aoffset(); + const uint b_offset = src1_idx(0, row, channel, samp) + get_boffset(); + const uint d_row = ((samp*p.ne02 + channel)*p.ne01 + row)*ncols; + const uint d_offset = d_row + get_doffset(); + + FLOAT_TYPE sum = FLOAT_TYPE(0.0f); + for (uint col = lane; col < ncols; col += sgsize) { + const FLOAT_TYPE xi = FLOAT_TYPE(data_a[a_offset + col]); + sum += xi * xi; + } + sum = subgroupAdd(sum); + + const FLOAT_TYPE mean = sum / FLOAT_TYPE(ncols); + const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1)); + + if (do_multiply2) { + for (uint col = lane; col < ncols; col += sgsize) { + const FLOAT_TYPE g = FLOAT_TYPE(data_b[b_offset + (ncols > p.ne10 ? fastmod(col, p.ne10) : col)]); + const FLOAT_TYPE y = scale * FLOAT_TYPE(data_a[a_offset + col]) * g; + data_d[d_offset + col] = D_TYPE(y * FLOAT_TYPE(data_c[d_row + col])); + } + } else if (do_multiply) { + if (ncols > p.ne10) { + for (uint col = lane; col < ncols; col += sgsize) { + data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + } + } else { + for (uint col = lane; col < ncols; col += sgsize) { + data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + } + } + } else if (do_post_scale) { + for (uint col = lane; col < ncols; col += sgsize) { + const float y = float(scale * FLOAT_TYPE(data_a[a_offset + col])); + data_d[d_offset + col] = D_TYPE(y * p.param2); + } + } else { + for (uint col = lane; col < ncols; col += sgsize) { + data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp index 4cd9b8da3596..e4e724b37232 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_conv.comp @@ -8,6 +8,14 @@ layout(constant_id = 0) const uint BLOCK_SIZE = 32; layout(constant_id = 1) const uint TOKENS_PER_WG = 16; layout(constant_id = 2) const bool APPLY_BIAS = false; layout(constant_id = 3) const bool APPLY_SILU = false; +// 0: SSM_CONV over the materialised concat (src0 = [nc - 1 + n_t, nr, n_s]). +// 1: direct conv + SiLU: src0 is x [nr, n_t, n_s] (nb01 = token stride, nb02 = sequence stride, bytes), binding 2 +// is the ncs-column state [ncs, nr, n_s]; the concat is never read. Kernel width must be 4 (ncs = 3). +// 2: concat tail: write concat columns [nc, ncs + n_t) of every channel row from the state and x (src0 = x, +// binding 2 = state, dst = concat with row stride dst_nb1 and sequence stride dst_nb2; nc = first column). +layout(constant_id = 4) const uint MODE = 0; +// MODE 1: tokens per thread. The 4-tap window slides in registers, so a token costs one x load and one store. +layout(constant_id = 5) const uint TOKENS_PER_THREAD = 1; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z = 1) in; @@ -24,6 +32,62 @@ layout(push_constant) uniform PushConstants { }; void main() { + if (MODE == 1) { + // one thread per (channel, TOKENS_PER_THREAD tokens); lanes run along channels, so every x load and + // dst store is a contiguous row segment + const uint c = gl_GlobalInvocationID.x; + const uint t0 = (gl_WorkGroupID.y * TOKENS_PER_WG + gl_LocalInvocationID.y) * TOKENS_PER_THREAD; + const uint s = gl_WorkGroupID.z; + if (c >= nr || t0 >= n_t || s >= n_s) { + return; + } + const uint xs = nb01 / 4; + const uint x_base = s * (nb02 / 4) + c; + const uint st_base = s * ncs * nr + c * ncs; + const uint k_base = c * (nb11 / 4); + const vec4 w = vec4(src1[k_base], src1[k_base + 1], src1[k_base + 2], src1[k_base + 3]); + + // window: times t0-3, t0-2, t0-1; a negative time reads state column t + ncs + float win[3]; + [[unroll]] for (uint i = 0; i < 3; i++) { + const int t = int(t0) + int(i) - 3; + if (t < 0) { + win[i] = bias[st_base + uint(t + int(ncs))]; + } else { + win[i] = src0[x_base + uint(t) * xs]; + } + } + const uint d_base = s * (dst_nb2 / 4) + c; + [[unroll]] for (uint i = 0; i < TOKENS_PER_THREAD; i++) { + const uint t = t0 + i; + if (t >= n_t) { + break; + } + const float v = src0[x_base + t * xs]; + // the same dot() as the nc == 4 path over the concat, so the result is bit-identical + float sum = dot(vec4(win[0], win[1], win[2], v), w); + sum = sum / (1.0f + exp(-sum)); + dst[d_base + t * (dst_nb1 / 4)] = sum; + win[0] = win[1]; + win[1] = win[2]; + win[2] = v; + } + return; + } else if (MODE == 2) { + const uint ncols = ncs + n_t - nc; + const uint idx = gl_GlobalInvocationID.x; + const uint s = gl_WorkGroupID.y; + if (idx >= ncols * nr || s >= n_s) { + return; + } + // lanes along channels: x reads coalesce, the (few) columns per channel are the outer index + const uint c = idx % nr; + const uint j = nc + idx / nr; + const float v = j < ncs ? bias[s * ncs * nr + c * ncs + j] : src0[s * (nb02 / 4) + (j - ncs) * (nb01 / 4) + c]; + dst[s * (dst_nb2 / 4) + c * (dst_nb1 / 4) + j] = v; + return; + } + const uint i1 = gl_GlobalInvocationID.x; const uint i2 = gl_WorkGroupID.y * TOKENS_PER_WG + gl_LocalInvocationID.y; const uint i3 = gl_WorkGroupID.z; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp index 0b757f38e186..25c1a0e0ba65 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_nary_search.comp @@ -77,7 +77,7 @@ void topk(const uint row) { ivec2 a = dst_row[tid]; ivec2 b = dst_row[tid + s]; if (a.x >= p.orig_ncols || - b.x < p.orig_ncols && b.y > a.y) { + b.x < p.orig_ncols && f2ui(intBitsToFloat(b.y)) > f2ui(intBitsToFloat(a.y))) { dst_row[tid] = b; } } @@ -95,9 +95,10 @@ void topk(const uint row) { int shift = 32 - SUBGROUP_SIZE_LOG2; uint mask = ((1 << SUBGROUP_SIZE_LOG2) - 1) << shift; - // The current range. + // The current range, inclusive at both ends (f2ui(+inf) = 0xFF800000, so an exclusive bound there + // could never count +inf). uint range_min = 0; - uint range_max = 0xFF800000; + uint range_hi = 0xFFFFFFFF; // How many are above the current range, and how many we need to find. uint total = 0; uint limit = min(p.k, p.ncols_input - gl_WorkGroupID.x * BLOCK_SIZE); @@ -113,7 +114,7 @@ void topk(const uint row) { if (tid < p.ncols_input) { float y = intBitsToFloat(dst_row[tid].y); uint fy = f2ui(y); - if (fy >= range_min && fy < range_max) { + if (fy >= range_min && fy <= range_hi) { uint bucket = (fy & mask) >> shift; atomicAdd(counts[bucket], 1); } @@ -138,8 +139,8 @@ void topk(const uint row) { total = sh_total; // Update the range, and break if we've found the K-th largest. - range_max = range_min + ((min_idx + 1) << shift); - range_min = range_min + (min_idx << shift); + range_min = range_min + (uint(min_idx) << shift); + range_hi = range_min + ((1u << shift) - 1u); if (total == p.k) { break; 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 7d62e92e3422..4938b218f04e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl @@ -1869,6 +1869,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), @@ -1909,6 +1990,71 @@ i32vec2 iq4nl_to_i8x8(uint32_t vui) { #endif #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) #if !defined(USE_OCP_FP4) const int8_t kvalues_mxfp4_const[16] = { 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 e7e303e50cf9..e7c931bd0df6 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,8 +251,13 @@ bool is_iq_quant(const std::string& type_name) { return string_starts_with(type_name, "iq"); } +bool is_rocmfp_quant(const std::string& type_name) { + return string_starts_with(type_name, "rocmfp"); +} + +// types compiled as their own matmul shader instead of the MmTypeA quant shader bool is_lut_quant(const std::string& type_name) { - return is_iq_quant(type_name) || type_name == "mxfp4" || type_name == "nvfp4"; + return is_iq_quant(type_name) || type_name == "mxfp4" || type_name == "nvfp4" || is_rocmfp_quant(type_name); } std::string lut_load_vec_a(const std::string& type_name) { @@ -625,7 +635,8 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c }; #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4" || tname == "iq3_s" || tname == "iq4_xs")) { + if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4" || tname == "iq3_s" || tname == "iq4_xs" || + 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 @@ -737,6 +748,11 @@ void process_shaders() { #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) string_to_spv("flash_attn_f32_f16", "flash_attn_cm1.comp", merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}), fp16, true, false, f16acc); + // multi-row prefill flash attention (GGML_VK_FA_MR): f32 accumulation only + if (!f16acc) { + string_to_spv("flash_attn_f32_f16_mr", "flash_attn_cm1_mr.comp", + merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}), fp16, true, false, f16acc); + } #endif } @@ -813,7 +829,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" || tname == "iq4_xs") { + 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" || tname == "iq4_xs") { 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"}})); @@ -829,9 +847,15 @@ 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", {}); + // f16 V -> V^T per head for the coopmat1 flash attention + string_to_spv("dequant_f16_transpose_vt", "dequant_f16_transpose_vt.comp", {}); + } shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp"; @@ -845,6 +869,37 @@ 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("gdn_chunk_prep_f32", "gdn_chunk_prep.comp", {}); + string_to_spv("gdn_chunk_scan_f32", "gdn_chunk_scan.comp", {}); + 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", {}); + string_to_spv("flash_attn_sel_f16", "flash_attn_sel.comp", {}); + string_to_spv("flash_attn_sel_prep", "flash_attn_sel_prep.comp", {}); +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + string_to_spv("flash_attn_sel_cm_f16", "flash_attn_sel_cm.comp", {}); +#endif +#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("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"}}); @@ -857,6 +912,7 @@ void process_shaders() { string_to_spv("rms_norm_mul_add_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); string_to_spv("rms_norm_set_rows_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); string_to_spv("rms_norm_set_rows_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); + string_to_spv("rms_norm_small_f32", "rms_norm_small.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}})); string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}})); @@ -885,13 +941,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"}}); } @@ -934,8 +998,6 @@ void process_shaders() { string_to_spv("fa_decode_ph2", "flash_attn_decode_phase_2.comp", {}, true, true, false, false); #endif - 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"}}); @@ -962,6 +1024,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"}}); @@ -1336,7 +1399,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" && tname != "iq4_xs") { + 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" && tname != "iq4_xs") { 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 cdc10a75970f..f245e5a0b9d3 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( @@ -8082,6 +8161,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/README.md b/gguf-py/README.md index ca7e09c68184..83817e3edad5 100644 --- a/gguf-py/README.md +++ b/gguf-py/README.md @@ -30,6 +30,8 @@ pip install gguf[gui] [gguf/scripts/gguf_new_metadata.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_new_metadata.py) — Copies a GGUF file with added/modified/removed metadata values. +[gguf/scripts/gguf_extract_ple.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_extract_ple.py) - Writes the per-layer embedding (PLE) tables of a model to their own GGUF, optionally in another type, for `--ple`. + [gguf/scripts/gguf_editor_gui.py](https://github.com/ggml-org/llama.cpp/blob/master/gguf-py/gguf/scripts/gguf_editor_gui.py) — Allows for viewing, editing, adding, or removing metadata values within a GGUF file as well as viewing its tensors with a Qt interface. ## Development diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 9a7a5e5bfaf6..67bdfaf37e1d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1198,6 +1198,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 @@ -1985,6 +1988,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", @@ -2976,12 +2982,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, @@ -5721,6 +5738,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 @@ -5779,6 +5802,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 @@ -5915,6 +5955,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/scripts/gguf_extract_ple.py b/gguf-py/gguf/scripts/gguf_extract_ple.py new file mode 100644 index 000000000000..04090b7f17ce --- /dev/null +++ b/gguf-py/gguf/scripts/gguf_extract_ple.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import logging +import os +import re +import sys +from pathlib import Path +from typing import BinaryIO + +import numpy as np +from tqdm import tqdm + +# Necessary to load the local gguf package +if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent.parent / 'gguf-py').exists(): + sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import gguf + +logger = logging.getLogger("gguf-extract-ple") + +DEFAULT_TENSORS = ["per_layer_token_embd.weight"] + +OUT_TYPES = { + "f32": gguf.GGMLQuantizationType.F32, + "f16": gguf.GGMLQuantizationType.F16, + "bf16": gguf.GGMLQuantizationType.BF16, + "q8_0": gguf.GGMLQuantizationType.Q8_0, + "q5_1": gguf.GGMLQuantizationType.Q5_1, + "q5_0": gguf.GGMLQuantizationType.Q5_0, + "q4_1": gguf.GGMLQuantizationType.Q4_1, + "q4_0": gguf.GGMLQuantizationType.Q4_0, +} + +LOSSLESS_SOURCES = (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, gguf.GGMLQuantizationType.BF16) + + +def model_files(path: Path) -> list[Path]: + # every shard of a split model, else just the file + m = re.match(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", path.name) + if m is None: + return [path] + n = int(m.group(3)) + return [path.with_name(f"{m.group(1)}-{i:05d}-of-{n:05d}.gguf") for i in range(1, n + 1)] + + +def write_rows(fout: BinaryIO, rows: np.ndarray, src_type: gguf.GGMLQuantizationType, out_type: gguf.GGMLQuantizationType, out_row_bytes: int, chunk_rows: int) -> None: + # converted and written a chunk at a time, so a table larger than RAM is never resident + n_rows = rows.shape[0] + bar = tqdm(desc="Writing", total=n_rows * out_row_bytes, unit="byte", unit_scale=True) + for r0 in range(0, n_rows, chunk_rows): + n = min(chunk_rows, n_rows - r0) + chunk = np.asarray(rows[r0:r0 + n]) + if out_type != src_type: + f32 = gguf.quants.dequantize(chunk, src_type).astype(np.float32, copy=False) + chunk = gguf.quants.quantize(f32, out_type) + chunk = np.ascontiguousarray(chunk).view(np.uint8) + assert chunk.nbytes == n * out_row_bytes + chunk.tofile(fout) + bar.update(chunk.nbytes) + bar.close() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Write the per-layer embedding (PLE) tables of a model to their own GGUF, for llama.cpp --ple") + parser.add_argument("model", type=Path, help="GGUF model, or any shard of a split model (all shards are searched)") + parser.add_argument("output", type=Path, help="GGUF file to write") + parser.add_argument("--tensor", action="append", help=f"tensor to extract, can be repeated (default: {', '.join(DEFAULT_TENSORS)})") + parser.add_argument("--type", choices=["keep", *OUT_TYPES], default="keep", help="type of the written tables (default: keep, a byte copy)") + parser.add_argument("--metadata-from", type=Path, help="take the architecture keys from this GGUF, for a source that has none (e.g. a lone shard)") + parser.add_argument("--chunk-rows", type=int, default=1 << 20, help="rows converted and written at a time (default: 1048576)") + parser.add_argument("--force", action="store_true", help="overwrite the output file") + parser.add_argument("--verbose", action="store_true", help="increase output verbosity") + args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"]) + + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) + + if args.output.exists() and not args.force: + logger.error(f"{args.output} exists, use --force to overwrite") + sys.exit(1) + + names = args.tensor or DEFAULT_TENSORS + + files = model_files(args.model) + readers = [gguf.GGUFReader(f, "r") for f in files] + + meta = gguf.GGUFReader(args.metadata_from, "r") if args.metadata_from else readers[0] + arch_field = meta.get_field(gguf.Keys.General.ARCHITECTURE) + if arch_field is None: + logger.error(f"{args.metadata_from or files[0]} has no {gguf.Keys.General.ARCHITECTURE}, pass --metadata-from with a GGUF of the model") + sys.exit(1) + arch = arch_field.contents() + + # the keys that describe the layout of the tables; llama.cpp checks them against the model + keys = [f for f in meta.fields.values() if f.name.startswith(f"{arch}.ple.") or f.name == f"{arch}.embedding_length_per_layer_input"] + + found: dict[str, gguf.ReaderTensor] = {} + for rd in readers: + for t in rd.tensors: + if t.name in names: + found[t.name] = t + missing = [n for n in names if n not in found] + if missing: + logger.error(f"not found in {args.model}: {', '.join(missing)}") + sys.exit(1) + + writer = gguf.GGUFWriter(args.output, arch=arch, endianess=readers[0].endianess) + + name_field = meta.get_field(gguf.Keys.General.NAME) + if name_field is not None: + writer.add_name(name_field.contents()) + for f in keys: + val_type = f.types[0] + sub_type = f.types[-1] if val_type == gguf.GGUFValueType.ARRAY else None + writer.add_key_value(f.name, f.contents(), val_type, sub_type=sub_type) + logger.info(f"key {f.name}") + if not keys: + logger.warning(f"no {arch}.ple.* keys found, llama.cpp cannot check the tables against the model") + + jobs: list[tuple[np.ndarray, gguf.GGMLQuantizationType, gguf.GGMLQuantizationType, int]] = [] + for name in names: + t = found[name] + src_type = t.tensor_type + out_type = src_type if args.type == "keep" else OUT_TYPES[args.type] + if out_type != src_type and src_type not in LOSSLESS_SOURCES: + logger.warning(f"{name}: converting {src_type.name} -> {out_type.name} adds a second rounding, prefer an F32/F16/BF16 source") + + ne = [int(x) for x in t.shape] + n_rows = int(np.prod(ne[1:])) + rows = t.data.reshape(n_rows, -1) + + block_size, type_size = gguf.GGML_QUANT_SIZES[out_type] + if ne[0] % block_size != 0: + logger.error(f"{name}: row size {ne[0]} is not a multiple of the {out_type.name} block size {block_size}") + sys.exit(1) + out_row_bytes = ne[0] // block_size * type_size + + byte_shape = (*reversed(ne[1:]), out_row_bytes) + writer.add_tensor_info(name, byte_shape, np.dtype(np.uint8), n_rows * out_row_bytes, raw_dtype=out_type) + jobs.append((rows, src_type, out_type, out_row_bytes)) + logger.info(f"tensor {name}: {src_type.name} -> {out_type.name}, shape {ne}, {n_rows * out_row_bytes / 1e9:.2f} GB") + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_ti_data_to_file() + writer.close() + + # the data section follows the tensor infos; each tensor starts on an alignment boundary, as their offsets assume + align = writer.data_alignment + with open(args.output, "r+b") as fout: + fout.seek(0, os.SEEK_END) + for rows, src_type, out_type, out_row_bytes in jobs: + fout.write(bytes(-fout.tell() % align)) + write_rows(fout, rows, src_type, out_type, out_row_bytes, args.chunk_rows) + fout.write(bytes(-fout.tell() % align)) + + logger.info(f"wrote {args.output}") + + +if __name__ == "__main__": + main() 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/gguf-py/pyproject.toml b/gguf-py/pyproject.toml index 07e6f7fee82d..f4ebb016b5d4 100644 --- a/gguf-py/pyproject.toml +++ b/gguf-py/pyproject.toml @@ -23,6 +23,7 @@ gguf-convert-endian = "gguf.scripts.gguf_convert_endian:main" gguf-dump = "gguf.scripts.gguf_dump:main" gguf-set-metadata = "gguf.scripts.gguf_set_metadata:main" gguf-new-metadata = "gguf.scripts.gguf_new_metadata:main" +gguf-extract-ple = "gguf.scripts.gguf_extract_ple:main" gguf-editor-gui = "gguf.scripts.gguf_editor_gui:main" [project.optional-dependencies] 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 dea4d46062b7..8039c144167c 100644 --- a/include/llama.h +++ b/include/llama.h @@ -158,6 +158,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 }; @@ -219,6 +237,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 { @@ -353,6 +372,9 @@ extern "C" { bool no_host; // bypass host buffer allowing extra buffers to be used bool no_alloc; // only load metadata and simulate memory allocations bool load_mtp; // whether to load MTP layers + + // optional GGUF with per-layer embedding (PLE) tables: its tensors replace the lookup tables of the same name in the model + const char * path_ple; }; struct llama_sampler_seq_config { @@ -372,6 +394,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` @@ -1122,6 +1145,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; @@ -7654,6 +7746,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(); @@ -7666,6 +7759,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..a9a9dd93d8b1 --- /dev/null +++ b/tests/test-qsa-fast.cpp @@ -0,0 +1,180 @@ +#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. + +static void set_qsa_fast_max(const char * value) { +#ifdef _WIN32 + _putenv_s("QSA_FAST_MAX", value ? value : ""); +#else + if (value) { + setenv("QSA_FAST_MAX", value, 1); + } else { + unsetenv("QSA_FAST_MAX"); + } +#endif +} + +// 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.) + set_qsa_fast_max("127"); + 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)); + set_qsa_fast_max("0"); + 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)); + set_qsa_fast_max("9999"); // 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)); + set_qsa_fast_max(nullptr); + } +} + +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-qsa-recover.cpp b/tests/test-qsa-recover.cpp new file mode 100644 index 000000000000..5c7e21f2338c --- /dev/null +++ b/tests/test-qsa-recover.cpp @@ -0,0 +1,70 @@ +// qsa_rebuild_prefix: rebuild the QSA prefix of one sequence from a unified KV cache that also holds other sequences +#include "../src/prefix.h" +#include "ggml.h" +#include + +static llama_kv_cells make_cells(uint32_t n) { llama_kv_cells c; c.resize(n); return c; } +static void put(llama_kv_cells & c, uint32_t cell, llama_seq_id seq, llama_pos pos) { c.pos_set(cell, pos); c.seq_add(cell, seq); } + +int main() { + // one sequence alone, in position order: what already worked + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s)); + GGML_ASSERT(s.sequence == 0 && s.cells.size() == 10 && s.block_positions.size() == 2); + for (int p = 0; p < 10; ++p) { GGML_ASSERT(s.cells[p] == p && s.positions[p] == p); } + } + // another conversation resident first (cells 0..14), ours after it, stored in reverse cell order + { + auto c = make_cells(64); + for (int p = 0; p < 15; ++p) { put(c, p, 1, p); } + for (int p = 0; p < 10; ++p) { put(c, 29 - p, 0, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s)); + GGML_ASSERT(s.sequence == 0 && s.cells.size() == 10 && s.block_positions.size() == 2); + for (int p = 0; p < 10; ++p) { GGML_ASSERT(s.cells[p] == 29 - p && s.positions[29 - p] == p); } + for (int i = 0; i < 15; ++i) { GGML_ASSERT(s.positions[i] == -1); } + qsa_prefix_state t(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 1, t) && t.cells.size() == 15 && t.positions[29] == -1); + } + // two conversations decoded together: their cells alternate + { + auto c = make_cells(64); + for (int p = 0; p < 20; ++p) { put(c, 2*p, 0, p); put(c, 2*p + 1, 1, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 1, s) && s.cells.size() == 20); + for (int p = 0; p < 20; ++p) { GGML_ASSERT(s.cells[p] == 2*p + 1); } + } + // a new conversation, nothing cached yet, others resident: an empty prefix its first ubatch extends + { + auto c = make_cells(64); + for (int p = 0; p < 15; ++p) { put(c, p, 1, p); } + qsa_prefix_state s(64); + GGML_ASSERT(qsa_rebuild_prefix(c, 0, s) && s.cells.empty() && s.sequence == 0); + GGML_ASSERT(s.apply(0, 0, {20, 21, 22})); + } + // still refused: a hole in the positions, a cell shared with another sequence, a duplicated position + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { if (p != 5) { put(c, p, 0, p); } } + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + c.seq_add(3, 1); + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + { + auto c = make_cells(64); + for (int p = 0; p < 10; ++p) { put(c, p, 0, p); } + put(c, 40, 0, 3); + qsa_prefix_state s(64); + GGML_ASSERT(!qsa_rebuild_prefix(c, 0, s)); + } + std::cout << "PASS: prefixes rebuilt next to other sequences; holes, shared cells and duplicates refused\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-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index 4ad0d6f9efca..5c58593994c4 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -1,6 +1,7 @@ #include "arg.h" #include "common.h" #include "ggml-backend.h" +#include "gguf.h" #include "log.h" #include "llama-cpp.h" #include "llama.h" @@ -12,9 +13,12 @@ #include #include #include +#include #include #include +#include #include +#include #include #include #include @@ -310,6 +314,272 @@ static test_status test_multi_seq_split_replay(const common_params & params, lla return test_status::PASS; } +static void set_tensor_data_scaled(ggml_tensor * tensor, void * userdata) { + const size_t seed = *(const size_t *) userdata ^ std::hash{}(tensor->name); + std::mt19937 gen(seed); + std::normal_distribution dis(0.0f, 1.0f); + + // norm weights and matrices at scale 1: at the 0.01 scale of the generated models the recurrent + // branch is too small to change the logits, so a wrong recurrent state would go unnoticed + const bool is_norm = strstr(tensor->name, "norm") != nullptr; + const float scale = ggml_n_dims(tensor) > 1 ? 1.0f : 1.0e-2f; + + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + std::vector tmp(ggml_nelements(tensor)); + for (auto & x : tmp) { + x = is_norm ? 1.0f : scale*dis(gen); + } + ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); +} + +// same hparams as the model file, new random weights +static llama_model * load_model_scaled(common_params & params) { + gguf_init_params gparams = { + /*.no_alloc =*/ true, + /*.ctx =*/ nullptr, + }; + gguf_context * meta_file = gguf_init_from_file(params.model.path.c_str(), gparams); + if (meta_file == nullptr) { + return nullptr; + } + gguf_context * meta = gguf_init_empty(); + gguf_set_kv(meta, meta_file); + gguf_free(meta_file); + + size_t seed = 1234; + llama_model * model = llama_model_init_from_user(meta, set_tensor_data_scaled, &seed, common_model_params_to_llama(params)); + gguf_free(meta); + return model; +} + +static void set_graph_reuse_disable(bool disable) { +#ifdef _WIN32 + _putenv_s("LLAMA_GRAPH_REUSE_DISABLE", disable ? "1" : ""); +#else + if (disable) { + setenv("LLAMA_GRAPH_REUSE_DISABLE", "1", 1); + } else { + unsetenv("LLAMA_GRAPH_REUSE_DISABLE"); + } +#endif +} + +// Two sequences share ubatches while seq_cp swaps them, forks one and rolls it back, with graph reuse on. +// With n_rs_seq = 0 the recurrent state can be read in place when no state is copied, so a seq_cp between +// two same-shape ubatches must stop the graph from being reused. Checked against a context without graph +// reuse (same ubatches, so bit-exact), a context with n_rs_seq = 1, which never reads the state in place, +// and a context that decodes one sequence per ubatch. +static bool test_seq_cp_graph_reuse(const common_params & params, llama_model * model) { + const char * func = __func__; + const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + + const auto make = [&](uint32_t n_rs_seq, bool reuse) { + auto cparams = common_context_params_to_llama(params); + cparams.n_seq_max = 3; + cparams.n_rs_seq = n_rs_seq; + cparams.n_ctx = 256; + cparams.n_batch = 64; + cparams.n_ubatch = 64; + cparams.kv_unified = true; + set_graph_reuse_disable(!reuse); + llama_context * ctx = llama_init_from_model(model, cparams); + set_graph_reuse_disable(false); + return ctx; + }; + + struct arm { + const char * name; + llama_context * ctx; + bool split; // one llama_decode per sequence + float eps; // max |logit diff| against the first arm, relative to its max |logit| + float diff_max; + }; + std::vector arms = { + { "reuse", make(0, true), false, 0.0f, 0.0f }, + { "no-reuse", make(0, false), false, 1e-6f, 0.0f }, + { "no-in-place", make(1, false), false, 1e-5f, 0.0f }, + { "seq-split", make(0, false), true, 1e-4f, 0.0f }, + }; + + const auto cleanup = [&]() { + for (auto & a : arms) { + llama_free(a.ctx); + } + }; + + for (const auto & a : arms) { + if (a.ctx == nullptr) { + fprintf(stderr, "%s : failed to init %s context\n", __func__, a.name); + cleanup(); + return false; + } + } + + llama_pos pos[3] = { 0, 0, 0 }; + uint32_t n_step = 0; + bool ok = true; + + const auto decode = [&](const std::vector & seqs, uint32_t n_tokens) { + if (!ok) { + return; + } + n_step++; + const auto token = [&](llama_seq_id s, uint32_t i) { + return (llama_token) ((13*n_step + 7*(uint32_t) s + 3*i + 1) % (uint32_t) n_vocab); + }; + + const size_t n_rows = seqs.size()*n_tokens; + std::vector> logits(arms.size()); + llama_batch batch = llama_batch_init(n_rows, 0, 1); + for (size_t k = 0; ok && k < arms.size(); ++k) { + logits[k].reserve(n_rows*n_vocab); + const size_t n_batches = arms[k].split ? seqs.size() : 1; + for (size_t b = 0; ok && b < n_batches; ++b) { + common_batch_clear(batch); + for (size_t j = 0; j < seqs.size(); ++j) { + if (arms[k].split && j != b) { + continue; + } + for (uint32_t i = 0; i < n_tokens; ++i) { + common_batch_add(batch, token(seqs[j], i), pos[seqs[j]] + (llama_pos) i, { seqs[j] }, true); + } + } + ok = llama_decode(arms[k].ctx, batch) == 0; + for (int32_t i = 0; ok && i < batch.n_tokens; ++i) { + const float * l = llama_get_logits_ith(arms[k].ctx, i); + logits[k].insert(logits[k].end(), l, l + n_vocab); + } + } + if (!ok) { + fprintf(stderr, "%s : %s decode failed at step %u\n", func, arms[k].name, n_step); + } + } + llama_batch_free(batch); + + for (size_t r = 0; ok && r < n_rows; ++r) { + const float * ref = logits[0].data() + r*n_vocab; + float ref_max = 0.0f; + for (int t = 0; t < n_vocab; ++t) { + ref_max = std::max(ref_max, std::fabs(ref[t])); + } + for (size_t k = 1; k < arms.size(); ++k) { + const float * cur = logits[k].data() + r*n_vocab; + float diff = 0.0f; + for (int t = 0; t < n_vocab; ++t) { + diff = std::max(diff, logit_diff(ref[t], cur[t])); + } + diff /= std::max(ref_max, std::numeric_limits::min()); + arms[k].diff_max = std::max(arms[k].diff_max, diff); + if (!(diff <= arms[k].eps)) { + fprintf(stderr, "%s : step %u, seq %d row %zu: %s differs from %s by %g (rel)\n", + func, n_step, seqs[r/n_tokens], r % n_tokens, arms[k].name, arms[0].name, (double) diff); + ok = false; + } + } + } + for (llama_seq_id s : seqs) { + pos[s] += (llama_pos) n_tokens; + } + }; + + const auto seq_cp = [&](llama_seq_id src, llama_seq_id dst) { + for (auto & a : arms) { + llama_memory_t mem = llama_get_memory(a.ctx); + ok = ok && llama_memory_seq_rm(mem, dst, -1, -1); + llama_memory_seq_cp(mem, src, dst, -1, -1); + } + pos[dst] = pos[src]; + }; + + const auto seq_rm = [&](llama_seq_id s) { + for (auto & a : arms) { + ok = ok && llama_memory_seq_rm(llama_get_memory(a.ctx), s, -1, -1); + } + pos[s] = 0; + }; + + const auto swap_01 = [&]() { + seq_cp(0, 2); + seq_cp(1, 0); + seq_cp(2, 1); + seq_rm(2); + }; + + // seq 1 is decoded first, so the first shared ubatch reorders the cells + decode({ 1 }, 5); + decode({ 0 }, 5); + for (uint32_t n_tokens : { 2u, 3u, 1u }) { + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + + // same ubatch shape, head and rs_z as the step before: only the state copy map changes + swap_01(); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + + decode({ 0 }, n_tokens); + decode({ 1 }, n_tokens); + decode({ 0 }, n_tokens); + decode({ 1 }, n_tokens); + + // fork seq 0, advance it, then roll it back to the fork + seq_cp(0, 2); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + seq_cp(2, 0); + decode({ 0, 1 }, n_tokens); + seq_rm(2); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + } + + const int32_t n_reused = llama_perf_context(arms[0].ctx).n_reused; + if (ok && n_reused == 0) { + fprintf(stderr, "%s : graph reuse was not exercised\n", __func__); + ok = false; + } + + if (ok) { + fprintf(stderr, "%s : %u steps matched, %d graphs reused (max rel diff: %s %g, %s %g, %s %g)\n", __func__, n_step, n_reused, + arms[1].name, (double) arms[1].diff_max, arms[2].name, (double) arms[2].diff_max, arms[3].name, (double) arms[3].diff_max); + } + + cleanup(); + return ok; +} + +static test_status test_seq_cp_graph_reuse_for_model(const std::string & model_path, const common_params & base_params) { + common_params params = base_params; + params.model.path = model_path; + + gguf_init_params gparams = { + /*.no_alloc =*/ true, + /*.ctx =*/ nullptr, + }; + gguf_context * meta = gguf_init_from_file(model_path.c_str(), gparams); + if (meta == nullptr) { + LOG_ERR("%s: failed to read model metadata\n", __func__); + return test_status::FAIL; + } + const int64_t key = gguf_find_key(meta, "general.architecture"); + const std::string arch = key >= 0 ? gguf_get_val_str(meta, key) : ""; + gguf_free(meta); + + // the in-place recurrent state is only enabled for qwen35 and qwen35moe + if (arch != "qwen35" && arch != "qwen35moe") { + LOG_INF("%s: skipping test_seq_cp_graph_reuse for %s\n", __func__, arch.c_str()); + return test_status::SKIP; + } + + llama_model_ptr model_scaled{load_model_scaled(params)}; + if (!model_scaled) { + LOG_ERR("%s: failed to create scaled model\n", __func__); + return test_status::FAIL; + } + return test_seq_cp_graph_reuse(params, model_scaled.get()) ? test_status::PASS : test_status::FAIL; +} + // Save a rolled-back single-seq state, restore it into fresh and dirty // contexts, and verify exact logit matches on replay. static test_status test_rollback(const common_params & params, llama_model * model, uint8_t fill) { @@ -608,27 +878,29 @@ int main(int argc, char ** argv) { // silence everything but the table itself (LOG has verbosity LOG_LEVEL_OUTPUT = 0) common_log_set_verbosity_thold(0); - LOG("%-*s %-8s %s\n", (int) name_width, "Model", "rollback", "split replay"); + LOG("%-*s %-8s %-12s %s\n", (int) name_width, "Model", "rollback", "split replay", "seq cp graph reuse"); common_log_flush(common_log_main()); - size_t n_pass[2] = { 0, 0 }; - size_t n_skip[2] = { 0, 0 }; - size_t n_fail[2] = { 0, 0 }; + size_t n_pass[3] = { 0, 0, 0 }; + size_t n_skip[3] = { 0, 0, 0 }; + size_t n_fail[3] = { 0, 0, 0 }; for (const auto & model_path : models) { const auto name = std::filesystem::path(model_path).filename().string(); LOG("%-*s", (int) name_width, name.c_str()); const test_results res = run_tests_for_model(model_path, params); + const test_status graph = (res.rollback == test_status::FAIL || res.replay == test_status::FAIL) + ? test_status::SKIP : test_seq_cp_graph_reuse_for_model(model_path, params); // all status strings have the same raw length, so the columns line up; // pad the first status to the width of the "rollback" header + separator - LOG(" %s %s", test_status_str(res.rollback), test_status_str(res.replay)); + LOG(" %s %s %s", test_status_str(res.rollback), test_status_str(res.replay), test_status_str(graph)); LOG("\n"); common_log_flush(common_log_main()); - const test_status all[2] = { res.rollback, res.replay }; - for (int t = 0; t < 2; ++t) { + const test_status all[3] = { res.rollback, res.replay, graph }; + for (int t = 0; t < 3; ++t) { switch (all[t]) { case test_status::PASS: n_pass[t]++; break; case test_status::FAIL: n_fail[t]++; break; @@ -645,11 +917,17 @@ int main(int argc, char ** argv) { LOG_INF("%s: split replay: %zu passed, %zu skipped, %zu failed (of %zu)\n", __func__, n_pass[1], n_skip[1], n_fail[1], models.size()); - return (n_fail[0] + n_fail[1]) == 0 ? 0 : 1; + LOG_INF("%s: seq cp graph reuse: %zu passed, %zu skipped, %zu failed (of %zu)\n", + __func__, n_pass[2], n_skip[2], n_fail[2], models.size()); + + return (n_fail[0] + n_fail[1] + n_fail[2]) == 0 ? 0 : 1; } // single-model mode const test_results res = run_tests_for_model(params.model.path, params); - return (res.rollback == test_status::FAIL || res.replay == test_status::FAIL) ? 1 : 0; + if (res.rollback == test_status::FAIL || res.replay == test_status::FAIL) { + return 1; + } + return test_seq_cp_graph_reuse_for_model(params.model.path, params) == test_status::FAIL ? 1 : 0; } diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index dee5e17be6dd..b29af241fd52 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -457,14 +457,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)}; - LOGV(LOG_LEVEL_INFO, "\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host"); + LOGV(LOG_LEVEL_INFO, "\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; @@ -522,6 +522,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()) { @@ -770,10 +775,10 @@ struct test_suite { // column headers for the --models table, one per test, in the order they are run static const std::vector test_names = { - "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf", + "baseline", "seq_rm", "state_load", "cp_h", "cp_d", "cp_h_s", "cp_d_s", "rt", "rf", "cp_d_c", }; -// Run the full save/load test suite (tests 1-9) for a single model. +// Run the full save/load test suite (tests 1-10) for a single model. // Returns the per-test results. static test_suite run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { test_suite suite; @@ -852,6 +857,9 @@ static test_suite run_save_load_tests_for_model(const std::string & model_path, // Test 9: state restore failure suite.results.push_back(test_state_restore_failure(model, params, tokens) ? test_status::PASS : test_status::FAIL); + // Test 10: on-device restore with different source and destination chunking. + suite.results.push_back(test_seq_cp_scatter(model, params, tokens, 10, true, true) ? test_status::PASS : test_status::FAIL); + return suite; } diff --git a/tools/cli/README.md b/tools/cli/README.md index acecc74a6d9d..4e66adb987a4 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -57,6 +57,7 @@ | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | +| `--ple FNAME` | GGUF with per-layer embedding (PLE) tables to use instead of the tables in the model file; each tensor in it replaces the lookup table of the same name in the model and must have the same shape (default: unused)
(env: LLAMA_ARG_PLE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index e2ac0668c647..efb5bc81c078 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -140,6 +140,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | +| `--ple FNAME` | GGUF with per-layer embedding (PLE) tables to use instead of the tables in the model file; each tensor in it replaces the lookup table of the same name in the model and must have the same shape (default: unused)
(env: LLAMA_ARG_PLE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index 512a2c024891..70da812f0424 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -68,6 +68,7 @@ test parameters: -fa, --flash-attn (default: auto) -dev, --device (default: auto) -lzm, --lazy-mode (default: auto) + --ple GGUF with PLE tables to use instead of the model's (default: none) -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -embd, --embeddings <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 70e15d044986..8dde3b73fa4c 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"); } @@ -355,6 +357,7 @@ struct cmd_params { std::vector split_mode; std::vector load_mode; std::vector lazy_mode; + std::vector ple; std::vector main_gpu; std::vector no_kv_offload; std::vector flash_attn; @@ -401,6 +404,7 @@ static const cmd_params cmd_params_defaults = { /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, + /* ple */ { "" }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -478,7 +482,8 @@ 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(" --ple GGUF with PLE tables to use instead of the model's (default: none)\n"); 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"); @@ -819,6 +824,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") { @@ -833,6 +840,18 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } params.lazy_mode.insert(params.lazy_mode.end(), modes.begin(), modes.end()); + } else if (arg == "--ple") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + for (auto & f : p) { + if (f == "none") { + f = ""; + } + } + params.ple.insert(params.ple.end(), p.begin(), p.end()); } else if (arg == "-mg" || arg == "--main-gpu") { if (++i >= argc) { invalid_param = true; @@ -1156,6 +1175,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.lazy_mode.empty()) { params.lazy_mode = cmd_params_defaults.lazy_mode; } + if (params.ple.empty()) { + params.ple = cmd_params_defaults.ple; + } if (params.main_gpu.empty()) { params.main_gpu = cmd_params_defaults.main_gpu; } @@ -1226,6 +1248,7 @@ struct cmd_params_instance { llama_split_mode split_mode; llama_load_mode load_mode; llama_lazy_mode lazy_mode; + std::string ple; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1249,6 +1272,7 @@ struct cmd_params_instance { mparams.split_mode = split_mode; mparams.load_mode = load_mode; mparams.lazy_mode = lazy_mode; + mparams.path_ple = ple.empty() ? nullptr : ple.c_str(); mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); mparams.no_host = no_host; @@ -1297,7 +1321,7 @@ struct cmd_params_instance { return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe && split_mode == other.split_mode && main_gpu == other.main_gpu && tensor_split == other.tensor_split && - load_mode == other.load_mode && lazy_mode == other.lazy_mode && + load_mode == other.load_mode && lazy_mode == other.lazy_mode && ple == other.ple && devices == other.devices && no_host == other.no_host && repack == other.repack && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); } @@ -1333,6 +1357,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & sm : params.split_mode) for (const auto & lm : params.load_mode) for (const auto & lzm : params.lazy_mode) + for (const auto & ple : params.ple) for (const auto & mg : params.main_gpu) for (const auto & devs : params.devices) for (const auto & ts : params.tensor_split) @@ -1374,6 +1399,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .split_mode = */ sm, /* .load_mode = */ lm, /* .lazy_mode = */ lzm, + /* .ple = */ ple, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1412,6 +1438,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .split_mode = */ sm, /* .load_mode = */ lm, /* .lazy_mode = */ lzm, + /* .ple = */ ple, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1450,6 +1477,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .split_mode = */ sm, /* .load_mode = */ lm, /* .lazy_mode = */ lzm, + /* .ple = */ ple, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1493,6 +1521,7 @@ struct test { llama_split_mode split_mode; llama_load_mode load_mode; llama_lazy_mode lazy_mode; + std::string ple; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1534,6 +1563,7 @@ struct test { split_mode = inst.split_mode; load_mode = inst.load_mode; lazy_mode = inst.lazy_mode; + ple = inst.ple; main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; flash_attn = inst.flash_attn; @@ -1602,7 +1632,7 @@ struct test { "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", - "tensor_buft_overrides", "load_mode", "lazy_mode", + "tensor_buft_overrides", "load_mode", "lazy_mode", "ple", "embeddings", "no_op_offload", "no_host", "repack", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", @@ -1699,6 +1729,7 @@ struct test { tensor_buft_overrides_str, llama_load_mode_name(load_mode), lazy_mode_str(lazy_mode), + ple.empty() ? "none" : ple, std::to_string(embeddings), std::to_string(no_op_offload), std::to_string(no_host), @@ -2023,6 +2054,9 @@ struct markdown_printer : public printer { if (params.lazy_mode.size() > 1 || params.lazy_mode != cmd_params_defaults.lazy_mode) { fields.emplace_back("lazy_mode"); } + if (params.ple.size() > 1 || params.ple != cmd_params_defaults.ple) { + fields.emplace_back("ple"); + } if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) { fields.emplace_back("embeddings"); } @@ -2098,6 +2132,8 @@ struct markdown_printer : public printer { } else if (field == "t/s") { snprintf(buf, sizeof(buf), "%.2f ± %.2f", t.avg_ts(), t.stdev_ts()); value = buf; + } else if (field == "ple") { + value = t.ple.empty() ? "none" : t.ple.substr(t.ple.find_last_of('/') + 1); } else if (vmap.find(field) != vmap.end()) { value = vmap.at(field); } else { 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/CMakeLists.txt b/tools/server/CMakeLists.txt index 4adaaceefd23..c9e240e3df5d 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -33,11 +33,29 @@ target_include_directories(${TARGET} PRIVATE ../mtmd) target_include_directories(${TARGET} PRIVATE ${PROJECT_SOURCE_DIR}) target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT}) +# stats dashboard page (GET /dashboard), embedded as a byte array + +set(DASHBOARD_HTML ${CMAKE_CURRENT_SOURCE_DIR}/dashboard/dashboard.html) +set(DASHBOARD_HTML_HPP ${CMAKE_CURRENT_BINARY_DIR}/dashboard.html.hpp) + +add_custom_command( + OUTPUT ${DASHBOARD_HTML_HPP} + COMMAND ${CMAKE_COMMAND} + "-DINPUT=${DASHBOARD_HTML}" + "-DOUTPUT=${DASHBOARD_HTML_HPP}" + "-DNAME=dashboard_html" + -P "${CMAKE_CURRENT_SOURCE_DIR}/dashboard/embed.cmake" + DEPENDS ${DASHBOARD_HTML} ${CMAKE_CURRENT_SOURCE_DIR}/dashboard/embed.cmake + COMMENT "Embedding dashboard.html" + VERBATIM +) + # llama-server-impl: server logic, reusable by app set(TARGET llama-server-impl) add_library(${TARGET} + ${DASHBOARD_HTML_HPP} server.cpp server-http.cpp server-http.h diff --git a/tools/server/README.md b/tools/server/README.md index 09fdb87e5daf..8a26ec3dc296 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -74,9 +74,14 @@ For the full list of features, please refer to [server's changelog](https://gith | `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | +| `--ple FNAME` | GGUF with per-layer embedding (PLE) tables to use instead of the tables in the model file; each tensor in it replaces the lookup table of the same name in the model and must have the same shape (default: unused)
(env: LLAMA_ARG_PLE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | + + +| Environment variable | Effect | +| -------------------- | ------ | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | @@ -232,8 +237,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) | @@ -402,6 +415,16 @@ Recommended `--cors-origins` setting, depending on where the server runs: Related flags: `--cors-origins`, `--cors-methods`, `--cors-headers`, `--cors-credentials` / `--no-cors-credentials`. Background and rationale: [#25655](https://github.com/ggml-org/llama.cpp/pull/25655). + +### Environment-only knobs (ggml/ggml-hip internals, no CLI argument) + +| Environment variable | Effect | +| -------------------- | ------ | +| `GGML_CUDA_POOL_FREE_MIN_MB=N` | device scratch blocks at or above N MiB (e.g. QSA prefill/decode intermediates that scale with the KV length) are freed on release instead of staying cached in the legacy/non-VMM pool | +| `QSA_DIRECT` | `-1/auto` (default) picks direct F16 reads when the batch is small against the cache, `0` always packs, `1` always reads direct; both paths are bit-identical | +| `GGML_QSA_RANKSORT=1` | force the O(n^2) rank sort in the QSA indexer path | +| `HC_NO_MIX_FUSE=1` | disable the hc_mix stream-mix fusion (pre-existing) | + ## Build `llama-server` is built alongside everything else from the root of the project @@ -1145,6 +1168,54 @@ In *router mode* the query param `?model={model_id}` has to be set. This endpoin | `llamacpp:spec_decode_num_drafts_total` | Counter | Total speculative decoding verification steps (0 when spec-decode is off). | | `llamacpp:spec_decode_num_accepted_tokens_per_pos_total` | Counter | Accepted tokens per draft position (labeled `position="N"`; absent when spec-decode is off or before the first completed speculative request). | +### GET `/dashboard`: Live stats dashboard + +A single self-contained page (no external resources, works offline, light/dark, phone-sized screens) that polls `/dashboard/stats` every 1.5 s and shows, per model: + +- decode and prefill tokens/s: live (running slots), last request, average since start, and a 10-minute throughput chart for each +- speculative / MTP draft acceptance rate and mean accepted length (when a `--spec-type` is set) +- slots and their state, prompt progress, KV/context usage per slot and in total, deferred (queued) requests +- the last finished requests with their timings, and the model/configuration (file, size, type, architecture, context, KV cache type, speculative settings, PLE on disk) + +It needs no flag and is independent of the chat UI (it is also served with `--no-ui`). The page itself is public and holds no data; the JSON it reads requires the API key when `--api-key` is set (the page asks for it and keeps it in the browser's localStorage). No prompts or generated text are exposed. + +In *router mode* (`--models-preset` / `--models-dir`) the dashboard lists every model with its status (`loaded`, `sleeping`, `loading`, `unloaded`, failed) and shows the live stats of the running ones. + +### GET `/dashboard/stats`: Dashboard data + +Returns the data behind `/dashboard`. Polling it does not reset the `--sleep-idle-seconds` timer and does not wake a sleeping model (the stats from before sleep are returned). In router mode the router queries each running child (in parallel, 1.5 s timeout) and adds the models that are not running. + +```json +{ + "mode": "single", // or "router" + "build_info": "b11336-2e63e11b6", + "t_now_ms": 1790946801345, + "models": [{ + "id": "qwen35-2b", + "status": "loaded", // loaded | sleeping | loading | unloaded | downloading + "info": { + "file": "Qwen3.5-2B-UD-Q4_K_XL.gguf", "arch": "qwen35", "ftype": "Q4_K - Medium", + "size_bytes": 1328790784, "n_params": 1881825088, "n_ctx_train": 262144, + "config": { "n_ctx": 8192, "n_ctx_slot": 4096, "n_parallel": 2, "cache_type_k": "f16", + "speculative": { "types": "ngram-simple", "n_max": 3, "n_min": 0 }, "lazy_mode": "auto", ... } + }, + "stats": { // null when the model is not running + "slots": [{ "id": 0, "state": "generating", "n_ctx": 4096, "n_tokens": 118, "n_gen": 45, + "prompt_tps": 126.0, "gen_tps": 20.9, "draft": { "acceptance": 0.06, "mean_accepted_len": 3.0, ... } }], + "kv": { "n_ctx": 8192, "n_ctx_slot": 4096, "n_used": 226, "unified": false }, + "queue": { "n_slots": 2, "n_processing": 1, "n_deferred": 0 }, + "totals": { "n_requests": 9, "n_prompt": 157, "n_gen": 779, "prompt_tps": 133.0, "gen_tps": 21.6, + "draft": { "n_draft": 66, "n_accepted": 4, "acceptance": 0.06, "mean_accepted_len": 3.0, "acceptance_per_pos": [1.0, 1.0] } }, + "recent": [{ "t_end_ms": 1790946821451, "id_slot": 0, "n_prompt": 20, "n_cached": 0, "n_gen": 111, + "prompt_tps": 166.1, "gen_tps": 26.1, "draft": { ... }, "stop": "eos", "truncated": false }], + "series": { "period_s": 2.0, "t_ms": [...], "gen_tps": [...], "prompt_tps": [...] } // last 10 minutes + } + }] +} +``` + +`gen_tps` follows the server's timings (decode steps per second, the first token is excluded); `series` is the aggregate throughput over all slots in 2 s buckets; `recent` keeps the last 64 requests. + ### POST `/slots/{id_slot}?action=save`: Save the prompt cache of the specified slot to a file. *Options:* @@ -1323,6 +1394,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 +1504,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 +1519,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 +1531,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. @@ -1794,6 +1875,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/dashboard/dashboard.html b/tools/server/dashboard/dashboard.html new file mode 100644 index 000000000000..55fceb08b5cb --- /dev/null +++ b/tools/server/dashboard/dashboard.html @@ -0,0 +1,617 @@ + + + + + + +llama-server dashboard + + + +
+ ← Chat +

llama-server

+ + + connecting + + +
+
+ + + + diff --git a/tools/server/dashboard/embed.cmake b/tools/server/dashboard/embed.cmake new file mode 100644 index 000000000000..c49c29be9fc8 --- /dev/null +++ b/tools/server/dashboard/embed.cmake @@ -0,0 +1,21 @@ +# Embed a file as a C byte array header. +# usage: cmake -DINPUT= -DOUTPUT=
-DNAME= -P embed.cmake + +file(READ "${INPUT}" hex HEX) +string(LENGTH "${hex}" hex_len) +math(EXPR n_bytes "${hex_len} / 2") + +# 0xNN, per byte, a line break every 32 bytes +string(REGEX REPLACE "([0-9a-f][0-9a-f])" "0x\\1," bytes "${hex}") +string(REGEX REPLACE "((0x[0-9a-f][0-9a-f],){32})" "\\1\n" bytes "${bytes}") + +file(WRITE "${OUTPUT}.tmp" +"// Generated by tools/server/dashboard/embed.cmake from ${NAME} - do not edit.\n" +"#pragma once\n" +"#include \n" +"static const unsigned char ${NAME}[] = {\n${bytes}\n};\n" +"static const size_t ${NAME}_len = ${n_bytes};\n") + +# only touch the header when the contents change, to avoid needless rebuilds +execute_process(COMMAND ${CMAKE_COMMAND} -E copy_if_different "${OUTPUT}.tmp" "${OUTPUT}") +file(REMOVE "${OUTPUT}.tmp") diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 8cf794f1e24f..14319681a086 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; } @@ -1417,12 +1417,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 8cb6b90da62f..f9a2510cb09e 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(); @@ -326,6 +328,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 2036319c1381..59a92ac61d54 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,10 +14,14 @@ #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" #include +#include +#include #include #include #include @@ -200,6 +204,7 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + llama_context * ctx_spf = nullptr; common_memory mem; @@ -217,6 +222,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; @@ -294,6 +302,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(); } @@ -363,6 +374,9 @@ struct server_slot { // clear multimodal state mbatch.reset(); + + spec_prefill_active = false; + spec_prefill_tokens.clear(); } void init_sampler() const { @@ -441,6 +455,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"); @@ -478,10 +500,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", @@ -497,7 +521,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) { @@ -521,6 +545,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(); @@ -592,11 +620,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 { @@ -667,7 +702,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); @@ -695,6 +730,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; @@ -853,6 +892,13 @@ struct server_context_impl { metrics.reset_bucket(); } + // live dashboard stats, see dash_to_json() + // note: not thread-safe, call from the main loop (or right before sleep) + // with_queue = false when the caller holds the task queue lock (sleeping state callback) + json get_dashboard_stats(bool with_queue = true) { + return dash_to_json(with_queue); + } + private: // note: accessing these fields outside of this class is not thread-safe // use server_context methods instead @@ -871,6 +917,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; @@ -903,6 +955,38 @@ struct server_context_impl { int64_t t_prompt_start = 0; // start of the oldest queued prompt decode uint64_t n_prompt_queued = 0; + // dashboard stats, see GET /dashboard/stats + // note: only touched from the main loop (and from update_cached_responses right before sleep) + struct dash_request { + int64_t t_end_ms; // unix time, ms + int32_t id_slot; + uint64_t n_prompt; // processed prompt tokens (excluding cached) + uint64_t n_cached; // prompt tokens reused from the cache + uint64_t n_gen; // generated tokens + double t_prompt_ms; + double t_gen_ms; + double prompt_tps; + double gen_tps; + uint64_t n_draft; + uint64_t n_draft_accepted; + uint64_t n_draft_verif_steps; + int stop; // stop_type + bool truncated; + }; + struct dash_sample { + int64_t t_us; // ggml_time_us() + int64_t t_unix_ms; + uint64_t n_gen; // cumulative generated tokens (finished + in flight) + uint64_t n_prompt; // cumulative processed prompt tokens + }; + static constexpr size_t DASH_MAX_REQUESTS = 64; + static constexpr size_t DASH_MAX_SAMPLES = 300; // 10 minutes at 2 s + static constexpr int64_t DASH_SAMPLE_PERIOD_US = 2000000; + + std::deque dash_requests; + std::deque dash_samples; + uint64_t dash_n_requests = 0; + json json_ui_settings = json::object(); // Necessary similarity of prompt for slot selection @@ -918,6 +1002,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; @@ -982,6 +1073,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) { @@ -992,6 +1130,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; @@ -1131,6 +1277,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"}}); @@ -1271,6 +1518,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(); @@ -1289,6 +1537,7 @@ struct server_context_impl { if (slot.stats.n_gen > 0) { metrics_on_prediction(slot); } + dash_on_request_done(slot); }; slot.reset(); @@ -1466,6 +1715,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 }; @@ -1686,6 +1942,9 @@ struct server_context_impl { } bool launch_slot_with_task(server_slot & slot, server_task && task) { + // anchor the throughput series before the counters start moving after an idle period + dash_sample(); + // process per-request lora adapters if (!task.params.lora.empty()) { auto task_loras = construct_lora_list(task.params.lora); @@ -1758,6 +2017,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 { @@ -2094,6 +2373,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(); @@ -2352,10 +2632,122 @@ struct server_context_impl { cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); } + // checkpoints are appended to the slot save file, after the llama state payload + // they cannot be recreated from the final state alone (a recurrent state cannot be rewound) + static constexpr uint32_t SLOT_CKPT_MAGIC = 0x504b4353; // "SCKP" + static constexpr uint32_t SLOT_CKPT_VERSION = 1; + + static bool ckpt_read(std::ifstream & ifs, void * dst, size_t size, size_t & n_read) { + if (!ifs.read((char *) dst, size)) { + return false; + } + n_read += size; + return true; + } + + static bool ckpt_read_buf(std::ifstream & ifs, std::vector & buf, size_t & n_read) { + uint64_t n = 0; + // 16 GiB cap, in case the size field itself is corrupted + if (!ckpt_read(ifs, &n, sizeof(n), n_read) || n > (1ull << 34)) { + return false; + } + buf.resize(n); + return n == 0 || ckpt_read(ifs, buf.data(), n, n_read); + } + + static void ckpt_write(std::ofstream & ofs, const void * src, size_t size, size_t & n_written) { + ofs.write((const char *) src, size); + n_written += size; + } + + static void ckpt_write_buf(std::ofstream & ofs, const std::vector & buf, size_t & n_written) { + const uint64_t n = buf.size(); + ckpt_write(ofs, &n, sizeof(n), n_written); + if (n > 0) { + ckpt_write(ofs, buf.data(), n, n_written); + } + } + + size_t save_slot_checkpoints(const std::string & filepath, const server_slot & slot) const { + if (slot.prompt.checkpoints.empty()) { + return 0; + } + std::ofstream ofs(filepath, std::ios::binary | std::ios::app); + if (!ofs) { + SRV_WRN("failed to append context checkpoints to '%s'\n", filepath.c_str()); + return 0; + } + size_t n_written = 0; + const uint32_t magic = SLOT_CKPT_MAGIC; + const uint32_t version = SLOT_CKPT_VERSION; + const uint32_t count = (uint32_t) slot.prompt.checkpoints.size(); + ckpt_write(ofs, &magic, sizeof(magic), n_written); + ckpt_write(ofs, &version, sizeof(version), n_written); + ckpt_write(ofs, &count, sizeof(count), n_written); + for (const auto & cur : slot.prompt.checkpoints) { + ckpt_write(ofs, &cur.n_tokens, sizeof(cur.n_tokens), n_written); + ckpt_write(ofs, &cur.pos_min, sizeof(cur.pos_min), n_written); + ckpt_write(ofs, &cur.pos_max, sizeof(cur.pos_max), n_written); + ckpt_write_buf(ofs, cur.data_tgt, n_written); + ckpt_write_buf(ofs, cur.data_dft, n_written); + ckpt_write_buf(ofs, cur.data_spec, n_written); + } + ofs.flush(); + if (!ofs) { + SRV_WRN("failed to append context checkpoints to '%s' - the appendix is incomplete\n", filepath.c_str()); + return 0; + } + SRV_INF("appended %u context checkpoint(s) (%.3f MiB) to '%s'\n", + count, (float) n_written / 1024 / 1024, filepath.c_str()); + return n_written; + } + + // returns the number of bytes consumed, 0 if there is no usable appendix + size_t load_slot_checkpoints(const std::string & filepath, size_t offset, server_slot & slot) const { + std::ifstream ifs(filepath, std::ios::binary); + if (!ifs || !ifs.seekg(offset)) { + return 0; + } + size_t n_read = 0; + uint32_t magic = 0; + uint32_t version = 0; + uint32_t count = 0; + if (!ckpt_read(ifs, &magic, sizeof(magic), n_read) || magic != SLOT_CKPT_MAGIC) { + return 0; + } + if (!ckpt_read(ifs, &version, sizeof(version), n_read) || version != SLOT_CKPT_VERSION || + !ckpt_read(ifs, &count, sizeof(count), n_read) || count > 1024) { + SRV_WRN("invalid context checkpoint appendix in '%s' - ignored\n", filepath.c_str()); + return 0; + } + std::list checkpoints; + for (uint32_t i = 0; i < count; ++i) { + common_prompt_checkpoint cur; + cur.id_task = -1; + if (!ckpt_read(ifs, &cur.n_tokens, sizeof(cur.n_tokens), n_read) || + !ckpt_read(ifs, &cur.pos_min, sizeof(cur.pos_min), n_read) || + !ckpt_read(ifs, &cur.pos_max, sizeof(cur.pos_max), n_read) || + !ckpt_read_buf(ifs, cur.data_tgt, n_read) || + !ckpt_read_buf(ifs, cur.data_dft, n_read) || + !ckpt_read_buf(ifs, cur.data_spec, n_read)) { + SRV_WRN("truncated context checkpoint appendix in '%s' - ignored\n", filepath.c_str()); + return 0; + } + checkpoints.push_back(std::move(cur)); + } + while (checkpoints.size() > (size_t) params_base.n_ctx_checkpoints) { + checkpoints.pop_front(); + } + slot.prompt.checkpoints = std::move(checkpoints); + SRV_INF("restored %zu context checkpoint(s) from '%s'\n", slot.prompt.checkpoints.size(), filepath.c_str()); + return n_read; + } + // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { // while yielding, an encode / decode is running and only reading the server state is safe - if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS && task.type != SERVER_TASK_TYPE_SLOT_GET) { + if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS && task.type != SERVER_TASK_TYPE_SLOT_GET && + task.type != SERVER_TASK_TYPE_DASHBOARD) { SRV_DBG("decoding, decline task, id_task = %d\n", task.id); return false; } @@ -2501,6 +2893,13 @@ struct server_context_impl { } queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_DASHBOARD: + { + auto res = std::make_unique(); + res->id = task.id; + res->data = dash_to_json(true); + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_SLOT_GET: { json slots_data = json::array(); @@ -2560,6 +2959,8 @@ struct server_context_impl { break; } + const size_t nwrite_ckpt = save_slot_checkpoints(filepath, *slot); + const int64_t t_end = ggml_time_us(); const double t_save_ms = (t_end - t_start) / 1000.0; @@ -2569,7 +2970,7 @@ struct server_context_impl { res->filename = filename; res->is_save = true; res->n_tokens = slot->prompt.tokens.size(); - res->n_bytes = nwrite; + res->n_bytes = nwrite + nwrite_ckpt; res->t_ms = t_save_ms; queue_results.send(std::move(res)); } break; @@ -2625,6 +3026,9 @@ struct server_context_impl { break; } + // nread is the end offset of the llama state payload within the file + const size_t nread_ckpt = load_slot_checkpoints(filepath, nread, *slot); + const int64_t t_end = ggml_time_us(); const double t_restore_ms = (t_end - t_start) / 1000.0; @@ -2634,7 +3038,7 @@ struct server_context_impl { res->filename = filename; res->is_save = false; res->n_tokens = slot->prompt.tokens.size(); - res->n_bytes = nread; + res->n_bytes = nread + nread_ckpt; res->t_ms = t_restore_ms; queue_results.send(std::move(res)); } break; @@ -2959,6 +3363,14 @@ struct server_context_impl { std::vector generating; std::vector drafting; + // --spec-draft-max-slots: with many slots generating, start no new drafts + int n_generating = 0; + iterate(slots, [&](server_slot & slot) { + n_generating += slot.state == SLOT_STATE_GENERATING; + }); + const int spec_max_slots = params_base.speculative.draft.max_slots; + const bool spec_new_ok = spec_max_slots <= 0 || n_generating <= spec_max_slots; + // determine which slots are generating and drafting iterate(slots, [&](server_slot & slot) { if (slot.state != SLOT_STATE_GENERATING) { @@ -2990,7 +3402,7 @@ struct server_context_impl { if (use_ckpt_tgt) { GGML_ASSERT(!slot.spec_ckpt.empty()); } - } else { + } else if (spec_new_ok) { GGML_ASSERT(slot.spec_i_batch.empty()); slot.spec_ckpt.update_pos( @@ -2999,7 +3411,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(); @@ -3038,7 +3450,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)) { @@ -3057,7 +3469,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); @@ -3069,7 +3481,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); } } }); @@ -3112,19 +3524,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) { @@ -3161,33 +3577,33 @@ 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; @@ -3195,7 +3611,7 @@ struct server_context_impl { const bool is_stateless_task = slot.task->type == SERVER_TASK_TYPE_EMBEDDING || slot.task->type == SERVER_TASK_TYPE_RERANK; - if (slot.task->params.cache_prompt && !is_stateless_task) { + if (slot.task->params.cache_prompt && !slot.spec_prefill_active && !is_stateless_task) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -3272,7 +3688,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)); @@ -3288,7 +3704,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; @@ -3313,7 +3729,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; @@ -3381,8 +3797,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); } @@ -3408,7 +3824,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; } } @@ -3459,7 +3875,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; @@ -3502,7 +3918,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) { @@ -3520,9 +3936,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); @@ -3548,7 +3965,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; } @@ -3564,15 +3981,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 @@ -3893,11 +4314,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); @@ -3923,10 +4357,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); @@ -4095,6 +4529,8 @@ struct server_context_impl { slot.stats.set_prompt_last(t_now); } } + + dash_sample(); } // flush any queued prompt metrics if all slots are now idle @@ -4129,6 +4565,236 @@ struct server_context_impl { dst[i] += src[i]; } } + + // + // dashboard + // + + static int64_t dash_unix_ms() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count(); + } + + static const char * dash_slot_state_str(slot_state state) { + switch (state) { + case SLOT_STATE_IDLE: return "idle"; + case SLOT_STATE_WAIT_OTHER: return "waiting"; + case SLOT_STATE_STARTED: return "started"; + case SLOT_STATE_PROCESSING_PROMPT: return "prompt"; + case SLOT_STATE_DONE_PROMPT: return "prompt_done"; + case SLOT_STATE_GENERATING: return "generating"; + } + return "unknown"; + } + + static const char * dash_stop_str(int stop) { + switch (stop) { + case STOP_TYPE_EOS: return "eos"; + case STOP_TYPE_WORD: return "word"; + case STOP_TYPE_LIMIT: return "limit"; + default: return "none"; + } + } + + // cumulative token counters, including the requests that are still running + uint64_t dash_gen_total() const { + uint64_t n = metrics.predict.count; + for (const auto & slot : slots) { + n += slot.stats.n_gen; + } + return n; + } + + uint64_t dash_prompt_total() const { + return metrics.prompt.count + n_prompt_queued; + } + + // take a throughput sample, at most once per DASH_SAMPLE_PERIOD_US + void dash_sample() { + const int64_t t_now = ggml_time_us(); + if (!dash_samples.empty() && t_now - dash_samples.back().t_us < DASH_SAMPLE_PERIOD_US) { + return; + } + dash_samples.push_back({ t_now, dash_unix_ms(), dash_gen_total(), dash_prompt_total() }); + while (dash_samples.size() > DASH_MAX_SAMPLES) { + dash_samples.pop_front(); + } + } + + void dash_on_request_done(const server_slot & slot) { + const auto & st = slot.stats; + if (!st.is_set() || (st.n_gen == 0 && st.n_prompt_processed == 0 && st.n_prompt_cached == 0)) { + return; + } + dash_n_requests++; + dash_requests.push_back({ + /* t_end_ms */ dash_unix_ms(), + /* id_slot */ slot.id, + /* n_prompt */ st.n_prompt_processed, + /* n_cached */ st.n_prompt_cached, + /* n_gen */ st.n_gen, + /* t_prompt_ms */ st.t_prompt_ms(), + /* t_gen_ms */ st.t_gen_ms(), + /* prompt_tps */ st.n_prompt_tps(), + /* gen_tps */ st.n_gen_tps(), + /* n_draft */ st.n_draft_tokens, + /* n_draft_accepted */ st.n_draft_accepted, + /* n_draft_verif_steps */ st.n_draft_verif_steps, + /* stop */ (int) slot.stop, + /* truncated */ slot.truncated, + }); + while (dash_requests.size() > DASH_MAX_REQUESTS) { + dash_requests.pop_front(); + } + } + + static json dash_draft_json(uint64_t n_draft, uint64_t n_accepted, uint64_t n_steps) { + return json { + {"n_draft", n_draft}, + {"n_accepted", n_accepted}, + {"n_steps", n_steps}, + {"acceptance", n_draft > 0 ? (double) n_accepted / (double) n_draft : 0.0}, + {"mean_accepted_len", n_steps > 0 ? 1.0 + (double) n_accepted / (double) n_steps : 0.0}, + }; + } + + // the live part of GET /dashboard/stats; no prompts or generated text in here + json dash_to_json(bool with_queue) { + dash_sample(); + + const int64_t t_now = ggml_time_us(); + + json jslots = json::array(); + int n_processing = 0; + uint64_t n_kv_used = 0; + for (const auto & slot : slots) { + const auto & st = slot.stats; + const bool busy = slot.is_processing(); + n_processing += busy ? 1 : 0; + n_kv_used += slot.prompt.n_tokens(); + + json js = { + {"id", slot.id}, + {"state", dash_slot_state_str(slot.state)}, + {"n_ctx", slot.n_ctx}, + {"n_tokens", slot.prompt.n_tokens()}, // tokens held in the KV cache by this slot + {"speculative", slot.can_speculate()}, + }; + if (busy && slot.task) { + js["id_task"] = slot.task->id; + js["n_prompt_tokens"] = slot.task->n_tokens(); + js["n_prompt_processed"] = st.n_prompt_processed; + js["n_prompt_cached"] = st.n_prompt_cached; + js["n_gen"] = st.n_gen; + js["n_remaining"] = slot.n_remaining(); + js["elapsed_ms"] = st.is_set() ? st.t_elapsed_us() / 1000.0 : 0.0; + if (slot.state == SLOT_STATE_PROCESSING_PROMPT && st.is_set()) { + // still processing: rate over the time spent so far + const double t_ms = (t_now - st.t_start) / 1000.0; + js["prompt_tps"] = t_ms > 0.0 ? 1e3 / t_ms * st.n_prompt_processed : 0.0; + } else { + js["prompt_tps"] = st.n_prompt_tps(); + } + js["gen_tps"] = st.n_gen_tps(); + js["draft"] = dash_draft_json(st.n_draft_tokens, st.n_draft_accepted, st.n_draft_verif_steps); + } + jslots.push_back(std::move(js)); + } + + // totals: finished requests (server_metrics) + the in-flight ones + uint64_t n_draft = metrics.n_draft_tokens, n_acc = metrics.n_draft_accepted, n_steps = metrics.n_draft_verif_steps; + for (const auto & slot : slots) { + n_draft += slot.stats.n_draft_tokens; + n_acc += slot.stats.n_draft_accepted; + n_steps += slot.stats.n_draft_verif_steps; + } + json jdraft = dash_draft_json(n_draft, n_acc, n_steps); + { + json per_pos = json::array(); + if (metrics.n_draft_verif_steps > 0) { + // drop the trailing never-accepted positions (the vector is sized for n_max) + size_t n = metrics.n_accepted_per_pos.size(); + while (n > 0 && metrics.n_accepted_per_pos[n - 1] == 0) { + n--; + } + for (size_t i = 0; i < n; ++i) { + per_pos.push_back((double) metrics.n_accepted_per_pos[i] / (double) metrics.n_draft_verif_steps); + } + } + jdraft["acceptance_per_pos"] = per_pos; + } + + json jrecent = json::array(); + for (auto it = dash_requests.rbegin(); it != dash_requests.rend(); ++it) { + const auto & r = *it; + jrecent.push_back({ + {"t_end_ms", r.t_end_ms}, + {"id_slot", r.id_slot}, + {"n_prompt", r.n_prompt}, + {"n_cached", r.n_cached}, + {"n_gen", r.n_gen}, + {"t_prompt_ms", r.t_prompt_ms}, + {"t_gen_ms", r.t_gen_ms}, + {"prompt_tps", r.prompt_tps}, + {"gen_tps", r.gen_tps}, + {"draft", dash_draft_json(r.n_draft, r.n_draft_accepted, r.n_draft_verif_steps)}, + {"stop", dash_stop_str(r.stop)}, + {"truncated", r.truncated}, + }); + } + + // throughput series: rate between consecutive samples + json s_t = json::array(), s_gen = json::array(), s_prompt = json::array(); + for (size_t i = 1; i < dash_samples.size(); ++i) { + const auto & a = dash_samples[i - 1]; + const auto & b = dash_samples[i]; + const double dt = (b.t_us - a.t_us) / 1e6; + if (dt <= 0.0) { + continue; + } + s_t .push_back(b.t_unix_ms); + s_gen .push_back(b.n_gen >= a.n_gen ? (b.n_gen - a.n_gen) / dt : 0.0); + s_prompt.push_back(b.n_prompt >= a.n_prompt ? (b.n_prompt - a.n_prompt) / dt : 0.0); + } + + const uint64_t t_start_unix_ms = dash_unix_ms() - (t_now - metrics.t_start) / 1000; + + return json { + {"t_now_ms", dash_unix_ms()}, + {"t_start_ms", t_start_unix_ms}, + {"slots", jslots}, + {"kv", { + {"n_ctx", n_ctx}, + {"n_ctx_slot", n_ctx_slot()}, + {"n_used", n_kv_used}, + {"unified", params_base.kv_unified}, + }}, + {"queue", { + {"n_slots", (int) slots.size()}, + {"n_processing", n_processing}, + {"n_deferred", with_queue ? queue_tasks.queue_tasks_deferred_size() : 0}, + }}, + {"totals", { + {"n_requests", dash_n_requests}, + {"n_prompt", metrics.prompt.count}, + {"n_prompt_cached", metrics.n_prompt_cached}, + {"t_prompt_s", metrics.prompt.time / 1e6}, + {"prompt_tps", metrics.prompt.n_per_second()}, + {"n_gen", dash_gen_total()}, + {"t_gen_s", metrics.predict.time / 1e6}, + {"gen_tps", metrics.predict.n_per_second()}, + {"n_decode", metrics.n_decode}, + {"draft", jdraft}, + }}, + {"recent", jrecent}, + {"series", { + {"period_s", DASH_SAMPLE_PERIOD_US / 1e6}, + {"t_ms", s_t}, + {"gen_tps", s_gen}, + {"prompt_tps", s_prompt}, + }}, + }; + } }; // @@ -4167,6 +4833,13 @@ server_context_meta server_context::get_meta() const { const char * ftype_name = llama_ftype_name(llama_model_ftype(impl->model_tgt)); + char model_arch[128] = {0}; + if (llama_model_meta_val_str(impl->model_tgt, "general.architecture", model_arch, sizeof(model_arch)) < 0) { + model_arch[0] = '\0'; + } + char model_desc[256] = {0}; + llama_model_desc(impl->model_tgt, model_desc, sizeof(model_desc)); + return server_context_meta { /* build_info */ std::string(llama_build_info()), /* model_name */ impl->model_name, @@ -4202,6 +4875,8 @@ server_context_meta server_context::get_meta() const { /* model_n_params */ llama_model_n_params(impl->model_tgt), /* model_size */ llama_model_size(impl->model_tgt), /* model_ftype */ ftype_name, + /* model_arch */ std::string(model_arch), + /* model_desc */ std::string(model_desc), }; } @@ -4624,6 +5299,69 @@ static json get_res_props(const server_context_meta & meta, const common_params return props; } +// static model / config info for GET /dashboard/stats +// note: do NOT use ctx_server here, this must be accessible during sleep +static json get_res_dashboard_info(const server_context_meta & meta, const common_params & params) { + const auto & spec = params.speculative; + const bool has_spec = !(spec.types.empty() || + (spec.types.size() == 1 && spec.types[0] == COMMON_SPECULATIVE_TYPE_NONE)); + + json jspec = nullptr; + if (has_spec) { + std::string types; + for (const auto t : spec.types) { + if (t == COMMON_SPECULATIVE_TYPE_NONE) { + continue; + } + types += (types.empty() ? "" : ",") + common_speculative_type_to_str(t); + } + jspec = { + {"types", types}, + {"n_max", spec.draft.n_max}, + {"n_min", spec.draft.n_min}, + }; + if (spec.has_dft()) { + jspec["draft_model"] = std::filesystem::path(spec.draft.mparams.path).filename().string(); + } + if (spec.draft.mtp_vocab > 0) { + jspec["mtp_vocab"] = spec.draft.mtp_vocab; + } + } + + return json { + {"name", meta.model_name}, + {"aliases", meta.model_aliases}, + {"file", std::filesystem::path(meta.model_path).filename().string()}, + {"arch", meta.model_arch}, + {"desc", meta.model_desc}, + {"ftype", meta.model_ftype}, + {"size_bytes", meta.model_size}, + {"n_params", meta.model_n_params}, + {"n_ctx_train", meta.model_n_ctx_train}, + {"n_vocab", meta.model_vocab_n_tokens}, + {"modalities", { + {"vision", meta.has_inp_image}, + {"audio", meta.has_inp_audio}, + {"video", meta.has_inp_video}, + }}, + {"config", { + {"n_ctx", params.n_ctx}, + {"n_ctx_slot", meta.slot_n_ctx}, + {"n_parallel", params.n_parallel}, + {"kv_unified", params.kv_unified}, + {"cache_type_k", ggml_type_name(params.cache_type_k)}, + {"cache_type_v", ggml_type_name(params.cache_type_v)}, + {"flash_attn", llama_flash_attn_type_name(params.flash_attn_type)}, + {"n_gpu_layers", params.n_gpu_layers}, + {"speculative", jspec}, + {"lazy_mode", params.lazy_mode == LLAMA_LAZY_MODE_OFF ? "off" + : params.lazy_mode == LLAMA_LAZY_MODE_AUTO ? "auto" + : params.lazy_mode == LLAMA_LAZY_MODE_ON ? "on" + : params.lazy_mode == LLAMA_LAZY_MODE_DIRECT ? "direct" : "unknown"}, + }}, + }; +} + json server_routes::get_model_info() const { return get_res_model_info(*meta); } @@ -4707,6 +5445,59 @@ void server_routes::init_routes() { return res; }; + this->get_dashboard_stats = [this](const server_http_req & req) { + // must not wake the model up, see task_resets_idle_timer() + auto res = create_response(true); + + json stats = nullptr; + bool is_sleeping = queue_tasks.is_sleeping(); + + if (!is_sleeping) { + { + server_task task(SERVER_TASK_TYPE_DASHBOARD); + task.id = res->rd.get_new_id(); + res->rd.post_task(std::move(task), true); // high-priority task + } + + // a task posted right before sleeping is never processed, do not wait for it + auto result = res->rd.next([&]{ + return req.should_stop() || queue_tasks.is_sleeping(); + }); + if (!result) { + if (req.should_stop()) { + return res; + } + is_sleeping = true; + } else if (result->is_error()) { + res->error(result->to_json()); + return res; + } else { + stats = result->to_json(); + } + } + + if (is_sleeping) { + std::unique_lock lock(mutex_cache); + stats = cached_dashboard; + } + + json model = { + {"id", meta->model_name}, + {"status", is_sleeping ? "sleeping" : "loaded"}, + {"info", get_res_dashboard_info(*meta, params)}, + {"stats", stats}, + }; + + res->ok({ + {"mode", "single"}, + {"build_info", meta->build_info}, + {"t_now_ms", std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count()}, + {"models", json::array({ model })}, + }); + return res; + }; + this->get_slots = [this](const server_http_req & req) { auto res = create_response(); if (!params.endpoint_slots) { @@ -5545,6 +6336,8 @@ void server_routes::update_cached_responses(bool is_sleeping) { cached_models = get_res_models(*meta); cached_props = get_res_props(*meta, params, true); cached_metrics = ctx_server.get_metrics(); + // note: called with the task queue locked, so the queue size must not be read + cached_dashboard = ctx_server.get_dashboard_stats(false); should_reset_buckets = false; diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 7265ccad156d..b94e97323c0a 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -52,6 +52,8 @@ struct server_context_meta { uint64_t model_n_params; uint64_t model_size; std::string model_ftype; + std::string model_arch; // general.architecture + std::string model_desc; // llama_model_desc() }; enum server_state { @@ -154,6 +156,7 @@ struct server_routes { server_http_context::handler_t post_rerank; server_http_context::handler_t get_lora_adapters; server_http_context::handler_t post_lora_adapters; + server_http_context::handler_t get_dashboard_stats; // to be used in router mode json get_model_info() const; @@ -186,6 +189,7 @@ struct server_routes { json cached_models = nullptr; json cached_props = nullptr; server_metrics cached_metrics; + json cached_dashboard = nullptr; // set when a scrape during sleep already reported the throughput buckets bool should_reset_buckets = false; // call right before sleep to update the cached responses diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index e46a1f6f1020..9856f191f116 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -1,6 +1,7 @@ #include "common.h" #include "http.h" #include "server-http.h" +#include "dashboard.html.hpp" #include "server-common.h" #include "ui.h" @@ -240,7 +241,7 @@ bool server_http_context::init_listener(const common_params & params) { // Frontend paths - all embedded UI assets static const std::unordered_set frontend_paths = []() { - std::unordered_set paths { "/" }; + std::unordered_set paths { "/", "/dashboard" }; for (const llama_ui_asset & a : llama_ui_get_assets()) { paths.insert("/" + a.name); } @@ -368,6 +369,15 @@ bool server_http_context::init_listener(const common_params & params) { // Web UI setup // + // stats dashboard: a single self-contained page, independent of the chat UI + // the page is public (it holds no data), the JSON it polls (/dashboard/stats) needs the API key + srv->Get(params.api_prefix + "/dashboard", [](const httplib::Request &, httplib::Response & res) { + res.set_header("Cache-Control", "no-cache"); + res.set_header("X-Content-Type-Options", "nosniff"); + res.set_content(reinterpret_cast(dashboard_html), dashboard_html_len, "text/html; charset=utf-8"); + return false; + }); + // Use new `params.ui` field (backed by old `params.webui` for compat) if (!params.ui) { SRV_INF("%s", "The UI is disabled\n"); diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index a8996b00c950..e9f4bf225a10 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -793,7 +794,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 @@ -2029,7 +2033,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)}, @@ -2094,6 +2098,99 @@ void server_models_routes::init_routes() { return res; }; + this->get_router_dashboard_stats = [this](const server_http_req &) { + // children have no API key (see unset_reserved_args), they only listen on CHILD_ADDR + static constexpr int DASHBOARD_CHILD_TIMEOUT_MS = 1500; + + auto res = std::make_unique(); + auto all_models = models.get_all_meta(); + + // query the running children in parallel, a busy one must not stall the others + std::vector> child_stats; + child_stats.reserve(all_models.size()); + for (const auto & meta : all_models) { + if (meta.hidden || !meta.is_ready_or_sleep() || meta.port <= 0) { + child_stats.push_back({}); + continue; + } + child_stats.push_back(std::async(std::launch::async, [port = meta.port]() -> json { + httplib::Client cli(CHILD_ADDR, port); + cli.set_connection_timeout(0, DASHBOARD_CHILD_TIMEOUT_MS * 1000); + cli.set_read_timeout(0, DASHBOARD_CHILD_TIMEOUT_MS * 1000); + cli.set_write_timeout(0, DASHBOARD_CHILD_TIMEOUT_MS * 1000); + auto resp = cli.Get("/dashboard/stats"); + if (!resp || resp->status != 200) { + return nullptr; + } + try { + json body = json::parse(resp->body); + const auto & arr = body.at("models"); + if (arr.is_array() && !arr.empty()) { + return arr.at(0); + } + } catch (const std::exception &) { + } + return nullptr; + })); + } + + json models_json = json::array(); + for (size_t i = 0; i < all_models.size(); ++i) { + const auto & meta = all_models[i]; + if (meta.hidden) { + continue; + } + + json child = nullptr; + if (child_stats[i].valid()) { + child = child_stats[i].get(); + } + + json entry = child.is_object() ? child : json::object(); + entry["id"] = meta.name; + entry["aliases"] = meta.aliases; + entry["source"] = server_model_source_to_string(meta.source); + // the router's view of the status wins, it also knows unloaded / loading / failed + entry["status"] = server_model_status_to_string(meta.status); + if (meta.is_failed()) { + entry["failed"] = true; + entry["exit_code"] = meta.exit_code; + } + if (meta.status == SERVER_MODEL_STATUS_LOADING && !meta.progress.is_null()) { + entry["progress"] = meta.progress; + } + if (!entry.contains("stats")) { + entry["stats"] = nullptr; + } + if (!entry.contains("info") || !entry["info"].is_object()) { + // not running: what the preset says about the file + json info = json::object(); + std::string model_path; + if (meta.preset.get_option("LLAMA_ARG_MODEL", model_path) && !model_path.empty()) { + info["file"] = std::filesystem::path(model_path).filename().string(); + std::error_code ec; + auto size = std::filesystem::file_size(model_path, ec); + if (!ec) { + info["size_bytes"] = (uint64_t) size; + } + } else if (std::string hf_repo; meta.preset.get_option("LLAMA_ARG_HF_REPO", hf_repo)) { + info["file"] = hf_repo; + } + entry["info"] = info; + } + models_json.push_back(entry); + } + + res_ok(res, { + {"mode", "router"}, + {"build_info", std::string(llama_build_info())}, + {"t_now_ms", std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count()}, + {"models", models_json}, + }); + return res; + }; + this->post_router_models_unload = [this](const server_http_req & req) { auto res = std::make_unique(); json body = json::parse(req.body); diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 90161bf34ad6..88e4df453437 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -360,6 +360,7 @@ struct server_models_routes { server_http_context::handler_t proxy_get; server_http_context::handler_t proxy_post; server_http_context::handler_t get_router_models; + server_http_context::handler_t get_router_dashboard_stats; // aggregates /dashboard/stats of the children server_http_context::handler_t post_router_models_load; server_http_context::handler_t post_router_models_unload; // management API diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..55672f396498 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -22,7 +22,8 @@ // static bool task_resets_idle_timer(server_task_type type) { - return type != SERVER_TASK_TYPE_METRICS; + // polling stats endpoints must not keep the model awake + return type != SERVER_TASK_TYPE_METRICS && type != SERVER_TASK_TYPE_DASHBOARD; } int server_queue::post(server_task && task, bool front) { 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..1d7e86e26b08 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -22,6 +22,7 @@ enum server_task_type { SERVER_TASK_TYPE_NEXT_RESPONSE, SERVER_TASK_TYPE_METRICS, SERVER_TASK_TYPE_SLOT_GET, + SERVER_TASK_TYPE_DASHBOARD, // used by /dashboard/stats SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE, @@ -320,6 +321,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 +378,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(); @@ -507,6 +513,15 @@ struct server_task_result_metrics : server_task_result { std::string to_metrics(); }; +// used by /dashboard/stats +struct server_task_result_dashboard : server_task_result { + json data; + + virtual json to_json() override { + return data; + } +}; + // used by /slots API struct server_task_result_slots : server_task_result { int n_idle_slots = 0; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 4568a11fb958..80dda38c6239 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -240,6 +240,7 @@ int llama_server(common_params & params, int argc, char ** argv) { // custom routes for router routes.get_props = models_routes->get_router_props; routes.get_models = models_routes->get_router_models; + routes.get_dashboard_stats = models_routes->get_router_dashboard_stats; ctx_http.post("/models", ex_wrapper(models_routes->post_router_models)); ctx_http.post("/models/load", ex_wrapper(models_routes->post_router_models_load)); @@ -288,6 +289,8 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters", ex_wrapper(routes.post_lora_adapters)); // Save & load slots ctx_http.get ("/slots", ex_wrapper(routes.get_slots)); + // dashboard (the page itself is served by server_http_context, see dashboard.html) + ctx_http.get ("/dashboard/stats", ex_wrapper(routes.get_dashboard_stats)); ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots)); // resumable streaming: a child binds the local session factories, the router binds diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index 5eca46cb292d..42a08817d8e6 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/tools/server/tests/unit/test_slot_save.py @@ -37,6 +37,8 @@ def test_slot_save_restore(): }) assert res.status_code == 200 assert res.body["n_saved"] == 84 + slot_file = os.path.join(server.slot_save_path, "slot1.bin") + assert res.body["n_written"] == os.path.getsize(slot_file) # Since we have cache, this should only process the last tokens res = server.make_request("POST", "/completion", data={ @@ -54,6 +56,7 @@ def test_slot_save_restore(): }) assert res.status_code == 200 assert res.body["n_restored"] == 84 + assert res.body["n_read"] == os.path.getsize(slot_file) # Since we have cache, slot 0 should only process the last tokens res = server.make_request("POST", "/completion", data={ @@ -546,3 +549,77 @@ def test_slot_restore_media_file_without_mmproj(mmproj_server): assert res.status_code == 200 assert res.body["timings"]["cache_n"] == 0 assert res.body["content"] == content + + +@pytest.fixture +def swa_server(): + swa = ServerPreset.tinygemma3() + swa.slot_save_path = "./tmp" + swa.temperature = 0.0 + swa.cache_ram = 0 + # Keep the first prompt checkpoint before the divergence point. + swa.n_ubatch = 32 + return swa + + +def test_slot_restore_preserves_context_checkpoints(swa_server): + server = swa_server + server.start() + + base = "The quick brown fox jumps over the lazy dog. " * 20 + + res = server.make_request("POST", "/completion", data={ + "prompt": base + "The first ending of this story is a happy one.", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + n_full = res.body["timings"]["prompt_n"] + + res = server.make_request("POST", "/completion", data={ + "prompt": base + "But the second ending was different and sad.", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + n_live = res.body["timings"]["prompt_n"] + assert n_live < n_full + + res = server.make_request("POST", "/slots/1?action=erase") + assert res.status_code == 200 + + res = server.make_request("POST", "/completion", data={ + "prompt": base + "The first ending of this story is a happy one.", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "ckpt_slot1.bin", + }) + assert res.status_code == 200 + assert res.body["n_saved"] > 0 + ckpt_file = os.path.join(server.slot_save_path, "ckpt_slot1.bin") + assert res.body["n_written"] == os.path.getsize(ckpt_file) + + res = server.make_request("POST", "/completion", data={ + "prompt": "Unrelated text with no common prefix occupies the slot now.", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/1?action=restore", data={ + "filename": "ckpt_slot1.bin", + }) + assert res.status_code == 200 + assert res.body["n_read"] == os.path.getsize(ckpt_file) + + res = server.make_request("POST", "/completion", data={ + "prompt": base + "But the second ending was different and sad.", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + assert res.body["timings"]["prompt_n"] == n_live diff --git a/tools/ui/README.md b/tools/ui/README.md index 99abfaa41fcf..f09d25476c9c 100644 --- a/tools/ui/README.md +++ b/tools/ui/README.md @@ -125,7 +125,8 @@ proxy: { '/models': SERVER_ORIGIN, '/tools': SERVER_ORIGIN, '/slots': SERVER_ORIGIN, - '/cors-proxy': SERVER_ORIGIN + '/cors-proxy': SERVER_ORIGIN, + '/dashboard': SERVER_ORIGIN }, ``` diff --git a/tools/ui/src/lib/components/app/actions/ActionIcon.svelte b/tools/ui/src/lib/components/app/actions/ActionIcon.svelte index 0ed22d932cc0..e4e23db25f43 100644 --- a/tools/ui/src/lib/components/app/actions/ActionIcon.svelte +++ b/tools/ui/src/lib/components/app/actions/ActionIcon.svelte @@ -14,6 +14,7 @@ onclick?: (e?: MouseEvent) => void; size?: ButtonSize; stopPropagationOnClick?: boolean; + target?: string; tooltip?: string; variant?: ButtonVariant; tooltipSide?: TooltipSide; @@ -29,6 +30,7 @@ onclick, size = 'sm', stopPropagationOnClick = false, + target, tooltip, tooltipSide = TooltipSide.TOP, variant = 'ghost' @@ -50,7 +52,9 @@ onclick?.(e); }} + rel={href && target === '_blank' ? 'noopener noreferrer' : undefined} {size} + {target} {variant} > {#if icon} diff --git a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte index 6bc2ca2d7fdd..40e20e852b49 100644 --- a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte +++ b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte @@ -1,5 +1,5 @@