diff --git a/.github/workflows/build-vulkan.yml b/.github/workflows/build-vulkan.yml new file mode 100644 index 0000000..c7f6cb7 --- /dev/null +++ b/.github/workflows/build-vulkan.yml @@ -0,0 +1,143 @@ +# Reusable build of moss-infer-vulkan, the standalone Vulkan program (vulkan/), called by ci.yml and release.yml. Per +# platform: the pinned Vulkan SDK (headers and glslc, checked against LunarG's SHA-256), configure and build vulkan/, +# the CPack archive. On Linux also the backend's checks, run on lavapipe (Mesa's CPU Vulkan driver: the plain kernels) +# with the Khronos validation layer from the SDK, since hosted runners have no GPU. No CUDA is involved. +name: build-vulkan + +on: + workflow_call: + inputs: + retention-days: + description: Days to keep the uploaded packages + type: number + default: 7 + +permissions: + contents: read + +env: + VULKAN_SDK_VERSION: 1.4.357.0 + +jobs: + vulkan: + name: ${{ matrix.platform }}-vulkan + runs-on: ${{ matrix.os }} + timeout-minutes: 120 + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-22.04 + platform: linux + sdk_sha256: 0f09bf6a0625e346bf004be70b92907e934a4c76606b323441b2baf3a5a0e66d + - os: windows-2022 + platform: windows + sdk_sha256: 81f474711e9042f4cd22b31b2f7a8870db2e428b21586fb43dd80150be97310d + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - uses: lukka/get-cmake@fffaaafeea488556c2c12dad60690008bc1caacb # v4.4.2 + + - uses: ilammy/msvc-dev-cmd@0b201ec74fa43914dc39ae48a89fd1d8cb592756 # v1.13.0 + if: matrix.platform == 'windows' + + - name: Toolchain and lavapipe (Linux) + if: matrix.platform == 'linux' + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends g++-12 zlib1g-dev mesa-vulkan-drivers libvulkan1 + sudo apt-get clean + echo "CXX=g++-12" >> "$GITHUB_ENV" + + # The SDK (330 MB on Linux, 290 MB on Windows) is cached per version: LunarG's server can be slow from the runners. + - uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + id: sdk-cache + with: + path: ${{ runner.temp }}/vulkan-sdk + key: vulkan-sdk-${{ matrix.os }}-${{ env.VULKAN_SDK_VERSION }}-v1 + + - name: Vulkan SDK download (Linux) + if: matrix.platform == 'linux' && steps.sdk-cache.outputs.cache-hit != 'true' + run: | + f="vulkansdk-linux-x86_64-$VULKAN_SDK_VERSION.tar.xz" + curl -fsSL --retry 5 --retry-all-errors --retry-delay 15 --connect-timeout 30 -o "$RUNNER_TEMP/$f" \ + "https://sdk.lunarg.com/sdk/download/$VULKAN_SDK_VERSION/linux/$f" + echo "${{ matrix.sdk_sha256 }} $RUNNER_TEMP/$f" | sha256sum -c - + mkdir -p "$RUNNER_TEMP/vulkan-sdk" + tar -xJf "$RUNNER_TEMP/$f" -C "$RUNNER_TEMP/vulkan-sdk" + rm "$RUNNER_TEMP/$f" + + - name: Vulkan SDK download (Windows) + if: matrix.platform == 'windows' && steps.sdk-cache.outputs.cache-hit != 'true' + shell: pwsh + run: | + $f = "vulkansdk-windows-X64-$env:VULKAN_SDK_VERSION.exe" + $exe = Join-Path $env:RUNNER_TEMP $f + curl.exe -fsSL --retry 5 --retry-all-errors --retry-delay 15 --connect-timeout 30 -o $exe ` + "https://sdk.lunarg.com/sdk/download/$env:VULKAN_SDK_VERSION/windows/$f" + if ($LASTEXITCODE -ne 0) { throw "download failed" } + $hash = (Get-FileHash -Algorithm SHA256 $exe).Hash.ToLower() + if ($hash -ne "${{ matrix.sdk_sha256 }}") { throw "SHA-256 mismatch: $hash" } + $root = Join-Path $env:RUNNER_TEMP "vulkan-sdk" + Start-Process -FilePath $exe -Wait -NoNewWindow -ArgumentList @("--root", $root, "--accept-licenses", "--default-answer", "--confirm-command", "install") + Remove-Item $exe + + - uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + if: steps.sdk-cache.outputs.cache-hit != 'true' + with: + path: ${{ runner.temp }}/vulkan-sdk + key: ${{ steps.sdk-cache.outputs.cache-primary-key }} + + - name: Vulkan SDK environment + shell: bash + run: | + if [ "${{ matrix.platform }}" = linux ]; then + sdk="$RUNNER_TEMP/vulkan-sdk/$VULKAN_SDK_VERSION/x86_64"; bin="$sdk/bin" + else + sdk="$RUNNER_TEMP/vulkan-sdk"; bin="$sdk/Bin" + fi + echo "VULKAN_SDK=$sdk" >> "$GITHUB_ENV" + echo "$bin" >> "$GITHUB_PATH" + "$bin/glslc" --version + + - name: Configure + shell: bash + run: | + tools=OFF + extra="" + if [ "${{ matrix.platform }}" = linux ]; then tools=ON; extra="-DMOSSINFER_STATIC_LIBSTDCXX=ON"; fi + cmake -S vulkan -B build-vulkan -G Ninja -DCMAKE_BUILD_TYPE=Release -DMOSSINFER_VULKAN_TOOLS=$tools $extra + + - name: Build + run: cmake --build build-vulkan + + - name: Checks on lavapipe (Linux) + if: matrix.platform == 'linux' + run: | + export VK_DRIVER_FILES="$(ls /usr/share/vulkan/icd.d/lvp_icd.*.json | head -n 1)" + export VK_LAYER_PATH="$VULKAN_SDK/share/vulkan/explicit_layer.d" + export LD_LIBRARY_PATH="$VULKAN_SDK/lib" + export VK_INSTANCE_LAYERS=VK_LAYER_KHRONOS_validation + build-vulkan/moss-infer-vulkan --list-devices + # Every check runs, so one run reports all failures. + status=0 + build-vulkan/vulkan_block_check --no-block || status=1 + build-vulkan/vulkan_te_kernel_check || status=1 + build-vulkan/vulkan_vae_check --kat-only || status=1 + exit $status + + - name: Package + shell: bash + run: | + cd build-vulkan + cpack + ls -l mossinfer-*-vulkan-* + + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 + with: + name: mossinfer-vulkan-${{ matrix.platform }} + path: | + build-vulkan/mossinfer-*-vulkan-*.zip + build-vulkan/mossinfer-*-vulkan-*.tar.gz + retention-days: ${{ inputs.retention-days }} + if-no-files-found: error diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6c7dee1..39d987a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,3 +28,7 @@ jobs: with: retention-days: 7 sccache-recache: ${{ inputs.recache == true }} + vulkan: + uses: ./.github/workflows/build-vulkan.yml + with: + retention-days: 7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 1cb6063..d614c7e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,7 +1,8 @@ # Tag vX.Y.Z (a release) or vX.Y.Z-rcN, N >= 1 (a release candidate, published as a pre-release that does not become # "latest"): build and check all four standard flavours, the Windows CUDA 13 SM75 flavour and the Linux CUDA 12 SM70 -# flavour (build.yml), then publish -# their archives (the full package and the -system-cuda variant of each) and a SHA256SUMS file as a GitHub release. +# flavour (build.yml) and moss-infer-vulkan for Linux and Windows (build-vulkan.yml), then publish +# their archives (the full package and the -system-cuda variant of each CUDA flavour, one archive per Vulkan platform) +# and a SHA256SUMS file as a GitHub release. # X.Y.Z must equal the MI_VERSION_* of # include/mossinfer/mossinfer.h and an rc must come before vX.Y.Z (.github/scripts/release_version.py), the tagged # commit must be on main with a successful ci run (.github/scripts/release_gate.sh), and the release must not exist @@ -92,9 +93,15 @@ jobs: with: retention-days: 30 + vulkan: + needs: version + uses: ./.github/workflows/build-vulkan.yml + with: + retention-days: 30 + # Checks and gathers the release assets with a read-only token; on a dispatch its artifact is the dry run's result. package: - needs: [version, build] + needs: [version, build, vulkan] runs-on: ubuntu-22.04 env: ARCHIVE_VERSION: ${{ needs.version.outputs.archive_version }} @@ -112,7 +119,7 @@ jobs: for full in mossinfer-*.zip mossinfer-*.tar.gz; do [ -f "$full" ] || continue case "$full" in "mossinfer-$ARCHIVE_VERSION-"*) ;; *) echo "$full is not version $ARCHIVE_VERSION"; exit 1 ;; esac - case "$full" in *-system-cuda.*) continue ;; esac + case "$full" in *-system-cuda.* | *-vulkan-*) continue ;; esac case "$full" in *.tar.gz) ext="tar.gz"; base="${full%.tar.gz}" ;; *) ext="zip"; base="${full%.zip}" ;; esac test -f "$base-system-cuda.$ext" || { echo "missing $base-system-cuda.$ext"; exit 1; } n=$((n + 1)) @@ -125,6 +132,9 @@ jobs: sm70="mossinfer-$ARCHIVE_VERSION-linux-cu12-x86_64-sm70.zip" test -f "$sm70" || { echo "missing SM70 package $sm70"; exit 1; } test -f "${sm70%.zip}-system-cuda.zip" || { echo "missing SM70 system-CUDA package"; exit 1; } + for vk in "mossinfer-$ARCHIVE_VERSION-linux-vulkan-x86_64.tar.gz" "mossinfer-$ARCHIVE_VERSION-windows-vulkan-x86_64.zip"; do + test -f "$vk" || { echo "missing Vulkan package $vk"; exit 1; } + done - name: SHA256SUMS shell: bash working-directory: dist diff --git a/README.md b/README.md index ccfc4fb..5304cb0 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A native C++/CUDA inference engine for the MOSS-Image text-to-image model (5B Di Models: MOSS-Image Base and SFT, in the official folders and common community formats; other models are refused ([supported models](#supported-models)). -Version 2.1.2. Supported: Linux x86-64 and Windows x64 (MSVC), each with a CUDA 12 and a CUDA 13 build, as a full package and a `-system-cuda` package. Turing (compute capability 7.5, e.g. RTX 2080) is supported by builds configured with `-DMOSSINFER_SM75=ON` (the `windows-msvc-cuda13-sm75` preset packages it): fp16 tensor-core GEMMs replace the bf16 ones, INT8 and INT4 ConvRot files run the INT8 GEMMs, the exact attention and the VAE run sm_75 kernels, all three quality tiers work, and an unset attention resolves to `exact`. See [docs/sm75.md](docs/sm75.md) for what runs, the measured speed on an RTX 2080 Ti and the remaining gaps. Volta (Tesla V100, sm_70) is supported by the CUDA 12 `linux-cuda12-sm70` preset; see [docs/sm70.md](docs/sm70.md) for its kernels, validation and performance. The full support matrix is in [docs/integration.md](docs/integration.md#support-status). +Version 2.1.2. Supported: Linux x86-64 and Windows x64 (MSVC), each with a CUDA 12 and a CUDA 13 build, as a full package and a `-system-cuda` package. Turing (compute capability 7.5, e.g. RTX 2080) is supported by builds configured with `-DMOSSINFER_SM75=ON` (the `windows-msvc-cuda13-sm75` preset packages it): fp16 tensor-core GEMMs replace the bf16 ones, INT8 and INT4 ConvRot files run the INT8 GEMMs, the exact attention and the VAE run sm_75 kernels, all three quality tiers work, and an unset attention resolves to `exact`. See [docs/sm75.md](docs/sm75.md) for what runs, the measured speed on an RTX 2080 Ti and the remaining gaps. Volta (Tesla V100, sm_70) is supported by the CUDA 12 `linux-cuda12-sm70` preset; see [docs/sm70.md](docs/sm70.md) for its kernels, validation and performance. Other GPUs (AMD, Intel, older NVIDIA) can run txt2img at the reference tier through `moss-infer-vulkan`, a separate and slower Vulkan program in its own package ([docs/vulkan.md](docs/vulkan.md)). The full support matrix is in [docs/integration.md](docs/integration.md#support-status). The caller picks the quality tier: `reference` (bf16, exact), `high` (the loaded DiT's GEMMs with 8-bit attention) or `turbo` (a 4-bit DiT). An INT8 DiT with `high` is the everyday choice: 2.5 times the bf16 speed on an RTX 5090, in 40 % less VRAM ([performance](docs/integration.md#performance)). Upgrading from 1.x: [migration](docs/integration.md#migrating-from-1x). @@ -30,6 +30,7 @@ Lowest VRAM against normal use at 1024×1024. `min_memory` streams the DiT and c - [docs/integration.md](docs/integration.md): building the engine into an application. Start here. - [docs/sm70.md](docs/sm70.md): Volta (Tesla V100) support: CUDA 12 build, exact INT8 accumulation, validation and performance. - [docs/sm75.md](docs/sm75.md): Turing (compute capability 7.5) support: build and package, what runs, measured speed on an RTX 2080 Ti, limitations. +- [docs/vulkan.md](docs/vulkan.md): `moss-infer-vulkan`, the Vulkan fallback for GPUs the CUDA engine does not support. - [docs/c-api.md](docs/c-api.md): the C library, `include/mossinfer/mossinfer.h`. - [docs/server-protocol.md](docs/server-protocol.md): `moss-server`, the JSON-lines sidecar process. diff --git a/docs/vulkan.md b/docs/vulkan.md new file mode 100644 index 0000000..49429f2 --- /dev/null +++ b/docs/vulkan.md @@ -0,0 +1,85 @@ +# Vulkan backend + +`moss-infer-vulkan` runs MOSS-Image on Vulkan. It is the fallback for GPUs the CUDA engine does not support (AMD, Intel, +older NVIDIA cards): a standalone program in its own package (`mossinfer-VERSION-linux-vulkan-x86_64.tar.gz`, +`mossinfer-VERSION-windows-vulkan-x86_64.zip`) that needs a Vulkan driver and neither CUDA nor an NVIDIA driver. It is +slower than CUDA; on a supported NVIDIA GPU, use `moss-infer`. + +It runs txt2img with the official Base and SFT bf16 files, the official euler sampler and schedule and classifier-free +guidance with the uncond skip layers. The tier is `reference`. Everything else is refused with a message that names +the option: other samplers and schedulers, LoRAs, img2img and inpainting, hires fix, control adapters, references, +the lossy accelerators, quantized or GGUF files, and the CUDA memory options. The tokenizer, the seed's initial noise +and the schedule are the engine's own, so a seed gives the same inputs as on CUDA. + +```sh +moss-infer-vulkan --model MOSS-Image/base --prompt "a red apple on a wooden table" --size 1024x1024 -o apple.png +moss-infer-vulkan --list-devices +``` + +The flags are those of `moss-infer generate` that the backend runs, with the same defaults (`--negative`, `--size`, +`--width`, `--height`, `--steps`, `--cfg`, `--shift`, `--seed`, `--skip-layers`, `--max-seq`, `--latent-out`, +`--no-metadata`, `--quiet`; `--tier` accepts `reference`). `--device N` picks a device (default: the first usable one, a +discrete GPU first) and `--kernels coop|coop-f16|plain` the kernels (default: the first in that order the device has); +both are printed at the start of every run. Progress goes to stderr one line per event: `text encoder`, `denoise / +`, `vae`, then a `done:` summary. The exit status is 0 on success, 2 for a usage error or an unsupported +request, 1 for any other failure. + +## Kernels + +| Kernels | Needs | Used for | +|---|---|---| +| `coop` | Vulkan 1.3, `VK_KHR_cooperative_matrix` and `VK_KHR_shader_bfloat16` with 16x16x16 bf16 fragments, a compute subgroup size fixable at 32, 40 KiB of shared memory | GEMMs, attention and the VAE's convolutions on matrix cores | +| `coop-f16` | the same with 16x16x16 f16 fragments instead of bf16 ones | devices with matrix cores but no bf16 fragments (Mesa's RADV exposes them only on RDNA4, for one): the GEMMs and attention on matrix cores, the VAE as `plain` | +| `plain` | Vulkan 1.2, buffer device addresses, 64-bit integers, 32 KiB of shared memory | every other device; no cooperative matrices, 16-bit arithmetic or particular subgroup size | + +`coop-f16` computes what `coop` computes. The weights and activations stay bf16; each GEMM operand row is scaled by a +power of two that brings it into f16's range as it is loaded, where a bf16 value is an f16 value, and the result is +scaled back, so the products are the same. Only attention's q, k and v convert unscaled: their rare values below 2^-14 +lose bits, and a value above 65504 would become infinite (the model's stay far below). On the golden case's text encoder and on real DiT blocks it gives the same figures as `coop`, and the +same step 0 accuracy below; whole images then differ as any two backends' do (see below). + +## Memory + +The text encoder streams its 32 layers from the model files, one layer on the device at a time. The DiT keeps as many +of its 42 blocks on the device as fit next to the run's activations, from the driver's memory budget +(`VK_EXT_memory_budget`). The other blocks stay in host memory and are copied on the transfer queue while the block +before them runs. So a GPU with less than the DiT's 9.5 GiB still runs it, at the cost of host memory and bus traffic. +The VAE decodes the whole image at once: about 1.6 GiB at 1024x1024, 3.1 GiB at 1456x1456 and 6 GiB at 2048x2048. A +decode the device cannot hold is refused before the text encoder runs. The text encoder and the DiT are on the device +one after the other, never together. + +Long operations are split into submits of bounded work, so no submit runs long enough to trip the Windows display +driver's watchdog on a slow GPU. + +## Accuracy + +Against the official PyTorch pipeline's golden data at 1024x1024 (`tests/golden/g_fridge_1024`), on the `coop` path: + +| Stage | Result | +|---|---| +| Text encoder, last hidden state | cosine 0.99988 (prompt), 0.99980 (negative prompt); torch itself with one kernel swapped: 0.99989, 0.99983 | +| DiT, step 0 prediction | cosine 0.999976 (cond), 0.999983 (uncond); the CUDA engine's test asks for 0.99988 | +| VAE decode of the golden latent | PSNR 69.4 dB against the golden PNG, 99.25 % identical bytes, max difference 1 | + +Whole images keep the composition and the text but not the pixels: 50 sampling steps magnify any difference. The +golden 1024x1024 case through `moss-infer` ends with a final latent at cosine 0.974 to the official pipeline's (the CUDA +backend: 0.992, being nearly bit-identical to torch op by op) and an image at 30.8 dB PSNR. That is the sampler's +sensitivity, not an error of the backend: the same Vulkan steps on torch's own text embeddings instead of the Vulkan +encoder's (cosine 0.9998 apart, about torch's own variation when one of its kernels changes) end 0.971 apart from each +other. A 512x512, 20-step image with seed 7 is 32.9 dB PSNR from the CUDA engine's image of the same command: the same apple, +details redrawn. + +## Building + +`vulkan/` is its own CMake project and needs no CUDA: Vulkan headers that define `VK_KHR_shader_bfloat16` and a `glslc` +that compiles `GL_EXT_bfloat16`, as in the Vulkan SDK 1.4.357 the release builds use. The program opens the Vulkan +loader at run time, so it links nothing of Vulkan. + +```sh +cmake -S vulkan -B build-vulkan -G Ninja -DMOSSINFER_VULKAN_TOOLS=OFF +cmake --build build-vulkan +cpack --config build-vulkan/CPackConfig.cmake -B build-vulkan +``` + +With the tools on (the default) it also builds the backend's checks and benchmarks; see +[vulkan/README.md](../vulkan/README.md). diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt new file mode 100644 index 0000000..3b2248e --- /dev/null +++ b/vulkan/CMakeLists.txt @@ -0,0 +1,175 @@ +# The Vulkan backend and moss-infer-vulkan, the standalone program for GPUs the CUDA engine does not support. Its own +# project, built without CUDA; the engine's build does not include it. See README.md. +cmake_minimum_required(VERSION 3.20) +include(${CMAKE_CURRENT_LIST_DIR}/../cmake/version.cmake) # MOSSINFER_VERSION, from the C header +project(mossinfer_vulkan VERSION ${MOSSINFER_VERSION} LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE Release) # the CPU reference is unusably slow unoptimized +endif() + +# The Vulkan headers and glslc, from $VULKAN_SDK first. No loader library: vk_api.cpp opens it at run time. +find_path(MI_VULKAN_INCLUDE_DIR vulkan/vulkan_core.h + HINTS "$ENV{VULKAN_SDK}/include" "$ENV{VULKAN_SDK}/Include" REQUIRED) +file(STRINGS ${MI_VULKAN_INCLUDE_DIR}/vulkan/vulkan_core.h _mi_bf16 REGEX "VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME") +if(NOT _mi_bf16) + message(FATAL_ERROR "${MI_VULKAN_INCLUDE_DIR}/vulkan/vulkan_core.h has no VK_KHR_shader_bfloat16: use a newer Vulkan SDK") +endif() +find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN_SDK}/bin" REQUIRED) +add_library(mi_vulkan_headers INTERFACE) +target_include_directories(mi_vulkan_headers SYSTEM INTERFACE ${MI_VULKAN_INCLUDE_DIR}) + +# Shaders become SPIR-V word lists compiled into the binary, so nothing beyond it ships. +set(_spirv_dir ${CMAKE_CURRENT_BINARY_DIR}/spirv) +file(MAKE_DIRECTORY ${_spirv_dir}) +set(_spirv_files "") +# The cooperative-matrix shader needs Vulkan 1.3; the rest must also load on Vulkan 1.2 devices, which reject newer SPIR-V. +# _mi_add_shader_as(source output target_env [-DNAME ...]) builds shaders/.comp into .spv.inc. +function(_mi_add_shader_as shader output target_env) + set(_source ${CMAKE_CURRENT_SOURCE_DIR}/shaders/${shader}.comp) + set(_output ${_spirv_dir}/${output}.spv.inc) + add_custom_command( + OUTPUT ${_output} + COMMAND ${GLSL_COMPILER} --target-env=${target_env} -fshader-stage=comp -O -mfmt=num ${ARGN} + -I${CMAKE_CURRENT_SOURCE_DIR}/shaders ${_source} -o ${_output} + DEPENDS ${_source} ${CMAKE_CURRENT_SOURCE_DIR}/shaders/common.glsl + COMMENT "glslc ${shader}.comp -> ${output}") + set(_spirv_files ${_spirv_files} ${_output} PARENT_SCOPE) +endfunction() +function(_mi_add_shader shader target_env) + _mi_add_shader_as(${shader} ${shader} ${target_env}) + set(_spirv_files ${_spirv_files} PARENT_SCOPE) +endfunction() +_mi_add_shader(gemm_coop vulkan1.3) +_mi_add_shader(gemm_bf16_coop vulkan1.3) +_mi_add_shader(attention_coop vulkan1.3) +_mi_add_shader_as(gemm_bf16_coop gemm_f16_coop vulkan1.3 -DF16) +_mi_add_shader_as(attention_coop attention_f16_coop vulkan1.3 -DF16) +foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain elementwise row_scale) + _mi_add_shader(${shader} vulkan1.2) +endforeach() + +add_library(mossinfer_vulkan STATIC vk_api.cpp runtime.cpp gemm.cpp block.cpp dit.cpp safetensors.cpp ${_spirv_files}) +target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) +# Hidden symbols: vk_api.cpp defines data objects named like the Vulkan entry points, which a shared library must not +# export (they would interpose with a loader in the same process). +set_target_properties(mossinfer_vulkan PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON) +# Headers only: vk_api.cpp opens the loader at run time, so nothing links against it. +find_package(Threads REQUIRED) +target_link_libraries(mossinfer_vulkan PUBLIC mi_vulkan_headers Threads::Threads ${CMAKE_DL_LIBS}) + +# The Qwen3.5 text encoder (text_encoder.h): its own kernels on top of the DiT block's GEMMs. +set(_te_spirv_files "") +foreach(shader IN ITEMS te_rms_norm te_conv_silu te_gates te_gdn_qk te_gdn_chunk te_gdn_solve te_gdn_state te_attention + te_gate_mul) + set(_source ${CMAKE_CURRENT_SOURCE_DIR}/shaders/${shader}.comp) + set(_output ${_spirv_dir}/${shader}.spv.inc) + add_custom_command( + OUTPUT ${_output} + COMMAND ${GLSL_COMPILER} --target-env=vulkan1.2 -fshader-stage=comp -O -mfmt=num + -I${CMAKE_CURRENT_SOURCE_DIR}/shaders ${_source} -o ${_output} + DEPENDS ${_source} ${CMAKE_CURRENT_SOURCE_DIR}/shaders/common.glsl ${CMAKE_CURRENT_SOURCE_DIR}/shaders/te_common.glsl + COMMENT "glslc ${shader}.comp") + list(APPEND _te_spirv_files ${_output}) +endforeach() +target_sources(mossinfer_vulkan PRIVATE te_kernels.cpp text_encoder.cpp ${_te_spirv_files}) + +# The FLUX.2 VAE decoder (vae.h): its cooperative-matrix kernel needs Vulkan 1.3, the rest loads on Vulkan 1.2. +set(_spirv_block ${_spirv_files}) +_mi_add_shader(vae_conv_coop vulkan1.3) +foreach(shader IN ITEMS vae_conv_plain vae_conv_narrow vae_gn_stats vae_gn_finalize vae_gn_apply vae_softmax vae_transpose + vae_latent_unpack vae_rgb8) + _mi_add_shader(${shader} vulkan1.2) +endforeach() +set(_spirv_vae ${_spirv_files}) +list(REMOVE_ITEM _spirv_vae ${_spirv_block}) +target_sources(mossinfer_vulkan PRIVATE vae.cpp generate.cpp ${_spirv_vae}) +# generate.cpp's CFG and Euler step round like the engine's kernels: no contraction. +set_source_files_properties(generate.cpp PROPERTIES COMPILE_OPTIONS "$,/fp:precise,-ffp-contract=off>") + +# moss-infer-vulkan, the standalone program: the backend plus the engine's host sources it shares with moss-infer (the +# tokenizer, the seed's initial noise, the schedule, the PNG writer), compiled here so that nothing needs CUDA. +set(_engine_src ${CMAKE_CURRENT_SOURCE_DIR}/../src) +add_library(mossinfer_vulkan_host STATIC + ${_engine_src}/tokenizer/qwen_bpe.cpp ${_engine_src}/tokenizer/unicode.cpp ${_engine_src}/tokenizer/prompt_tokens.cpp + ${_engine_src}/pipeline/samplers/initial_noise.cpp ${_engine_src}/pipeline/rng_torch_compat.cpp + ${_engine_src}/pipeline/samplers/schedules.cpp ${_engine_src}/pipeline/png_writer.cpp ${_engine_src}/util/json.cpp + ${_engine_src}/util/regular_file.cpp ${_engine_src}/util/replaceable_file.cpp) +target_include_directories(mossinfer_vulkan_host PUBLIC ${_engine_src} ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_executable(moss-infer-vulkan tools/moss_infer_vulkan.cpp) +target_link_libraries(moss-infer-vulkan PRIVATE mossinfer_vulkan mossinfer_vulkan_host) +if(MSVC) + target_compile_options(mossinfer_vulkan_host PRIVATE /utf-8) + target_compile_options(moss-infer-vulkan PRIVATE /utf-8) +endif() +# Release packages run on any glibc distribution of the build's age or newer, without its libstdc++. +option(MOSSINFER_STATIC_LIBSTDCXX "Link libstdc++ and libgcc statically into moss-infer-vulkan (Linux packages)" OFF) +if(MOSSINFER_STATIC_LIBSTDCXX AND NOT MSVC) + target_link_options(moss-infer-vulkan PRIVATE -static-libstdc++ -static-libgcc) +endif() + +# Package: the program and the license files. No Vulkan loader ships: the GPU driver provides it. +install(TARGETS moss-infer-vulkan RUNTIME DESTINATION bin) +install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../LICENSE ${CMAKE_CURRENT_SOURCE_DIR}/../NOTICE + ${CMAKE_CURRENT_SOURCE_DIR}/../THIRD_PARTY_NOTICES.md ${CMAKE_CURRENT_SOURCE_DIR}/../docs/vulkan.md + DESTINATION share/doc/mossinfer) +if(WIN32) + set(_mi_platform windows) +elseif(APPLE) + set(_mi_platform macos) +else() + set(_mi_platform linux) +endif() +string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _mi_arch) +if(_mi_arch STREQUAL "amd64") + set(_mi_arch x86_64) +endif() +set(CPACK_PACKAGE_NAME mossinfer-vulkan) +set(CPACK_PACKAGE_VERSION ${MOSSINFER_VERSION}) +set(CPACK_PACKAGE_FILE_NAME "mossinfer-${MOSSINFER_VERSION}-${_mi_platform}-vulkan-${_mi_arch}") +if(WIN32) + set(CPACK_GENERATOR ZIP) +else() + set(CPACK_GENERATOR TGZ) +endif() +include(CPack) + +# The checks and benchmarks; a release build turns them off. +option(MOSSINFER_VULKAN_TOOLS "Build the Vulkan backend's checks and benchmarks" ON) +if(NOT MOSSINFER_VULKAN_TOOLS) + return() +endif() + +add_executable(vulkan_check tools/vulkan_check.cpp) +target_link_libraries(vulkan_check PRIVATE mossinfer_vulkan) + +# The host replays the block is checked against must round where the kernels do: no contraction. +add_executable(vulkan_block_check tools/vulkan_block_check.cpp) +target_compile_options(vulkan_block_check PRIVATE $,/fp:precise,-ffp-contract=off>) +target_link_libraries(vulkan_block_check PRIVATE mossinfer_vulkan) + +# One real block against the engine's captured output. +add_executable(vulkan_block_real tools/vulkan_block_real.cpp) +target_link_libraries(vulkan_block_real PRIVATE mossinfer_vulkan) + +# The DiT's sampling loop from a golden case; CFG and Euler on the host round like the engine's kernels: no contraction. +add_executable(vulkan_dit_run tools/vulkan_dit_run.cpp) +target_compile_options(vulkan_dit_run PRIVATE $,/fp:precise,-ffp-contract=off>) +target_link_libraries(vulkan_dit_run PRIVATE mossinfer_vulkan) + +# The text encoder's kernels against host replays of their op order; the replays must round where the kernels do. +add_executable(vulkan_te_kernel_check tools/vulkan_te_kernel_check.cpp) +target_compile_options(vulkan_te_kernel_check PRIVATE $,/fp:precise,-ffp-contract=off>) +target_link_libraries(vulkan_te_kernel_check PRIVATE mossinfer_vulkan) + +# The whole text encoder against the official pipeline's golden hidden states. +add_executable(vulkan_te_check tools/vulkan_te_check.cpp) +target_link_libraries(vulkan_te_check PRIVATE mossinfer_vulkan) + +# Kernels against host replays, then the real decoder against the reference pipeline's PNG (zlib reads and writes PNGs). +find_package(ZLIB REQUIRED) +add_executable(vulkan_vae_check tools/vulkan_vae_check.cpp) +target_compile_options(vulkan_vae_check PRIVATE $,/fp:precise,-ffp-contract=off>) +target_link_libraries(vulkan_vae_check PRIVATE mossinfer_vulkan ZLIB::ZLIB) diff --git a/vulkan/README.md b/vulkan/README.md new file mode 100644 index 0000000..1e27f48 --- /dev/null +++ b/vulkan/README.md @@ -0,0 +1,139 @@ +# Vulkan backend + +The fallback for GPUs the CUDA engine does not support: `moss-infer-vulkan`, a standalone program built from this +directory without CUDA (user guide: [docs/vulkan.md](../docs/vulkan.md)). CUDA stays the path on every GPU it supports. +The program compiles in the engine's host sources it shares with `moss-infer` (tokenizer, initial noise, schedule, PNG +writer); the engine's own build does not include this directory. + +Every model runs on one of three kernel paths, chosen per device (the first the device has) and reported: + +| Path | Needs | | +|---|---|---| +| `coop` | Vulkan 1.3, `VK_KHR_cooperative_matrix` with `VK_KHR_shader_bfloat16` (16x16x16 bf16 x bf16 -> f32 in subgroup scope), compute subgroup size 32 with vote and shuffle, 40 KiB of shared memory | `Capabilities::coop_bf16` | +| `coop-f16` | the same with 16x16x16 f16 x f16 -> f32 instead of bf16 fragments; the VAE runs `plain` | `Capabilities::coop_f16` | +| `plain` | Vulkan 1.2, buffer device addresses, 64-bit integers, 32 KiB of shared memory | every other device; no cooperative matrices, 16-bit arithmetic or particular subgroup size | + +The Vulkan loader is opened at run time (`vk_api.h`), so nothing links against it. + +## Build and check + +```sh +cmake -S vulkan -B build-vulkan -G Ninja +cmake --build build-vulkan +build-vulkan/vulkan_check --list # each device, and why a device gets the plain path +build-vulkan/vulkan_block_check [--path P] # the DiT block's kernels against host replays, then one block's timing +build-vulkan/vulkan_block_real DIR # real blocks against the engine's output (files captured from its forward) +build-vulkan/vulkan_dit_run MODEL/transformer tests/golden g_fridge_1024 [--steps N] [--dit-budget GiB] +build-vulkan/vulkan_te_kernel_check # the text encoder's kernels against host replays +build-vulkan/vulkan_te_check --golden tests/golden +build-vulkan/vulkan_vae_check # the VAE's kernels, then the golden latent against the golden PNG +``` + +It needs the Vulkan headers with `VK_KHR_shader_bfloat16`, a `glslc` that compiles `GL_EXT_bfloat16`, and zlib for the +VAE check. Shaders compile into the binary. Every check exits non-zero on a failure, and passes under the Khronos +validation layer and on llvmpipe (Mesa's CPU driver: 8-wide subgroups, 32 KiB of shared memory, the plain path). +`MI_VK_PIPELINE_STATS=1` prints each pipeline's registers and spills where the driver reports them. + +## DiT + +`block.h` holds one DiT block's kernels on every path, with the engine's rounding points: RMSNorm with modulation, the +GEMMs with bias, gated-residual and SiTU-GLU epilogues, QK norm with RoPE, and flash attention (an online softmax over +32-key steps; the `coop` one rescales its output only when a row's maximum grew). `dit.h` runs the whole +MOSSImageTransformer on them: the time embedding and shared modulation, the text refiner, the joint sequences of the CFG +rows with per-row skipped blocks, and the head. + +The blocks that do not fit next to the activations (from `VK_EXT_memory_budget`) stay in host memory. Two device slots +alternate: each streamed block is copied on the transfer queue while the block before it runs. Every op over many +tokens runs in token ranges of at most 200 GFLOP (`coop`) or 20 GFLOP (`plain`) per submit, attention in query ranges, +so a slow device does not trip a display driver's watchdog. Neither changes a bit of the result. + +Against the engine's own output on real blocks of the 1024x1024 golden case (captured from its forward): + +| Block | Path | Same bits | Max scaled ULPs | Cosine | +|---:|---|---:|---:|---:| +| 1 | coop | 98.14 % | 1.00 | 0.999999999 | +| 22 | coop | 96.16 % | 2.00 | 0.999999989 | +| 41 | coop | 98.56 % | 32.00 | 0.999999919 | + +The `coop` GEMMs round like the tensor cores, hence so many identical bits. Step 0 of g_fridge_1024 against the official +pipeline: cosine 0.999976 (cond) and 0.999983 (uncond, block 22 skipped); the CUDA engine's test asks for 0.99988. + +`coop-f16` runs the `coop` kernels built with `F16`. The operands stay bf16 in memory and become f16 as they are +staged: `row_scale.comp` finds, per row of each GEMM's A and B, the power of two that puts the row's largest value in +[2^15, 2^16), where every bf16 value is an f16 value; the GEMM scales the rows by it on load and the accumulators back +before the epilogue, all exactly, so every product is the bf16 one while the values are above 2^-29 of their row's +largest. Attention converts q, k and v unscaled (bits lost below 2^-14, infinite above 65504) and feeds P as +bf16(p) * 2^15. On the 4090 it gives the table above to +the digit (98.14 %, 96.16 %, 98.56 %), step 0 at 0.999976 and 0.999985, and the text encoder's figures of `coop`; +`vulkan_block_check` adds a GEMM case whose rows span 2^-90 to 2^100. Its staging conversions and row scans cost +about 7 % of a 512x512 step there, nothing measurable at 1024x1024. + +## VAE + +`vae.h` is the FLUX.2 decoder from the final latents to RGB8. Its convolutions round both operands to TF32 as cuDNN does +(on `coop` each TF32 operand is two exact bf16 pieces), so every product is torch's and only the fp32 accumulation order +differs; GroupNorm, the mid attention and the upsampling follow the engine. It decodes in one piece: about 1.6 GiB at +1024x1024. The golden latent against the official pipeline's PNG: 69.39 dB (`coop`), 69.35 dB (`plain`), 99.25 % of the +bytes identical and none off by more than 1, as the CUDA engine's reference decode. + +## Text encoder + +`text_encoder.h` runs the Qwen3.5 text encoder (`Qwen3_5Model(input_ids).last_hidden_state`, bf16 [S, 4096] per prompt) +for several prompts at once. Its GEMMs are `BlockKernels`' on either path. Its own kernels (`te_kernels.h`) need no +cooperative matrices, 16-bit arithmetic or particular subgroup size, so every path shares them: the Qwen RMSNorms (with +torch's reduction order for the row count, and RoPE fused into the q/k norm), conv1d with SiLU, the GatedDeltaNet gates, +the engine's chunked gated delta rule in fp32 (four kernels, each within 32 KiB of shared memory), causal grouped-query +attention at head_dim 256, and the gated products. Weights stream from the safetensors files: four threads read layer +k + 2 into host staging while layer k runs and layer k + 1 is copied to the device on the transfer queue +(`Device::copy_async`). The embedding table stays on the host. Device memory is two layer buffers (833 MiB) plus about +130 KiB per token. + +```sh +build-vulkan/vulkan_te_kernel_check # each kernel against a host replay of its op order +build-vulkan/vulkan_te_check --golden tests/golden # the encoder against the official pipeline's hidden states +build-vulkan/vulkan_te_check --golden tests/golden --case g_long4096_1024 --profile +``` + +Against the goldens of the official PyTorch pipeline on an RTX 4090 Laptop, last_hidden_state, with torch's own +deviation from the same goldens when one of its kernels changes (the worse of SDPA math and cuBLASLt, +`tests/data/qwen35_text/noise_floor.json`): + +| Prompt | Tokens | Cosine, coop | Cosine, plain | Cosine, torch | Mean row rel. L2, coop | plain | torch | +|---|---:|---:|---:|---:|---:|---:|---:| +| g_fridge cond | 62 | 0.9998786 | 0.9999012 | 0.9998949 | 1.32e-2 | 1.27e-2 | 1.29e-2 | +| g_fridge uncond | 10 | 0.9998012 | 0.9997889 | 0.9998333 | 1.84e-2 | 1.90e-2 | 1.60e-2 | +| g_short cond | 29 | 0.9998910 | 0.9997967 | 0.9999026 | 1.40e-2 | 1.65e-2 | 1.33e-2 | +| g_layout cond | 172 | 0.9998827 | 0.9998885 | 0.9998857 | 1.43e-2 | 1.40e-2 | 1.44e-2 | +| g_layout uncond | 16 | 0.9998486 | 0.9998133 | 0.9998634 | 1.64e-2 | 1.81e-2 | 1.52e-2 | +| g_long4096 cond | 4096 | 0.9985149 | 0.9984682 | 0.9992133 | 2.29e-2 | 2.32e-2 | 1.89e-2 | + +Differences grow over the 32 layers until they reach torch's own sensitivity to a kernel change; after the first layer +they are at the level of one bf16 rounding (layer 0, row relative L2 at most 7.1e-3 on both paths). + +Times on that GPU with nothing else running, the weights read from a drive that held little of them in the page cache +(about 2.5 GB/s with four readers), which bounds the short prompts: + +| Prompts | Tokens | Coop: encode, of it compute | Plain: encode, of it compute | +|---|---:|---:|---:| +| g_fridge cond + uncond | 72 | 6.7 s, 0.6 s | 6.8 s, 2.8 s | +| g_long4096 cond + uncond | 4112 | 8.4 s, 7.2 s | 186 s, 185 s | + +At 4112 tokens on the cooperative-matrix path the GEMMs take 5.6 s of the 7.2 s, attention 0.7 s and the gated delta rule +0.8 s (`--profile`). + +## Files + +| File | Purpose | +|---|---| +| `vk_api.h`, `vk_api.cpp` | the Vulkan functions, resolved from the loader at run time | +| `runtime.h`, `runtime.cpp` | device selection, capabilities, buffers, pipelines, barriers, transfer-queue copies | +| `block.h`, `block.cpp` | the DiT block's kernels on every path | +| `dit.h`, `dit.cpp` | the DiT: text refiner, forward, block streaming | +| `text_encoder.h`, `te_kernels.h` and their `.cpp` | the Qwen3.5 text encoder | +| `vae.h`, `vae.cpp` | the FLUX.2 VAE decoder | +| `generate.h`, `generate.cpp` | text to image: the three models one after another, CFG and the official Euler step | +| `gemm.h`, `gemm.cpp` | the standalone f16 GEMM of the first bring-up (`vulkan_check`) | +| `safetensors.h`, `npy.h`, `half.h` | readers and conversions | +| `shaders/` | compute shaders | +| `tools/moss_infer_vulkan.cpp` | the program | +| `tools/` | checks and benchmarks | diff --git a/vulkan/block.cpp b/vulkan/block.cpp new file mode 100644 index 0000000..ae4d04c --- /dev/null +++ b/vulkan/block.cpp @@ -0,0 +1,334 @@ +#include "block.h" + +#include +#include +#include +#include + +namespace mi::vk { +namespace { + +const std::uint32_t kGemmSpirv[] = { +#include "gemm_plain.spv.inc" +}; +const std::uint32_t kGemmCoopSpirv[] = { +#include "gemm_bf16_coop.spv.inc" +}; +const std::uint32_t kNormSpirv[] = { +#include "norm_modulate.spv.inc" +}; +const std::uint32_t kQkSpirv[] = { +#include "qk_norm_rope.spv.inc" +}; +const std::uint32_t kElementwiseSpirv[] = { +#include "elementwise.spv.inc" +}; +const std::uint32_t kAttentionSpirv[] = { +#include "attention_plain.spv.inc" +}; +const std::uint32_t kAttentionCoopSpirv[] = { +#include "attention_coop.spv.inc" +}; +const std::uint32_t kGemmCoopF16Spirv[] = { +#include "gemm_f16_coop.spv.inc" +}; +const std::uint32_t kAttentionCoopF16Spirv[] = { +#include "attention_f16_coop.spv.inc" +}; +const std::uint32_t kRowScaleSpirv[] = { +#include "row_scale.spv.inc" +}; + +struct GemmPush { + std::uint64_t a, b, out, gate; + std::uint32_t m, n, k, lda, ldb, ldo; + std::uint64_t a_exponents, b_exponents; // kCoopF16 only: per row of A and of B +}; +struct RowScalePush { + std::uint64_t a, exponents; + std::uint32_t m, k, lda; + std::uint32_t pad = 0; // sizeof is pushed: no indeterminate bytes +}; +constexpr std::uint32_t kRowScaleRows = 8; // rows per row_scale.comp workgroup +constexpr std::uint32_t kMaxF16GemmRows = 1u << 20; // rows of A, and of B, that row_exponents_ holds +struct NormPush { + std::uint64_t x, weight, scale, shift, out; + std::uint32_t ldx, ldo, d; + float factor; +}; +struct QkPush { + std::uint64_t qkv, q_norm, k_norm, table; + std::uint32_t ld, table_ld, heads; + float factor; +}; + +struct AttentionPush { // attention_coop.comp and attention_plain.comp + std::uint64_t q, k, v, o; + std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_row0, k_row0, seqlen_q, seqlen_k; + float log2_scale; +}; +constexpr std::uint32_t kCoopAttentionRows = 64, kPlainAttentionRows = 32; + +enum GemmMode : std::uint32_t { kF32 = 0, kBF16 = 1, kGated = 2, kSitu = 3, kBias = 4, kBiasResidual = 5 }; + +struct ElementwisePush { + std::uint64_t a, b, out; + std::uint32_t n; +}; + +constexpr std::uint32_t kCoopSubgroup = 32; +constexpr std::uint32_t kCoopTileRows = 128, kCoopTileCols = 128; + +std::unique_ptr gemm_pipeline(Device& device, BlockPath path, GemmMode mode) { + const auto m = static_cast(mode); + if (path == BlockPath::kCoop) { + return std::make_unique(device, kGemmCoopSpirv, sizeof(kGemmCoopSpirv), sizeof(GemmPush), kCoopSubgroup, + std::vector{m}); + } + if (path == BlockPath::kCoopF16) { + return std::make_unique(device, kGemmCoopF16Spirv, sizeof(kGemmCoopF16Spirv), sizeof(GemmPush), + kCoopSubgroup, std::vector{m}); + } + return std::make_unique(device, kGemmSpirv, sizeof(kGemmSpirv), sizeof(GemmPush), 0, + std::vector{1u, m}); +} + +void require(bool condition, const std::string& message) { + if (!condition) throw std::invalid_argument("vulkan block: " + message); +} + +constexpr std::uint32_t kGemmTileRows = 128, kGemmTileCols = 128; + +bool is_coop(BlockPath path) { return path != BlockPath::kPlain; } + +} // namespace + +const char* block_path_name(BlockPath path) { + switch (path) { + case BlockPath::kCoop: return "coop"; + case BlockPath::kCoopF16: return "coop-f16"; + case BlockPath::kPlain: return "plain"; + } + return "?"; +} + +BlockPath block_path_from_name(const std::string& name) { + if (name == "coop") return BlockPath::kCoop; + if (name == "coop-f16") return BlockPath::kCoopF16; + if (name == "plain") return BlockPath::kPlain; + throw std::invalid_argument("kernels must be coop, coop-f16 or plain, not \"" + name + "\""); +} + +BlockPath default_block_path(const Capabilities& caps) { + return caps.coop_bf16 ? BlockPath::kCoop : caps.coop_f16 ? BlockPath::kCoopF16 : BlockPath::kPlain; +} + +std::string no_coop_reason(const Capabilities& caps) { + if (caps.coop_bf16 || caps.coop_f16) return {}; + if (caps.coop_bf16_reason == caps.coop_f16_reason) return caps.coop_bf16_reason; + return "bf16: " + caps.coop_bf16_reason + "; f16: " + caps.coop_f16_reason; +} + +void BlockKernels::dispatch_gemm(const Pipeline& pipeline, std::uint32_t mode, VkCommandBuffer cb, VkDeviceAddress a, + VkDeviceAddress b, VkDeviceAddress out, VkDeviceAddress gate, std::uint32_t m, + std::uint32_t n, std::uint32_t k, std::uint32_t lda, std::uint32_t ldb, + std::uint32_t ldo) const { + const std::uint32_t align = is_coop(path_) ? 8 : 4; + require(m > 0 && n > 0 && k > 0, "bad GEMM shape"); + require(n % align == 0 && k % align == 0 && lda % align == 0 && ldb % align == 0 && ldo % align == 0, + "GEMM N, K and strides must be multiples of " + std::to_string(align)); + GemmPush push{a, b, out, gate, m, n, k, lda, ldb, ldo, 0, 0}; + if (path_ == BlockPath::kCoopF16) { + const std::uint32_t b_rows = mode == kSitu ? 2 * n : n; + require(m <= kMaxF16GemmRows && b_rows <= kMaxF16GemmRows, + "GEMM M and N must be at most " + std::to_string(kMaxF16GemmRows) + " on coop-f16"); + push.a_exponents = row_exponents_->address(); + push.b_exponents = push.a_exponents + VkDeviceAddress{kMaxF16GemmRows} * 4; + for (const RowScalePush& scale : {RowScalePush{a, push.a_exponents, m, k, lda}, RowScalePush{b, push.b_exponents, b_rows, k, ldb}}) { + const std::uint32_t groups = (scale.m + kRowScaleRows - 1) / kRowScaleRows, x = std::min(groups, 65535); + row_scale_->dispatch(cb, &scale, x, (groups + x - 1) / x); + } + compute_barrier(cb); + } + if (is_coop(path_)) { + // The shader numbers tiles y * x_count + x and walks them in bands of N tiles. + const std::uint32_t cols = mode == kSitu ? kCoopTileCols / 2 : kCoopTileCols; + pipeline.dispatch(cb, &push, (n + cols - 1) / cols, (m + kCoopTileRows - 1) / kCoopTileRows); + } else { + const std::uint32_t cols = mode == kSitu ? kGemmTileCols / 2 : kGemmTileCols; + pipeline.dispatch(cb, &push, (n + cols - 1) / cols, (m + kGemmTileRows - 1) / kGemmTileRows); + } + compute_barrier(cb); +} + +BlockKernels::BlockKernels(Device& device) : BlockKernels(device, default_block_path(device.info().caps)) {} + +BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { + const Capabilities& caps = device.info().caps; + if (path == BlockPath::kCoop && !caps.coop_bf16) { + throw std::runtime_error(device.info().name + " cannot run the bf16 cooperative-matrix kernels: " + + caps.coop_bf16_reason); + } + if (path == BlockPath::kCoopF16 && !caps.coop_f16) { + throw std::runtime_error(device.info().name + " cannot run the f16 cooperative-matrix kernels: " + + caps.coop_f16_reason); + } + gemm_bf16_ = gemm_pipeline(device, path, kBF16); + gemm_gated_ = gemm_pipeline(device, path, kGated); + gemm_situ_ = gemm_pipeline(device, path, kSitu); + gemm_bias_ = gemm_pipeline(device, path, kBias); + gemm_bias_residual_ = gemm_pipeline(device, path, kBiasResidual); + silu_ = std::make_unique(device, kElementwiseSpirv, sizeof(kElementwiseSpirv), sizeof(ElementwisePush), 0, + std::vector{0u}); + situ_ = std::make_unique(device, kElementwiseSpirv, sizeof(kElementwiseSpirv), sizeof(ElementwisePush), 0, + std::vector{1u}); + norm_modulate_ = std::make_unique(device, kNormSpirv, sizeof(kNormSpirv), sizeof(NormPush)); + qk_norm_rope_ = std::make_unique(device, kQkSpirv, sizeof(kQkSpirv), sizeof(QkPush)); + if (path == BlockPath::kCoop) { + attention_ = std::make_unique(device, kAttentionCoopSpirv, sizeof(kAttentionCoopSpirv), + sizeof(AttentionPush), kCoopSubgroup); + } else if (path == BlockPath::kCoopF16) { + attention_ = std::make_unique(device, kAttentionCoopF16Spirv, sizeof(kAttentionCoopF16Spirv), + sizeof(AttentionPush), kCoopSubgroup); + row_scale_ = std::make_unique(device, kRowScaleSpirv, sizeof(kRowScaleSpirv), sizeof(RowScalePush)); + row_exponents_ = std::make_unique(device, VkDeviceSize{kMaxF16GemmRows} * 4 * 2); + } else { + attention_ = std::make_unique(device, kAttentionSpirv, sizeof(kAttentionSpirv), sizeof(AttentionPush)); + } +} + +BlockKernels::~BlockKernels() = default; + +void BlockKernels::record_norm_modulate(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, + VkDeviceAddress scale, VkDeviceAddress shift, VkDeviceAddress out, + std::uint32_t ldo, std::uint32_t rows, std::uint32_t d) const { + require(d % 4 == 0 && ldx % 4 == 0 && ldo % 4 == 0, "norm: d and strides must be multiples of 4"); + if (rows == 0) return; + // torch's MeanOps factor. + const float factor = static_cast(rows) / static_cast(static_cast(rows) * d); + const NormPush push{x, weight, scale, shift, out, ldx, ldo, d, factor}; + norm_modulate_->dispatch(cb, &push, rows, 1); + compute_barrier(cb); +} + +void BlockKernels::record_gemm(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, std::uint32_t m, std::uint32_t n, + std::uint32_t k) const { + dispatch_gemm(*gemm_bf16_, kBF16, cb, x, w, out, 0, m, n, k, ldx, ldw, ldo); +} + +void BlockKernels::record_gemm_gated_residual(VkCommandBuffer cb, VkDeviceAddress residual, std::uint32_t ldr, + VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, + VkDeviceAddress gate, std::uint32_t m, std::uint32_t n, + std::uint32_t k) const { + dispatch_gemm(*gemm_gated_, kGated, cb, x, w, residual, gate, m, n, k, ldx, ldw, ldr); +} + +void BlockKernels::record_gemm_situ_glu(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w_up_gate, std::uint32_t ldw, std::uint32_t m, + std::uint32_t f, std::uint32_t k) const { + dispatch_gemm(*gemm_situ_, kSitu, cb, x, w_up_gate, out, 0, m, f, k, ldx, ldw, ldo); +} + +void BlockKernels::record_gemm_bias(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, VkDeviceAddress bias, + std::uint32_t m, std::uint32_t n, std::uint32_t k) const { + dispatch_gemm(*gemm_bias_, kBias, cb, x, w, out, bias, m, n, k, ldx, ldw, ldo); +} + +void BlockKernels::record_gemm_bias_residual(VkCommandBuffer cb, VkDeviceAddress residual, std::uint32_t ldr, + VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, + VkDeviceAddress bias, std::uint32_t m, std::uint32_t n, + std::uint32_t k) const { + dispatch_gemm(*gemm_bias_residual_, kBiasResidual, cb, x, w, residual, bias, m, n, k, ldx, ldw, ldr); +} + +void BlockKernels::record_norm(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, + VkDeviceAddress out, std::uint32_t ldo, std::uint32_t rows, std::uint32_t d) const { + record_norm_modulate(cb, x, ldx, weight, 0, 0, out, ldo, rows, d); +} + +namespace { +void dispatch_elementwise(const Pipeline& p, VkCommandBuffer cb, const ElementwisePush& push) { + require(push.n % 4 == 0, "elementwise: n must be a multiple of 4"); + if (push.n == 0) return; + const std::uint32_t groups = (push.n / 4 + 63) / 64, x = std::min(groups, 65535); + p.dispatch(cb, &push, x, (groups + x - 1) / x); + compute_barrier(cb); +} +} // namespace + +void BlockKernels::record_silu(VkCommandBuffer cb, VkDeviceAddress a, VkDeviceAddress out, std::uint32_t n) const { + dispatch_elementwise(*silu_, cb, ElementwisePush{a, 0, out, n}); +} + +void BlockKernels::record_situ(VkCommandBuffer cb, VkDeviceAddress hidden, VkDeviceAddress gate, VkDeviceAddress out, + std::uint32_t n) const { + dispatch_elementwise(*situ_, cb, ElementwisePush{hidden, gate, out, n}); +} + +void BlockKernels::record_qk_norm_rope(VkCommandBuffer cb, VkDeviceAddress qkv, std::uint32_t ld, std::uint32_t tokens, + std::uint32_t heads, VkDeviceAddress q_norm, VkDeviceAddress k_norm, + VkDeviceAddress table, std::uint32_t table_ld) const { + require(ld % 4 == 0 && table_ld % 4 == 0, "qk_norm_rope: strides must be multiples of 4"); + if (tokens == 0) return; + // rms_norm_moss on a [L, heads, 128] head view sees L * heads rows. + const std::uint64_t rows = static_cast(tokens) * heads; + const float factor = static_cast(rows) / static_cast(rows * kDiTHeadDim); + const QkPush push{qkv, q_norm, k_norm, table, ld, table_ld, heads, factor}; + qk_norm_rope_->dispatch(cb, &push, 2 * heads, tokens); + compute_barrier(cb); +} + +void BlockKernels::record_attention(VkCommandBuffer cb, const AttentionArgs& a) const { + require(a.cu_q.size() == a.cu_k.size() && a.cu_q.size() >= 2, "attention: cu_q and cu_k must describe the same batch"); + // In double then rounded, as FlashAttention (float * M_LOG2E): keeps the exp2 argument bit-identical. + const float log2_scale = static_cast(static_cast(a.scale) * 1.4426950408889634074); + if (is_coop(path_)) { + require(a.ldq % 8 == 0 && a.ldk % 8 == 0 && a.ldv % 8 == 0 && a.ldo % 8 == 0 && a.hsq % 8 == 0 && a.hsk % 8 == 0 && + a.hsv % 8 == 0 && a.hso % 8 == 0, + "attention: strides must be multiples of 8 on the cooperative-matrix path"); + } + for (std::size_t b = 0; b + 1 < a.cu_q.size(); ++b) { + const std::int32_t seqlen_q = a.cu_q[b + 1] - a.cu_q[b], seqlen_k = a.cu_k[b + 1] - a.cu_k[b]; + if (seqlen_q <= 0) continue; + require(seqlen_k > 0, "attention: a sequence with queries has no keys"); + const AttentionPush push{a.q, a.k, a.v, a.o, a.ldq, a.ldk, a.ldv, a.ldo, a.hsq, a.hsk, a.hsv, a.hso, + static_cast(a.cu_q[b]), static_cast(a.cu_k[b]), + static_cast(seqlen_q), static_cast(seqlen_k), log2_scale}; + const std::uint32_t rows = is_coop(path_) ? kCoopAttentionRows : kPlainAttentionRows; + attention_->dispatch(cb, &push, a.heads, (static_cast(seqlen_q) + rows - 1) / rows); + } + compute_barrier(cb); +} + +const char* block_op_name(int op) { + static constexpr const char* kNames[kBlockOps] = {"norm1+modulate", "qkv gemm", "qk_norm_rope", "attention", + "o gemm+gated_res", "norm2+modulate", "up_gate gemm+situ", "down gemm+gated_res"}; + return kNames[op]; +} + +void record_block_op(const BlockKernels& k, VkCommandBuffer cb, int op, const BlockWeights& w, const BlockBuffers& b) { + constexpr std::uint32_t ldx = kDiTDim, ldq = 3 * kDiTDim, ldf = kDiTFfnDim, dim = kDiTDim, head = kDiTHeadDim; + const auto mod = [&](int i) { return b.mod + static_cast(i) * kDiTDim * 2; }; + const float scale = static_cast(1.0 / std::sqrt(static_cast(kDiTHeadDim))); + switch (op) { + case 0: k.record_norm_modulate(cb, b.x, ldx, w.norm1, mod(1), mod(0), b.y, ldx, b.tokens, dim); break; + case 1: k.record_gemm(cb, b.qkv, ldq, b.y, ldx, w.qkv, dim, b.tokens, 3 * dim, dim); break; + case 2: k.record_qk_norm_rope(cb, b.qkv, ldq, b.tokens, kDiTHeads, w.q_norm, w.k_norm, b.rope, head); break; + case 3: { + const VkDeviceAddress kv = b.qkv + static_cast(dim) * 2; + const AttentionArgs args{b.qkv, kv, kv + static_cast(dim) * 2, b.attn, ldq, ldq, ldq, ldx, + head, head, head, head, b.cu_seqlens, b.cu_seqlens, kDiTHeads, scale}; + k.record_attention(cb, args); + break; + } + case 4: k.record_gemm_gated_residual(cb, b.x, ldx, b.attn, ldx, w.o, dim, mod(2), b.tokens, dim, dim); break; + case 5: k.record_norm_modulate(cb, b.x, ldx, w.norm2, mod(4), mod(3), b.y, ldx, b.tokens, dim); break; + case 6: k.record_gemm_situ_glu(cb, b.ffn, ldf, b.y, ldx, w.up_gate, dim, b.tokens, kDiTFfnDim, dim); break; + case 7: k.record_gemm_gated_residual(cb, b.x, ldx, b.ffn, ldf, w.down, kDiTFfnDim, mod(5), b.tokens, dim, kDiTFfnDim); break; + default: throw std::invalid_argument("vulkan block: no such op"); + } +} + +} // namespace mi::vk diff --git a/vulkan/block.h b/vulkan/block.h new file mode 100644 index 0000000..2f6e820 --- /dev/null +++ b/vulkan/block.h @@ -0,0 +1,120 @@ +// The plain compute path of one MOSS DiT block, bf16 throughout (src/models/moss_dit/dit_block_fused.h): +// y = rms_norm_modulate(x, norm1, scale_a, shift_a) +// qkv = y * Wqkv^T +// q,k = rope(rms_norm(q|k)) per head, in place +// attn = attention(q, k, v) +// x += bf16(bf16(attn * Wo^T) * gate_a) +// y = rms_norm_modulate(x, norm2, scale_f, shift_f) +// ffn = situ(bf16(y * Wup^T), bf16(y * Wgate^T)) +// x += bf16(bf16(ffn * Wdown^T) * gate_f) +// Three paths. kPlain runs on any device the runtime accepts, needs no cooperative matrices, and does not depend on the +// subgroup size; its operands need row strides that are multiples of 4 elements. kCoop runs the GEMMs and the attention on +// bf16 cooperative matrices (Capabilities::coop_bf16) and needs strides, N and K that are multiples of 8. kCoopF16 runs the +// same kernels on f16 fragments (Capabilities::coop_f16) for devices without bf16 ones: a bf16 value is an f16 value +// while it is inside f16's normal range, so the operands convert exactly as they are staged, with every row of a GEMM's +// A and B scaled by a power of two that brings its largest value to [2^15, 2^16) and the output scaled back (all exact); +// q, k and v convert unscaled, and P enters P V as bf16(p) * 2^15. Each record_* adds its own barrier. +#pragma once + +#include +#include +#include +#include + +#include "runtime.h" + +namespace mi::vk { + +inline constexpr int kDiTDim = 3072; +inline constexpr int kDiTHeads = 24; +inline constexpr int kDiTHeadDim = 128; +inline constexpr int kDiTFfnDim = 8192; + +struct AttentionArgs { + VkDeviceAddress q, k, v, o; + std::uint32_t ldq, ldk, ldv, ldo; + std::uint32_t hsq, hsk, hsv, hso; + std::vector cu_q, cu_k; // batch + 1 row offsets, on the host: sequences are dispatched one by one + std::uint32_t heads; + float scale; // SDPA's 1 / sqrt(head_dim) for the DiT +}; + +enum class BlockPath { kCoop, kCoopF16, kPlain }; +const char* block_path_name(BlockPath path); // "coop", "coop-f16", "plain" +// The path named so; throws std::invalid_argument for any other name. +BlockPath block_path_from_name(const std::string& name); +// The fastest path the device has: kCoop, then kCoopF16, then kPlain. +BlockPath default_block_path(const Capabilities& caps); +// Why the device has no cooperative-matrix path (empty when it has one). +std::string no_coop_reason(const Capabilities& caps); + +class BlockKernels { + public: + // kCoop where the device has it, unless `path` names one. Throws if the device cannot run the path asked for. + explicit BlockKernels(Device& device); + BlockKernels(Device& device, BlockPath path); + ~BlockKernels(); + + BlockPath path() const { return path_; } + + void record_norm_modulate(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, + VkDeviceAddress scale, VkDeviceAddress shift, VkDeviceAddress out, std::uint32_t ldo, + std::uint32_t rows, std::uint32_t d) const; + void record_gemm(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, std::uint32_t ldx, + VkDeviceAddress w, std::uint32_t ldw, std::uint32_t m, std::uint32_t n, std::uint32_t k) const; + void record_gemm_gated_residual(VkCommandBuffer cb, VkDeviceAddress residual, std::uint32_t ldr, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, VkDeviceAddress gate, + std::uint32_t m, std::uint32_t n, std::uint32_t k) const; + void record_gemm_situ_glu(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w_up_gate, std::uint32_t ldw, std::uint32_t m, + std::uint32_t f, std::uint32_t k) const; + // out = bf16(x W^T + bias), and residual += bf16(x W^T + bias): nn.Linear with a bias, alone or before a residual add. + void record_gemm_bias(VkCommandBuffer cb, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress x, std::uint32_t ldx, + VkDeviceAddress w, std::uint32_t ldw, VkDeviceAddress bias, std::uint32_t m, std::uint32_t n, + std::uint32_t k) const; + void record_gemm_bias_residual(VkCommandBuffer cb, VkDeviceAddress residual, std::uint32_t ldr, VkDeviceAddress x, + std::uint32_t ldx, VkDeviceAddress w, std::uint32_t ldw, VkDeviceAddress bias, + std::uint32_t m, std::uint32_t n, std::uint32_t k) const; + // out = bf16(bf16(rms_norm(x)) * weight), without modulation. + void record_norm(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, VkDeviceAddress out, + std::uint32_t ldo, std::uint32_t rows, std::uint32_t d) const; + // Contiguous bf16 vectors of n elements, n a multiple of 4: out = silu(a); out = situ(hidden, gate). + void record_silu(VkCommandBuffer cb, VkDeviceAddress a, VkDeviceAddress out, std::uint32_t n) const; + void record_situ(VkCommandBuffer cb, VkDeviceAddress hidden, VkDeviceAddress gate, VkDeviceAddress out, + std::uint32_t n) const; + void record_qk_norm_rope(VkCommandBuffer cb, VkDeviceAddress qkv, std::uint32_t ld, std::uint32_t tokens, + std::uint32_t heads, VkDeviceAddress q_norm, VkDeviceAddress k_norm, VkDeviceAddress table, + std::uint32_t table_ld) const; + void record_attention(VkCommandBuffer cb, const AttentionArgs& args) const; + + private: + void dispatch_gemm(const Pipeline& pipeline, std::uint32_t mode, VkCommandBuffer cb, VkDeviceAddress a, + VkDeviceAddress b, VkDeviceAddress out, VkDeviceAddress gate, std::uint32_t m, std::uint32_t n, + std::uint32_t k, std::uint32_t lda, std::uint32_t ldb, std::uint32_t ldo) const; + + BlockPath path_; + std::unique_ptr gemm_bf16_, gemm_gated_, gemm_situ_, gemm_bias_, gemm_bias_residual_, norm_modulate_, + qk_norm_rope_, attention_, silu_, situ_; + // kCoopF16: the per-row power-of-two exponents of the GEMM's A and B, written by row_scale_ before each GEMM (the GEMMs + // run one after another, each behind a barrier, so one buffer serves them all). + std::unique_ptr row_scale_; + std::unique_ptr row_exponents_; +}; + +struct BlockWeights { + VkDeviceAddress norm1, norm2, q_norm, k_norm, qkv, o, up_gate, down; +}; + +// One run of `tokens` tokens: x, y [tokens, kDiTDim]; qkv [tokens, 3 kDiTDim]; attn [tokens, kDiTDim]; +// ffn [tokens, kDiTFfnDim]; rope [tokens, 128]; mod [6 kDiTDim]: shift_a, scale_a, gate_a, shift_f, scale_f, gate_f. +struct BlockBuffers { + VkDeviceAddress x, y, qkv, attn, ffn, rope, mod; + std::vector cu_seqlens; + std::uint32_t tokens; +}; + +inline constexpr int kBlockOps = 8; +const char* block_op_name(int op); +void record_block_op(const BlockKernels& kernels, VkCommandBuffer cb, int op, const BlockWeights& w, const BlockBuffers& b); + +} // namespace mi::vk diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp new file mode 100644 index 0000000..5932a62 --- /dev/null +++ b/vulkan/dit.cpp @@ -0,0 +1,474 @@ +#include "dit.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "safetensors.h" + +namespace mi::vk { +namespace { + +constexpr std::uint32_t kDim = kDiTDim, kFfn = kDiTFfnDim, kHead = kDiTHeadDim, kQkv = 3 * kDiTDim; +constexpr VkDeviceSize kBf16 = 2; + +// The RoPE band frequencies 1 / 10000^(i / 16), as the reference builds them in fp32 on the CPU (kernels/basic/rope.h). +constexpr std::uint32_t kBandBits[16] = {0x3f800000, 0x3f0ff59a, 0x3ea1e89b, 0x3e361887, 0x3dcccccd, 0x3d6655c2, + 0x3d0186e3, 0x3c91ad39, 0x3c23d70a, 0x3bb8449c, 0x3b4f3e38, 0x3ae91528, + 0x3a83126f, 0x3a136a16, 0x39a5cb60, 0x393a7753}; + +std::uint16_t to_bf16(float v) { + const auto u = std::bit_cast(v); + if ((u & 0x7fffffffu) > 0x7f800000u) return static_cast((u >> 16) | 0x0040u); + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} +float from_bf16(std::uint16_t b) { return std::bit_cast(static_cast(b) << 16); } + +[[noreturn]] void fail(const std::string& what) { throw std::runtime_error("vulkan dit: " + what); } + +double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +VkDeviceAddress at(const Buffer& b, VkDeviceSize byte_offset) { return b.address() + byte_offset; } + +} // namespace + +std::vector rope_table(std::uint32_t height, std::uint32_t width, std::uint32_t text_length) { + const std::uint64_t hw = std::uint64_t{height} * width, n = hw + text_length; + const float h_mid = static_cast((static_cast(height) - 1.0) / 2.0); + const float w_mid = static_cast((static_cast(width) - 1.0) / 2.0); + std::vector out(n * kHead); + for (std::uint64_t r = 0; r < n; ++r) { + for (std::uint32_t p = 0; p < 64; ++p) { + const std::uint32_t axis = p >> 4, band = p & 15; + float coord = 0.0f; + if (r < hw) { + if (axis == 1) coord = static_cast(text_length); + if (axis == 2) coord = static_cast(r / width) - h_mid; + if (axis == 3) coord = static_cast(r % width) - w_mid; + } else if (axis == 1) { + coord = static_cast(r - hw); + } + const float phase = coord * std::bit_cast(kBandBits[band]); + // sinf/cosf to fp32 as torch's CUDA kernel; libm's are correctly rounded or within an ulp, and bf16 hides that. + out[r * kHead + p] = to_bf16(std::sin(phase)); + out[r * kHead + 64 + p] = to_bf16(std::cos(phase)); + } + } + return out; +} + +Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::path& transformer_dir, const DitOptions& options) + : device_(device), kernels_(kernels) { + max_gflop_ = options.submit_gflop > 0 ? options.submit_gflop : kernels.path() != BlockPath::kPlain ? 200.0 : 20.0; + const st::Index index(transformer_dir); + std::vector host; + // Places tensors in one new buffer, each at a 256-byte aligned offset, and uploads them (or, for a host buffer, reads + // them straight into its mapping); returns their offsets. + const auto place = [&](const std::vector& names, const std::vector>& shapes, + Buffer::Kind kind) { + std::vector offsets; + VkDeviceSize total = 0; + for (std::size_t i = 0; i < names.size(); ++i) { + const st::TensorInfo& t = index.at(names[i]); + if (t.dtype != "BF16") fail(names[i] + " is " + t.dtype + ", expected BF16"); + if (t.shape != shapes[i]) fail(names[i] + " has an unexpected shape"); + offsets.push_back(total); + total += (t.bytes + 255) / 256 * 256; + } + // Host copies of streamed blocks are read by the transfer queue. + owned_.push_back(std::make_unique(device_, total, kind, kind == Buffer::Kind::kStaging)); + Buffer& buf = *owned_.back(); + for (std::size_t i = 0; i < names.size(); ++i) { + const std::uint64_t bytes = index.at(names[i]).bytes; + if (kind == Buffer::Kind::kStaging) { + index.read_into(names[i], static_cast(buf.mapped()) + offsets[i], bytes); + } else { + host.resize(bytes); + index.read_into(names[i], host.data(), host.size()); + buf.upload(host.data(), host.size(), offsets[i]); + } + } + if (kind == Buffer::Kind::kDevice) weight_bytes_ += total; + else host_weight_bytes_ += total; + return offsets; + }; + const auto load = [&](const std::vector& names, const std::vector>& shapes) { + const std::vector offsets = place(names, shapes, Buffer::Kind::kDevice); + std::vector addresses; + for (const VkDeviceSize o : offsets) addresses.push_back(owned_.back()->address() + o); + return addresses; + }; + const std::int64_t d = kDim, f = kFfn; + + auto a = load({"video_embedder.weight", "video_embedder.bias", "time_embedding.0.weight", "time_embedding.0.bias", + "time_embedding.2.weight", "time_embedding.2.bias", "shared_block_modulation.weight", + "shared_block_modulation.bias", "video_head_modulation.weight", "video_head_modulation.bias", + "final_norm_video.weight", "final_linear_video.weight", "text_embedder.proj.weight", + "text_embedder.proj.bias"}, + {{d, kDiTLatentChannels}, {d}, {d, kDiTFreqDim}, {d}, {d, d}, {d}, {6 * d, d}, {6 * d}, {2 * d, d}, {2 * d}, + {d}, {kDiTLatentChannels, d}, {d, kDiTTextDim}, {d}}); + g_ = {a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7], a[8], a[9], a[10], a[11], a[12], a[13]}; + + for (int b = 0; b < 2; ++b) { + const std::string p = "text_embedder.blocks." + std::to_string(b) + "."; + auto t = load({p + "norm1.weight", p + "norm2.weight", p + "q_norm.weight", p + "k_norm.weight", p + "qkv.weight", + p + "qkv.bias", p + "proj.weight", p + "proj.bias", p + "ffn_up.weight", p + "ffn_up.bias", + p + "ffn_gate.weight", p + "ffn_gate.bias", p + "ffn_down.weight", p + "ffn_down.bias"}, + {{d}, {d}, {kHead}, {kHead}, {3 * d, d}, {3 * d}, {d, d}, {d}, {f, d}, {f}, {f, d}, {f}, {d, f}, {d}}); + text_blocks_.push_back({t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7], t[8], t[9], t[10], t[11], t[12], t[13]}); + } + + // Blocks: the first `resident` in device memory, the rest in host memory, each copied into one device slot before it + // runs. q, k and v sit back to back (256-aligned sizes of 18 MiB each), so they form the [9216, 3072] QKV weight. + const std::vector> block_shapes = {{d}, {d}, {kHead}, {kHead}, {d, d}, {d, d}, {d, d}, {d, d}, + {2 * f, d}, {d, f}}; + const auto block_names = [](int l) { + const std::string p = "layers." + std::to_string(l) + "."; + return std::vector{p + "norm1.weight", p + "norm2.weight", p + "q_norm.weight", p + "k_norm.weight", + p + "proj.q.weight", p + "proj.k.weight", p + "proj.v.weight", p + "proj.o.weight", + p + "proj.ffn_up_gate.weight", p + "proj.ffn_down.weight"}; + }; + const auto block_weights = [&](VkDeviceAddress base, const std::vector& o) { + if (o[5] != o[4] + kDim * kDim * kBf16 || o[6] != o[5] + kDim * kDim * kBf16) fail("q, k and v are not contiguous"); + return BlockWeights{base + o[0], base + o[1], base + o[2], base + o[3], base + o[4], base + o[7], base + o[8], base + o[9]}; + }; + VkDeviceSize block_bytes = 0; + for (const std::string& n : block_names(0)) block_bytes += (index.at(n).bytes + 255) / 256 * 256; + int resident = kDiTLayers; + if (options.device_budget > 0) { + const std::uint64_t free = options.device_budget > weight_bytes_ ? options.device_budget - weight_bytes_ : 0; + // Streamed blocks need two slots: the next one copies while the current one runs. + resident = free >= kDiTLayers * block_bytes ? kDiTLayers + : static_cast(free > 2 * block_bytes ? free / block_bytes - 2 : 0); + } + for (int l = 0; l < kDiTLayers; ++l) { + const bool on_device = l < resident; + const std::vector o = place(block_names(l), block_shapes, on_device ? Buffer::Kind::kDevice : Buffer::Kind::kStaging); + if (on_device) { + blocks_.push_back(block_weights(owned_.back()->address(), o)); + host_blocks_.push_back(nullptr); + slot_of_.push_back(-1); + } else { + const int s = (l - resident) % 2; + if (!slots_[s]) { + slots_[s] = std::make_unique(device_, block_bytes, Buffer::Kind::kDevice, true); + weight_bytes_ += block_bytes; + } + blocks_.push_back(block_weights(slots_[s]->address(), o)); + host_blocks_.push_back(owned_.back().get()); + slot_of_.push_back(s); + } + } + resident_blocks_ = resident; +} + +Dit::~Dit() { + // Copies still in flight read host buffers this destructor frees. A lost device fails the wait; nothing is in flight + // then. + for (Slot& s : slot_state_) { + if (s.ticket == 0) continue; + try { + device_.host_wait(s.ticket); + } catch (const std::exception&) { + } + } +} + +void Dit::issue_copy(int layer) { + if (layer < 0 || host_blocks_[static_cast(layer)] == nullptr) return; + Slot& s = slot_state_[slot_of_[static_cast(layer)]]; + if (s.layer == layer) return; + if (s.ticket != 0) { // a copy nobody ran from (the next forward skipped it): retire it + device_.host_wait(s.ticket); + if (!s.gpu_waited) { + device_.gpu_wait(s.ticket); + device_.run([](VkCommandBuffer) {}); + } + } + const Buffer& host = *host_blocks_[static_cast(layer)]; + s = {layer, device_.copy_async(host, 0, *slots_[slot_of_[static_cast(layer)]], 0, host.bytes()), false}; +} + +int Dit::next_streamed(int layer) const { + for (int k = 1; k <= kDiTLayers; ++k) { + const int l = (layer + k) % kDiTLayers; + if (host_blocks_[static_cast(l)] != nullptr) return l; + } + return -1; +} + +void Dit::run_block_op(int op, const BlockWeights& w, const BlockBuffers& b) { + // Multiply-adds per token, x2 for FLOP; attention's per query token over its sequence's keys. + static constexpr double kFlopPerToken[kBlockOps] = {0, 2.0 * kDim * kQkv, 0, 0, 2.0 * kDim * kDim, 0, + 2.0 * kDim * 2 * kFfn, 2.0 * kFfn * kDim}; + const double limit = max_gflop_ * 1e9; + if (op == 3) { + // Attention splits along each sequence's queries: every query row reads all of its sequence's keys. + for (std::size_t s = 0; s + 1 < b.cu_seqlens.size(); ++s) { + const auto start = static_cast(b.cu_seqlens[s]), len = static_cast(b.cu_seqlens[s + 1]) - start; + const double per_query = 4.0 * len * kDim; + std::uint32_t chunk = len; + if (limit > 0 && per_query * len > limit) chunk = std::max(64, static_cast(limit / per_query) / 64 * 64); + for (std::uint32_t q0 = 0; q0 < len; q0 += chunk) { + const std::uint32_t n = std::min(chunk, len - q0); + const VkDeviceAddress base = b.qkv + VkDeviceAddress{start} * kQkv * kBf16; + const VkDeviceAddress q = base + VkDeviceAddress{q0} * kQkv * kBf16; + const AttentionArgs args{q, base + kDim * kBf16, base + 2 * kDim * kBf16, b.attn + (VkDeviceAddress{start} + q0) * kDim * kBf16, + kQkv, kQkv, kQkv, kDim, kHead, kHead, kHead, kHead, {0, static_cast(n)}, + {0, static_cast(len)}, kDiTHeads, + static_cast(1.0 / std::sqrt(static_cast(kHead)))}; + device_.run([&](VkCommandBuffer cb) { kernels_.record_attention(cb, args); }); + } + } + return; + } + // Row-local ops split along the tokens, in multiples of the GEMM tiles' 128 rows. + std::uint32_t chunk = b.tokens; + if (limit > 0 && kFlopPerToken[op] * b.tokens > limit) { + chunk = std::max(128, static_cast(limit / kFlopPerToken[op]) / 128 * 128); + } + for (std::uint32_t t0 = 0; t0 < b.tokens; t0 += chunk) { + BlockBuffers part = b; + const VkDeviceAddress t = t0; + part.x += t * kDim * kBf16; + part.y += t * kDim * kBf16; + part.qkv += t * kQkv * kBf16; + part.attn += t * kDim * kBf16; + part.ffn += t * kFfn * kBf16; + part.rope += t * kHead * kBf16; + part.tokens = std::min(chunk, b.tokens - t0); + device_.run([&](VkCommandBuffer cb) { record_block_op(kernels_, cb, op, w, part); }); + } +} + +std::uint64_t Dit::activation_bytes(std::uint64_t tokens, std::uint64_t hw) { + // forward(): x, y and attn [tokens, 3072]; qkv [tokens, 9216]; ffn [tokens, 8192]; rope [tokens, 128]; the video tokens + // and the head's projection [hw, 128]; the modulation and time scratch. encode_text() reuses them for a prompt no longer + // than the joint sequence, with its own up, gate [S, 8192] and embeddings [S, 4096]. + const std::uint64_t per_token = (3 * kDim + kQkv + kFfn + kHead) * kBf16; + return tokens * per_token + 2 * hw * kDiTLatentChannels * kBf16 + 2 * tokens * kFfn * kBf16 + + tokens * kDiTTextDim * kBf16 + (8 * kDim + 2 * kDim + kDiTFreqDim) * kBf16; +} + +Buffer& Dit::scratch(std::unique_ptr& slot, VkDeviceSize bytes) { + if (!slot || slot->bytes() < bytes) { + slot.reset(); + slot = std::make_unique(device_, bytes); + } + return *slot; +} + +std::unique_ptr Dit::encode_text(const std::vector& embeds, std::uint32_t length) { + if (length == 0 || embeds.size() != std::size_t{length} * kDiTTextDim) fail("encode_text: embeds must be [length, 4096]"); + const VkDeviceSize s = length; + Buffer& e = scratch(embeds_, s * kDiTTextDim * kBf16); + e.upload(embeds.data(), embeds.size() * kBf16); + Buffer& y = scratch(y_, s * kDim * kBf16); + Buffer& qkv = scratch(qkv_, s * kQkv * kBf16); + Buffer& attn = scratch(attn_, s * kDim * kBf16); + Buffer& up = scratch(up_, s * kFfn * kBf16); + Buffer& gate = scratch(gate_, s * kFfn * kBf16); + Buffer& ffn = scratch(ffn_, s * kFfn * kBf16); + Buffer& rope = scratch(rope_, s * kHead * kBf16); + const std::vector table = rope_table(0, 0, length); + rope.upload(table.data(), table.size() * kBf16); + auto out = std::make_unique(device_, s * kDim * kBf16); + const VkDeviceAddress x = out->address(); + const VkDeviceAddress ya = y.address(), qa = qkv.address(), aa = attn.address(), ua = up.address(), ga = gate.address(), + fa = ffn.address(), ea = e.address(); + // Each linear in token ranges of at most max_gflop_ (see run_block_op); norms and elementwise ops run whole. + const auto rows = [&](double flop_per_token, const std::function& op) { + const double limit = max_gflop_ * 1e9; + std::uint32_t chunk = length; + if (limit > 0 && flop_per_token * length > limit) { + chunk = std::max(128, static_cast(limit / flop_per_token) / 128 * 128); + } + for (std::uint32_t t0 = 0; t0 < length; t0 += chunk) { + device_.run([&](VkCommandBuffer cb) { op(cb, t0, std::min(chunk, length - t0)); }); + } + }; + const auto off = [](VkDeviceAddress base, std::uint64_t t, std::uint64_t width) { return base + t * width * kBf16; }; + + rows(2.0 * kDiTTextDim * kDim, [&](VkCommandBuffer cb, VkDeviceAddress t, std::uint32_t n) { + kernels_.record_gemm_bias(cb, off(x, t, kDim), kDim, off(ea, t, kDiTTextDim), kDiTTextDim, g_.text_w, kDiTTextDim, g_.text_b, + n, kDim, kDiTTextDim); + }); + for (const TextWeights& b : text_blocks_) { + device_.run([&](VkCommandBuffer cb) { kernels_.record_norm(cb, x, kDim, b.norm1, ya, kDim, length, kDim); }); + rows(2.0 * kDim * kQkv, [&](VkCommandBuffer cb, VkDeviceAddress t, std::uint32_t n) { + kernels_.record_gemm_bias(cb, off(qa, t, kQkv), kQkv, off(ya, t, kDim), kDim, b.qkv_w, kDim, b.qkv_b, n, kQkv, kDim); + }); + device_.run([&](VkCommandBuffer cb) { + kernels_.record_qk_norm_rope(cb, qa, kQkv, length, kDiTHeads, b.q_norm, b.k_norm, rope.address(), kHead); + }); + run_block_op(3, BlockWeights{}, BlockBuffers{0, 0, qa, aa, 0, 0, 0, {0, static_cast(length)}, length}); + rows(2.0 * kDim * kDim, [&](VkCommandBuffer cb, VkDeviceAddress t, std::uint32_t n) { + kernels_.record_gemm_bias_residual(cb, off(x, t, kDim), kDim, off(aa, t, kDim), kDim, b.proj_w, kDim, b.proj_b, n, kDim, + kDim); + }); + device_.run([&](VkCommandBuffer cb) { kernels_.record_norm(cb, x, kDim, b.norm2, ya, kDim, length, kDim); }); + rows(4.0 * kDim * kFfn, [&](VkCommandBuffer cb, VkDeviceAddress t, std::uint32_t n) { + kernels_.record_gemm_bias(cb, off(ua, t, kFfn), kFfn, off(ya, t, kDim), kDim, b.up_w, kDim, b.up_b, n, kFfn, kDim); + kernels_.record_gemm_bias(cb, off(ga, t, kFfn), kFfn, off(ya, t, kDim), kDim, b.gate_w, kDim, b.gate_b, n, kFfn, kDim); + }); + device_.run([&](VkCommandBuffer cb) { kernels_.record_situ(cb, ua, ga, fa, length * kFfn); }); + rows(2.0 * kFfn * kDim, [&](VkCommandBuffer cb, VkDeviceAddress t, std::uint32_t n) { + kernels_.record_gemm_bias_residual(cb, off(x, t, kDim), kDim, off(fa, t, kFfn), kFfn, b.down_w, kFfn, b.down_b, n, kDim, + kFfn); + }); + } + return out; +} + +void Dit::time_modulation(float timestep) { + // P:232-241: a float64 sinusoid rounded to fp32 then bf16 (the .float() and back are exact), then the bf16 MLP. + constexpr int kHalf = kDiTFreqDim / 2; + std::vector emb(kDiTFreqDim); + for (int i = 0; i < kHalf; ++i) { + const double freq = std::pow(10000.0, -static_cast(i) / kHalf); + const double phase = static_cast(timestep) * freq; + emb[static_cast(i)] = to_bf16(static_cast(std::cos(phase))); + emb[static_cast(kHalf + i)] = to_bf16(static_cast(std::sin(phase))); + } + Buffer& t_in = scratch(t_in_, kDiTFreqDim * kBf16); + t_in.upload(emb.data(), emb.size() * kBf16); + Buffer& hidden = scratch(hidden_, kDim * kBf16); + Buffer& time = scratch(time_, kDim * kBf16); + Buffer& mod = scratch(mod_, 6 * kDim * kBf16); + Buffer& head = scratch(head_, 2 * kDim * kBf16); + device_.run([&](VkCommandBuffer cb) { + kernels_.record_gemm_bias(cb, hidden.address(), kDim, t_in.address(), kDiTFreqDim, g_.time0_w, kDiTFreqDim, g_.time0_b, + 1, kDim, kDiTFreqDim); + kernels_.record_silu(cb, hidden.address(), hidden.address(), kDim); + kernels_.record_gemm_bias(cb, time.address(), kDim, hidden.address(), kDim, g_.time2_w, kDim, g_.time2_b, 1, kDim, kDim); + kernels_.record_gemm_bias(cb, mod.address(), 6 * kDim, time.address(), kDim, g_.mod_w, kDim, g_.mod_b, 1, 6 * kDim, kDim); + kernels_.record_gemm_bias(cb, head.address(), 2 * kDim, time.address(), kDim, g_.head_w, kDim, g_.head_b, 1, 2 * kDim, + kDim); + }); +} + +std::vector> Dit::forward(const std::vector& latents, std::uint32_t height, std::uint32_t width, + float timestep, const std::vector& rows) { + const std::uint64_t hw = std::uint64_t{height} * width; + if (rows.empty() || hw == 0 || latents.size() != hw * kDiTLatentChannels) fail("forward: latents must be [128, h, w]"); + std::vector off{0}; + for (const Row& r : rows) { + if (r.text == nullptr || r.length == 0 || r.text->bytes() < std::uint64_t{r.length} * kDim * kBf16) { + fail("forward: a row has no refined text"); + } + off.push_back(off.back() + hw + r.length); + } + const std::uint64_t total = off.back(); + op_ms_.assign(kBlockOps + 1, 0.0); + const double t_start = now_ms(); + double copy_ms = 0.0; + + Buffer& x = scratch(x_, total * kDim * kBf16); + scratch(y_, total * kDim * kBf16); + scratch(qkv_, total * kQkv * kBf16); + scratch(attn_, total * kDim * kBf16); + scratch(ffn_, total * kFfn * kBf16); + Buffer& rope = scratch(rope_, total * kHead * kBf16); + Buffer& tokens = scratch(tokens_, hw * kDiTLatentChannels * kBf16); + Buffer& proj = scratch(proj_, hw * kDiTLatentChannels * kBf16); + + time_modulation(timestep); + + // The video tokens [hw, 128] in bf16 (the reference's .to(bf16) of the transposed latents), the same for every row. + std::vector host_tokens(hw * kDiTLatentChannels); + for (std::uint64_t p = 0; p < hw; ++p) { + for (std::uint32_t c = 0; c < kDiTLatentChannels; ++c) host_tokens[p * kDiTLatentChannels + c] = to_bf16(latents[c * hw + p]); + } + tokens.upload(host_tokens.data(), host_tokens.size() * kBf16); + for (std::size_t r = 0; r < rows.size(); ++r) { + const std::vector table = rope_table(height, width, rows[r].length); + rope.upload(table.data(), table.size() * kBf16, off[r] * kHead * kBf16); + } + device_.run([&](VkCommandBuffer cb) { + for (std::size_t r = 0; r < rows.size(); ++r) { + kernels_.record_gemm_bias(cb, at(x, off[r] * kDim * kBf16), kDim, tokens.address(), kDiTLatentChannels, g_.video_w, + kDiTLatentChannels, g_.video_b, static_cast(hw), kDim, kDiTLatentChannels); + record_copy(cb, *rows[r].text, 0, x, (off[r] + hw) * kDim * kBf16, std::uint64_t{rows[r].length} * kDim * kBf16); + } + }); + + for (int l = 0; l < kDiTLayers; ++l) { + Slot* slot = nullptr; + int prefetch_after = -1; + if (host_blocks_[static_cast(l)] != nullptr) { + // This block's copy (issued while the previous streamed block ran, or now), waited for on the device by its first + // op. The next streamed block's copy goes into the other slot now, on the transfer queue while this block computes; + // when it shares this slot (an odd number of streamed blocks wraps around to the first), only after this block. + const double t0 = now_ms(); + issue_copy(l); + slot = &slot_state_[slot_of_[static_cast(l)]]; + if (!slot->gpu_waited) { + device_.gpu_wait(slot->ticket); + slot->gpu_waited = true; + } + prefetch_after = next_streamed(l); + if (prefetch_after == l) prefetch_after = -1; + if (prefetch_after >= 0 && slot_of_[static_cast(prefetch_after)] != slot_of_[static_cast(l)]) { + issue_copy(prefetch_after); + prefetch_after = -1; + } + copy_ms += now_ms() - t0; + } + // Each maximal run of consecutive rows that executes block l is one varlen batch over its contiguous tokens. + for (std::size_t r0 = 0; r0 < rows.size();) { + if (rows[r0].skip.test(static_cast(l))) { + ++r0; + continue; + } + std::size_t r1 = r0; + std::vector cu{0}; + while (r1 < rows.size() && !rows[r1].skip.test(static_cast(l))) { + cu.push_back(static_cast(off[r1 + 1] - off[r0])); + ++r1; + } + const BlockBuffers b{at(x, off[r0] * kDim * kBf16), y_->address(), qkv_->address(), attn_->address(), ffn_->address(), + at(rope, off[r0] * kHead * kBf16), mod_->address(), cu, + static_cast(off[r1] - off[r0])}; + for (int op = 0; op < kBlockOps; ++op) { + const double t0 = now_ms(); + run_block_op(op, blocks_[static_cast(l)], b); + op_ms_[static_cast(op)] += now_ms() - t0; + } + r0 = r1; + } + if (slot != nullptr && slot->ticket != 0) { + device_.host_wait(slot->ticket); // every run of this block has completed (Device::run waits) + slot->ticket = 0; + } + if (prefetch_after >= 0) issue_copy(prefetch_after); // this block's slot is free now + } + + // Head (P:256-259): final norm with the head modulation (shift = head[0], scale = head[1]), linear, transpose. + std::vector> out(rows.size(), std::vector(hw * kDiTLatentChannels)); + std::vector host_proj(hw * kDiTLatentChannels); + for (std::size_t r = 0; r < rows.size(); ++r) { + device_.run([&](VkCommandBuffer cb) { + kernels_.record_norm_modulate(cb, at(x, off[r] * kDim * kBf16), kDim, g_.final_norm, at(*head_, kDim * kBf16), + head_->address(), y_->address(), kDim, static_cast(hw), kDim); + kernels_.record_gemm(cb, proj.address(), kDiTLatentChannels, y_->address(), kDim, g_.final_w, kDim, + static_cast(hw), kDiTLatentChannels, kDim); + }); + proj.download(host_proj.data(), host_proj.size() * kBf16); + for (std::uint64_t p = 0; p < hw; ++p) { + for (std::uint32_t c = 0; c < kDiTLatentChannels; ++c) out[r][c * hw + p] = from_bf16(host_proj[p * kDiTLatentChannels + c]); + } + } + double blocks = 0.0; + for (int op = 0; op < kBlockOps; ++op) blocks += op_ms_[static_cast(op)]; + op_ms_[kBlockOps] = now_ms() - t_start - blocks; + last_copy_ms_ = copy_ms; + return out; +} + +} // namespace mi::vk diff --git a/vulkan/dit.h b/vulkan/dit.h new file mode 100644 index 0000000..46c64b8 --- /dev/null +++ b/vulkan/dit.h @@ -0,0 +1,118 @@ +// The MOSS-Image DiT (MOSSImageTransformer, transformer/pipeline_moss_image.py) on the Vulkan block kernels, every weight +// resident in device memory (bf16, about 9.5 GiB): +// - encode_text(): the text refiner (text_embedder: a linear, then 2 un-modulated blocks), once per prompt. +// - forward(): one sampling step for several rows (cond, uncond, ...). The time embedding and the shared modulation run +// once; the rows' joint sequences [video_r | text_r] are concatenated, so each linear runs once over all rows and +// attention is one varlen call per run of rows. A block that some rows skip runs on the other rows' token ranges only. +// Each op follows the reference's order and bf16/fp32 rounding, as the engine's CUDA forward (src/models/moss_dit). +// Work is submitted op by op, and an op over many tokens in parts (DitOptions::submit_gflop). +#pragma once + +#include +#include +#include +#include +#include + +#include "block.h" +#include "runtime.h" + +namespace mi::vk { + +inline constexpr int kDiTLayers = 42; +inline constexpr int kDiTTextDim = 4096; +inline constexpr int kDiTFreqDim = 256; +inline constexpr int kDiTLatentChannels = 128; + +using SkipSet = std::bitset; + +struct DitOptions { + // Device memory the weights may take; 0: every block resident. Blocks past the budget stay in host memory (mapped, + // host-coherent) and are copied into one device slot right before they run, so a device with too little memory for + // the 9.5 GiB of weights still runs the DiT, at the cost of moving those blocks over the bus every step. + std::uint64_t device_budget = 0; + // Estimated work per submit, in GFLOP; 0 picks the path's default (200 cooperative matrix, 20 plain). A block op + // larger than that runs in token ranges (attention in query ranges), each its own submit, so no submit runs long + // enough to trip a display driver's watchdog (Windows' TDR, 2 s by default) on a slow device. + double submit_gflop = 0; +}; + +class Dit { + public: + // transformer_dir: the model's transformer/ directory (config.json and the sharded safetensors). + Dit(Device& device, const BlockKernels& kernels, const std::filesystem::path& transformer_dir, + const DitOptions& options = {}); + ~Dit(); + Dit(const Dit&) = delete; + Dit& operator=(const Dit&) = delete; + + std::uint64_t weight_bytes() const { return weight_bytes_; } // device memory, the streaming slot included + std::uint64_t host_weight_bytes() const { return host_weight_bytes_; } + int resident_blocks() const { return resident_blocks_; } + // Device memory of the activations of one forward() over `tokens` joint tokens and an h x w latent grid. + static std::uint64_t activation_bytes(std::uint64_t tokens, std::uint64_t hw); + + // The text refiner on one prompt's text-encoder output, bf16 bits [length, 4096]: device bf16 [length, 3072]. + std::unique_ptr encode_text(const std::vector& embeds, std::uint32_t length); + + struct Row { + const Buffer* text = nullptr; // from encode_text() + std::uint32_t length = 0; + SkipSet skip; + }; + // One step: latents f32 [128, height, width] (the sampler's), timestep as the scheduler passes it. Returns each row's + // prediction, f32 [128, height, width]. + std::vector> forward(const std::vector& latents, std::uint32_t height, std::uint32_t width, + float timestep, const std::vector& rows); + + // Milliseconds of the last forward(), summed per op kind over its blocks (block_op_name order), then the rest (which + // includes last_copy_ms(), the time spent issuing streamed blocks' copies; they then run alongside the compute). + const std::vector& last_op_ms() const { return op_ms_; } + double last_copy_ms() const { return last_copy_ms_; } + + private: + Buffer& scratch(std::unique_ptr& slot, VkDeviceSize bytes); + void run_block_op(int op, const BlockWeights& w, const BlockBuffers& b); + void time_modulation(float timestep); + + Device& device_; + const BlockKernels& kernels_; + std::uint64_t weight_bytes_ = 0, host_weight_bytes_ = 0; + int resident_blocks_ = 0; + double max_gflop_ = 0; + std::vector host_blocks_; // per block: its host copy when streamed, else null + // Streamed blocks alternate between two device slots: block k's copy runs on the transfer queue while block k - 1 + // computes. + struct Slot { + int layer = -1; // the block the slot holds or is receiving + std::uint64_t ticket = 0; // its copy, until retired + bool gpu_waited = false; // a run has waited for that copy on the device + }; + std::unique_ptr slots_[2]; + Slot slot_state_[2]; + std::vector slot_of_; // per block: its slot when streamed, else -1 + void issue_copy(int layer); + int next_streamed(int layer) const; + double last_copy_ms_ = 0.0; + std::vector> owned_; + std::vector blocks_; + struct Globals { + VkDeviceAddress video_w, video_b, time0_w, time0_b, time2_w, time2_b, mod_w, mod_b, head_w, head_b, final_norm, + final_w, text_w, text_b; + } g_{}; + struct TextWeights { + VkDeviceAddress norm1, norm2, q_norm, k_norm, qkv_w, qkv_b, proj_w, proj_b, up_w, up_b, gate_w, gate_b, down_w, + down_b; + }; + std::vector text_blocks_; + + std::unique_ptr x_, y_, qkv_, attn_, ffn_, rope_, tokens_, proj_, up_, gate_, embeds_; + std::unique_ptr mod_, head_, t_in_, hidden_, time_; + std::vector op_ms_; +}; + +// The RoPE table rows of one joint sequence (P:178-185), bf16 bits [height * width + text_length, 128]: sin of the 64 +// phases, then their cos. height = width = 0 gives the text rows alone. +std::vector rope_table(std::uint32_t height, std::uint32_t width, std::uint32_t text_length); + +} // namespace mi::vk diff --git a/vulkan/gemm.cpp b/vulkan/gemm.cpp new file mode 100644 index 0000000..2b69096 --- /dev/null +++ b/vulkan/gemm.cpp @@ -0,0 +1,84 @@ +#include "gemm.h" + +#include +#include + +namespace mi::vk { +namespace { + +const std::uint32_t kCoopSpirv[] = { +#include "gemm_coop.spv.inc" +}; +const std::uint32_t kPlainSpirv[] = { +#include "gemm_plain.spv.inc" +}; + +constexpr std::uint32_t kCoopTile = 64; +constexpr std::uint32_t kPlainTileRows = 128, kPlainTileCols = 128; +constexpr std::uint32_t kCoopSubgroup = 32; + +struct Push { + std::uint64_t a, b, c; + std::uint32_t m, n, k, lda, ldb, ldc; +}; +struct PlainPush { + std::uint64_t a, b, out, gate; + std::uint32_t m, n, k, lda, ldb, ldo; +}; + +void require(bool condition, const std::string& message) { + if (!condition) throw std::invalid_argument("vulkan gemm: " + message); +} + +std::uint32_t tiles(std::uint32_t size, std::uint32_t tile) { return (size + tile - 1) / tile; } + +} // namespace + +const char* gemm_path_name(GemmPath path) { return path == GemmPath::kCoop ? "coop" : "plain"; } + +Gemm::Gemm(Device& device) : Gemm(device, device.info().caps.coop ? GemmPath::kCoop : GemmPath::kPlain) {} + +Gemm::Gemm(Device& device, GemmPath path) : path_(path) { + if (path == GemmPath::kCoop) { + if (!device.info().caps.coop) { + throw std::runtime_error(device.info().name + " cannot run the cooperative-matrix GEMM: " + + device.info().caps.coop_reason); + } + coop_ = std::make_unique(device, kCoopSpirv, sizeof(kCoopSpirv), sizeof(Push), kCoopSubgroup); + } else { + plain_f16_ = std::make_unique(device, kPlainSpirv, sizeof(kPlainSpirv), sizeof(PlainPush), 0, + std::vector{0, 0}); + plain_bf16_ = std::make_unique(device, kPlainSpirv, sizeof(kPlainSpirv), sizeof(PlainPush), 0, + std::vector{1, 0}); + } +} + +Gemm::~Gemm() = default; + +void Gemm::record(VkCommandBuffer cb, const GemmArgs& args, InputType input) const { + require(args.a != nullptr && args.b != nullptr && args.c != nullptr, "null buffer"); + require(args.m > 0 && args.n > 0 && args.k > 0, "M, N and K must be positive"); + require(args.lda >= args.k && args.ldb >= args.k && args.ldc >= args.n, "a stride is smaller than its row"); + require(args.lda % 4 == 0 && args.ldb % 4 == 0, "lda and ldb must be multiples of 4"); + require(args.k % (path_ == GemmPath::kCoop ? 16 : 4) == 0, + std::string("K must be a multiple of ") + (path_ == GemmPath::kCoop ? "16" : "4")); + require((args.m - 1ull) * args.lda + args.k <= args.a->bytes() / 2, "A is too small"); + require((args.n - 1ull) * args.ldb + args.k <= args.b->bytes() / 2, "B is too small"); + require((args.m - 1ull) * args.ldc + args.n <= args.c->bytes() / 4, "C is too small"); + require(path_ == GemmPath::kPlain || input == InputType::kF16, "the cooperative-matrix path takes f16 inputs"); + require(path_ == GemmPath::kCoop || args.n % 4 == 0, "the plain path needs N to be a multiple of 4"); + + if (path_ == GemmPath::kCoop) { + const Push push{args.a->address(), args.b->address(), args.c->address(), args.m, args.n, args.k, args.lda, args.ldb, + args.ldc}; + coop_->dispatch(cb, &push, tiles(args.n, kCoopTile), tiles(args.m, kCoopTile)); + } else { + const PlainPush push{args.a->address(), args.b->address(), args.c->address(), 0, args.m, args.n, args.k, args.lda, + args.ldb, args.ldc}; + const Pipeline& pipeline = input == InputType::kBF16 ? *plain_bf16_ : *plain_f16_; + pipeline.dispatch(cb, &push, tiles(args.n, kPlainTileCols), tiles(args.m, kPlainTileRows)); + } + compute_barrier(cb); +} + +} // namespace mi::vk diff --git a/vulkan/gemm.h b/vulkan/gemm.h new file mode 100644 index 0000000..8e10b10 --- /dev/null +++ b/vulkan/gemm.h @@ -0,0 +1,45 @@ +// The GEMM behind every linear: C[M,N] = A[M,K] * B[N,K]^T, with B in the engine's row-major weight layout, f32 +// accumulation and f32 output. The fragments cannot run everywhere, so a plain compute path takes every other device. +#pragma once + +#include +#include + +#include "runtime.h" + +namespace mi::vk { + +enum class GemmPath { kCoop, kPlain }; +enum class InputType { kF16, kBF16 }; + +const char* gemm_path_name(GemmPath path); + +struct GemmArgs { + const Buffer* a = nullptr; + const Buffer* b = nullptr; + Buffer* c = nullptr; + std::uint32_t m = 0, n = 0, k = 0; + std::uint32_t lda = 0, ldb = 0, ldc = 0; // in elements; lda and ldb are multiples of 4 +}; + +class Gemm { + public: + // The cooperative-matrix path where the device has it, unless `path` names one. Throws if the device cannot run it. + explicit Gemm(Device& device); + Gemm(Device& device, GemmPath path); + ~Gemm(); + + GemmPath path() const { return path_; } + + // Records the GEMM and a barrier. Checks the arguments. kBF16 needs the plain path; the coop path takes kF16 only, + // because bf16 fragments are not widely available and the fallback must not depend on them. + void record(VkCommandBuffer cb, const GemmArgs& args, InputType input = InputType::kF16) const; + + private: + GemmPath path_; + std::unique_ptr coop_; + std::unique_ptr plain_f16_; + std::unique_ptr plain_bf16_; +}; + +} // namespace mi::vk diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp new file mode 100644 index 0000000..6ef986c --- /dev/null +++ b/vulkan/generate.cpp @@ -0,0 +1,135 @@ +#include "generate.h" + +#include +#include +#include +#include +#include + +#include "text_encoder.h" +#include "vae.h" + +namespace mi::vk { +namespace { + +double now_s() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +[[noreturn]] void fail(const std::string& what) { throw std::invalid_argument("vulkan generate: " + what); } + +} // namespace + +GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r) { + if (r.width == 0 || r.height == 0 || r.width % 16 != 0 || r.height % 16 != 0) fail("width and height must be multiples of 16"); + const std::uint32_t lh = r.height / 16, lw = r.width / 16; + if (r.noise.size() != std::size_t{kDiTLatentChannels} * lh * lw) fail("noise must be [128, height / 16, width / 16]"); + if (r.timesteps.empty() || r.sigmas.size() != r.timesteps.size() + 1) fail("sigmas must have one more entry than timesteps"); + if (r.cond_ids.empty()) fail("the prompt has no tokens"); + const bool cfg = r.cfg != 1.0f; + if (cfg && r.uncond_ids.empty()) fail("the negative prompt has no tokens"); + const auto report = [&](const char* stage, int done, int total) { + if (r.progress) r.progress(stage, done, total); + }; + + GenerateResult out; + out.width = r.width; + out.height = r.height; + + // The VAE first (about 0.1 GiB of weights, kept through the run), so a decode the device cannot hold fails before the + // text encoder and the steps run rather than after them. + std::optional vae; + if (r.decode) { + vae.emplace(device, r.model_dir / "vae" / "diffusion_pytorch_model.safetensors", + path == BlockPath::kCoop ? VaePath::kCoop : VaePath::kPlain); + } + if (r.decode) { + const Capabilities& caps = device.info().caps; + const std::uint64_t need = vae->workspace_bytes(lh, lw), largest = vae->largest_buffer_bytes(lh, lw); + const std::uint64_t limit = std::min(caps.max_buffer_bytes, caps.max_allocation_bytes); + if (limit != 0 && largest > limit) { + throw std::runtime_error("the VAE decode at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs a " + + std::to_string(largest >> 20) + " MiB buffer; " + device.info().name + " takes at most " + + std::to_string(limit >> 20) + " MiB"); + } + if (need > device.free_device_memory()) { + throw std::runtime_error("the VAE decode at " + std::to_string(r.width) + "x" + std::to_string(r.height) + " needs " + + std::to_string(need >> 20) + " MiB of device memory; " + device.info().name + " has " + + std::to_string(device.free_device_memory() >> 20) + " MiB free"); + } + } + + // Text encoder: streamed layer by layer, freed before the DiT loads. + std::vector> embeds; + { + double t0 = now_s(); + report("text encoder", 0, 1); + TextEncoder encoder(device, r.model_dir / "text_encoder", path); + std::vector> sequences{r.cond_ids}; + if (cfg) sequences.push_back(r.uncond_ids); + embeds = encoder.encode(sequences); + out.text_encoder_s = now_s() - t0; + report("text encoder", 1, 1); + } + + // DiT: resident, the text refiner once per prompt, then the steps. + { + double t0 = now_s(); + const BlockKernels kernels(device, path); + // Blocks past what the device can hold with this run's activations stream from host memory. + std::uint64_t tokens = 0; + for (const auto& e : embeds) tokens += std::uint64_t{lh} * lw + e.size() / kDiTTextDim; + constexpr std::uint64_t kReserve = std::uint64_t{512} << 20; // pipelines, command buffers, the driver + const std::uint64_t free = device.free_device_memory(), need = Dit::activation_bytes(tokens, std::uint64_t{lh} * lw) + kReserve; + DitOptions options; + options.device_budget = free > need ? free - need : 1; + Dit dit(device, kernels, r.model_dir / "transformer", options); + out.dit_resident_blocks = dit.resident_blocks(); + std::vector> texts; + std::vector rows; + for (std::size_t i = 0; i < embeds.size(); ++i) { + const auto length = static_cast(embeds[i].size() / kDiTTextDim); + texts.push_back(dit.encode_text(embeds[i], length)); + rows.push_back({texts.back().get(), length, i == 0 ? SkipSet{} : r.uncond_skip}); + } + out.dit_load_s = now_s() - t0; + + t0 = now_s(); + std::vector latents = r.noise; + const int steps = static_cast(r.timesteps.size()); + report("denoise", 0, steps); + for (int s = 0; s < steps; ++s) { + const auto si = static_cast(s); + const std::vector> pred = dit.forward(latents, lh, lw, r.timesteps[si], rows); + // CFG then the official Euler step in the engine's op order (sampler_kernels.cu): u + g (c - u); x + h v, with + // h = sigma_next - sigma rounded once from double. The fp32 expressions below compile without contraction. + const float h = static_cast(static_cast(r.sigmas[si + 1]) - static_cast(r.sigmas[si])); + for (std::size_t i = 0; i < latents.size(); ++i) { + float v = pred[0][i]; + if (cfg) { + const float diff = pred[0][i] - pred[1][i]; + const float scaled = r.cfg * diff; + v = pred[1][i] + scaled; + } + const float step = h * v; + latents[i] = latents[i] + step; + } + report("denoise", s + 1, steps); + } + out.denoise_s = now_s() - t0; + out.latents = std::move(latents); + } + + // VAE decode, after handing the final latents to the caller (a failed decode keeps them). + if (r.on_latents) r.on_latents(out.latents); + if (r.decode) { + const double t0 = now_s(); + report("vae", 0, 1); + vae->decode(out.latents.data(), lh, lw, out.rgb); + out.vae_s = now_s() - t0; + report("vae", 1, 1); + } + return out; +} + +} // namespace mi::vk diff --git a/vulkan/generate.h b/vulkan/generate.h new file mode 100644 index 0000000..d770166 --- /dev/null +++ b/vulkan/generate.h @@ -0,0 +1,49 @@ +// Text to image on the Vulkan backend: the text encoder, the DiT's sampling loop with CFG and the official Euler step, and +// the VAE decode, one model at a time on the device (the text encoder streams its layers and is freed before the DiT +// loads; the DiT is freed before the VAE loads). The caller supplies what the engine computes on the host for any +// backend: token ids (tokenizer and chat template), the initial noise (the seed's torch-compatible draw) and the +// schedule, so a seed gives the same inputs as on CUDA. +#pragma once + +#include +#include +#include +#include +#include + +#include "block.h" +#include "dit.h" +#include "runtime.h" + +namespace mi::vk { + +struct GenerateRequest { + std::filesystem::path model_dir; // MOSS-Image directory: text_encoder/, transformer/, vae/ + std::vector cond_ids; + std::vector uncond_ids; // unused when cfg == 1 + std::uint32_t width = 0, height = 0; // pixels, multiples of 16 + std::vector noise; // f32 [128, height / 16, width / 16] + std::vector sigmas; // steps + 1 + std::vector timesteps; // steps + float cfg = 3.0f; + SkipSet uncond_skip; + bool decode = true; // false: stop at the final latents + // Called with the final latents before the decode. + std::function& latents)> on_latents; + // stage: "text encoder", "denoise", "vae"; done of total within it. + std::function progress; +}; + +struct GenerateResult { + std::uint32_t width = 0, height = 0; + std::vector rgb; // [height, width, 3]; empty without the decode + std::vector latents; // the final latents, f32 [128, height / 16, width / 16] + double text_encoder_s = 0, dit_load_s = 0, denoise_s = 0, vae_s = 0; + int dit_resident_blocks = 0; // of 42; the rest streamed from host memory every step +}; + +// path: the DiT's and text encoder's kernels (kCoop needs Capabilities::coop_bf16, kCoopF16 coop_f16); the VAE runs +// kCoop on kCoop and kPlain otherwise (its cooperative-matrix kernels need bf16 fragments). +GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& request); + +} // namespace mi::vk diff --git a/vulkan/half.h b/vulkan/half.h new file mode 100644 index 0000000..a068d7d --- /dev/null +++ b/vulkan/half.h @@ -0,0 +1,59 @@ +// Host conversions between fp32 and the engine's 16-bit types. +#pragma once + +#include +#include +#include + +namespace mi::vk { + +inline float half_to_float(std::uint16_t h) { + const std::uint32_t exp = (h >> 10) & 0x1f; + const std::uint32_t mant = h & 0x3ff; + float f = 0.0f; + if (exp == 0) { + f = std::ldexp(static_cast(mant), -24); // zero or subnormal: mant * 2^-24, exact in fp32 + } else { + const std::uint32_t bits = exp == 31 ? 0x7f800000u | (mant << 13) : ((exp + 127 - 15) << 23) | (mant << 13); + std::memcpy(&f, &bits, 4); + } + return (h & 0x8000u) != 0 ? -f : f; +} + +// Rounds to nearest, ties to even, like the hardware conversion. +inline std::uint16_t float_to_half(float f) { + std::uint32_t x = 0; + std::memcpy(&x, &f, 4); + const auto sign = static_cast((x >> 16) & 0x8000u); + const std::uint32_t magnitude = x & 0x7fffffffu; + if (magnitude > 0x7f800000u) return static_cast(sign | 0x7e00u); // NaN + const std::int32_t exp = static_cast(magnitude >> 23) - 127 + 15; + if (exp >= 31) return static_cast(sign | 0x7c00u); // overflow and infinity + if (exp < -10) return sign; // below half of the smallest subnormal + const std::uint32_t mant = magnitude & 0x7fffffu; + const std::uint32_t significand = exp > 0 ? mant : mant | 0x800000u; + const int shift = exp > 0 ? 13 : 14 - exp; + std::uint32_t rounded = significand >> shift; + const std::uint32_t rest = significand & ((1u << shift) - 1u); + const std::uint32_t half = 1u << (shift - 1); + if (rest > half || (rest == half && (rounded & 1u) != 0)) ++rounded; + // An addition, not an OR: rounding can carry out of the mantissa into the exponent. + return static_cast(sign | ((exp > 0 ? static_cast(exp) << 10 : 0u) + rounded)); +} + +inline float bf16_to_float(std::uint16_t b) { + const std::uint32_t bits = static_cast(b) << 16; + float f = 0.0f; + std::memcpy(&f, &bits, 4); + return f; +} + +// Rounds to nearest, ties to even. +inline std::uint16_t float_to_bf16(float f) { + std::uint32_t x = 0; + std::memcpy(&x, &f, 4); + if ((x & 0x7fffffffu) > 0x7f800000u) return static_cast(x >> 16 | 0x7fc0u); // NaN + return static_cast((x + 0x7fffu + ((x >> 16) & 1u)) >> 16); +} + +} // namespace mi::vk diff --git a/vulkan/npy.h b/vulkan/npy.h new file mode 100644 index 0000000..02a81d9 --- /dev/null +++ b/vulkan/npy.h @@ -0,0 +1,180 @@ +// Reader for NumPy .npy files (format versions 1-3): little-endian, C-order arrays of fixed-size numeric dtypes. A copy +// of the engine's util/npy, header-only and without its helpers, for the checks that compare against golden data. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mi::vk::npy { + +struct Array { + std::string dtype; // numpy descr without byte order, e.g. "f4", "i8", "u2" + std::vector shape; + std::vector data; // raw element bytes, C order + + std::int64_t numel() const { + std::int64_t n = 1; + for (const std::int64_t d : shape) n *= d; + return n; + } + + // Copies the elements out as T. Throws unless dtype is exactly T's kind and size, e.g. "f4" for float. + template + std::vector to_vector() const { + static_assert(std::is_arithmetic_v); + constexpr char kind = std::is_same_v ? 'b' + : std::is_floating_point_v ? 'f' + : std::is_signed_v ? 'i' + : 'u'; + if (dtype != kind + std::to_string(sizeof(T))) { + throw std::runtime_error("npy: array of dtype " + dtype + " read as " + kind + std::to_string(sizeof(T))); + } + std::vector out(data.size() / sizeof(T)); + if (!out.empty()) std::memcpy(out.data(), data.data(), out.size() * sizeof(T)); + return out; + } +}; + +namespace detail { + +[[noreturn]] inline void fail(const std::filesystem::path& path, const std::string& what) { + throw std::runtime_error("npy: " + path.string() + ": " + what); +} + +// The value text following 'key': in the header dict literal. +inline std::string_view dict_value(std::string_view header, std::string_view key) { + const std::string quoted = "'" + std::string(key) + "'"; + const std::size_t at = header.find(quoted); + if (at == std::string_view::npos) return {}; + const std::size_t colon = header.find(':', at + quoted.size()); + if (colon == std::string_view::npos) return {}; + const std::size_t start = header.find_first_not_of(' ', colon + 1); + if (start == std::string_view::npos) return {}; + const char open = header[start]; + if (open == '\'') { + const std::size_t end = header.find('\'', start + 1); + return end == std::string_view::npos ? std::string_view{} : header.substr(start + 1, end - start - 1); + } + if (open == '(') { + const std::size_t end = header.find(')', start); + return end == std::string_view::npos ? std::string_view{} : header.substr(start + 1, end - start - 1); + } + const std::size_t end = header.find_first_of(",}", start); + return header.substr(start, end - start); +} + +template +std::optional parse_uint(std::string_view text) { + T v{}; + const auto [ptr, ec] = std::from_chars(text.data(), text.data() + text.size(), v); + if (ec != std::errc() || ptr != text.data() + text.size()) return std::nullopt; + if constexpr (std::is_signed_v) { + if (v < 0) return std::nullopt; + } + return v; +} + +} // namespace detail + +// Throws std::runtime_error on unsupported or malformed files. +inline Array load(const std::filesystem::path& path) { + using detail::fail; + std::ifstream in(path, std::ios::binary); + if (!in) fail(path, "cannot open"); + in.seekg(0, std::ios::end); + const auto file_size = static_cast(in.tellg()); + in.seekg(0); + char magic[8]; + if (!in.read(magic, 8) || std::string_view(magic, 6) != "\x93NUMPY") fail(path, "bad magic"); + const int major = static_cast(magic[6]); + std::uint32_t header_len = 0; + if (major == 1) { + unsigned char b[2]; + if (!in.read(reinterpret_cast(b), 2)) fail(path, "truncated header length"); + header_len = static_cast(b[0] | (b[1] << 8)); + } else if (major == 2 || major == 3) { + unsigned char b[4]; + if (!in.read(reinterpret_cast(b), 4)) fail(path, "truncated header length"); + header_len = static_cast(b[0]) | (static_cast(b[1]) << 8) | + (static_cast(b[2]) << 16) | (static_cast(b[3]) << 24); + } else { + fail(path, "unsupported format version " + std::to_string(major)); + } + if (header_len > file_size - static_cast(in.tellg())) fail(path, "truncated header"); + std::string header(header_len, '\0'); + if (!in.read(header.data(), header_len)) fail(path, "truncated header"); + + const std::string_view descr = detail::dict_value(header, "descr"); + if (descr.size() < 3) fail(path, "missing descr"); + if (descr[0] != '<' && descr[0] != '|') fail(path, "only little-endian arrays are supported"); + if (detail::dict_value(header, "fortran_order").starts_with("True")) fail(path, "fortran order is not supported"); + + Array arr; + arr.dtype = std::string(descr.substr(1)); + const char kind = arr.dtype[0]; + if (kind != 'f' && kind != 'i' && kind != 'u' && kind != 'b') fail(path, "unsupported dtype " + arr.dtype); + const std::optional itemsize = detail::parse_uint(std::string_view(arr.dtype).substr(1)); + if (!itemsize || (*itemsize != 1 && *itemsize != 2 && *itemsize != 4 && *itemsize != 8)) { + fail(path, "unsupported dtype " + arr.dtype); + } + + const std::string_view shape = detail::dict_value(header, "shape"); + std::size_t pos = 0; + while (pos < shape.size()) { + const std::size_t next = shape.find(',', pos); + const std::string_view item = + shape.substr(pos, next == std::string_view::npos ? std::string_view::npos : next - pos); + const std::size_t first = item.find_first_not_of(' '); + if (first != std::string_view::npos) { + const std::size_t last = item.find_last_not_of(' '); + const std::optional dim = detail::parse_uint(item.substr(first, last - first + 1)); + if (!dim) fail(path, "malformed shape '" + std::string(shape) + "'"); + arr.shape.push_back(*dim); + } + if (next == std::string_view::npos) break; + pos = next + 1; + } + + std::uint64_t bytes = *itemsize; + for (const std::int64_t d : arr.shape) { + if (d != 0 && bytes > std::numeric_limits::max() / static_cast(d)) { + fail(path, "element count overflows"); + } + bytes *= static_cast(d); + } + if (bytes > file_size - static_cast(in.tellg())) fail(path, "truncated data"); + arr.data.resize(static_cast(bytes)); + if (!in.read(reinterpret_cast(arr.data.data()), static_cast(bytes))) + fail(path, "truncated data"); + return arr; +} + +// Writes a little-endian C-order array (format 1.0); descr e.g. "& shape, + const void* data, std::size_t bytes) { + std::string header = std::string("{'descr': '") + descr + "', 'fortran_order': False, 'shape': ("; + for (const std::int64_t d : shape) header += std::to_string(d) + ", "; + header += "), }"; + while ((10 + header.size() + 1) % 64 != 0) header += ' '; + header += '\n'; + std::ofstream out(path, std::ios::binary); + out.write("\x93NUMPY\x01\x00", 8); + const unsigned char len[2] = {static_cast(header.size() & 0xff), static_cast(header.size() >> 8)}; + out.write(reinterpret_cast(len), 2); + out.write(header.data(), static_cast(header.size())); + out.write(static_cast(data), static_cast(bytes)); + if (!out) throw std::runtime_error("npy: cannot write " + path.string()); +} + +} // namespace mi::vk::npy diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp new file mode 100644 index 0000000..4c7007c --- /dev/null +++ b/vulkan/runtime.cpp @@ -0,0 +1,668 @@ +#include "runtime.h" + +#include +#include +#include +#include +#include + +namespace mi::vk { +namespace { + +// MI_VK_PIPELINE_STATS=1 prints the driver's statistics of every pipeline (registers, spills, ...), where the device has +// VK_KHR_pipeline_executable_properties. A debugging aid; off by default. +bool pipeline_stats_requested() { + const char* v = std::getenv("MI_VK_PIPELINE_STATS"); + return v != nullptr && v[0] == '1'; +} + +// gemm_coop.comp's shared memory: an f16 A and B tile of 64x32, and an f32 output tile of 64x64. +constexpr std::uint32_t kCoopSharedBytes = 2 * 64 * 32 * 2 + 64 * 64 * 4; +constexpr std::uint32_t kCoopSubgroup = 32; +constexpr std::uint32_t kCoopSubgroupsPerWorkgroup = 8; +// The DiT kernels' largest shared memory (bf16 and f16 fragments alike): gemm_bf16_coop.comp's two stages of 128x32 A and +// B tiles, rows padded to 80 B. +constexpr std::uint32_t kCoopBf16SharedBytes = 4 * 128 * 80; + +void check(VkResult result, const char* what) { + if (result != VK_SUCCESS) throw std::runtime_error(std::string(what) + " failed: VkResult " + std::to_string(result)); +} + +VkInstance create_instance() { + std::string error; + if (!load_vulkan(&error)) throw std::runtime_error(error); + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; + app.pApplicationName = "mossinfer"; + app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo info{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; + info.pApplicationInfo = &app; + VkInstance instance = VK_NULL_HANDLE; + check(vkCreateInstance(&info, nullptr, &instance), "vkCreateInstance"); + return instance; +} + +std::vector physical_devices(VkInstance instance) { + std::uint32_t count = 0; + check(vkEnumeratePhysicalDevices(instance, &count, nullptr), "vkEnumeratePhysicalDevices"); + std::vector devices(count); + check(vkEnumeratePhysicalDevices(instance, &count, devices.data()), "vkEnumeratePhysicalDevices"); + return devices; +} + +bool has_extension(VkPhysicalDevice physical, const char* name) { + std::uint32_t count = 0; + vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, nullptr); + std::vector extensions(count); + vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, extensions.data()); + return std::any_of(extensions.begin(), extensions.end(), + [&](const VkExtensionProperties& e) { return std::strcmp(e.extensionName, name) == 0; }); +} + +// What both cooperative-matrix paths need: Vulkan 1.3, the extension and feature, 16-bit storage, the memory model, and a +// compute subgroup size fixable at 32 with 8 per workgroup. Empty when the device has all of it. +std::string why_no_coop_base(VkPhysicalDevice physical, const VkPhysicalDeviceProperties& properties, Capabilities& caps) { + if (properties.apiVersion < VK_API_VERSION_1_3) return "needs Vulkan 1.3"; + VkPhysicalDeviceVulkan13Properties p13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_PROPERTIES}; + VkPhysicalDeviceProperties2 p2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, &p13}; + vkGetPhysicalDeviceProperties2(physical, &p2); + caps.min_subgroup = p13.minSubgroupSize; + caps.max_subgroup = p13.maxSubgroupSize; + if (!has_extension(physical, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME)) return "no VK_KHR_cooperative_matrix"; + + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coop{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &coop}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11}; + vkGetPhysicalDeviceFeatures2(physical, &features); + if (coop.cooperativeMatrix == 0) return "the cooperativeMatrix feature is off"; + if (v11.storageBuffer16BitAccess == 0) return "no storageBuffer16BitAccess"; + if (v12.shaderFloat16 == 0) return "no shaderFloat16"; + if (v12.vulkanMemoryModel == 0) return "no vulkanMemoryModel"; + if (v13.subgroupSizeControl == 0 || v13.computeFullSubgroups == 0) return "no subgroup size control"; + if (p13.minSubgroupSize > kCoopSubgroup || p13.maxSubgroupSize < kCoopSubgroup || + (p13.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT) == 0) { + return "cannot fix the compute subgroup size at 32 (device range " + std::to_string(p13.minSubgroupSize) + "-" + + std::to_string(p13.maxSubgroupSize) + ")"; + } + if (p13.maxComputeWorkgroupSubgroups < kCoopSubgroupsPerWorkgroup) return "needs 8 subgroups per workgroup"; + VkPhysicalDeviceSubgroupProperties subgroup{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_PROPERTIES}; + VkPhysicalDeviceProperties2 sp{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, &subgroup}; + vkGetPhysicalDeviceProperties2(physical, &sp); + constexpr VkSubgroupFeatureFlags kNeeded = + VK_SUBGROUP_FEATURE_BASIC_BIT | VK_SUBGROUP_FEATURE_VOTE_BIT | VK_SUBGROUP_FEATURE_SHUFFLE_BIT; + if ((subgroup.supportedStages & VK_SHADER_STAGE_COMPUTE_BIT) == 0 || (subgroup.supportedOperations & kNeeded) != kNeeded) { + return "no compute subgroup vote and shuffle operations"; + } + return {}; +} + +std::vector coop_shapes(VkInstance instance, VkPhysicalDevice physical) { + auto query = reinterpret_cast( + vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); + if (query == nullptr) return {}; + std::uint32_t count = 0; + if (query(physical, &count, nullptr) != VK_SUCCESS) return {}; + std::vector shapes(count); + for (auto& s : shapes) s.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; + if (query(physical, &count, shapes.data()) != VK_SUCCESS) return {}; + return shapes; +} + +bool has_shape(const std::vector& shapes, VkComponentTypeKHR input) { + return std::any_of(shapes.begin(), shapes.end(), [&](const VkCooperativeMatrixPropertiesKHR& s) { + return s.MSize == 16 && s.NSize == 16 && s.KSize == 16 && s.AType == input && s.BType == input && + s.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && s.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR && + s.saturatingAccumulation == VK_FALSE && s.scope == VK_SCOPE_SUBGROUP_KHR; + }); +} + +void describe_coop(VkInstance instance, VkPhysicalDevice physical, const VkPhysicalDeviceProperties& properties, + Capabilities& caps) { + const std::string base = why_no_coop_base(physical, properties, caps); + const std::vector shapes = + base.empty() ? coop_shapes(instance, physical) : std::vector{}; + const std::uint32_t shared = properties.limits.maxComputeSharedMemorySize; + + caps.coop_reason = base; + if (caps.coop_reason.empty() && shared < kCoopSharedBytes) { + caps.coop_reason = "needs " + std::to_string(kCoopSharedBytes) + " bytes of shared memory"; + } + if (caps.coop_reason.empty() && !has_shape(shapes, VK_COMPONENT_TYPE_FLOAT16_KHR)) { + caps.coop_reason = "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; + } + caps.coop = caps.coop_reason.empty(); + + caps.coop_bf16_reason = base; + if (caps.coop_bf16_reason.empty() && !has_extension(physical, VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME)) { + caps.coop_bf16_reason = "no VK_KHR_shader_bfloat16"; + } + if (caps.coop_bf16_reason.empty()) { + VkPhysicalDeviceShaderBfloat16FeaturesKHR bf16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &bf16}; + vkGetPhysicalDeviceFeatures2(physical, &features); + if (bf16.shaderBFloat16Type == 0 || bf16.shaderBFloat16CooperativeMatrix == 0) { + caps.coop_bf16_reason = "no bf16 cooperative matrices"; + } + } + if (caps.coop_bf16_reason.empty() && shared < kCoopBf16SharedBytes) { + caps.coop_bf16_reason = "needs " + std::to_string(kCoopBf16SharedBytes) + " bytes of shared memory"; + } + if (caps.coop_bf16_reason.empty() && !has_shape(shapes, VK_COMPONENT_TYPE_BFLOAT16_KHR)) { + caps.coop_bf16_reason = "no 16x16x16 bf16 x bf16 -> f32 cooperative matrix in subgroup scope"; + } + caps.coop_bf16 = caps.coop_bf16_reason.empty(); + + caps.coop_f16_reason = base; + if (caps.coop_f16_reason.empty() && shared < kCoopBf16SharedBytes) { + caps.coop_f16_reason = "needs " + std::to_string(kCoopBf16SharedBytes) + " bytes of shared memory"; + } + if (caps.coop_f16_reason.empty() && !has_shape(shapes, VK_COMPONENT_TYPE_FLOAT16_KHR)) { + caps.coop_f16_reason = "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; + } + caps.coop_f16 = caps.coop_f16_reason.empty(); +} + +DeviceInfo describe(VkInstance instance, VkPhysicalDevice physical, std::uint32_t index) { + VkPhysicalDeviceProperties properties{}; + vkGetPhysicalDeviceProperties(physical, &properties); + DeviceInfo info; + info.index = index; + info.name = properties.deviceName; + info.vendor_id = properties.vendorID; + info.device_id = properties.deviceID; + info.discrete = properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU; + info.caps.shared_memory_bytes = properties.limits.maxComputeSharedMemorySize; + { + VkPhysicalDeviceVulkan13Properties p13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_PROPERTIES}; + VkPhysicalDeviceVulkan11Properties p11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_PROPERTIES}; + if (properties.apiVersion >= VK_API_VERSION_1_3) p11.pNext = &p13; + VkPhysicalDeviceProperties2 p2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2, &p11}; + if (properties.apiVersion >= VK_API_VERSION_1_2) vkGetPhysicalDeviceProperties2(physical, &p2); + info.caps.max_allocation_bytes = p11.maxMemoryAllocationSize; + // Before Vulkan 1.3 the limit is not reported; the allocation limit bounds a buffer here anyway. + info.caps.max_buffer_bytes = properties.apiVersion >= VK_API_VERSION_1_3 ? p13.maxBufferSize : p11.maxMemoryAllocationSize; + } + if (properties.apiVersion < VK_API_VERSION_1_2) { + info.unusable_reason = "needs Vulkan 1.2"; + return info; + } + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v12}; + vkGetPhysicalDeviceFeatures2(physical, &features); + if (v12.bufferDeviceAddress == 0) info.unusable_reason = "no bufferDeviceAddress"; + if (features.features.shaderInt64 == 0) info.unusable_reason = "no shaderInt64"; + describe_coop(instance, physical, properties, info.caps); + return info; +} + +} // namespace + +std::vector list_devices() { + VkInstance instance = create_instance(); + std::vector infos; + std::uint32_t index = 0; + for (VkPhysicalDevice physical : physical_devices(instance)) infos.push_back(describe(instance, physical, index++)); + vkDestroyInstance(instance, nullptr); + return infos; +} + +Device::Device(int index) { + if (index < -1) throw std::runtime_error("Vulkan device index " + std::to_string(index) + " is negative"); + instance_ = create_instance(); + const std::vector devices = physical_devices(instance_); + if (index >= 0 && static_cast(index) >= devices.size()) { + const std::size_t count = devices.size(); + vkDestroyInstance(instance_, nullptr); + throw std::runtime_error("no Vulkan device " + std::to_string(index) + ": there are " + std::to_string(count)); + } + bool chosen = false; + for (std::uint32_t i = 0; i < devices.size(); ++i) { + DeviceInfo candidate = describe(instance_, devices[i], i); + if (index >= 0) { + if (static_cast(index) != i) continue; + if (!candidate.unusable_reason.empty()) { + throw std::runtime_error("Vulkan device " + std::to_string(i) + " (" + candidate.name + ") cannot be used: " + + candidate.unusable_reason); + } + } else if (!candidate.unusable_reason.empty() || (chosen && (info_.discrete || !candidate.discrete))) { + continue; + } + physical_ = devices[i]; + info_ = std::move(candidate); + chosen = true; + } + if (!chosen) throw std::runtime_error("no usable Vulkan device"); + vkGetPhysicalDeviceMemoryProperties(physical_, &memory_); + + std::uint32_t families = 0; + vkGetPhysicalDeviceQueueFamilyProperties(physical_, &families, nullptr); + std::vector queue_families(families); + vkGetPhysicalDeviceQueueFamilyProperties(physical_, &families, queue_families.data()); + std::uint32_t family = families; + for (std::uint32_t i = 0; i < families && family == families; ++i) { + if ((queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) != 0) family = i; + } + if (family == families) throw std::runtime_error("no compute queue on " + info_.name); + // A family with transfer and neither graphics nor compute is a copy engine that runs beside compute work. + std::uint32_t transfer = families; + for (std::uint32_t i = 0; i < families && transfer == families; ++i) { + const VkQueueFlags flags = queue_families[i].queueFlags; + if ((flags & VK_QUEUE_TRANSFER_BIT) != 0 && (flags & (VK_QUEUE_GRAPHICS_BIT | VK_QUEUE_COMPUTE_BIT)) == 0) transfer = i; + } + family_ = family; + transfer_family_ = transfer == families ? family : transfer; + + // Only what the shaders declare, so a shader that needs more fails at pipeline creation instead of working by luck. + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coop{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}; + coop.cooperativeMatrix = VK_TRUE; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &coop}; + v13.subgroupSizeControl = VK_TRUE; + v13.computeFullSubgroups = VK_TRUE; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; + v12.bufferDeviceAddress = VK_TRUE; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; + features.features.shaderInt64 = VK_TRUE; + VkPhysicalDeviceShaderBfloat16FeaturesKHR bf16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR}; + bf16.shaderBFloat16Type = VK_TRUE; + bf16.shaderBFloat16CooperativeMatrix = VK_TRUE; + std::vector extensions; + if (info_.caps.coop_bf16) { + coop.pNext = &bf16; + extensions.push_back(VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + } + if (info_.caps.coop || info_.caps.coop_bf16 || info_.caps.coop_f16) { + v12.pNext = &v13; + v12.shaderFloat16 = VK_TRUE; + v12.vulkanMemoryModel = VK_TRUE; + v11.storageBuffer16BitAccess = VK_TRUE; + features.pNext = &v11; + extensions.push_back(VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME); + } else { + features.pNext = &v12; // a Vulkan 1.2 device need not know the 1.3 structure + } + + VkPhysicalDevicePipelineExecutablePropertiesFeaturesKHR executable{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_EXECUTABLE_PROPERTIES_FEATURES_KHR}; + if (pipeline_stats_requested() && has_extension(physical_, VK_KHR_PIPELINE_EXECUTABLE_PROPERTIES_EXTENSION_NAME)) { + executable.pipelineExecutableInfo = VK_TRUE; + executable.pNext = features.pNext; + features.pNext = &executable; + extensions.push_back(VK_KHR_PIPELINE_EXECUTABLE_PROPERTIES_EXTENSION_NAME); + pipeline_stats_ = true; + } + + const float priority = 1.0f; + VkDeviceQueueCreateInfo queue_infos[2] = {{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}, + {VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}}; + for (VkDeviceQueueCreateInfo& q : queue_infos) { + q.queueCount = 1; + q.pQueuePriorities = &priority; + } + queue_infos[0].queueFamilyIndex = family; + queue_infos[1].queueFamilyIndex = transfer_family_; + VkDeviceCreateInfo device_info{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO}; + device_info.pNext = &features; + device_info.queueCreateInfoCount = transfer_family_ != family ? 2 : 1; + device_info.pQueueCreateInfos = queue_infos; + device_info.enabledExtensionCount = static_cast(extensions.size()); + device_info.ppEnabledExtensionNames = extensions.data(); + check(vkCreateDevice(physical_, &device_info, nullptr, &device_), "vkCreateDevice"); + vkGetDeviceQueue(device_, family, 0, &queue_); + + VkCommandPoolCreateInfo pool_info{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; + pool_info.queueFamilyIndex = family; + check(vkCreateCommandPool(device_, &pool_info, nullptr, &pool_), "vkCreateCommandPool"); + if (transfer_family_ != family) { + vkGetDeviceQueue(device_, transfer_family_, 0, &transfer_queue_); + pool_info.queueFamilyIndex = transfer_family_; + check(vkCreateCommandPool(device_, &pool_info, nullptr, &transfer_pool_), "vkCreateCommandPool"); + } +} + +Device::~Device() { + if (transfer_queue_ != VK_NULL_HANDLE) vkQueueWaitIdle(transfer_queue_); + if (queue_ != VK_NULL_HANDLE) vkQueueWaitIdle(queue_); + for (auto& [ticket, c] : copies_) { + if (c.fence != VK_NULL_HANDLE) vkDestroyFence(device_, c.fence, nullptr); + vkDestroySemaphore(device_, c.done, nullptr); + } + if (transfer_pool_ != VK_NULL_HANDLE) vkDestroyCommandPool(device_, transfer_pool_, nullptr); + if (pool_ != VK_NULL_HANDLE) vkDestroyCommandPool(device_, pool_, nullptr); + if (device_ != VK_NULL_HANDLE) vkDestroyDevice(device_, nullptr); + if (instance_ != VK_NULL_HANDLE) vkDestroyInstance(instance_, nullptr); +} + +std::uint64_t Device::free_device_memory() const { + VkPhysicalDeviceMemoryBudgetPropertiesEXT budget{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_BUDGET_PROPERTIES_EXT}; + VkPhysicalDeviceMemoryProperties2 properties{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MEMORY_PROPERTIES_2}; + const bool has_budget = has_extension(physical_, VK_EXT_MEMORY_BUDGET_EXTENSION_NAME); + if (has_budget) properties.pNext = &budget; + vkGetPhysicalDeviceMemoryProperties2(physical_, &properties); + std::uint64_t best = 0; + for (std::uint32_t i = 0; i < properties.memoryProperties.memoryHeapCount; ++i) { + const VkMemoryHeap& heap = properties.memoryProperties.memoryHeaps[i]; + if ((heap.flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) == 0) continue; + std::uint64_t free = heap.size; + if (has_budget) free = budget.heapBudget[i] > budget.heapUsage[i] ? budget.heapBudget[i] - budget.heapUsage[i] : 0; + best = std::max(best, free); + } + return best; +} + +std::uint32_t Device::find_memory_type(std::uint32_t type_bits, VkMemoryPropertyFlags flags) const { + for (std::uint32_t i = 0; i < memory_.memoryTypeCount; ++i) { + if ((type_bits & (1u << i)) != 0 && (memory_.memoryTypes[i].propertyFlags & flags) == flags) return i; + } + throw std::runtime_error("no Vulkan memory type with flags " + std::to_string(flags)); +} + +void Device::run(const std::function& record) { + VkCommandBufferAllocateInfo allocate{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; + allocate.commandPool = pool_; + allocate.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + allocate.commandBufferCount = 1; + VkCommandBuffer cb = VK_NULL_HANDLE; + check(vkAllocateCommandBuffers(device_, &allocate, &cb), "vkAllocateCommandBuffers"); + VkCommandBufferBeginInfo begin{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + begin.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + try { + check(vkBeginCommandBuffer(cb, &begin), "vkBeginCommandBuffer"); + record(cb); + check(vkEndCommandBuffer(cb), "vkEndCommandBuffer"); + } catch (...) { + vkFreeCommandBuffers(device_, pool_, 1, &cb); // an op's argument check threw while recording + throw; + } + VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO}; + submit.commandBufferCount = 1; + submit.pCommandBuffers = &cb; + std::vector waits; + for (const std::uint64_t ticket : run_waits_) waits.push_back(copies_.at(ticket).done); + const std::vector wait_stages(waits.size(), + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT); + submit.waitSemaphoreCount = static_cast(waits.size()); + submit.pWaitSemaphores = waits.data(); + submit.pWaitDstStageMask = wait_stages.data(); + const VkResult submitted = vkQueueSubmit(queue_, 1, &submit, VK_NULL_HANDLE); + const std::vector waited_for = submitted == VK_SUCCESS ? std::move(run_waits_) : std::vector{}; + if (submitted == VK_SUCCESS) run_waits_.clear(); + const VkResult waited = submitted == VK_SUCCESS ? vkQueueWaitIdle(queue_) : submitted; + vkFreeCommandBuffers(device_, pool_, 1, &cb); + check(submitted, "vkQueueSubmit"); + check(waited, "vkQueueWaitIdle"); + for (const std::uint64_t ticket : waited_for) { + copies_.at(ticket).run_done = true; + release_if_done(ticket); + } +} + +Buffer::Buffer(Device& device, VkDeviceSize bytes, Kind kind, bool transfer) : device_(device), bytes_(bytes) { + const Capabilities& caps = device.info().caps; + if ((caps.max_buffer_bytes != 0 && bytes > caps.max_buffer_bytes) || + (caps.max_allocation_bytes != 0 && bytes > caps.max_allocation_bytes)) { + throw std::runtime_error("a buffer of " + std::to_string(bytes >> 20) + " MiB exceeds " + device.info().name + + "'s largest buffer or allocation (" + + std::to_string(std::min(caps.max_buffer_bytes, caps.max_allocation_bytes) >> 20) + " MiB)"); + } + const bool on_device = kind == Kind::kDevice; + VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; + info.size = bytes; + info.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + if (on_device) info.usage |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; + const std::uint32_t families[2] = {device_.queue_family(), device_.transfer_family()}; + if (transfer && families[0] != families[1]) { + info.sharingMode = VK_SHARING_MODE_CONCURRENT; // no ownership transfers between the two queues + info.queueFamilyIndexCount = 2; + info.pQueueFamilyIndices = families; + } + check(vkCreateBuffer(device_.handle(), &info, nullptr, &buffer_), "vkCreateBuffer"); + + VkMemoryRequirements requirements{}; + vkGetBufferMemoryRequirements(device_.handle(), buffer_, &requirements); + VkMemoryAllocateFlagsInfo flags{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO}; + flags.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT; + VkMemoryAllocateInfo allocate{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + if (on_device) allocate.pNext = &flags; + allocate.allocationSize = requirements.size; + allocate.memoryTypeIndex = device_.find_memory_type( + requirements.memoryTypeBits, on_device ? VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT + : VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + const VkResult allocated = vkAllocateMemory(device_.handle(), &allocate, nullptr, &memory_); + if (allocated != VK_SUCCESS) { + vkDestroyBuffer(device_.handle(), buffer_, nullptr); + check(allocated, "vkAllocateMemory"); + } + check(vkBindBufferMemory(device_.handle(), buffer_, memory_, 0), "vkBindBufferMemory"); + if (on_device) { + VkBufferDeviceAddressInfo address{VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO}; + address.buffer = buffer_; + address_ = vkGetBufferDeviceAddress(device_.handle(), &address); + } else { + check(vkMapMemory(device_.handle(), memory_, 0, VK_WHOLE_SIZE, 0, &mapped_), "vkMapMemory"); + } +} + +Buffer::~Buffer() { + vkDestroyBuffer(device_.handle(), buffer_, nullptr); + vkFreeMemory(device_.handle(), memory_, nullptr); +} + +void Buffer::upload(const void* data, VkDeviceSize bytes, VkDeviceSize offset) { + if (offset > bytes_ || bytes > bytes_ - offset) throw std::runtime_error("Buffer::upload past the end of the buffer"); + Buffer staging(device_, bytes, Kind::kStaging); + std::memcpy(staging.mapped(), data, static_cast(bytes)); + device_.run([&](VkCommandBuffer cb) { + const VkBufferCopy copy{0, offset, bytes}; + vkCmdCopyBuffer(cb, staging.handle(), buffer_, 1, ©); + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, + nullptr, 0, nullptr); + }); +} + +void Buffer::download(void* data, VkDeviceSize bytes, VkDeviceSize offset) { + if (offset > bytes_ || bytes > bytes_ - offset) throw std::runtime_error("Buffer::download past the end of the buffer"); + Buffer staging(device_, bytes, Kind::kStaging); + device_.run([&](VkCommandBuffer cb) { + const VkBufferCopy copy{offset, 0, bytes}; + vkCmdCopyBuffer(cb, buffer_, staging.handle(), 1, ©); + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_HOST_BIT, 0, 1, &barrier, 0, nullptr, 0, + nullptr); + }); + std::memcpy(data, staging.mapped(), static_cast(bytes)); +} + +Pipeline::Pipeline(Device& device, const std::uint32_t* spirv, std::size_t spirv_bytes, std::uint32_t push_bytes, + std::uint32_t required_subgroup_size, const std::vector& specialization) + : device_(device), push_bytes_(push_bytes) { + VkShaderModuleCreateInfo module_info{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; + module_info.codeSize = spirv_bytes; + module_info.pCode = spirv; + check(vkCreateShaderModule(device_.handle(), &module_info, nullptr, &module_), "vkCreateShaderModule"); + + VkPushConstantRange range{}; + range.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT; + range.size = push_bytes; + VkPipelineLayoutCreateInfo layout_info{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; + layout_info.pushConstantRangeCount = 1; + layout_info.pPushConstantRanges = ⦥ + check(vkCreatePipelineLayout(device_.handle(), &layout_info, nullptr, &layout_), "vkCreatePipelineLayout"); + + std::vector entries; + for (std::uint32_t i = 0; i < specialization.size(); ++i) { + entries.push_back({i, i * static_cast(sizeof(std::uint32_t)), sizeof(std::uint32_t)}); + } + VkSpecializationInfo specialization_info{}; + specialization_info.mapEntryCount = static_cast(entries.size()); + specialization_info.pMapEntries = entries.data(); + specialization_info.dataSize = specialization.size() * sizeof(std::uint32_t); + specialization_info.pData = specialization.data(); + + VkPipelineShaderStageRequiredSubgroupSizeCreateInfo subgroup{ + VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; + subgroup.requiredSubgroupSize = required_subgroup_size; + VkComputePipelineCreateInfo info{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; + info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT; + info.stage.module = module_; + info.stage.pName = "main"; + if (!specialization.empty()) info.stage.pSpecializationInfo = &specialization_info; + if (required_subgroup_size != 0) { + info.stage.pNext = &subgroup; + info.stage.flags = VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT; + } + info.layout = layout_; + if (device_.pipeline_stats()) info.flags |= VK_PIPELINE_CREATE_CAPTURE_STATISTICS_BIT_KHR; + check(vkCreateComputePipelines(device_.handle(), VK_NULL_HANDLE, 1, &info, nullptr, &pipeline_), + "vkCreateComputePipelines"); + if (device_.pipeline_stats()) print_statistics(spirv_bytes); +} + +void Pipeline::print_statistics(std::size_t spirv_bytes) const { + auto properties = reinterpret_cast( + vkGetDeviceProcAddr(device_.handle(), "vkGetPipelineExecutablePropertiesKHR")); + auto statistics = reinterpret_cast( + vkGetDeviceProcAddr(device_.handle(), "vkGetPipelineExecutableStatisticsKHR")); + if (properties == nullptr || statistics == nullptr) return; + VkPipelineInfoKHR pipeline_info{VK_STRUCTURE_TYPE_PIPELINE_INFO_KHR}; + pipeline_info.pipeline = pipeline_; + std::uint32_t count = 0; + properties(device_.handle(), &pipeline_info, &count, nullptr); + std::vector executables(count, {VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_PROPERTIES_KHR}); + properties(device_.handle(), &pipeline_info, &count, executables.data()); + for (std::uint32_t e = 0; e < count; ++e) { + VkPipelineExecutableInfoKHR executable{VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_INFO_KHR}; + executable.pipeline = pipeline_; + executable.executableIndex = e; + std::uint32_t n = 0; + statistics(device_.handle(), &executable, &n, nullptr); + std::vector stats(n, {VK_STRUCTURE_TYPE_PIPELINE_EXECUTABLE_STATISTIC_KHR}); + statistics(device_.handle(), &executable, &n, stats.data()); + std::fprintf(stderr, "pipeline (SPIR-V %zu bytes) %s:", spirv_bytes, executables[e].name); + for (const auto& s : stats) { + if (s.format == VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_UINT64_KHR) { + std::fprintf(stderr, " %s=%llu", s.name, static_cast(s.value.u64)); + } else if (s.format == VK_PIPELINE_EXECUTABLE_STATISTIC_FORMAT_INT64_KHR) { + std::fprintf(stderr, " %s=%lld", s.name, static_cast(s.value.i64)); + } + } + std::fprintf(stderr, "\n"); + } +} + +Pipeline::~Pipeline() { + vkDestroyPipeline(device_.handle(), pipeline_, nullptr); + vkDestroyPipelineLayout(device_.handle(), layout_, nullptr); + vkDestroyShaderModule(device_.handle(), module_, nullptr); +} + +void Pipeline::dispatch(VkCommandBuffer cb, const void* push, std::uint32_t x, std::uint32_t y) const { + vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline_); + vkCmdPushConstants(cb, layout_, VK_SHADER_STAGE_COMPUTE_BIT, 0, push_bytes_, push); + vkCmdDispatch(cb, x, y, 1); +} + +std::uint64_t Device::copy_async(const Buffer& src, VkDeviceSize src_offset, const Buffer& dst, VkDeviceSize dst_offset, + VkDeviceSize bytes) { + if (src_offset > src.bytes() || bytes > src.bytes() - src_offset || dst_offset > dst.bytes() || + bytes > dst.bytes() - dst_offset) { + throw std::runtime_error("copy_async past the end of a buffer"); + } + const bool own_queue = transfer_queue_ != VK_NULL_HANDLE; + Copy c; + VkCommandBufferAllocateInfo allocate{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; + allocate.commandPool = own_queue ? transfer_pool_ : pool_; + allocate.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + allocate.commandBufferCount = 1; + check(vkAllocateCommandBuffers(device_, &allocate, &c.cb), "vkAllocateCommandBuffers"); + VkCommandBufferBeginInfo begin{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + begin.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + check(vkBeginCommandBuffer(c.cb, &begin), "vkBeginCommandBuffer"); + const VkBufferCopy region{src_offset, dst_offset, bytes}; + vkCmdCopyBuffer(c.cb, src.handle(), dst.handle(), 1, ®ion); + check(vkEndCommandBuffer(c.cb), "vkEndCommandBuffer"); + const VkFenceCreateInfo fence_info{VK_STRUCTURE_TYPE_FENCE_CREATE_INFO}; + check(vkCreateFence(device_, &fence_info, nullptr, &c.fence), "vkCreateFence"); + const VkSemaphoreCreateInfo semaphore_info{VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO}; + check(vkCreateSemaphore(device_, &semaphore_info, nullptr, &c.done), "vkCreateSemaphore"); + VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO}; + submit.commandBufferCount = 1; + submit.pCommandBuffers = &c.cb; + submit.signalSemaphoreCount = 1; + submit.pSignalSemaphores = &c.done; + // On the compute queue, earlier runs completed before this was called; a run's own waits order nothing here. + check(vkQueueSubmit(own_queue ? transfer_queue_ : queue_, 1, &submit, c.fence), "vkQueueSubmit"); + const std::uint64_t ticket = next_ticket_++; + copies_[ticket] = c; + return ticket; +} + +void Device::gpu_wait(std::uint64_t ticket) { + const auto it = copies_.find(ticket); + if (it == copies_.end() || it->second.run_done || + std::find(run_waits_.begin(), run_waits_.end(), ticket) != run_waits_.end()) { + throw std::runtime_error("gpu_wait on an unknown copy or one already waited for"); + } + run_waits_.push_back(ticket); +} + +void Device::host_wait(std::uint64_t ticket) { + const auto it = copies_.find(ticket); + if (it == copies_.end() || it->second.host_done) throw std::runtime_error("host_wait on an unknown copy"); + Copy& c = it->second; + check(vkWaitForFences(device_, 1, &c.fence, VK_TRUE, UINT64_MAX), "vkWaitForFences"); + vkFreeCommandBuffers(device_, transfer_queue_ != VK_NULL_HANDLE ? transfer_pool_ : pool_, 1, &c.cb); + vkDestroyFence(device_, c.fence, nullptr); + c.cb = VK_NULL_HANDLE; + c.fence = VK_NULL_HANDLE; + c.host_done = true; + release_if_done(ticket); +} + +// The semaphore lives until the copy completed on the host's side and a run waited for it: a binary semaphore is +// destroyed only when no signal or wait on it is pending. +void Device::release_if_done(std::uint64_t ticket) { + const auto it = copies_.find(ticket); + if (it == copies_.end() || !it->second.host_done || !it->second.run_done) return; + vkDestroySemaphore(device_, it->second.done, nullptr); + copies_.erase(it); +} + +void compute_barrier(VkCommandBuffer cb) { + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 1, &barrier, 0, nullptr, + 0, nullptr); +} + +void record_copy(VkCommandBuffer cb, const Buffer& src, VkDeviceSize src_offset, const Buffer& dst, VkDeviceSize dst_offset, + VkDeviceSize bytes) { + if (src_offset > src.bytes() || bytes > src.bytes() - src_offset || dst_offset > dst.bytes() || + bytes > dst.bytes() - dst_offset) { + throw std::runtime_error("record_copy past the end of a buffer"); + } + VkMemoryBarrier before{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + before.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT; + before.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 1, &before, 0, nullptr, 0, nullptr); + const VkBufferCopy copy{src_offset, dst_offset, bytes}; + vkCmdCopyBuffer(cb, src.handle(), dst.handle(), 1, ©); + VkMemoryBarrier after{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + after.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + after.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &after, 0, nullptr, + 0, nullptr); +} + +} // namespace mi::vk diff --git a/vulkan/runtime.h b/vulkan/runtime.h new file mode 100644 index 0000000..1d4dd73 --- /dev/null +++ b/vulkan/runtime.h @@ -0,0 +1,165 @@ +// The Vulkan device runtime: one device, one compute queue, buffers with device addresses, and pipelines from SPIR-V +// embedded in the binary. Failures throw std::runtime_error. +#pragma once + +#include "vk_api.h" + +#include +#include +#include +#include +#include +#include + +namespace mi::vk { + +// What a device offers beyond the base requirements (Vulkan 1.2, buffer device addresses, 64-bit integers). +struct Capabilities { + // The cooperative-matrix GEMM can run: 16x16x16 f16 x f16 -> f32 in subgroup scope, a compute subgroup size fixable at + // 32 with 8 per workgroup, and its shared memory. Otherwise `coop_reason` says why not. + bool coop = false; + std::string coop_reason; + // The DiT block's cooperative-matrix kernels can run: the same subgroup and memory-model requirements, with + // VK_KHR_shader_bfloat16 fragments and a 16x16x16 bf16 x bf16 -> f32 shape. Otherwise `coop_bf16_reason` says why not. + bool coop_bf16 = false; + std::string coop_bf16_reason; + // The same kernels on f16 fragments (16x16x16 f16 x f16 -> f32), for devices without bf16 ones: the bf16 operands + // convert to f16 as they are staged. Otherwise `coop_f16_reason` says why not. + bool coop_f16 = false; + std::string coop_f16_reason; + std::uint32_t shared_memory_bytes = 0; + std::uint32_t min_subgroup = 0, max_subgroup = 0; + // The largest single buffer and allocation the device takes (maxBufferSize on Vulkan 1.3, maxMemoryAllocationSize). + std::uint64_t max_buffer_bytes = 0, max_allocation_bytes = 0; +}; + +struct DeviceInfo { + std::uint32_t index = 0; + std::string name; + std::uint32_t vendor_id = 0, device_id = 0; + bool discrete = false; + std::string unusable_reason; // empty when the runtime can use the device + Capabilities caps; +}; + +// Enumerates every physical device, usable or not. Does not create a logical device. +std::vector list_devices(); + +class Buffer; + +class Device { + public: + // index < 0: the first usable device, preferring a discrete one. + explicit Device(int index = -1); + ~Device(); + Device(const Device&) = delete; + Device& operator=(const Device&) = delete; + + const DeviceInfo& info() const { return info_; } + + // Device-local memory this process can still allocate: the driver's budget minus its usage where the device has + // VK_EXT_memory_budget, else the size of the largest device-local heap. + std::uint64_t free_device_memory() const; + + // Records into one command buffer, submits it and waits. + void run(const std::function& record); + + // Asynchronous buffer copies on a transfer-only queue (a DMA engine) where the device has one, else on the compute + // queue. Buffers they read or write must be created with `transfer` set. copy_async returns at once with a ticket; + // gpu_wait(ticket) makes the next run() wait on the device for that copy, which makes its writes visible to the run + // (once per ticket); host_wait(ticket) blocks until the copy has completed, after which its source may be reused. The + // two may come in either order; a ticket is released once both happened (or with the device). + bool has_transfer_queue() const { return transfer_queue_ != VK_NULL_HANDLE; } + std::uint64_t copy_async(const Buffer& src, VkDeviceSize src_offset, const Buffer& dst, VkDeviceSize dst_offset, + VkDeviceSize bytes); + void gpu_wait(std::uint64_t ticket); + void host_wait(std::uint64_t ticket); + std::uint32_t queue_family() const { return family_; } + std::uint32_t transfer_family() const { return transfer_family_; } + + VkDevice handle() const { return device_; } + bool pipeline_stats() const { return pipeline_stats_; } + std::uint32_t find_memory_type(std::uint32_t type_bits, VkMemoryPropertyFlags flags) const; + + private: + VkInstance instance_ = VK_NULL_HANDLE; + VkPhysicalDevice physical_ = VK_NULL_HANDLE; + VkDevice device_ = VK_NULL_HANDLE; + VkQueue queue_ = VK_NULL_HANDLE; + VkCommandPool pool_ = VK_NULL_HANDLE; + VkPhysicalDeviceMemoryProperties memory_{}; + DeviceInfo info_; + bool pipeline_stats_ = false; + std::uint32_t family_ = 0, transfer_family_ = 0; + VkQueue transfer_queue_ = VK_NULL_HANDLE; + VkCommandPool transfer_pool_ = VK_NULL_HANDLE; + struct Copy { + VkCommandBuffer cb = VK_NULL_HANDLE; + VkFence fence = VK_NULL_HANDLE; + VkSemaphore done = VK_NULL_HANDLE; + bool host_done = false; // host_wait returned + bool run_done = false; // a run() waited for it on the device + }; + void release_if_done(std::uint64_t ticket); + std::map copies_; + std::uint64_t next_ticket_ = 1; + std::vector run_waits_; // tickets the next run() waits for +}; + +// A device-local storage buffer with a device address, or a host-visible staging buffer. +class Buffer { + public: + enum class Kind { kDevice, kStaging }; + // `transfer`: also used by Device::copy_async (shared with the transfer queue's family). + Buffer(Device& device, VkDeviceSize bytes, Kind kind = Kind::kDevice, bool transfer = false); + ~Buffer(); + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + + VkBuffer handle() const { return buffer_; } + VkDeviceAddress address() const { return address_; } + VkDeviceSize bytes() const { return bytes_; } + void* mapped() const { return mapped_; } + + // Copies through a staging buffer and waits; `offset` is in bytes from the start of this buffer. + void upload(const void* data, VkDeviceSize bytes, VkDeviceSize offset = 0); + void download(void* data, VkDeviceSize bytes, VkDeviceSize offset = 0); + + private: + Device& device_; + VkBuffer buffer_ = VK_NULL_HANDLE; + VkDeviceMemory memory_ = VK_NULL_HANDLE; + VkDeviceAddress address_ = 0; + VkDeviceSize bytes_ = 0; + void* mapped_ = nullptr; +}; + +// A compute pipeline with push constants only. required_subgroup_size != 0 also requires full subgroups. +class Pipeline { + public: + Pipeline(Device& device, const std::uint32_t* spirv, std::size_t spirv_bytes, std::uint32_t push_bytes, + std::uint32_t required_subgroup_size = 0, const std::vector& specialization = {}); + ~Pipeline(); + Pipeline(const Pipeline&) = delete; + Pipeline& operator=(const Pipeline&) = delete; + + void dispatch(VkCommandBuffer cb, const void* push, std::uint32_t x, std::uint32_t y) const; + + private: + void print_statistics(std::size_t spirv_bytes) const; + Device& device_; + VkShaderModule module_ = VK_NULL_HANDLE; + VkPipelineLayout layout_ = VK_NULL_HANDLE; + VkPipeline pipeline_ = VK_NULL_HANDLE; + std::uint32_t push_bytes_ = 0; +}; + +// Makes earlier compute writes visible to later compute and transfer work. A queue wait only reaches the host. +void compute_barrier(VkCommandBuffer cb); + +// Records a device copy of `bytes` from src at src_offset to dst at dst_offset, after earlier compute and transfer writes +// and before later compute work. +void record_copy(VkCommandBuffer cb, const Buffer& src, VkDeviceSize src_offset, const Buffer& dst, VkDeviceSize dst_offset, + VkDeviceSize bytes); + +} // namespace mi::vk diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp new file mode 100644 index 0000000..56c524c --- /dev/null +++ b/vulkan/safetensors.cpp @@ -0,0 +1,408 @@ +#include "safetensors.h" + +#include + +#include +#include +#include +#include + +#ifdef _WIN32 +#include + +#include +#else +#include +#endif + +namespace mi::vk::st { +namespace { + +[[noreturn]] void fail(const std::string& what) { throw std::runtime_error("safetensors: " + what); } + +// The JSON subset safetensors headers and index files use: objects, arrays, strings, numbers, true/false/null. +struct Json { + enum class Kind { kNull, kBool, kNumber, kString, kArray, kObject } kind = Kind::kNull; + double number = 0; + std::int64_t integer = 0; // exact for integral numbers + std::string text; + std::vector items; + std::vector> members; + + const Json* find(const std::string& key) const { + for (const auto& [k, v] : members) { + if (k == key) return &v; + } + return nullptr; + } +}; + +class Parser { + public: + explicit Parser(const std::string& s) : s_(s) {} + + Json parse_document() { + Json v = value(); + skip(); + if (i_ != s_.size()) fail("trailing bytes after the JSON value"); + return v; + } + + private: + void skip() { + while (i_ < s_.size() && (s_[i_] == ' ' || s_[i_] == '\n' || s_[i_] == '\r' || s_[i_] == '\t')) ++i_; + } + char peek() { + skip(); + if (i_ >= s_.size()) fail("unexpected end of JSON"); + return s_[i_]; + } + void expect(char c) { + if (peek() != c) fail(std::string("expected '") + c + "' in JSON"); + ++i_; + } + + Json value() { + const char c = peek(); + if (c == '{') return object(); + if (c == '[') return array(); + if (c == '"') { + Json v; + v.kind = Json::Kind::kString; + v.text = string(); + return v; + } + if (s_.compare(i_, 4, "true") == 0 || s_.compare(i_, 5, "false") == 0) { + Json v; + v.kind = Json::Kind::kBool; + v.integer = s_[i_] == 't'; + i_ += s_[i_] == 't' ? 4 : 5; + return v; + } + if (s_.compare(i_, 4, "null") == 0) { + i_ += 4; + return {}; + } + return number(); + } + + Json object() { + Json v; + v.kind = Json::Kind::kObject; + expect('{'); + if (peek() == '}') { + ++i_; + return v; + } + for (;;) { + std::string key = (peek(), string()); + expect(':'); + v.members.emplace_back(std::move(key), value()); + if (peek() == ',') { + ++i_; + continue; + } + expect('}'); + return v; + } + } + + Json array() { + Json v; + v.kind = Json::Kind::kArray; + expect('['); + if (peek() == ']') { + ++i_; + return v; + } + for (;;) { + v.items.push_back(value()); + if (peek() == ',') { + ++i_; + continue; + } + expect(']'); + return v; + } + } + + std::string string() { + expect('"'); + std::string out; + while (i_ < s_.size() && s_[i_] != '"') { + char c = s_[i_++]; + if (c == '\\') { + if (i_ >= s_.size()) break; + const char e = s_[i_++]; + switch (e) { + case 'n': + c = '\n'; + break; + case 't': + c = '\t'; + break; + case 'r': + c = '\r'; + break; + case 'b': + c = '\b'; + break; + case 'f': + c = '\f'; + break; + case 'u': { + // Names here are ASCII; other code points are kept as UTF-8 of the BMP value. + if (i_ + 4 > s_.size()) fail("bad \\u escape"); + const unsigned cp = static_cast(std::stoul(s_.substr(i_, 4), nullptr, 16)); + i_ += 4; + if (cp < 0x80) { + out += static_cast(cp); + } else if (cp < 0x800) { + out += static_cast(0xC0 | (cp >> 6)); + out += static_cast(0x80 | (cp & 0x3F)); + } else { + out += static_cast(0xE0 | (cp >> 12)); + out += static_cast(0x80 | ((cp >> 6) & 0x3F)); + out += static_cast(0x80 | (cp & 0x3F)); + } + continue; + } + default: + c = e; + break; + } + } + out += c; + } + if (i_ >= s_.size()) fail("unterminated JSON string"); + ++i_; + return out; + } + + Json number() { + const std::size_t start = i_; + bool integral = true; + while (i_ < s_.size() && (std::strchr("+-0123456789.eE", s_[i_]) != nullptr)) { + if (s_[i_] == '.' || s_[i_] == 'e' || s_[i_] == 'E') integral = false; + ++i_; + } + if (start == i_) fail("unexpected character in JSON"); + Json v; + v.kind = Json::Kind::kNumber; + const std::string text = s_.substr(start, i_ - start); + v.number = std::stod(text); + v.integer = integral ? std::stoll(text) : static_cast(v.number); + return v; + } + + const std::string& s_; + std::size_t i_ = 0; +}; + +std::uint64_t dtype_bytes(const std::string& dtype) { + if (dtype == "BF16" || dtype == "F16" || dtype == "I16" || dtype == "U16") return 2; + if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; + if (dtype == "F64" || dtype == "I64" || dtype == "U64") return 8; + if (dtype == "I8" || dtype == "U8" || dtype == "BOOL" || dtype == "F8_E4M3" || dtype == "F8_E5M2") return 1; + return 0; +} + +#ifdef _WIN32 +std::mutex g_read_mutex; // _lseeki64 + _read share the descriptor's position +#endif + +std::string utf8(const std::filesystem::path& p) { + const std::u8string u = p.u8string(); + return std::string(u.begin(), u.end()); +} + +void read_at(int fd, void* dst, std::uint64_t offset, std::uint64_t bytes, const std::string& path) { + auto* out = static_cast(dst); + while (bytes > 0) { + const std::uint64_t want = bytes < (std::uint64_t{1} << 30) ? bytes : (std::uint64_t{1} << 30); +#ifdef _WIN32 + std::lock_guard lock(g_read_mutex); + if (_lseeki64(fd, static_cast<__int64>(offset), SEEK_SET) < 0) fail("seek failed in " + path); + const int got = _read(fd, out, static_cast(want)); +#else + const ssize_t got = ::pread(fd, out, static_cast(want), static_cast(offset)); +#endif + if (got <= 0) fail("read failed or hit the end of " + path); + out += got; + offset += static_cast(got); + bytes -= static_cast(got); + } +} + +void close_fd(int fd) { + if (fd < 0) return; +#ifdef _WIN32 + _close(fd); +#else + ::close(fd); +#endif +} + +} // namespace + +File::File(const std::filesystem::path& file) : path_(utf8(file)) { + const std::string& path = path_; +#ifdef _WIN32 + fd_ = _wopen(file.c_str(), _O_RDONLY | _O_BINARY); +#else + fd_ = ::open(file.c_str(), O_RDONLY | O_CLOEXEC); +#endif + if (fd_ < 0) fail("cannot open " + path); + try { + std::uint64_t header_bytes = 0; + unsigned char le[8]; + read_at(fd_, le, 0, 8, path); + for (int i = 7; i >= 0; --i) header_bytes = (header_bytes << 8) | le[i]; + if (header_bytes == 0 || header_bytes > (std::uint64_t{100} << 20)) fail("implausible header size in " + path); + std::string header(static_cast(header_bytes), '\0'); + read_at(fd_, header.data(), 8, header_bytes, path); + const Json root = Parser(header).parse_document(); + if (root.kind != Json::Kind::kObject) fail("header of " + path + " is not an object"); + const std::uint64_t data_start = 8 + header_bytes; + for (const auto& [name, v] : root.members) { + if (name == "__metadata__") continue; + const Json* dtype = v.find("dtype"); + const Json* shape = v.find("shape"); + const Json* offsets = v.find("data_offsets"); + if (dtype == nullptr || shape == nullptr || offsets == nullptr || dtype->kind != Json::Kind::kString || + shape->kind != Json::Kind::kArray || offsets->kind != Json::Kind::kArray || offsets->items.size() != 2) { + fail("malformed entry " + name + " in " + path); + } + TensorInfo t; + t.name = name; + t.dtype = dtype->text; + std::uint64_t numel = 1; + for (const Json& d : shape->items) { + if (d.kind != Json::Kind::kNumber || d.integer < 0) fail("bad shape of " + name); + t.shape.push_back(d.integer); + numel *= static_cast(d.integer); + } + const std::int64_t begin = offsets->items[0].integer, end = offsets->items[1].integer; + if (begin < 0 || end < begin) fail("bad data_offsets of " + name); + t.offset = data_start + static_cast(begin); + t.bytes = static_cast(end - begin); + const std::uint64_t item = dtype_bytes(t.dtype); + if (item != 0 && numel * item != t.bytes) fail("size of " + name + " does not match its dtype and shape"); + by_name_[t.name] = tensors_.size(); + tensors_.push_back(std::move(t)); + } + } catch (...) { + close_fd(fd_); // the destructor does not run for a constructor that throws + throw; + } +} + +File::~File() { close_fd(fd_); } + +const TensorInfo* File::find(const std::string& name) const { + const auto it = by_name_.find(name); + return it == by_name_.end() ? nullptr : &tensors_[it->second]; +} + +const TensorInfo& File::at(const std::string& name) const { + const TensorInfo* t = find(name); + if (t == nullptr) fail(path_ + " has no tensor " + name); + return *t; +} + +void File::read(const TensorInfo& t, void* dst, std::uint64_t rel, std::uint64_t bytes) const { + if (rel > t.bytes || bytes > t.bytes - rel) fail("read past the end of tensor " + t.name); + read_at(fd_, dst, t.offset + rel, bytes, path_); +} + +std::map read_index(const std::filesystem::path& file) { + const std::string path = utf8(file); + std::ifstream in(file, std::ios::binary); + if (!in) fail("cannot open " + path); + std::ostringstream ss; + ss << in.rdbuf(); + const std::string text = ss.str(); + const Json root = Parser(text).parse_document(); + const Json* map = root.find("weight_map"); + if (map == nullptr || map->kind != Json::Kind::kObject) fail(path + " has no weight_map object"); + std::map out; + for (const auto& [name, file] : map->members) { + if (file.kind != Json::Kind::kString) fail("weight_map entry " + name + " is not a string"); + out[name] = file.text; + } + return out; +} + +std::vector File::read_f32(const std::string& name) const { + const TensorInfo& t = at(name); + if (t.dtype == "F32") { + std::vector out(static_cast(t.bytes / 4)); + read(t, out.data()); + return out; + } + if (t.dtype == "BF16") { + std::vector raw(static_cast(t.bytes / 2)); + read(t, raw.data()); + std::vector out(raw.size()); + for (std::size_t i = 0; i < raw.size(); ++i) { + const std::uint32_t bits = static_cast(raw[i]) << 16; + std::memcpy(&out[i], &bits, 4); + } + return out; + } + throw std::runtime_error("safetensors " + path_ + ": tensor '" + name + "' is " + t.dtype + ", not F32 or BF16"); +} + +Index::Index(const std::filesystem::path& path) { + const auto add = [&](const std::filesystem::path& file) { + files_.push_back(std::make_unique(file)); + const File& f = *files_.back(); + for (const TensorInfo& t : f.tensors()) { + if (!tensors_.emplace(t.name, std::make_pair(&f, &t)).second) { + throw std::runtime_error("safetensors: " + t.name + " appears in more than one file under " + utf8(path)); + } + } + }; + if (!std::filesystem::is_directory(path)) { + add(path); + return; + } + std::vector indexes, files; + for (const auto& entry : std::filesystem::directory_iterator(path)) { + const std::string name = utf8(entry.path().filename()); + if (name.ends_with(".safetensors.index.json")) indexes.push_back(entry.path()); + else if (name.ends_with(".safetensors")) files.push_back(entry.path()); + } + if (indexes.size() > 1) throw std::runtime_error("safetensors: " + utf8(path) + " holds more than one index file"); + if (indexes.empty()) { + if (files.size() != 1) { + throw std::runtime_error("safetensors: " + utf8(path) + " holds no index file and not exactly one .safetensors file"); + } + add(files.front()); + return; + } + std::map shards; + const std::map weight_map = read_index(indexes.front()); + for (const auto& [name, shard] : weight_map) shards[shard] = true; + for (const auto& [shard, unused] : shards) add(path / shard); + for (const auto& [name, shard] : weight_map) { + if (!has(name)) throw std::runtime_error("safetensors: " + name + " is in the index but in no shard"); + } +} + +const TensorInfo& Index::at(const std::string& name) const { + const auto it = tensors_.find(name); + if (it == tensors_.end()) throw std::runtime_error("safetensors: no tensor " + name); + return *it->second.second; +} + +void Index::read_into(const std::string& name, void* out, std::uint64_t bytes) const { + const TensorInfo& t = at(name); + if (bytes != t.bytes) { + throw std::runtime_error("safetensors: read of " + std::to_string(bytes) + " bytes of " + name + ", which has " + + std::to_string(t.bytes)); + } + tensors_.at(name).first->read(t, out); +} + +} // namespace mi::vk::st diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h new file mode 100644 index 0000000..3752d89 --- /dev/null +++ b/vulkan/safetensors.h @@ -0,0 +1,67 @@ +// A minimal safetensors reader: the header of one file, tensor bytes read with pread, and the shard map of a +// model.safetensors.index.json. Failures throw std::runtime_error. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace mi::vk::st { + +struct TensorInfo { + std::string name; + std::string dtype; // as stored, e.g. "BF16", "F32" + std::vector shape; + std::uint64_t offset = 0; // of the first byte, from the start of the file + std::uint64_t bytes = 0; +}; + +class File { + public: + explicit File(const std::filesystem::path& path); // opened by its wide name on Windows + ~File(); + File(const File&) = delete; + File& operator=(const File&) = delete; + + const std::string& path() const { return path_; } // UTF-8, for messages + const std::vector& tensors() const { return tensors_; } + const TensorInfo* find(const std::string& name) const; // null if absent + const TensorInfo& at(const std::string& name) const; // throws if absent + + // `bytes` bytes of tensor `t` from byte `rel` of it into `dst`. + void read(const TensorInfo& t, void* dst, std::uint64_t rel, std::uint64_t bytes) const; + void read(const TensorInfo& t, void* dst) const { read(t, dst, 0, t.bytes); } + // A whole F32 or BF16 tensor as fp32 values (bf16 widened exactly), or throws. + std::vector read_f32(const std::string& name) const; + + private: + std::string path_; + int fd_ = -1; + std::vector tensors_; + std::map by_name_; +}; + +// model.safetensors.index.json's weight_map: tensor name -> shard file name (relative to the index's directory). +std::map read_index(const std::filesystem::path& path); + +// The tensors of one file, or of every shard of a model directory, by name. +class Index { + public: + // A .safetensors file, or a directory holding one *.safetensors.index.json (its shards) or exactly one .safetensors. + explicit Index(const std::filesystem::path& path); + + bool has(const std::string& name) const { return tensors_.count(name) != 0; } + const TensorInfo& at(const std::string& name) const; // throws if absent + // Reads a whole tensor into caller memory of exactly its size, or throws. + void read_into(const std::string& name, void* out, std::uint64_t bytes) const; + + private: + std::vector> files_; + std::map> tensors_; +}; + +} // namespace mi::vk::st diff --git a/vulkan/shaders/attention_coop.comp b/vulkan/shaders/attention_coop.comp new file mode 100644 index 0000000..03e96c2 --- /dev/null +++ b/vulkan/shaders/attention_coop.comp @@ -0,0 +1,268 @@ +// Scaled-dot-product attention over one sequence on bf16 cooperative matrices: head_dim 128, non-causal, bf16 in and out, +// fp32 online softmax, the same steps as attention_plain.comp and FlashAttention: +// S = q * K_step^T (fragments, fp32 accumulators) +// max = max(max, rowmax(S)); rescale = exp2((old max - max) * log2_scale); row_sum *= rescale; acc_o *= rescale +// p = exp2(S * log2_scale - max * log2_scale); row_sum += p; P = bf16(p) +// acc_o += P * V_step (fragments, fp32 accumulators) +// and at the end out = bf16(acc_o * (1 / row_sum)). Key steps walk from the last down to the first, as FlashAttention's. +// +// One workgroup per 64 query rows of one head: 128 invocations, 4 full subgroups of 32, which the pipeline must require. +// Subgroup s owns query rows 16 s .. +16; lane l owns row l / 2 and keys or columns of half l % 2 of it. A key step is 32 +// keys, staged in shared memory with the V step. Accumulators have no row access in cooperative matrices, so S and the +// output pass through the subgroup's staging, and acc_o goes there for a rescale only when a row's maximum grew. +// +// Built with F16 (attention_f16_coop.spv) for devices with f16 fragments only: q, k and v convert to f16 unscaled as they +// are staged, and P enters as bf16(p) * 2^15 (p <= 1, so it stays an f16 normal down to 2^-29), undone with 1 / row_sum. +#version 460 +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#ifdef F16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define ELEM float16_t +#define TO_FRAG(v, scale) bf8_to_f16(v, scale) +const float P_SCALE = 32768.0; +#else +#extension GL_EXT_bfloat16 : require +#define ELEM bfloat16_t +#define TO_FRAG(v, scale) (v) +const float P_SCALE = 1.0; +#endif +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_vote : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { uvec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t q; + uint64_t k; + uint64_t v; + uint64_t o; + uint ldq; + uint ldk; + uint ldv; + uint ldo; + uint hsq; + uint hsk; + uint hsv; + uint hso; + uint q_row0; + uint k_row0; + uint seqlen_q; + uint seqlen_k; + float log2_scale; +} pc; + +const uint D = 128u; +const uint ROWS = 64u; // query rows per workgroup +const uint STEP = 32u; // keys per step +const uint KV_ROW = 17u; // uvec4 per staged K or V row: 16 of data, 1 of padding +const uint S_ROW = 9u; // uvec4 per staging row: 32 floats and 4 of padding +const uint P_ROW = 5u; // uvec4 per P row: 32 bf16 and 8 of padding +const uint SK = 0u; +const uint SV = SK + STEP * KV_ROW; +const uint SS = SV + STEP * KV_ROW; +const uint SP = SS + 4u * 16u * S_ROW; +const uint TOTAL = SP + 4u * 16u * P_ROW; + +// Q is staged once through the K and V region before the first step: 64 rows of 17 uvec4 fit in 2 * 32 rows. +shared uvec4 smem[TOTAL]; + +#define FRAG_A coopmat +#define FRAG_B coopmat +#define ACC coopmat + +void sync_subgroup() { + subgroupMemoryBarrierShared(); + subgroupBarrier(); +} + +// Invocation i loads chunks i % 16 of rows i / 16 + 8 j, j = 0..3, of the K and V steps at key n0. +void load_kv(uint64_t k_head, uint64_t v_head, uint n0, uint cols, out uvec4 rk[4], out uvec4 rv[4]) { + const uint i = gl_LocalInvocationIndex; + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint r = i / 16u + 8u * j; + rk[j] = uvec4(0u); + rv[j] = uvec4(0u); + if (r < cols) { + rk[j] = Vec4s(k_head + (uint64_t(n0 + r) * pc.ldk) * 2ul).d[i % 16u]; + rv[j] = Vec4s(v_head + (uint64_t(n0 + r) * pc.ldv) * 2ul).d[i % 16u]; + } + } +} + +void store_kv(const uvec4 rk[4], const uvec4 rv[4]) { + const uint i = gl_LocalInvocationIndex; + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint at = (i / 16u + 8u * j) * KV_ROW + i % 16u; + smem[SK + at] = TO_FRAG(rk[j], 1.0); + smem[SV + at] = TO_FRAG(rv[j], 1.0); + } +} + +void main() { + const uint head = gl_WorkGroupID.x; + const uint q_first = gl_WorkGroupID.y * ROWS; + const uint sg = gl_SubgroupID; + const uint lane = gl_SubgroupInvocationID; + const uint frow = lane / 2u, half_ = lane % 2u; + const uint i = gl_LocalInvocationIndex; + + // Q rows q_first .. +64 through the K and V region, zero past the sequence. + [[unroll]] for (uint j = 0u; j < 8u; ++j) { + const uint r = i / 16u + 8u * j; + uvec4 v = uvec4(0u); + if (q_first + r < pc.seqlen_q) { + const uint64_t row = uint64_t(pc.q_row0 + q_first + r); + v = Vec4s(pc.q + (row * pc.ldq + uint64_t(head) * pc.hsq) * 2ul).d[i % 16u]; + } + smem[SK + r * KV_ROW + i % 16u] = TO_FRAG(v, 1.0); + } + barrier(); + FRAG_A fq[8]; + [[unroll]] for (uint d = 0u; d < 8u; ++d) { + coopMatLoad(fq[d], smem, SK + sg * 16u * KV_ROW + d * 2u, KV_ROW, gl_CooperativeMatrixLayoutRowMajor); + } + barrier(); + + const uint64_t k_head = pc.k + (uint64_t(pc.k_row0) * pc.ldk + uint64_t(head) * pc.hsk) * 2ul; + const uint64_t v_head = pc.v + (uint64_t(pc.k_row0) * pc.ldv + uint64_t(head) * pc.hsv) * 2ul; + const uint steps = (pc.seqlen_k + STEP - 1u) / STEP; + + ACC acc[8]; + [[unroll]] for (uint d = 0u; d < 8u; ++d) acc[d] = ACC(0.0); + float row_max = 0.0, row_sum = 0.0; // of row frow, the same in both lanes of the row + + const uint ss = SS + sg * 16u * S_ROW, sp = SP + sg * 16u * P_ROW; + uvec4 rk[4], rv[4]; + load_kv(k_head, v_head, (steps - 1u) * STEP, pc.seqlen_k - (steps - 1u) * STEP, rk, rv); + store_kv(rk, rv); + barrier(); + + for (uint step_plus_1 = steps; step_plus_1 > 0u; --step_plus_1) { + const uint step_index = step_plus_1 - 1u; + const bool first = step_plus_1 == steps; + const uint n0 = step_index * STEP; + const uint cols = min(STEP, pc.seqlen_k - n0); + if (step_index > 0u) load_kv(k_head, v_head, n0 - STEP, STEP, rk, rv); + + // S = q K^T: 16 rows by 32 keys. + ACC s[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + s[c] = ACC(0.0); + [[unroll]] for (uint d = 0u; d < 8u; ++d) { + FRAG_B fk; + coopMatLoad(fk, smem, SK + c * 16u * KV_ROW + d * 2u, KV_ROW, gl_CooperativeMatrixLayoutColumnMajor); + s[c] = coopMatMulAdd(fq[d], fk, s[c]); + } + coopMatStore(s[c], smem, ss + c * 4u, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + } + sync_subgroup(); + + // Lane: row frow, keys 16 half_ .. +16. + float sv[16]; + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const vec4 v4 = uintBitsToFloat(smem[ss + frow * S_ROW + half_ * 4u + j]); + sv[4u * j] = v4.x; + sv[4u * j + 1u] = v4.y; + sv[4u * j + 2u] = v4.z; + sv[4u * j + 3u] = v4.w; + } + float cur = first ? -1.0 / 0.0 : row_max; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + if (16u * half_ + j < cols) cur = max(cur, sv[j]); + } + cur = max(cur, subgroupShuffleXor(cur, 1u)); + float rescale = 1.0; + if (!first) { + precise const float diff = (row_max - cur) * pc.log2_scale; + rescale = exp2(diff); + precise const float scaled = row_sum * rescale; + row_sum = scaled; + } + row_max = cur; + precise const float max_scaled = isinf(cur) && cur < 0.0 ? 0.0 : cur * pc.log2_scale; + precise float part = 0.0; + uint p_bits[8]; + [[unroll]] for (uint j = 0u; j < 8u; ++j) p_bits[j] = 0u; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + float w = 0.0; + if (16u * half_ + j < cols) w = exp2(fma(sv[j], pc.log2_scale, -max_scaled)); + part = part + w; + p_bits[j / 2u] |= bf(w) << (16u * (j % 2u)); + } + precise const float total = row_sum + (part + subgroupShuffleXor(part, 1u)); + row_sum = first ? part + subgroupShuffleXor(part, 1u) : total; + smem[sp + frow * P_ROW + half_ * 2u] = TO_FRAG(uvec4(p_bits[0], p_bits[1], p_bits[2], p_bits[3]), P_SCALE); + smem[sp + frow * P_ROW + half_ * 2u + 1u] = TO_FRAG(uvec4(p_bits[4], p_bits[5], p_bits[6], p_bits[7]), P_SCALE); + + // Every lane has read S from the staging before the rescale overwrites it. + sync_subgroup(); + // acc_o *= rescale, through the staging two fragments at a time, when some row's maximum grew. + if (subgroupAny(rescale != 1.0)) { + [[unroll]] for (uint pair = 0u; pair < 4u; ++pair) { + coopMatStore(acc[2u * pair], smem, ss, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc[2u * pair + 1u], smem, ss + 4u, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + sync_subgroup(); + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint at = ss + frow * S_ROW + half_ * 4u + j; + precise const vec4 scaled = uintBitsToFloat(smem[at]) * rescale; + smem[at] = floatBitsToUint(scaled); + } + sync_subgroup(); + coopMatLoad(acc[2u * pair], smem, ss, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(acc[2u * pair + 1u], smem, ss + 4u, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + sync_subgroup(); + } + } + sync_subgroup(); + + // acc_o += P V: P is 16 rows by 32 keys, V 32 keys by 128 columns. + FRAG_A fp[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + coopMatLoad(fp[c], smem, sp + c * 2u, P_ROW, gl_CooperativeMatrixLayoutRowMajor); + } + [[unroll]] for (uint d = 0u; d < 8u; ++d) { + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + FRAG_B fv; + coopMatLoad(fv, smem, SV + c * 16u * KV_ROW + d * 2u, KV_ROW, gl_CooperativeMatrixLayoutRowMajor); + acc[d] = coopMatMulAdd(fp[c], fv, acc[d]); + } + } + barrier(); + if (step_index > 0u) { + store_kv(rk, rv); + barrier(); + } + } + + // out = bf16(acc_o * (1 / row_sum)), through the staging two fragments at a time. + // F16: times 2^-15 for P's scale; a power of two, so (acc_o * inv_sum) rounds as without it. + const float inv_sum = ((row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum) * (1.0 / P_SCALE); + const uint q_local = q_first + sg * 16u + frow; + const bool write = q_local < pc.seqlen_q; + const uint64_t out_row = pc.o + ((uint64_t(pc.q_row0) + q_local) * pc.ldo + uint64_t(head) * pc.hso) * 2ul; + [[unroll]] for (uint pair = 0u; pair < 4u; ++pair) { + coopMatStore(acc[2u * pair], smem, ss, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc[2u * pair + 1u], smem, ss + 4u, S_ROW, gl_CooperativeMatrixLayoutRowMajor); + sync_subgroup(); + uvec4 packed[2]; + [[unroll]] for (uint j = 0u; j < 2u; ++j) { + precise const vec4 a = uintBitsToFloat(smem[ss + frow * S_ROW + half_ * 4u + 2u * j]) * inv_sum; + precise const vec4 b = uintBitsToFloat(smem[ss + frow * S_ROW + half_ * 4u + 2u * j + 1u]) * inv_sum; + packed[j] = uvec4(bf4(a), bf4(b)); + } + // Columns 32 pair + 16 half_ .. +16 of this head: two uvec4 of 8 bf16. + if (write) { + Vec4s out_v = Vec4s(out_row + uint64_t(32u * pair + 16u * half_) * 2ul); + out_v.d[0] = packed[0]; + out_v.d[1] = packed[1]; + } + sync_subgroup(); + } +} diff --git a/vulkan/shaders/attention_plain.comp b/vulkan/shaders/attention_plain.comp new file mode 100644 index 0000000..973c026 --- /dev/null +++ b/vulkan/shaders/attention_plain.comp @@ -0,0 +1,209 @@ +// Scaled-dot-product attention over one sequence without cooperative matrices, head_dim 128, non-causal, bf16 in and out, +// fp32 online softmax, FlashAttention's steps: +// S = q * K_step^T +// max = max(max, rowmax(S)); rescale = exp2((old max - max) * log2_scale); row_sum *= rescale; acc_o *= rescale +// p = exp2(S * log2_scale - max * log2_scale); row_sum += p; P = bf16(p) +// acc_o += P * V_step +// and at the end out = bf16(acc_o * (1 / row_sum)). Key steps of 32 walk from the last down to the first, as +// FlashAttention's reverse iteration. It needs no subgroup operations or 16-bit arithmetic, and 30 KiB of shared memory. +// +// One workgroup of 128 invocations per 32 query rows of one head. Invocation t computes S for rows 2 (t / 8) .. +2 and +// keys 4 (t % 8) .. +4; holds acc_o for the same two rows and columns 16 (t % 8) .. +16; and runs the softmax of row t / 4 +// over keys 8 (t % 4) .. +8, the four parts of a row meeting through shared memory. +#version 460 +#extension GL_GOOGLE_include_directive : require +#extension GL_EXT_control_flow_attributes : require +#include "common.glsl" + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer U32s { uint d[]; }; + +layout(push_constant) uniform PC { + uint64_t q; + uint64_t k; + uint64_t v; + uint64_t o; + uint ldq; + uint ldk; + uint ldv; + uint ldo; + uint hsq; + uint hsk; + uint hsv; + uint hso; + uint q_row0; + uint k_row0; + uint seqlen_q; + uint seqlen_k; + float log2_scale; +} pc; + +const uint ROWS = 32u; +const uint STEP = 32u; +const uint WORDS = 64u; // uints per 128-dim bf16 row +const uint STRIDE = 65u; // staged row stride in uints: odd, so 8 rows at one column fall in 8 banks + +shared uint sq[ROWS * STRIDE]; +shared uint sk[STEP * STRIDE]; +shared uint sv[STEP * STRIDE]; +shared float ss[ROWS * (STEP + 1u)]; +shared float part_max[ROWS * 4u]; +shared float part_sum[ROWS * 4u]; +shared float row_scale[ROWS]; +shared float row_inv[ROWS]; + +void sync_shared() { barrier(); } + +// Invocation t stages words 16 (t % 4) .. +16 of key rows t / 4 of the K and V steps at key n0. +void load_kv(uint64_t k_head, uint64_t v_head, uint n0, uint cols, out uint rk[16], out uint rv[16]) { + const uint t = gl_LocalInvocationIndex; + const uint r = t / 4u, w0 = (t % 4u) * 16u; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + rk[j] = 0u; + rv[j] = 0u; + } + if (r < cols) { + U32s krow = U32s(k_head + uint64_t(n0 + r) * pc.ldk * 2ul), vrow = U32s(v_head + uint64_t(n0 + r) * pc.ldv * 2ul); + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + rk[j] = krow.d[w0 + j]; + rv[j] = vrow.d[w0 + j]; + } + } +} + +void store_kv(const uint rk[16], const uint rv[16]) { + const uint t = gl_LocalInvocationIndex; + const uint at = (t / 4u) * STRIDE + (t % 4u) * 16u; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + sk[at + j] = rk[j]; + sv[at + j] = rv[j]; + } +} + +vec2 pair(uint w) { return uintBitsToFloat(uvec2(w << 16, w & 0xffff0000u)); } + +void main() { + const uint t = gl_LocalInvocationIndex; + const uint head = gl_WorkGroupID.x; + const uint q_first = gl_WorkGroupID.y * ROWS; + const uint pr = t / 8u, quad = t % 8u; // S and acc_o: rows 2 pr, 2 pr + 1 + const uint srow = t / 4u, spart = t % 4u; // softmax: row srow, keys 8 spart .. +8 + + // Q rows q_first .. +32, zero past the sequence. + { + const uint r = t / 4u, w0 = (t % 4u) * 16u; + const bool valid = q_first + r < pc.seqlen_q; + U32s qrow = U32s(pc.q + ((uint64_t(pc.q_row0) + q_first + r) * pc.ldq + uint64_t(head) * pc.hsq) * 2ul); + [[unroll]] for (uint j = 0u; j < 16u; ++j) sq[r * STRIDE + w0 + j] = valid ? qrow.d[w0 + j] : 0u; + } + + const uint64_t k_head = pc.k + (uint64_t(pc.k_row0) * pc.ldk + uint64_t(head) * pc.hsk) * 2ul; + const uint64_t v_head = pc.v + (uint64_t(pc.k_row0) * pc.ldv + uint64_t(head) * pc.hsv) * 2ul; + const uint steps = (pc.seqlen_k + STEP - 1u) / STEP; + + float acc[2][16]; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint c = 0u; c < 16u; ++c) acc[r][c] = 0.0; + } + float row_max = 0.0, row_sum = 0.0; // of row srow, kept alike by its four softmax invocations + + uint rk[16], rv[16]; + load_kv(k_head, v_head, (steps - 1u) * STEP, pc.seqlen_k - (steps - 1u) * STEP, rk, rv); + store_kv(rk, rv); + sync_shared(); + + for (uint step_plus_1 = steps; step_plus_1 > 0u; --step_plus_1) { + const uint step_index = step_plus_1 - 1u; + const bool first = step_plus_1 == steps; + const uint n0 = step_index * STEP; + const uint cols = min(STEP, pc.seqlen_k - n0); + if (step_index > 0u) load_kv(k_head, v_head, n0 - STEP, STEP, rk, rv); + + // S: rows 2 pr + i, keys 4 quad + j. + float s[2][4]; + [[unroll]] for (uint i = 0u; i < 2u; ++i) { + [[unroll]] for (uint j = 0u; j < 4u; ++j) s[i][j] = 0.0; + } + for (uint w = 0u; w < WORDS; ++w) { + const vec2 q0 = pair(sq[(2u * pr) * STRIDE + w]), q1 = pair(sq[(2u * pr + 1u) * STRIDE + w]); + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const vec2 kk = pair(sk[(4u * quad + j) * STRIDE + w]); + s[0][j] = fma(q0.y, kk.y, fma(q0.x, kk.x, s[0][j])); + s[1][j] = fma(q1.y, kk.y, fma(q1.x, kk.x, s[1][j])); + } + } + [[unroll]] for (uint i = 0u; i < 2u; ++i) { + [[unroll]] for (uint j = 0u; j < 4u; ++j) ss[(2u * pr + i) * (STEP + 1u) + 4u * quad + j] = s[i][j]; + } + sync_shared(); + + // Softmax of row srow over keys 8 spart .. +8. + float mine = -1.0 / 0.0; + [[unroll]] for (uint j = 0u; j < 8u; ++j) { + if (8u * spart + j < cols) mine = max(mine, ss[srow * (STEP + 1u) + 8u * spart + j]); + } + part_max[srow * 4u + spart] = mine; + sync_shared(); + float cur = first ? -1.0 / 0.0 : row_max; + [[unroll]] for (uint p = 0u; p < 4u; ++p) cur = max(cur, part_max[srow * 4u + p]); + float rescale = 1.0; + if (!first) { + precise const float diff = (row_max - cur) * pc.log2_scale; + rescale = exp2(diff); + } + row_max = cur; + precise const float max_scaled = isinf(cur) && cur < 0.0 ? 0.0 : cur * pc.log2_scale; + precise float part = 0.0; + [[unroll]] for (uint j = 0u; j < 8u; ++j) { + const uint at = srow * (STEP + 1u) + 8u * spart + j; + float w = 0.0; + if (8u * spart + j < cols) w = exp2(fma(ss[at], pc.log2_scale, -max_scaled)); + part = part + w; + ss[at] = uintBitsToFloat(bf(w) << 16); + } + part_sum[srow * 4u + spart] = part; + if (spart == 0u) row_scale[srow] = rescale; + sync_shared(); + precise const float step_sum = (part_sum[srow * 4u] + part_sum[srow * 4u + 1u]) + + (part_sum[srow * 4u + 2u] + part_sum[srow * 4u + 3u]); + precise const float scaled_sum = row_sum * rescale; + row_sum = first ? step_sum : scaled_sum + step_sum; + + // acc_o = acc_o * rescale + P V for rows 2 pr + i, columns 16 quad .. +16. + if (!first) { + [[unroll]] for (uint i = 0u; i < 2u; ++i) { + const float r = row_scale[2u * pr + i]; + [[unroll]] for (uint c = 0u; c < 16u; ++c) acc[i][c] *= r; + } + } + for (uint j = 0u; j < cols; ++j) { + const float p0 = ss[(2u * pr) * (STEP + 1u) + j], p1 = ss[(2u * pr + 1u) * (STEP + 1u) + j]; + [[unroll]] for (uint c = 0u; c < 8u; ++c) { + const vec2 vv = pair(sv[j * STRIDE + 8u * quad + c]); + acc[0][2u * c] = fma(p0, vv.x, acc[0][2u * c]); + acc[0][2u * c + 1u] = fma(p0, vv.y, acc[0][2u * c + 1u]); + acc[1][2u * c] = fma(p1, vv.x, acc[1][2u * c]); + acc[1][2u * c + 1u] = fma(p1, vv.y, acc[1][2u * c + 1u]); + } + } + sync_shared(); + if (step_index > 0u) { + store_kv(rk, rv); + sync_shared(); + } + } + + if (spart == 0u) row_inv[srow] = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + sync_shared(); + [[unroll]] for (uint i = 0u; i < 2u; ++i) { + const uint q_local = q_first + 2u * pr + i; + if (q_local >= pc.seqlen_q) continue; + const float inv = row_inv[2u * pr + i]; + U32s orow = U32s(pc.o + ((uint64_t(pc.q_row0) + q_local) * pc.ldo + uint64_t(head) * pc.hso) * 2ul); + [[unroll]] for (uint c = 0u; c < 8u; ++c) { + precise const float lo = acc[i][2u * c] * inv, hi = acc[i][2u * c + 1u] * inv; + orow.d[8u * quad + c] = bf(lo) | (bf(hi) << 16); + } + } +} diff --git a/vulkan/shaders/common.glsl b/vulkan/shaders/common.glsl new file mode 100644 index 0000000..6b00802 --- /dev/null +++ b/vulkan/shaders/common.glsl @@ -0,0 +1,64 @@ +// bf16 travels as its bits in 32-bit words, so these shaders need neither 16-bit storage nor 16-bit arithmetic. As in the +// engine's kernels every step is one fp32 operation rounded once, so results are declared precise: contraction into a +// fused multiply-add would change which way a value rounds. +#ifndef MI_VK_COMMON_GLSL +#define MI_VK_COMMON_GLSL + +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#extension GL_EXT_buffer_reference : require + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer Words { uvec2 d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 4) buffer Floats { float d[]; }; + +vec4 f4(uvec2 w) { + return uintBitsToFloat(uvec4(w.x << 16, w.x & 0xffff0000u, w.y << 16, w.y & 0xffff0000u)); +} + +// fp32 to bf16 bits, round to nearest even. A NaN stays a NaN: the rounding carry could otherwise wrap it to -0. +uint bf(float v) { + const uint u = floatBitsToUint(v); + if ((u & 0x7fffffffu) > 0x7f800000u) return (u >> 16) | 0x0040u; + return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16; +} +uvec2 bf4(vec4 v) { return uvec2(bf(v.x) | (bf(v.y) << 16), bf(v.z) | (bf(v.w) << 16)); } +vec4 rbf4(vec4 v) { return f4(bf4(v)); } + +// 2^e as an fp32, exact for e in [-126, 127] (exp2 need not be). +float pow2(int e) { return uintBitsToFloat(uint(127 + e) << 23); } + +// Eight bf16 (a uvec4 of pairs) times `scale`, a power of two, as eight f16 in the same packing. Exact while each product +// is an f16 normal or zero; the coop-f16 kernels pick the scales so that the large values are. +uvec4 bf8_to_f16(uvec4 v, float scale) { + uvec4 r; + for (uint j = 0u; j < 4u; ++j) { + r[j] = packHalf2x16(vec2(uintBitsToFloat(v[j] << 16), uintBitsToFloat(v[j] & 0xffff0000u)) * scale); + } + return r; +} + +// exp(-2 |x|) stays in range where GLSL's tanh may divide infinities; beyond 12 the result rounds to 1 in fp32 anyway. +// Below 1/16 the subtraction 1 - e cancels (to 0 for |x| under 2^-25), so there the series x - x^3/3 + 2x^5/15 takes +// over, whose first omitted term is under 2^-27 of the result. +vec4 tanh_safe(vec4 x) { + const vec4 e = exp(-2.0 * min(abs(x), vec4(12.0))); + precise const vec4 t = (1.0 - e) / (1.0 + e); + precise const vec4 x2 = x * x; + precise const vec4 series = x * (1.0 + x2 * (-1.0 / 3.0 + x2 * (2.0 / 15.0))); + return mix(sign(x) * t, series, lessThan(abs(x), vec4(0.0625))); +} + +vec4 situ(vec4 h, vec4 g) { + const float inv25 = 1.0 / 25.0; // torch divides by a CPU scalar as a multiply by its fp32 reciprocal + precise const vec4 g_quarter = g * 0.25; + const vec4 tanh_g = rbf4(tanh_safe(rbf4(g_quarter))); + precise const vec4 four_tanh = 4.0 * tanh_g; + precise const vec4 sigmoid = 1.0 / (1.0 + exp(-g)); + precise const vec4 left = rbf4(four_tanh) * rbf4(sigmoid); + precise const vec4 h_scaled = h * inv25; + const vec4 tanh_h = rbf4(tanh_safe(rbf4(h_scaled))); + precise const vec4 right = 25.0 * tanh_h; + precise const vec4 product = rbf4(left) * rbf4(right); + return product; +} + +#endif diff --git a/vulkan/shaders/elementwise.comp b/vulkan/shaders/elementwise.comp new file mode 100644 index 0000000..e35d985 --- /dev/null +++ b/vulkan/shaders/elementwise.comp @@ -0,0 +1,30 @@ +// Row-local bf16 ops the DiT needs besides the block's, n elements, n a multiple of 4. MODE: +// 0: out = bf16(silu(a)), silu(x) = x / (1 + exp(-x)) in fp32 (torch's SiLU) +// 1: out = bf16(situ(a, b)), the text refiner's SiTU over separate hidden (a) and gate (b) tensors +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const int MODE = 0; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint n; +} pc; + +void main() { + // A 2-D grid where one dimension would pass 65535 workgroups (the guaranteed limit). + const uint i = gl_GlobalInvocationID.y * (gl_NumWorkGroups.x * 64u) + gl_GlobalInvocationID.x; + if (i * 4u >= pc.n) return; + const vec4 a = f4(Words(pc.a).d[i]); + if (MODE == 0) { + precise const vec4 denominator = 1.0 + exp(-a); + precise const vec4 y = a / denominator; + Words(pc.out_).d[i] = bf4(y); + } else { + Words(pc.out_).d[i] = bf4(situ(a, f4(Words(pc.b).d[i]))); + } +} diff --git a/vulkan/shaders/gemm_bf16_coop.comp b/vulkan/shaders/gemm_bf16_coop.comp new file mode 100644 index 0000000..7ca452e --- /dev/null +++ b/vulkan/shaders/gemm_bf16_coop.comp @@ -0,0 +1,259 @@ +// out[M,N] = A[M,K] * B[N,K]^T on bf16 cooperative matrices with fp32 accumulators: the DiT's linears on devices with +// VK_KHR_shader_bfloat16 fragments. A and B are row-major bf16; lda, ldb and ldo are multiples of 8, and N and K are +// multiples of 8. M is free. MODE, as in gemm_plain.comp: +// 0: out is f32 +// 1: out is bf16 +// 2: out[m,n] += bf16(bf16(ab) * gate[n]), gated residual, bf16 in place +// 3: out[m,n] = situ(bf16(A * B[n]), bf16(A * B[N + n])), SiTU-GLU over the stored [hidden | gate] weight, bf16 +// 4: out[m,n] = bf16(ab + bias[n]), nn.Linear with a bias (bias in `gate`), bf16 +// 5: out[m,n] += bf16(ab + bias[n]), the same then a residual add, bf16 in place +// +// One workgroup per 128x128 tile of the product: 256 invocations, 8 full subgroups of 32, which the pipeline must require. +// Subgroup s owns rows 32 (s % 4) .. +32 and the columns 32 (s / 4) .. +32 of each half of the tile, so in MODE 3, where the +// first half holds 64 hidden columns and the second the same 64 gate columns, each subgroup owns both factors of its +// outputs. K walks in steps of 32 through two shared-memory stages: the next step's tiles load into registers while the +// current one multiplies. Rows of the stages are padded to 80 bytes, so 8 rows of 16 bytes fall in distinct banks. +// +// Built with F16 (gemm_f16_coop.spv) for devices with f16 fragments only: A and B are still bf16 in memory and become f16 +// as they are staged. Row m of A is scaled by 2^-a_exponents[m] and row j of B by 2^-b_exponents[j] (row_scale.comp, which +// puts each row's largest value in [2^15, 2^16)), and the accumulators are scaled back by both before the epilogue, so +// every product is the bf16 path's exactly while the scaled operands are f16 normals (values above 2^-29 of their row's +// largest). The scaling back is one combined exponent, applied in two exact halves (scaled_back). +#version 460 +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#ifdef F16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#define ELEM float16_t +#else +#extension GL_EXT_bfloat16 : require +#define ELEM bfloat16_t +#endif +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const int MODE = 1; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { uvec4 d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 4) readonly buffer Ints { int d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 16) readonly buffer Int4s { ivec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint64_t gate; + uint M; + uint N; + uint K; + uint lda; + uint ldb; + uint ldo; + uint64_t a_exponents; // F16 only + uint64_t b_exponents; // F16 only +} pc; + +const uint BM = 128u; +const uint BN = 128u; // rows of B per tile; MODE 3 covers BN / 2 output columns +const uint BK = 32u; +const uint ROW = 5u; // uvec4 per staged row: 4 of data, 1 of padding +const uint STAGE = BM * ROW; // uvec4 per staged operand +const uint BAND = 8u; // output tiles walk bands of N tiles, so a band of B stays in cache while every M tile passes +const uint STAGING = 64u; // uvec4 per 16x16 fp32 fragment in the epilogue + +// Stage s of A at s * 2 * STAGE, of B at (s * 2 + 1) * STAGE. The epilogue reuses the start: 2 fragments per subgroup. +shared uvec4 smem[4u * STAGE]; + +#define FRAG_A coopmat +#define FRAG_B coopmat +#define ACC coopmat + +uint b_row(uint n_base, uint r) { + if (MODE != 3) return n_base + r; + return r < BN / 2u ? n_base + r : pc.N + n_base + (r - BN / 2u); +} +bool b_valid(uint n_base, uint r) { return (MODE == 3 ? n_base + r % (BN / 2u) : n_base + r) < pc.N; } + +// F16: 2^-exponent of the A and B rows this invocation loads. +float a_scale[2], b_scale[2]; + +// Invocation i loads chunk i % 4 of rows i / 4 and i / 4 + 64 of the A and B tiles at k0. +void load_tiles(uint m_base, uint n_base, uint k0, out uvec4 ra[2], out uvec4 rb[2]) { + const uint i = gl_LocalInvocationIndex; + const uint chunk = i % 4u; + const uint k = k0 + chunk * 8u; + const bool k_ok = k < pc.K; + [[unroll]] for (uint h = 0u; h < 2u; ++h) { + const uint r = i / 4u + h * 64u; + const uint m = m_base + r; + ra[h] = uvec4(0u); + if (k_ok && m < pc.M) ra[h] = Vec4s(pc.a + (uint64_t(m) * pc.lda + k) * 2ul).d[0]; + rb[h] = uvec4(0u); + if (k_ok && b_valid(n_base, r)) rb[h] = Vec4s(pc.b + (uint64_t(b_row(n_base, r)) * pc.ldb + k) * 2ul).d[0]; +#ifdef F16 + ra[h] = bf8_to_f16(ra[h], a_scale[h]); + rb[h] = bf8_to_f16(rb[h], b_scale[h]); +#endif + } +} + +// F16: the exponents that undo A's scaling in output row m and B's in output columns j .. j + 3 (j a multiple of 4 below +// the row count of B). +int row_exponent(uint m) { +#ifdef F16 + return m < pc.M ? Ints(pc.a_exponents).d[m] : 0; +#else + return 0; +#endif +} +ivec4 col_exponents(uint j) { +#ifdef F16 + return Int4s(pc.b_exponents + uint64_t(j) * 4ul).d[0]; +#else + return ivec4(0); +#endif +} + +// v * 2^e, as two multiplies by 2^(e/2) and 2^(e - e/2): e lies in [-252, 224] and each half in fp32's range, and the +// intermediate lies between v and the result, so it overflows or underflows only where the result does. Exact. +vec4 scaled_back(vec4 v, ivec4 e) { +#ifdef F16 + const ivec4 h = e / 2; + precise const vec4 once = v * uintBitsToFloat(uvec4(127 + h) << 23); + precise const vec4 twice = once * uintBitsToFloat(uvec4(127 + (e - h)) << 23); + return twice; +#else + return v; +#endif +} + +void store_tiles(uint stage, const uvec4 ra[2], const uvec4 rb[2]) { + const uint i = gl_LocalInvocationIndex; + [[unroll]] for (uint h = 0u; h < 2u; ++h) { + const uint at = (i / 4u + h * 64u) * ROW + i % 4u; + smem[stage * 2u * STAGE + at] = ra[h]; + smem[(stage * 2u + 1u) * STAGE + at] = rb[h]; + } +} + +// Within a subgroup only: the staging fragments are private to it. +void sync_subgroup() { + subgroupMemoryBarrierShared(); + subgroupBarrier(); +} + +vec4 staged(uint base, uint row, uint col4) { return uintBitsToFloat(smem[base + row * 4u + col4]); } + +void main() { + const uint num_m = (pc.M + BM - 1u) / BM; + const uint cols_per_tile = MODE == 3 ? BN / 2u : BN; + const uint num_n = (pc.N + cols_per_tile - 1u) / cols_per_tile; + const uint id = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; + const uint per_band = num_m * BAND; + const uint band = id / per_band; + const uint pos = id - band * per_band; + const uint width = min(BAND, num_n - band * BAND); + const uint m_base = (pos / width) * BM; + const uint n_base = (band * BAND + pos % width) * cols_per_tile; + + const uint sg = gl_SubgroupID; + const uint row0 = (sg % 4u) * 32u; + const uint col0 = (sg / 4u) * 32u; + + ACC acc[2][4]; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) acc[r][c] = ACC(0.0); + } + + [[unroll]] for (uint h = 0u; h < 2u; ++h) { + const uint r = gl_LocalInvocationIndex / 4u + h * 64u; + const uint m = m_base + r; +#ifdef F16 + a_scale[h] = m < pc.M ? pow2(-Ints(pc.a_exponents).d[m]) : 1.0; + b_scale[h] = b_valid(n_base, r) ? pow2(-Ints(pc.b_exponents).d[b_row(n_base, r)]) : 1.0; +#else + a_scale[h] = 1.0; + b_scale[h] = 1.0; +#endif + } + + const uint steps = (pc.K + BK - 1u) / BK; + uvec4 ra[2], rb[2]; + load_tiles(m_base, n_base, 0u, ra, rb); + store_tiles(0u, ra, rb); + barrier(); + for (uint step = 0u; step < steps; ++step) { + const bool more = step + 1u < steps; + if (more) load_tiles(m_base, n_base, (step + 1u) * BK, ra, rb); + const uint a_base = (step & 1u) * 2u * STAGE, b_base = a_base + STAGE; + [[unroll]] for (uint kk = 0u; kk < BK / 16u; ++kk) { + FRAG_A fa[2]; + FRAG_B fb[4]; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + coopMatLoad(fa[r], smem, a_base + (row0 + r * 16u) * ROW + kk * 2u, ROW, gl_CooperativeMatrixLayoutRowMajor); + } + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + const uint col = (c / 2u) * (BN / 2u) + col0 + (c % 2u) * 16u; + coopMatLoad(fb[c], smem, b_base + col * ROW + kk * 2u, ROW, gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) acc[r][c] = coopMatMulAdd(fa[r], fb[c], acc[r][c]); + } + } + if (more) store_tiles((step + 1u) & 1u, ra, rb); + barrier(); + } + + // Epilogue: each 16x16 fragment goes through the subgroup's staging, then lane l writes 8 outputs of row l / 2. + const uint base0 = sg * 2u * STAGING, base1 = base0 + STAGING; + const uint lane = gl_SubgroupInvocationID; + const uint frow = lane / 2u, fcol = (lane % 2u) * 8u; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + const uint m = m_base + row0 + r * 16u + frow; + const int e_row = row_exponent(m); + [[unroll]] for (uint c = 0u; c < (MODE == 3 ? 2u : 4u); ++c) { + coopMatStore(acc[r][c], smem, base0, 4u, gl_CooperativeMatrixLayoutRowMajor); + if (MODE == 3) coopMatStore(acc[r][c + 2u], smem, base1, 4u, gl_CooperativeMatrixLayoutRowMajor); + sync_subgroup(); + const uint n = n_base + (MODE == 3 ? 0u : (c / 2u) * (BN / 2u)) + col0 + (c % 2u) * 16u + fcol; + if (m < pc.M && n < pc.N) { + const vec4 lo = scaled_back(staged(base0, frow, fcol / 4u), e_row + col_exponents(n)); + const vec4 hi = scaled_back(staged(base0, frow, fcol / 4u + 1u), e_row + col_exponents(n + 4u)); + const uint64_t at = uint64_t(m) * pc.ldo + n; + if (MODE == 0) { + Vec4s(pc.out_ + at * 4ul).d[0] = floatBitsToUint(lo); + Vec4s(pc.out_ + at * 4ul + 16ul).d[0] = floatBitsToUint(hi); + } else if (MODE == 1) { + Vec4s(pc.out_ + at * 2ul).d[0] = uvec4(bf4(lo), bf4(hi)); + } else if (MODE == 2) { + Vec4s target = Vec4s(pc.out_ + at * 2ul); + const uvec4 old = target.d[0]; + const uvec4 g = Vec4s(pc.gate + uint64_t(n) * 2ul).d[0]; + precise const vec4 gated_lo = rbf4(lo) * f4(g.xy), gated_hi = rbf4(hi) * f4(g.zw); + precise const vec4 sum_lo = f4(old.xy) + rbf4(gated_lo), sum_hi = f4(old.zw) + rbf4(gated_hi); + target.d[0] = uvec4(bf4(sum_lo), bf4(sum_hi)); + } else if (MODE == 3) { + const vec4 g_lo = scaled_back(staged(base1, frow, fcol / 4u), e_row + col_exponents(pc.N + n)); + const vec4 g_hi = scaled_back(staged(base1, frow, fcol / 4u + 1u), e_row + col_exponents(pc.N + n + 4u)); + Vec4s(pc.out_ + at * 2ul).d[0] = + uvec4(bf4(situ(rbf4(lo), rbf4(g_lo))), bf4(situ(rbf4(hi), rbf4(g_hi)))); + } else { + const uvec4 bias = Vec4s(pc.gate + uint64_t(n) * 2ul).d[0]; + precise const vec4 biased_lo = lo + f4(bias.xy), biased_hi = hi + f4(bias.zw); + Vec4s target = Vec4s(pc.out_ + at * 2ul); + if (MODE == 4) { + target.d[0] = uvec4(bf4(biased_lo), bf4(biased_hi)); + } else { + const uvec4 old = target.d[0]; + precise const vec4 sum_lo = f4(old.xy) + rbf4(biased_lo), sum_hi = f4(old.zw) + rbf4(biased_hi); + target.d[0] = uvec4(bf4(sum_lo), bf4(sum_hi)); + } + } + } + sync_subgroup(); + } + } +} diff --git a/vulkan/shaders/gemm_coop.comp b/vulkan/shaders/gemm_coop.comp new file mode 100644 index 0000000..dbb17d8 --- /dev/null +++ b/vulkan/shaders/gemm_coop.comp @@ -0,0 +1,120 @@ +// C[M,N] = A[M,K] * B[N,K]^T on cooperative matrices: f16 inputs, f32 accumulate, f32 output. A and B are row-major with +// strides lda and ldb, both multiples of 4; C has stride ldc. M and N are free, K is a multiple of 16. +// +// One workgroup per 64x64 output tile: 128 invocations, exactly 4 full subgroups of 32, which the pipeline must require. +// Each subgroup owns a 32x32 quarter. A and B tiles of 32 K are staged in shared memory, so each element is read from +// global memory once per workgroup. Anything past M or N stages as zero and is never written, so no padding is needed. +#version 460 +#extension GL_KHR_cooperative_matrix : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#extension GL_EXT_buffer_reference : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_memory_scope_semantics : require + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 8) readonly buffer F16Vec { f16vec4 d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 4) writeonly buffer F32 { float d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t c; + uint M; + uint N; + uint K; + uint lda; + uint ldb; + uint ldc; +} pc; + +const uint BM = 64u; +const uint BN = 64u; +const uint BK = 32u; +const uint BAND = 8u; // output tiles walk bands of N tiles, so a band of B stays in cache while every M tile passes +const uint INVOCATIONS = 128u; + +shared float16_t sa[BM * BK]; +shared float16_t sb[BN * BK]; +shared float sc[BM * BN]; + +void sync_shared() { + controlBarrier(gl_ScopeWorkgroup, gl_ScopeWorkgroup, gl_StorageSemanticsShared, + gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible); +} + +void main() { + const uint num_m = (pc.M + BM - 1u) / BM; + const uint num_n = (pc.N + BN - 1u) / BN; + const uint id = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; + const uint per_band = num_m * BAND; + const uint band = id / per_band; + const uint pos = id - band * per_band; + const uint width = min(BAND, num_n - band * BAND); + const uint m_base = (pos / width) * BM; + const uint n_base = (band * BAND + pos % width) * BN; + + const uint invocation = gl_LocalInvocationIndex; + const uint row0 = (gl_SubgroupID >> 1) * 32u; + const uint col0 = (gl_SubgroupID & 1u) * 32u; + + coopmat acc00, acc01, acc10, acc11; + const float zero = 0.0; + acc00 = coopmat(zero); + acc01 = acc00; + acc10 = acc00; + acc11 = acc00; + + for (uint k0 = 0u; k0 < pc.K; k0 += BK) { + const uint fragments = min(BK, pc.K - k0) / 16u; + const uint vectors = fragments * 4u; + const f16vec4 none = f16vec4(0.0hf); + for (uint i = invocation; i < BM * BK / 4u; i += INVOCATIONS) { + const uint line = i / 8u; + const uint element = i % 8u; + const uint at = line * BK + element * 4u; + const uint m = m_base + line; + f16vec4 v = none; + if (m < pc.M && element < vectors) v = F16Vec(pc.a + (uint64_t(m) * pc.lda + k0) * 2ul).d[element]; + sa[at] = v.x; + sa[at + 1u] = v.y; + sa[at + 2u] = v.z; + sa[at + 3u] = v.w; + const uint n = n_base + line; + f16vec4 w = none; + if (n < pc.N && element < vectors) w = F16Vec(pc.b + (uint64_t(n) * pc.ldb + k0) * 2ul).d[element]; + sb[at] = w.x; + sb[at + 1u] = w.y; + sb[at + 2u] = w.z; + sb[at + 3u] = w.w; + } + sync_shared(); + for (uint f = 0u; f < fragments; ++f) { + const uint k = f * 16u; + coopmat fa0, fa1; + coopmat fb0, fb1; + coopMatLoad(fa0, sa, row0 * BK + k, BK, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(fa1, sa, (row0 + 16u) * BK + k, BK, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(fb0, sb, col0 * BK + k, BK, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(fb1, sb, (col0 + 16u) * BK + k, BK, gl_CooperativeMatrixLayoutColumnMajor); + acc00 = coopMatMulAdd(fa0, fb0, acc00); + acc01 = coopMatMulAdd(fa0, fb1, acc01); + acc10 = coopMatMulAdd(fa1, fb0, acc10); + acc11 = coopMatMulAdd(fa1, fb1, acc11); + } + sync_shared(); + } + + coopMatStore(acc00, sc, row0 * BN + col0, BN, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc01, sc, row0 * BN + col0 + 16u, BN, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc10, sc, (row0 + 16u) * BN + col0, BN, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc11, sc, (row0 + 16u) * BN + col0 + 16u, BN, gl_CooperativeMatrixLayoutRowMajor); + sync_shared(); + for (uint i = invocation; i < BM * BN; i += INVOCATIONS) { + const uint m = m_base + i / BN; + const uint n = n_base + i % BN; + if (m < pc.M && n < pc.N) F32(pc.c + (uint64_t(m) * pc.ldc + n) * 4ul).d[0] = sc[i]; + } +} diff --git a/vulkan/shaders/gemm_plain.comp b/vulkan/shaders/gemm_plain.comp new file mode 100644 index 0000000..f344fff --- /dev/null +++ b/vulkan/shaders/gemm_plain.comp @@ -0,0 +1,156 @@ +// out[M,N] = A[M,K] * B[N,K]^T without cooperative matrices, for any device with buffer device addresses and 64-bit +// integers: fp32 fused multiply-adds with K ascending, inputs f16 or bf16. A and B are row-major with strides lda and ldb, +// both multiples of 4. M is free; N and K are multiples of 4. It does not depend on the subgroup size. +// +// One workgroup of 256 invocations per 128x128 tile of the product. Invocation (tx, ty) = (i % 16, i / 16) owns rows +// 4 ty .. +4 and 64 + 4 ty .. +4 and columns 4 tx .. +4 and 64 + 4 tx .. +4 of the tile, so in MODE 3, where the first half +// of the tile's B rows holds 64 hidden columns and the second the same 64 gate columns, each invocation owns both factors +// of its outputs. K walks in steps of 8 through two shared-memory stages of fp32 (17 KiB): the next step loads into +// registers while the current one multiplies. MODE: +// 0: out is f32 +// 1: out is bf16 +// 2: out[m,n] += bf16(bf16(ab) * gate[n]), gated residual, bf16 in place +// 3: out[m,n] = situ(bf16(A * B[n]), bf16(A * B[N + n])), SiTU-GLU over the stored [hidden | gate] weight, bf16 +// 4: out[m,n] = bf16(ab + bias[n]), nn.Linear with a bias (bias in `gate`), bf16 +// 5: out[m,n] += bf16(ab + bias[n]), the same then a residual add, bf16 in place +#version 460 +#extension GL_GOOGLE_include_directive : require +#extension GL_EXT_control_flow_attributes : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const bool INPUT_BF16 = true; +layout(constant_id = 1) const int MODE = 0; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint64_t gate; + uint M; + uint N; + uint K; + uint lda; + uint ldb; + uint ldo; +} pc; + +const uint BM = 128u; +const uint BN = 128u; // rows of B per tile; MODE 3 covers BN / 2 output columns +const uint BK = 8u; +const uint LINE = 132u; // floats per staged k: 128 and 4 of padding +const uint STAGE = BK * LINE; + +// [stage][k][m] for A and [stage][k][n] for B: k-major, so 4 rows or columns at one k are 4 adjacent floats. +shared float sa[2u * STAGE]; +shared float sb[2u * STAGE]; + +vec4 decode(uvec2 w) { return INPUT_BF16 ? f4(w) : vec4(unpackHalf2x16(w.x), unpackHalf2x16(w.y)); } + +uint b_row(uint n_base, uint r) { + if (MODE != 3) return n_base + r; + return r < BN / 2u ? n_base + r : pc.N + n_base + (r - BN / 2u); +} +bool b_valid(uint n_base, uint r) { return (MODE == 3 ? n_base + r % (BN / 2u) : n_base + r) < pc.N; } + +// Invocation i loads 4 elements, k = k0 + 4 (i % 2) .. +4, of row i / 2 of the A tile and of the B tile. +void load(uint m_base, uint n_base, uint k0, out vec4 ra, out vec4 rb) { + const uint i = gl_LocalInvocationIndex; + const uint r = i / 2u, k = k0 + (i % 2u) * 4u; + ra = vec4(0.0); + rb = vec4(0.0); + if (k < pc.K) { + const uint m = m_base + r; + if (m < pc.M) ra = decode(Words(pc.a + (uint64_t(m) * pc.lda + k) * 2ul).d[0]); + if (b_valid(n_base, r)) rb = decode(Words(pc.b + (uint64_t(b_row(n_base, r)) * pc.ldb + k) * 2ul).d[0]); + } +} + +void store(uint stage, vec4 ra, vec4 rb) { + const uint i = gl_LocalInvocationIndex; + const uint r = i / 2u, k = (i % 2u) * 4u; + // Element (k + j, r): a transposing store, conflict-free since the two k halves of a subgroup sit 16 banks apart. + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint f = stage * STAGE + (k + j) * LINE + r; + sa[f] = ra[j]; + sb[f] = rb[j]; + } +} + +void main() { + const uint i = gl_LocalInvocationIndex; + const uint tx = i % 16u, ty = i / 16u; + const uint cols_per_tile = MODE == 3 ? BN / 2u : BN; + const uint m_base = gl_WorkGroupID.y * BM, n_base = gl_WorkGroupID.x * cols_per_tile; + + vec4 acc[8][2]; // [row r: 4 ty + r % 4 + 64 (r / 4)][column group g: 4 tx + 64 g] + [[unroll]] for (uint r = 0u; r < 8u; ++r) { + acc[r][0] = vec4(0.0); + acc[r][1] = vec4(0.0); + } + + const uint steps = (pc.K + BK - 1u) / BK; + vec4 ra, rb; + load(m_base, n_base, 0u, ra, rb); + store(0u, ra, rb); + barrier(); + for (uint step = 0u; step < steps; ++step) { + const bool more = step + 1u < steps; + if (more) load(m_base, n_base, (step + 1u) * BK, ra, rb); + const uint base = (step & 1u) * STAGE; + [[unroll]] for (uint k = 0u; k < BK; ++k) { + const uint line = base + k * LINE; + const uint ar = line + 4u * ty, bc = line + 4u * tx; + const vec4 a_lo = vec4(sa[ar], sa[ar + 1u], sa[ar + 2u], sa[ar + 3u]); + const vec4 a_hi = vec4(sa[ar + 64u], sa[ar + 65u], sa[ar + 66u], sa[ar + 67u]); + const vec4 b_lo = vec4(sb[bc], sb[bc + 1u], sb[bc + 2u], sb[bc + 3u]); + const vec4 b_hi = vec4(sb[bc + 64u], sb[bc + 65u], sb[bc + 66u], sb[bc + 67u]); + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + acc[r][0] = fma(vec4(a_lo[r]), b_lo, acc[r][0]); + acc[r][1] = fma(vec4(a_lo[r]), b_hi, acc[r][1]); + acc[r + 4u][0] = fma(vec4(a_hi[r]), b_lo, acc[r + 4u][0]); + acc[r + 4u][1] = fma(vec4(a_hi[r]), b_hi, acc[r + 4u][1]); + } + } + if (more) store((step + 1u) & 1u, ra, rb); + barrier(); + } + + [[unroll]] for (uint r = 0u; r < 8u; ++r) { + const uint m = m_base + 4u * ty + r % 4u + 64u * (r / 4u); + if (m >= pc.M) continue; + [[unroll]] for (uint g = 0u; g < (MODE == 3 ? 1u : 2u); ++g) { + const uint n = n_base + 4u * tx + 64u * g; + if (n >= pc.N) continue; + const vec4 v = acc[r][g]; + const uint64_t at = uint64_t(m) * pc.ldo + n; + if (MODE == 0) { + Floats out_f = Floats(pc.out_ + at * 4ul); + out_f.d[0] = v.x; + out_f.d[1] = v.y; + out_f.d[2] = v.z; + out_f.d[3] = v.w; + continue; + } + Words out_w = Words(pc.out_ + at * 2ul); + if (MODE == 1) { + out_w.d[0] = bf4(v); + } else if (MODE == 2) { + const vec4 gate = f4(Words(pc.gate + uint64_t(n) * 2ul).d[0]); + precise const vec4 gated = rbf4(v) * gate; + precise const vec4 sum = f4(out_w.d[0]) + rbf4(gated); + out_w.d[0] = bf4(sum); + } else if (MODE == 3) { + out_w.d[0] = bf4(situ(rbf4(v), rbf4(acc[r][1]))); + } else { + precise const vec4 biased = v + f4(Words(pc.gate + uint64_t(n) * 2ul).d[0]); + if (MODE == 4) { + out_w.d[0] = bf4(biased); + } else { + precise const vec4 sum = f4(out_w.d[0]) + rbf4(biased); + out_w.d[0] = bf4(sum); + } + } + } + } +} diff --git a/vulkan/shaders/norm_modulate.comp b/vulkan/shaders/norm_modulate.comp new file mode 100644 index 0000000..32dd1d3 --- /dev/null +++ b/vulkan/shaders/norm_modulate.comp @@ -0,0 +1,68 @@ +// out[rows,d] = bf16(bf16(bf16(bf16(rms_norm(x)) * w) * bf16(1 + scale)) + shift); w, scale, shift: [d]. With a null scale +// address, out = bf16(bf16(rms_norm(x)) * w). x and out are +// bf16 with strides that are multiples of 4, and d is a multiple of 4. One row per workgroup. The fp32 mean of squares is +// summed in the order of torch's reduce kernel: 32 lanes of 4 running values, then a halving tree. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t weight; + uint64_t scale; + uint64_t shift; + uint64_t out_; + uint ldx; + uint ldo; + uint d; + float factor; +} pc; + +shared float red[32]; + +// In Vulkan GLSL barrier() also makes shared memory available and visible, at workgroup scope. memoryBarrierShared() would +// add device scope, which needs a feature this shader otherwise has no use for. +void sync_shared() { barrier(); } + +void main() { + const uint t = gl_LocalInvocationIndex; + const uint64_t row = gl_WorkGroupID.x; + Words X = Words(pc.x + row * pc.ldx * 2ul); + const uint groups = pc.d / 4u; + + precise vec4 acc = vec4(0.0); + for (uint g = t; g < groups; g += 32u) { + const vec4 v = f4(X.d[g]); + acc = acc + v * v; + } + precise const float total = acc.x + acc.y + acc.z + acc.w; + red[t] = total; + sync_shared(); + for (uint off = 16u; off > 0u; off >>= 1) { + if (t < off) { + precise const float sum = red[t] + red[t + off]; + red[t] = sum; + } + sync_shared(); + } + precise const float mean = red[0] * pc.factor; + precise const float mean_eps = mean + 1e-6; + const float r = inversesqrt(mean_eps); + + Words O = Words(pc.out_ + row * pc.ldo * 2ul); + Words W = Words(pc.weight), S = Words(pc.scale), H = Words(pc.shift); + for (uint g = t; g < groups; g += 32u) { + precise const vec4 normed = f4(X.d[g]) * r; + precise const vec4 y = rbf4(normed) * f4(W.d[g]); + if (pc.scale == 0ul) { // rms_norm_moss alone, as the text refiner's norms + O.d[g] = bf4(y); + continue; + } + precise const vec4 one_plus = 1.0 + f4(S.d[g]); + precise const vec4 scaled = rbf4(y) * rbf4(one_plus); + precise const vec4 shifted = rbf4(scaled) + f4(H.d[g]); + O.d[g] = bf4(shifted); + } +} diff --git a/vulkan/shaders/qk_norm_rope.comp b/vulkan/shaders/qk_norm_rope.comp new file mode 100644 index 0000000..07a9a9e --- /dev/null +++ b/vulkan/shaders/qk_norm_rope.comp @@ -0,0 +1,68 @@ +// In place, per head: q and k = rope(rms_norm(.)). qkv: [tokens, >= 2 * heads * 128] bf16 with q heads then k heads; +// q_norm, k_norm: [128]; table: [tokens, 128] RoPE rows, [sin(64) | cos(64)]. One (token, q-or-k head) per workgroup: lane l +// holds dims 4l..4l+3, and RoPE pairs dim d < 64 with d + 64 through shared memory. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t qkv; + uint64_t q_norm; + uint64_t k_norm; + uint64_t table; + uint ld; + uint table_ld; + uint heads; + float factor; +} pc; + +shared float red[32]; +shared vec4 norm_shared[32]; + +// In Vulkan GLSL barrier() also makes shared memory available and visible, at workgroup scope. memoryBarrierShared() would +// add device scope, which needs a feature this shader otherwise has no use for. +void sync_shared() { barrier(); } + +void main() { + const uint lane = gl_LocalInvocationIndex; + const uint head = gl_WorkGroupID.x; + const uint64_t token = gl_WorkGroupID.y; + Words row = Words(pc.qkv + (token * pc.ld + uint64_t(head) * 128ul) * 2ul); + Words weight = Words(head < pc.heads ? pc.q_norm : pc.k_norm); + + const vec4 x = f4(row.d[lane]); + precise const vec4 squares = x * x; + precise const float total = squares.x + squares.y + squares.z + squares.w; + red[lane] = total; + sync_shared(); + for (uint off = 16u; off > 0u; off >>= 1) { + if (lane < off) { + precise const float sum = red[lane] + red[lane + off]; + red[lane] = sum; + } + sync_shared(); + } + precise const float mean = red[0] * pc.factor; + precise const float mean_eps = mean + 1e-6; + const float r = inversesqrt(mean_eps); + + precise const vec4 normed = x * r; + precise const vec4 weighted = rbf4(normed) * f4(weight.d[lane]); + const vec4 y = rbf4(weighted); + norm_shared[lane] = y; + sync_shared(); + + const vec4 other = norm_shared[lane ^ 16u]; + Words tab = Words(pc.table + token * pc.table_ld * 2ul); + const uint pair = lane & 15u; + const vec4 si = f4(tab.d[pair]), co = f4(tab.d[pair + 16u]); + const bool lo_half = lane < 16u; + const vec4 lo = lo_half ? y : other, hi = lo_half ? other : y; + // moss_rope_apply's rotate_pair: lo' = lo*cos + (-hi)*sin, hi' = hi*cos + lo*sin, each op rounded. + precise const vec4 a = (lo_half ? lo : hi) * co; + precise const vec4 b = lo_half ? -hi * si : lo * si; + precise const vec4 rotated = rbf4(a) + rbf4(b); + row.d[lane] = bf4(rotated); +} diff --git a/vulkan/shaders/row_scale.comp b/vulkan/shaders/row_scale.comp new file mode 100644 index 0000000..fb5c7d8 --- /dev/null +++ b/vulkan/shaders/row_scale.comp @@ -0,0 +1,57 @@ +// The power-of-two exponent of each bf16 GEMM input row for the coop-f16 kernels: exponents[m] = e such that the row's +// largest magnitude times 2^-e lies in [2^15, 2^16), where every bf16 value is an f16 value (the largest bf16 below 2^16 +// is 65280, f16's largest is 65504), from the row's finite values (an infinity or NaN converts as itself); zero for a +// row without a finite nonzero value. In [-126, 112], so 2^-e and the GEMM's two halves of a scaling back are fp32 values. +// K is a multiple of 8, lda of 8. +// +// A workgroup of 256 invocations takes 8 rows, 32 invocations per row; no subgroup operations. +#version 460 +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#extension GL_EXT_buffer_reference : require +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 16) readonly buffer Vec4s { uvec4 d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 4) writeonly buffer Ints { int d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t exponents; + uint M; + uint K; + uint lda; +} pc; + +shared uint largest[256]; + +// The larger finite bf16 magnitude (its bits without the sign, which order as the magnitudes do) of the two in a word. +uint finite(uint mag) { return mag >= 0x7f80u ? 0u : mag; } +uint word_max(uint w) { return max(finite(w & 0x7fffu), finite((w >> 16) & 0x7fffu)); } + +void main() { + const uint i = gl_LocalInvocationIndex; + const uint row = (gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x) * 8u + i / 32u; + const uint lane = i % 32u; + uint m = 0u; + if (row < pc.M) { + Vec4s r = Vec4s(pc.a + uint64_t(row) * pc.lda * 2ul); + for (uint c = lane; c < pc.K / 8u; c += 32u) { + const uvec4 v = r.d[c]; + m = max(m, max(max(word_max(v.x), word_max(v.y)), max(word_max(v.z), word_max(v.w)))); + } + } + largest[i] = m; + barrier(); + [[unroll]] for (uint width = 16u; width > 0u; width /= 2u) { + if (lane < width) largest[i] = max(largest[i], largest[i + width]); + barrier(); + } + if (lane == 0u && row < pc.M) { + const uint top = largest[i]; + const int biased = int(top >> 7); // bf16 exponent field + int e = 0; + if (top != 0u) e = max(biased - 127 - 15, -126); + Ints(pc.exponents).d[row] = e; + } +} diff --git a/vulkan/shaders/te_attention.comp b/vulkan/shaders/te_attention.comp new file mode 100644 index 0000000..e05d966 --- /dev/null +++ b/vulkan/shaders/te_attention.comp @@ -0,0 +1,151 @@ +// Causal scaled-dot-product attention over one sequence for the text encoder's full-attention layers: head_dim 256, query +// head h reading key/value head h / group, bf16 in and out, fp32 online softmax, any device. A workgroup of 64 takes 16 +// query rows of one head and walks the key tiles of 16 that they see, from the last down to the first, as +// FlashAttention's reverse iteration; tiles align with the rows, so each row's first tile holds its own key: +// S = q * K_tile^T (dims ascending, each product and sum rounded; keys past the row masked) +// max, row_sum and acc_o rescaled to the new maximum +// P = bf16(exp2(S * log2_scale - max * log2_scale)), row_sum += P before rounding, keys ascending +// acc_o += P * V_tile (keys ascending) +// and out = bf16(acc_o / row_sum). Invocation t computes S for row t / 4, keys t % 4 + 4 j, and holds output dims +// 4t .. 4t + 3 of all 16 rows. Q, K and V tiles sit in shared memory with rows padded by 8 bytes against bank conflicts. +// Rows q_first .. row_end - 1 of the sequence, q_first a multiple of 16; q, k, v and o point at its first row. +#version 460 +#extension GL_GOOGLE_include_directive : require +#extension GL_EXT_control_flow_attributes : require +#include "te_common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t q; + uint64_t k; + uint64_t v; + uint64_t o; + uint ldq; + uint ldk; + uint ldv; + uint ldo; + uint hsq; + uint hsk; + uint hsv; + uint hso; + uint q_first; + uint row_end; + uint group; // query heads per key/value head + float log2_scale; +} pc; + +const uint BR = 16u; // query rows per workgroup +const uint BC = 16u; // keys per tile +const uint W = 64u; // uvec2 (4 bf16) per 256-dim row +const uint WP = 65u; // padded +const uint SP = 17u; // padded S row + +shared uvec2 q_tile[BR * WP]; +shared uvec2 k_tile[BC * WP]; +shared uvec2 v_tile[BC * WP]; +shared float s_tile[BR * SP]; +shared float rescale_shared[BR]; +shared float inv_sum_shared[BR]; + +#define NEG_INF uintBitsToFloat(0xff800000u) + +void main() { + const uint t = gl_LocalInvocationIndex; + const uint head = gl_WorkGroupID.x, kv_head = head / pc.group; + const uint i0 = pc.q_first + gl_WorkGroupID.y * BR; + const uint n_keys = min(i0 + BR, pc.row_end); // keys any valid row of the block sees + + for (uint e = t; e < BR * W; e += 64u) { + const uint r = e / W, w = e % W; + uvec2 val = uvec2(0u); + if (i0 + r < pc.row_end) val = Words(pc.q + (uint64_t(i0 + r) * pc.ldq + uint64_t(head) * pc.hsq) * 2ul).d[w]; + q_tile[r * WP + w] = val; + } + + const uint64_t k_head = pc.k + uint64_t(kv_head) * pc.hsk * 2ul; + const uint64_t v_head = pc.v + uint64_t(kv_head) * pc.hsv * 2ul; + const uint my_row = t / 4u, my_cols = t % 4u; + + precise vec4 acc[BR]; + [[unroll]] for (uint r = 0u; r < BR; ++r) acc[r] = vec4(0.0); + float row_max = 0.0, row_sum = 0.0; // of row t, for t < 16 + + const uint tiles = (n_keys + BC - 1u) / BC; + for (uint tile_plus_1 = tiles; tile_plus_1 > 0u; --tile_plus_1) { + const uint j0 = (tile_plus_1 - 1u) * BC; + const bool first = tile_plus_1 == tiles; + barrier(); // the previous tile's readers are done + for (uint e = t; e < BC * W; e += 64u) { + const uint c = e / W, w = e % W; + uvec2 kv = uvec2(0u), vv = uvec2(0u); + if (j0 + c < n_keys) { + kv = Words(k_head + uint64_t(j0 + c) * pc.ldk * 2ul).d[w]; + vv = Words(v_head + uint64_t(j0 + c) * pc.ldv * 2ul).d[w]; + } + k_tile[c * WP + w] = kv; + v_tile[c * WP + w] = vv; + } + barrier(); + + [[unroll]] for (uint jj = 0u; jj < 4u; ++jj) { + const uint c = my_cols + 4u * jj; + precise float s = 0.0; + for (uint w = 0u; w < W; ++w) { + const vec4 q4 = f4(q_tile[my_row * WP + w]); + const vec4 k4 = f4(k_tile[c * WP + w]); + s = s + q4.x * k4.x; + s = s + q4.y * k4.y; + s = s + q4.z * k4.z; + s = s + q4.w * k4.w; + } + const bool masked = j0 + c > i0 + my_row || j0 + c >= n_keys; + s_tile[my_row * SP + c] = masked ? NEG_INF : s; + } + barrier(); + + if (t < BR) { + float cur = first ? NEG_INF : row_max; + for (uint c = 0u; c < BC; ++c) cur = max(cur, s_tile[t * SP + c]); + float rescale = 1.0; + precise float sum = 0.0; + if (!first) { + precise const float diff = (row_max - cur) * pc.log2_scale; + rescale = exp2(diff); + precise const float scaled = row_sum * rescale; + sum = scaled; + } + row_max = cur; + precise const float max_scaled = isinf(cur) ? 0.0 : cur * pc.log2_scale; + for (uint c = 0u; c < BC; ++c) { + const float s = s_tile[t * SP + c]; + // FlashAttention's one fused step; a device without a fused fma rounds it twice, within the bf16 weight's error. + const float w = isinf(s) ? 0.0 : exp2(fma(s, pc.log2_scale, -max_scaled)); + sum = sum + w; + s_tile[t * SP + c] = uintBitsToFloat(bf(w) << 16); + } + row_sum = sum; + rescale_shared[t] = rescale; + } + barrier(); + + if (!first) { + [[unroll]] for (uint r = 0u; r < BR; ++r) { + precise const vec4 scaled = acc[r] * rescale_shared[r]; + acc[r] = scaled; + } + } + for (uint c = 0u; c < BC; ++c) { + const vec4 v4 = f4(v_tile[c * WP + t]); + [[unroll]] for (uint r = 0u; r < BR; ++r) acc[r] = acc[r] + s_tile[r * SP + c] * v4; + } + } + + if (t < BR) inv_sum_shared[t] = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + barrier(); + [[unroll]] for (uint r = 0u; r < BR; ++r) { + if (i0 + r >= pc.row_end) continue; + precise const vec4 result = acc[r] * inv_sum_shared[r]; + Words(pc.o + (uint64_t(i0 + r) * pc.ldo + uint64_t(head) * pc.hso) * 2ul).d[t] = bf4(result); + } +} diff --git a/vulkan/shaders/te_common.glsl b/vulkan/shaders/te_common.glsl new file mode 100644 index 0000000..14162ab --- /dev/null +++ b/vulkan/shaders/te_common.glsl @@ -0,0 +1,57 @@ +// What the text encoder's shaders add to common.glsl: 32-bit words of two bf16, single bf16 reads, and fp32 vectors. +#ifndef MI_VK_TE_COMMON_GLSL +#define MI_VK_TE_COMMON_GLSL + +#include "common.glsl" + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer Uints { uint d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4f { vec4 d[]; }; + +vec2 f2(uint w) { return uintBitsToFloat(uvec2(w << 16, w & 0xffff0000u)); } +uint bf2(vec2 v) { return bf(v.x) | (bf(v.y) << 16); } + +// Element `index` of a bf16 array at a 4-byte aligned address. +float bf16_at(uint64_t base, uint64_t index) { + const uint w = Uints(base + (index >> 1) * 4ul).d[0]; + return uintBitsToFloat((index & 1ul) == 0ul ? w << 16 : w & 0xffff0000u); +} + +// fp32 element `index` of an array: GLSL indexes with 32 bits, so 64-bit offsets go into the address. +float f32_at(uint64_t base, uint64_t index) { return Floats(base + index * 4ul).d[0]; } +void f32_store(uint64_t base, uint64_t index, float v) { Floats(base + index * 4ul).d[0] = v; } + +// exp(x) within about 1 fp32 ULP, where Vulkan's exp may be off by 3 + 2|x| ULP: x = n ln2 + r with ln2 in two parts, +// |r| <= ln2 / 2, exp(r) by its degree-7 Taylor polynomial (truncation below 0.1 ULP), times 2^n. Results below 2^-126 +// flush to zero. +float exp_f(float x) { + if (x < -87.3) return 0.0; + if (x > 88.7) return uintBitsToFloat(0x7f800000u); + const float n = roundEven(x * 1.4426950408889634); + precise const float r_hi = fma(n, -0.693145751953125, x); // ln2's leading 16 bits: n * them is exact + precise const float r = fma(n, -1.4286068203094172e-06, r_hi); + precise float p = 1.0 / 5040.0; + p = fma(p, r, 1.0 / 720.0); + p = fma(p, r, 1.0 / 120.0); + p = fma(p, r, 1.0 / 24.0); + p = fma(p, r, 1.0 / 6.0); + p = fma(p, r, 0.5); + p = fma(p, r, 1.0); + p = fma(p, r, 1.0); + const int e = int(n); + // 2^e as two factors, so that e = -126 .. -127 and 128 stay representable. + return (p * uintBitsToFloat(uint(127 + e / 2) << 23)) * uintBitsToFloat(uint(127 + e - e / 2) << 23); +} + +// torch's fp32 silu and sigmoid: x / (1 + exp(-x)) and 1 / (1 + exp(-x)). +float silu_f(float x) { + precise const float den = 1.0 + exp(-x); + precise const float r = x / den; + return r; +} +float sigmoid_f(float x) { + precise const float den = 1.0 + exp(-x); + precise const float r = 1.0 / den; + return r; +} + +#endif diff --git a/vulkan/shaders/te_conv_silu.comp b/vulkan/shaders/te_conv_silu.comp new file mode 100644 index 0000000..c9d819f --- /dev/null +++ b/vulkan/shaders/te_conv_silu.comp @@ -0,0 +1,38 @@ +// GatedDeltaNet's causal depthwise conv1d (kernel 4) and SiLU over one sequence, as gdn::causal_conv1d_silu: +// acc = sum_k w[c][k] * x[t - 3 + k][c] (fused multiply-adds, k ascending, rows before the sequence are zero) +// y = bf16(v / (1 + exp(-v))), v = bf16(acc) +// x: bf16 rows of the sequence, stride ldx; w: bf16 [channels][4]; y: bf16 rows, stride ldy. Invocation i handles channels +// 4i .. 4i + 3 of token gl_WorkGroupID.y. x and y must not overlap. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t w; + uint64_t y; + uint ldx; + uint ldy; + uint channels; +} pc; + +void main() { + const uint group = gl_GlobalInvocationID.x; + if (group * 4u >= pc.channels) return; + const int t = int(gl_WorkGroupID.y); + Words W = Words(pc.w + uint64_t(group) * 32ul); // 4 channels x 4 taps + const vec4 w0 = f4(W.d[0]), w1 = f4(W.d[1]), w2 = f4(W.d[2]), w3 = f4(W.d[3]); + precise vec4 acc = vec4(0.0); + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + if (src < 0) continue; + const vec4 xv = f4(Words(pc.x + uint64_t(src) * pc.ldx * 2ul).d[group]); + const vec4 wk = vec4(w0[k], w1[k], w2[k], w3[k]); + acc = fma(wk, xv, acc); + } + const vec4 v = rbf4(acc); + const vec4 y = vec4(silu_f(v.x), silu_f(v.y), silu_f(v.z), silu_f(v.w)); + Words(pc.y + uint64_t(t) * pc.ldy * 2ul).d[group] = bf4(y); +} diff --git a/vulkan/shaders/te_gate_mul.comp b/vulkan/shaders/te_gate_mul.comp new file mode 100644 index 0000000..41e3644 --- /dev/null +++ b/vulkan/shaders/te_gate_mul.comp @@ -0,0 +1,44 @@ +// Gated products of the text encoder, bf16, per element as torch eager: +// MODE 0, Qwen3_5MLP: out = bf16(bf16(silu(a)) * b) +// MODE 1, the attention gate: out = bf16(a * bf16(sigmoid(b))) +// Column c of a row reads a[c], b[(c / block) * block_stride + c % block] and writes out[c]; rows have strides lda, ldb and +// ldo. cols, block, block_stride and the strides are multiples of 4. out may alias a. Invocation: 4 columns of row +// gl_WorkGroupID.y. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const int MODE = 0; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint lda; + uint ldb; + uint ldo; + uint cols; + uint block; + uint block_stride; +} pc; + +void main() { + const uint c = gl_GlobalInvocationID.x * 4u; + if (c >= pc.cols) return; + const uint64_t r = gl_WorkGroupID.y; + const uint cb = (c / pc.block) * pc.block_stride + c % pc.block; + const vec4 a = f4(Words(pc.a + (r * pc.lda + c) * 2ul).d[0]); + const vec4 b = f4(Words(pc.b + (r * pc.ldb + cb) * 2ul).d[0]); + vec4 y; + if (MODE == 0) { + const vec4 act = rbf4(vec4(silu_f(a.x), silu_f(a.y), silu_f(a.z), silu_f(a.w))); + precise const vec4 product = act * b; + y = product; + } else { + const vec4 gate = rbf4(vec4(sigmoid_f(b.x), sigmoid_f(b.y), sigmoid_f(b.z), sigmoid_f(b.w))); + precise const vec4 product = a * gate; + y = product; + } + Words(pc.out_ + (r * pc.ldo + c) * 2ul).d[0] = bf4(y); +} diff --git a/vulkan/shaders/te_gates.comp b/vulkan/shaders/te_gates.comp new file mode 100644 index 0000000..c7fcf9c --- /dev/null +++ b/vulkan/shaders/te_gates.comp @@ -0,0 +1,55 @@ +// GatedDeltaNet's gates, as gdn::gates: beta = bf16(sigmoid(b)), g = -exp(A_log) * softplus(a + dt_bias) in fp32 (softplus +// threshold 20). ba: bf16 rows [b(32) | a(32)] with stride ldba; a_log, dt_bias: bf16 [32]; g: f32 [tokens, 32]; beta: bf16 +// [tokens, 32]. Invocation i handles heads 2 (i % 16) and 2 (i % 16) + 1 of token i / 16. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t ba; + uint64_t a_log; + uint64_t dt_bias; + uint64_t g; + uint64_t beta; + uint ldba; + uint tokens; +} pc; + +// log(1 + y) for y >= 0. Vulkan's log has only an absolute error bound (2^-21) on [0.5, 2], so below y = 1 this sums +// 2 atanh(s) = 2 (s + s^3 / 3 + ... + s^17 / 17), s = y / (2 + y) <= 1/3, whose dropped terms are below 2^-27 of it. +float log1p_f(float y) { + if (y > 1.0) return log(1.0 + y); + precise const float s = y / (2.0 + y); + precise const float s2 = s * s; + precise float p = 1.0 / 17.0; + p = p * s2 + 1.0 / 15.0; + p = p * s2 + 1.0 / 13.0; + p = p * s2 + 1.0 / 11.0; + p = p * s2 + 1.0 / 9.0; + p = p * s2 + 1.0 / 7.0; + p = p * s2 + 1.0 / 5.0; + p = p * s2 + 1.0 / 3.0; + p = p * s2 + 1.0; + precise const float r = 2.0 * s * p; + return r; +} + +float softplus_f(float x) { return x > 20.0 ? x : log1p_f(exp_f(x)); } + +void main() { + const uint i = gl_GlobalInvocationID.x; + if (i >= pc.tokens * 16u) return; + const uint64_t t = i / 16u; + const uint pair = i % 16u; + Uints row = Uints(pc.ba + t * pc.ldba * 2ul); + const vec2 b = f2(row.d[pair]), a = f2(row.d[16u + pair]); + const vec2 a_log = f2(Uints(pc.a_log).d[pair]), dt = f2(Uints(pc.dt_bias).d[pair]); + Uints(pc.beta + (t * 32ul + 2ul * pair) * 2ul).d[0] = bf2(vec2(sigmoid_f(b.x), sigmoid_f(b.y))); + precise const vec2 x = a + dt; + precise const vec2 g = -vec2(exp_f(a_log.x), exp_f(a_log.y)) * vec2(softplus_f(x.x), softplus_f(x.y)); + Floats G = Floats(pc.g + (t * 32ul + 2ul * pair) * 4ul); + G.d[0] = g.x; + G.d[1] = g.y; +} diff --git a/vulkan/shaders/te_gdn_chunk.comp b/vulkan/shaders/te_gdn_chunk.comp new file mode 100644 index 0000000..82af22d --- /dev/null +++ b/vulkan/shaders/te_gdn_chunk.comp @@ -0,0 +1,116 @@ +// The intra-chunk part of the chunked gated delta rule (chunk 64) for one chunk of one sequence and one value head h, as +// the engine's prep kernel: +// gcum = cumsum(g) over the chunk, eg = exp(gcum), el = exp(gcum[n - 1] - gcum) +// a[i][j] = (k_i . k_j) * beta_i * exp(gcum_i - gcum_j) for j < i, else 0 (the UT system's strict lower part) +// attn[i][j] = (q_i . k_j) * exp(gcum_i - gcum_j) for j <= i, else 0 +// with q and k of key head h / 2 from te_gdn_qk, dot products as fused multiply-adds over d ascending, and rows past the +// sequence zero. Writes gcum, eg, el [32][tokens] and a, attn [32][tokens][64] for the chunk's rows. Workgroup (chunk, h): +// invocation (ty, tx) = (i / 16, i % 16) owns rows ty + 16 r and columns tx + 16 c, r and c < 4. +#version 460 +#extension GL_GOOGLE_include_directive : require +#extension GL_EXT_control_flow_attributes : require +#include "te_common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t qn; + uint64_t kn; + uint64_t g; + uint64_t beta; + uint64_t gcum; + uint64_t eg; + uint64_t el; + uint64_t a; + uint64_t attn; + uint tokens; + uint row0; // the sequence's first token + uint len; +} pc; + +const uint C = 64u; + +shared float gc[C]; +shared float bt[C]; + +vec4 row4(uint64_t base, uint kh, uint t, uint d4) { + return Vec4f(base + ((uint64_t(kh) * pc.tokens + t) * 128ul) * 4ul).d[d4]; +} + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint h = gl_WorkGroupID.y, kh = h / 2u; + const uint c0 = gl_WorkGroupID.x * C; + const uint t0 = pc.row0 + c0; + const uint n = min(C, pc.len - c0); + + if (tid < C) { + const bool valid = tid < n; + bt[tid] = valid ? bf16_at(pc.beta, uint64_t(t0 + tid) * 32ul + h) : 0.0; + gc[tid] = valid ? Floats(pc.g).d[(t0 + tid) * 32u + h] : 0.0; // the padding's decay is 0 in log space + } + barrier(); + if (tid == 0u) { + precise float acc = 0.0; + for (uint i = 0u; i < C; ++i) { + acc = acc + gc[i]; + gc[i] = acc; + } + } + barrier(); + if (tid < n) { + const uint64_t at = uint64_t(h) * pc.tokens + t0 + tid; + precise const float tail = gc[n - 1u] - gc[tid]; + f32_store(pc.gcum, at, gc[tid]); + f32_store(pc.eg, at, exp_f(gc[tid])); + f32_store(pc.el, at, exp_f(tail)); + } + + const uint ty = tid / 16u, tx = tid % 16u; + precise float kk[4][4], qk[4][4]; + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + kk[r][c] = 0.0; + qk[r][c] = 0.0; + } + } + for (uint d4 = 0u; d4 < 32u; ++d4) { + vec4 qi[4], ki[4], kj[4]; + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + const uint i = ty + 16u * r, j = tx + 16u * r; + qi[r] = i < n ? row4(pc.qn, kh, t0 + i, d4) : vec4(0.0); + ki[r] = i < n ? row4(pc.kn, kh, t0 + i, d4) : vec4(0.0); + kj[r] = j < n ? row4(pc.kn, kh, t0 + j, d4) : vec4(0.0); + } + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + float s = kk[r][c], p = qk[r][c]; + s = fma(ki[r].x, kj[c].x, s); + s = fma(ki[r].y, kj[c].y, s); + s = fma(ki[r].z, kj[c].z, s); + s = fma(ki[r].w, kj[c].w, s); + p = fma(qi[r].x, kj[c].x, p); + p = fma(qi[r].y, kj[c].y, p); + p = fma(qi[r].z, kj[c].z, p); + p = fma(qi[r].w, kj[c].w, p); + kk[r][c] = s; + qk[r][c] = p; + } + } + } + + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + const uint i = ty + 16u * r; + if (i >= n) continue; + const uint64_t row = (uint64_t(h) * pc.tokens + t0 + i) * C; + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + const uint j = tx + 16u * c; + precise const float diff = gc[i] - gc[j]; + const float decay = j <= i ? exp_f(diff) : 0.0; + precise const float lower = (kk[r][c] * bt[i]) * decay; + precise const float masked = qk[r][c] * decay; + f32_store(pc.a, row + j, j < i ? lower : 0.0); + f32_store(pc.attn, row + j, masked); + } + } +} diff --git a/vulkan/shaders/te_gdn_qk.comp b/vulkan/shaders/te_gdn_qk.comp new file mode 100644 index 0000000..8122c9c --- /dev/null +++ b/vulkan/shaders/te_gdn_qk.comp @@ -0,0 +1,61 @@ +// The gated delta rule's q and k in fp32, as the engine's chunk prep: q = x * rsqrt(sum(x^2) + 1e-6) * 128^-0.5 and +// k = x * rsqrt(sum(x^2) + 1e-6) per 128-wide key head of the conv output (q heads at column 0, k heads at column 2048). +// Lane l sums its 4 squares with fused multiply-adds, then the 32 partial sums take a halving tree (the same sum as a +// butterfly of shuffles). qn, kn: f32 [16][tokens][128]. Workgroup (key head, token). +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t qkv; + uint64_t qn; + uint64_t kn; + uint ldqkv; + uint tokens; +} pc; + +const float Q_SCALE = 0.08838834764831845; // 128^-0.5 in fp32 +const uint KEY_DIM = 2048u; + +shared float red_q[32]; +shared float red_k[32]; + +void main() { + const uint lane = gl_LocalInvocationIndex; + const uint head = gl_WorkGroupID.x; + const uint64_t t = gl_WorkGroupID.y; + Words row = Words(pc.qkv + t * pc.ldqkv * 2ul); + const vec4 q = f4(row.d[head * 32u + lane]); + const vec4 k = f4(row.d[(KEY_DIM + head * 128u) / 4u + lane]); + precise float sq = 0.0, sk = 0.0; + sq = fma(q.x, q.x, sq); + sq = fma(q.y, q.y, sq); + sq = fma(q.z, q.z, sq); + sq = fma(q.w, q.w, sq); + sk = fma(k.x, k.x, sk); + sk = fma(k.y, k.y, sk); + sk = fma(k.z, k.z, sk); + sk = fma(k.w, k.w, sk); + red_q[lane] = sq; + red_k[lane] = sk; + barrier(); + for (uint off = 16u; off > 0u; off >>= 1) { + if (lane < off) { + precise const float a = red_q[lane] + red_q[lane + off]; + precise const float b = red_k[lane] + red_k[lane + off]; + red_q[lane] = a; + red_k[lane] = b; + } + barrier(); + } + precise const float q_sum = red_q[0] + 1e-6; + precise const float k_sum = red_k[0] + 1e-6; + const float iq = inversesqrt(q_sum), ik = inversesqrt(k_sum); + precise const vec4 qo = (q * iq) * Q_SCALE; + precise const vec4 ko = k * ik; + const uint64_t at = (uint64_t(head) * pc.tokens + t) * 128ul + lane * 4u; + Vec4f(pc.qn + at * 4ul).d[0] = qo; + Vec4f(pc.kn + at * 4ul).d[0] = ko; +} diff --git a/vulkan/shaders/te_gdn_solve.comp b/vulkan/shaders/te_gdn_solve.comp new file mode 100644 index 0000000..a9a29cc --- /dev/null +++ b/vulkan/shaders/te_gdn_solve.comp @@ -0,0 +1,62 @@ +// The UT transform of one chunk of one sequence and value head h, as the engine's prep kernel: forward substitution +// x_i = b_i - sum_{j u [32][tokens][128] +// blocks 2, 3: b = (k * beta) * exp(gcum) -> w [32][tokens][128] +// with v the value head h of the conv output (column 4096 + 128 h) and k of key head h / 2. Workgroup (4 chunk + block, h): +// invocation l solves column l of its block; its solution so far lives in column l of shared memory, which only it uses. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t qkv; // conv output, bf16 rows of stride ldqkv + uint64_t kn; + uint64_t beta; + uint64_t eg; + uint64_t a; + uint64_t u; + uint64_t w; + uint ldqkv; + uint tokens; + uint row0; + uint len; +} pc; + +const uint C = 64u; +const uint VALUE_COL = 4096u; + +shared float xs[C * C]; + +void main() { + const uint lane = gl_LocalInvocationIndex; + const uint h = gl_WorkGroupID.y, kh = h / 2u; + const uint block = gl_WorkGroupID.x % 4u; + const bool is_u = block < 2u; + const uint col = (block & 1u) * C + lane; + const uint c0 = (gl_WorkGroupID.x / 4u) * C; + const uint t0 = pc.row0 + c0; + const uint n = min(C, pc.len - c0); + + for (uint i = 0u; i < n; ++i) { + const uint t = t0 + i; + const float beta = bf16_at(pc.beta, uint64_t(t) * 32ul + h); + float b; + if (is_u) { + precise const float v = bf16_at(pc.qkv + uint64_t(t) * pc.ldqkv * 2ul, VALUE_COL + h * 128u + col) * beta; + b = v; + } else { + const float k = f32_at(pc.kn, (uint64_t(kh) * pc.tokens + t) * 128ul + col); + precise const float kb = (k * beta) * f32_at(pc.eg, uint64_t(h) * pc.tokens + t); + b = kb; + } + Floats A = Floats(pc.a + ((uint64_t(h) * pc.tokens + t) * C) * 4ul); + precise float acc = b; + for (uint j = 0u; j < i; ++j) acc = fma(-A.d[j], xs[j * C + lane], acc); + xs[i * C + lane] = acc; + } + const uint64_t x = is_u ? pc.u : pc.w; + for (uint i = 0u; i < n; ++i) f32_store(x, (uint64_t(h) * pc.tokens + t0 + i) * 128ul + col, xs[i * C + lane]); +} diff --git a/vulkan/shaders/te_gdn_state.comp b/vulkan/shaders/te_gdn_state.comp new file mode 100644 index 0000000..5e8be53 --- /dev/null +++ b/vulkan/shaders/te_gdn_state.comp @@ -0,0 +1,110 @@ +// The inter-chunk scan of the chunked gated delta rule for one sequence, value head h and 16 value columns vc0 .. +16, +// sequential over the chunks, as the engine's state kernel with S [128][16] the state slice (zero at the start): +// v_new = u - w S +// o = (q exp(gcum)) S + attn v_new -> out bf16 [tokens][32][128] +// S = S exp(gcum[n - 1]) + (k exp(gcum[n - 1] - gcum))^T v_new +// Products are fused multiply-adds: over d ascending for w S and q S, then on over j ascending for attn v_new starting from +// q S, over j ascending for the update, and S exp(...) + update as one more. Workgroup (slice, h): invocation +// (rg, cp) = (i / 8, i % 8) owns columns 2 cp and 2 cp + 1 of rows rg + 16 r (r < 4) of the chunk and of state rows +// rg + 16 r (r < 8). +#version 460 +#extension GL_GOOGLE_include_directive : require +#extension GL_EXT_control_flow_attributes : require +#include "te_common.glsl" + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t qn; + uint64_t kn; + uint64_t gcum; + uint64_t eg; + uint64_t el; + uint64_t attn; + uint64_t u; + uint64_t w; + uint64_t out_; + uint tokens; + uint row0; + uint len; +} pc; + +const uint C = 64u; +const uint BV = 16u; + +shared vec2 s_state[128 * (BV / 2u)]; +shared vec2 s_vn[C * (BV / 2u)]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint h = gl_WorkGroupID.y, kh = h / 2u; + const uint cp = tid % 8u, rg = tid / 8u; + const uint vc = gl_WorkGroupID.x * BV + 2u * cp; + const uint64_t ht = uint64_t(h) * pc.tokens, kt = uint64_t(kh) * pc.tokens; + + [[unroll]] for (uint r = 0u; r < 8u; ++r) s_state[(rg + 16u * r) * 8u + cp] = vec2(0.0); + barrier(); + + for (uint c0 = 0u; c0 < pc.len; c0 += C) { + const uint t0 = pc.row0 + c0; + const uint n = min(C, pc.len - c0); + + // w S and q S for rows rg + 16 r. + precise vec2 pw[4], pq[4]; + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + pw[r] = vec2(0.0); + pq[r] = vec2(0.0); + const uint i = rg + 16u * r; + if (i < n) { + Floats W = Floats(pc.w + ((ht + t0 + i) * 128ul) * 4ul); + Floats Q = Floats(pc.qn + ((kt + t0 + i) * 128ul) * 4ul); + const float eg = f32_at(pc.eg, ht + t0 + i); + precise vec2 accw = vec2(0.0), accq = vec2(0.0); + for (uint d = 0u; d < 128u; ++d) { + const vec2 s = s_state[d * 8u + cp]; + precise const float qd = Q.d[d] * eg; + accw = fma(vec2(W.d[d]), s, accw); + accq = fma(vec2(qd), s, accq); + } + pw[r] = accw; + pq[r] = accq; + } + } + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + const uint i = rg + 16u * r; + vec2 vn = vec2(0.0); + if (i < n) { + Floats U = Floats(pc.u + ((ht + t0 + i) * 128ul + vc) * 4ul); + precise const vec2 diff = vec2(U.d[0], U.d[1]) - pw[r]; + vn = diff; + } + s_vn[i * 8u + cp] = vn; + } + barrier(); + + // o = q S + attn v_new. + [[unroll]] for (uint r = 0u; r < 4u; ++r) { + const uint i = rg + 16u * r; + if (i >= n) continue; + Floats A = Floats(pc.attn + ((ht + t0 + i) * C) * 4ul); + precise vec2 o = pq[r]; + for (uint j = 0u; j <= i; ++j) o = fma(vec2(A.d[j]), s_vn[j * 8u + cp], o); + Uints(pc.out_ + ((uint64_t(t0 + i) * 32ul + h) * 128ul + vc) * 2ul).d[0] = bf2(o); + } + + // S = S exp(gcum[n - 1]) + kd^T v_new, kd = k exp(gcum[n - 1] - gcum). + const float cd = exp_f(f32_at(pc.gcum, ht + t0 + n - 1u)); + [[unroll]] for (uint r = 0u; r < 8u; ++r) { + const uint d = rg + 16u * r; + precise vec2 upd = vec2(0.0); + for (uint j = 0u; j < n; ++j) { + precise const float kd = f32_at(pc.kn, (kt + t0 + j) * 128ul + d) * f32_at(pc.el, ht + t0 + j); + upd = fma(vec2(kd), s_vn[j * 8u + cp], upd); + } + const uint at = d * 8u + cp; + precise const vec2 next = fma(s_state[at], vec2(cd), upd); + s_state[at] = next; + } + barrier(); + } +} diff --git a/vulkan/shaders/te_rms_norm.comp b/vulkan/shaders/te_rms_norm.comp new file mode 100644 index 0000000..52948fc --- /dev/null +++ b/vulkan/shaders/te_rms_norm.comp @@ -0,0 +1,112 @@ +// The Qwen3.5 RMSNorms over rows of d bf16 values, one row per workgroup of 32 invocations, any device: +// MODE 0, Qwen3_5RMSNorm: out = bf16(x * rsqrt(mean(x^2) + eps) * (1 + w)) +// MODE 1, Qwen3_5RMSNormGated: out = bf16(bf16(w * bf16(x * rsqrt(mean(x^2) + eps))) * silu(gate)) +// MODE 0 with a table also applies the text RoPE to the first 64 dims of the row: pairs (c, c + 32), c < 32, rotated by the +// table row [cos(32) | sin(32)] of the row's token, each product and the sum rounded to bf16 (apply_rotary_pos_emb). +// +// Rows are (token, head) pairs: workgroup (head, token) reads x + token * ldx + head * hsx, and likewise out and gate. out +// may alias x. The mean of squares is summed in the order of torch's reduce kernel for `bw` threads per row (a power of +// two from 32 to 512): virtual thread v sums 4 running squares over the vec4 groups v, v + bw, ..., adds them left to +// right, the bw partial sums halve down to 32 and then take a shuffle-down tree. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "te_common.glsl" + +layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const int MODE = 0; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t weight; + uint64_t gate; + uint64_t out_; + uint64_t table; // 0: no RoPE + uint ldx; + uint hsx; + uint ldo; + uint hso; + uint ldg; + uint hsg; + uint table_ld; + uint d; + uint bw; + float factor; // torch's MeanOps factor: rows / (rows * d) in fp32 + float eps; +} pc; + +shared float red[512]; +shared vec4 rope_in[16]; + +void main() { + const uint lane = gl_LocalInvocationIndex; + const uint64_t head = gl_WorkGroupID.x; + const uint64_t token = gl_WorkGroupID.y; + Words X = Words(pc.x + (token * pc.ldx + head * pc.hsx) * 2ul); + const uint groups = pc.d / 4u; + + for (uint v = lane; v < pc.bw; v += 32u) { + precise vec4 acc = vec4(0.0); + for (uint g = v; g < groups; g += pc.bw) { + const vec4 e = f4(X.d[g]); + acc = acc + e * e; + } + precise const float total = ((acc.x + acc.y) + acc.z) + acc.w; + red[v] = total; + } + barrier(); + for (uint off = pc.bw / 2u; off >= 32u; off >>= 1) { + for (uint v = lane; v < off; v += 32u) { + precise const float sum = red[v] + red[v + off]; + red[v] = sum; + } + barrier(); + } + for (uint off = 16u; off > 0u; off >>= 1) { + if (lane < off) { + precise const float sum = red[lane] + red[lane + off]; + red[lane] = sum; + } + barrier(); + } + precise const float mean = red[0] * pc.factor; + precise const float mean_eps = mean + pc.eps; + const float r = inversesqrt(mean_eps); + + Words O = Words(pc.out_ + (token * pc.ldo + head * pc.hso) * 2ul); + Words W = Words(pc.weight); + const bool rope = MODE == 0 && pc.table != 0ul; + for (uint g = lane; g < groups; g += 32u) { + precise const vec4 xn = f4(X.d[g]) * r; + vec4 y; + if (MODE == 0) { + precise const vec4 one_plus = 1.0 + f4(W.d[g]); + precise const vec4 scaled = xn * one_plus; + y = scaled; + } else { + const vec4 gv = f4(Words(pc.gate + (token * pc.ldg + head * pc.hsg) * 2ul).d[g]); + precise const vec4 weighted = f4(W.d[g]) * rbf4(xn); + const vec4 act = vec4(silu_f(gv.x), silu_f(gv.y), silu_f(gv.z), silu_f(gv.w)); + precise const vec4 gated = rbf4(weighted) * act; + y = gated; + } + if (rope && g < 16u) { + rope_in[g] = rbf4(y); + } else { + O.d[g] = bf4(y); + } + } + if (rope) { + barrier(); + if (lane < 16u) { + Words T = Words(pc.table + token * pc.table_ld * 2ul); + const uint pair = lane & 7u; + const vec4 co = f4(T.d[pair]), si = f4(T.d[pair + 8u]); + const vec4 lo = rope_in[pair], hi = rope_in[pair + 8u]; + const bool first_half = lane < 8u; + precise const vec4 a = (first_half ? lo : hi) * co; + precise const vec4 b = first_half ? -hi * si : lo * si; + precise const vec4 rotated = rbf4(a) + rbf4(b); + O.d[lane] = bf4(rotated); + } + } +} diff --git a/vulkan/shaders/vae_conv_coop.comp b/vulkan/shaders/vae_conv_coop.comp new file mode 100644 index 0000000..cacef84 --- /dev/null +++ b/vulkan/shaders/vae_conv_coop.comp @@ -0,0 +1,265 @@ +// The VAE's convolutions and fp32 GEMMs on bf16 cooperative matrices with fp32 accumulators: +// out[m, n] = (sum_k A[m, k] * B[n, k] + bias[n]) + residual[m, n] (each + rounded separately, as the reference) +// with fp32 A and B in memory. Each operand is split on load into A_TERMS (B_TERMS) bf16 pieces whose sum is the operand +// (or its TF32 rounding), and the products of pieces i, j with i + j < max(A_TERMS, B_TERMS) are accumulated: +// TF32 with A_TERMS 2, B_TERMS 1: cuDNN's TF32 convolution for weights that are bf16 values (the MOSS bundle's are). A is +// rounded to TF32 (half away from zero, as cvt.rna), whose 11 significant bits are exactly hi + lo in bf16, so every +// product is the exact product of the TF32 operands; only the fp32 accumulation order differs from cuDNN's. +// no TF32, A_TERMS 3, B_TERMS 1: an fp32 GEMM with exact products for bf16-valued weights (fp32 = 3 bf16 pieces). +// no TF32, A_TERMS 2, B_TERMS 2: 16 significant bits per operand, for attention scores and P V. +// TAPS selects A: 9 is a 3x3 convolution with zero padding 1 and 1 a 1x1 convolution over an NHWC input of C channels, for +// an H x W output whose pixels are the rows m (UPSAMPLE: the input is H/2 x W/2 and read through a nearest 2x upsample); +// 0 is a row-major A with row stride lda. B is row-major [N, K] (a convolution's weights in KRSC order, K = TAPS C). +// C is a multiple of 8; K, lda, ldb and ldo are multiples of 4. M and N are free. +// +// A workgroup computes rows m0 + 128 y .. +128 of the output and 32 FN columns from 32 FN x: 256 invocations, 8 full +// subgroups of 32, which the pipeline must require. Subgroup s owns rows 32 (s % 4) .. +32 and columns 16 FN (s / 4) .. +// +16 FN. K walks in steps of 16 through two shared-memory stages; the next step loads into registers while the current +// one multiplies. Staged rows are 32 bytes (16 bf16), so a 16 x 16 fragment is one contiguous 512-byte block. +#version 460 +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_bfloat16 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const uint TAPS = 9; +layout(constant_id = 1) const bool UPSAMPLE = false; +layout(constant_id = 2) const uint A_TERMS = 2; // 1..3 +layout(constant_id = 3) const uint B_TERMS = 1; // 1..3 +layout(constant_id = 4) const bool TF32 = true; +layout(constant_id = 5) const uint FN = 4; // 1..4 + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { uvec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint64_t bias; // 0: none + uint64_t residual; // 0: none; [M, N] with stride ldo, may be out_ + uint M; // rows end (exclusive) + uint N; + uint K; + uint lda; + uint ldb; + uint ldo; + uint H; + uint W; + uint C; + uint m0; // first row +} pc; + +const uint BM = 128u; +const uint BN = 32u * FN; +const uint BK = 16u; +const uint TERMS = A_TERMS > B_TERMS ? A_TERMS : B_TERMS; +// uvec4 per stage: A pieces of BM rows, then B pieces of BN rows, 2 uvec4 per row. +const uint A_PIECE = BM * 2u; +const uint B_PIECE = BN * 2u; +const uint STAGE = A_TERMS * A_PIECE + B_TERMS * B_PIECE; +const uint STAGING = 64u; // uvec4 per 16x16 fp32 fragment in the epilogue: 8 subgroups need 512, below any STAGE + +shared uvec4 smem[2u * STAGE]; + +#define FRAG_A coopmat +#define FRAG_B coopmat +#define ACC coopmat + +// Round half away from zero to 10 mantissa bits; infinities and NaNs pass. +float tf32(float v) { + const uint u = floatBitsToUint(v); + if ((u & 0x7f800000u) == 0x7f800000u) return v; + return uintBitsToFloat((u + 0x1000u) & 0xffffe000u); +} + +// 8 fp32 values -> TERMS pieces of 8 bf16 (round to nearest even each; the remainders are exact in fp32). +void split8(vec4 lo, vec4 hi, bool round_tf32, uint terms, out uvec4 piece[3]) { + vec4 r0 = lo, r1 = hi; + if (round_tf32) { + r0 = vec4(tf32(r0.x), tf32(r0.y), tf32(r0.z), tf32(r0.w)); + r1 = vec4(tf32(r1.x), tf32(r1.y), tf32(r1.z), tf32(r1.w)); + } + [[unroll]] for (uint t = 0u; t < 3u; ++t) { + piece[t] = uvec4(0u); + if (t < terms) { + const uvec2 p0 = bf4(r0), p1 = bf4(r1); + piece[t] = uvec4(p0, p1); + precise const vec4 n0 = r0 - f4(p0), n1 = r1 - f4(p1); + r0 = n0; + r1 = n1; + } + } +} + +// The 8 values of A at row m, columns k .. k + 7; zero outside the tensor or the padding. +void load_a(uint m, uint k, out vec4 lo, out vec4 hi) { + lo = vec4(0.0); + hi = vec4(0.0); + if (m >= pc.M || k >= pc.K) return; + uint64_t at; + if (TAPS == 0u) { + at = uint64_t(m) * pc.lda + k; + } else { + const uint tap = k / pc.C, c = k - tap * pc.C; + const uint oy = m / pc.W, ox = m - oy * pc.W; + if (TAPS == 1u) { + at = uint64_t(m) * pc.C + c; + } else { + const int iy = int(oy) + int(tap / 3u) - 1, ix = int(ox) + int(tap % 3u) - 1; + if (iy < 0 || ix < 0 || iy >= int(pc.H) || ix >= int(pc.W)) return; + if (UPSAMPLE) { + at = (uint64_t(uint(iy) >> 1) * (pc.W >> 1) + (uint(ix) >> 1)) * pc.C + c; + } else { + at = (uint64_t(iy) * pc.W + uint(ix)) * pc.C + c; + } + } + } + Vec4s p = Vec4s(pc.a + at * 4ul); + lo = uintBitsToFloat(p.d[0]); + if (k + 4u < pc.K) hi = uintBitsToFloat(p.d[1]); // K is a multiple of 4: the chunk's second half may lie past it +} + +void load_b(uint n, uint k, out vec4 lo, out vec4 hi) { + lo = vec4(0.0); + hi = vec4(0.0); + if (n >= pc.N || k >= pc.K) return; + Vec4s p = Vec4s(pc.b + (uint64_t(n) * pc.ldb + k) * 4ul); + lo = uintBitsToFloat(p.d[0]); + if (k + 4u < pc.K) hi = uintBitsToFloat(p.d[1]); +} + +// Invocation i loads 8 values (columns 8 (i % 2) ..) of tile row i / 2 of A and, while i / 2 < BN, of B. +struct Regs { + vec4 a_lo, a_hi, b_lo, b_hi; +}; + +void load_step(uint m_base, uint n_base, uint k0, out Regs r) { + const uint i = gl_LocalInvocationIndex; + const uint row = i / 2u, k = k0 + (i % 2u) * 8u; + load_a(m_base + row, k, r.a_lo, r.a_hi); + r.b_lo = vec4(0.0); + r.b_hi = vec4(0.0); + if (row < BN) load_b(n_base + row, k, r.b_lo, r.b_hi); +} + +void store_step(uint stage, const Regs r) { + const uint i = gl_LocalInvocationIndex; + const uint row = i / 2u, half_ = i % 2u; + const uint base = stage * STAGE; + uvec4 piece[3]; + split8(r.a_lo, r.a_hi, TF32, A_TERMS, piece); + [[unroll]] for (uint t = 0u; t < 3u; ++t) { + if (t < A_TERMS) smem[base + t * A_PIECE + row * 2u + half_] = piece[t]; + } + if (row < BN) { + split8(r.b_lo, r.b_hi, TF32, B_TERMS, piece); + [[unroll]] for (uint t = 0u; t < 3u; ++t) { + if (t < B_TERMS) smem[base + A_TERMS * A_PIECE + t * B_PIECE + row * 2u + half_] = piece[t]; + } + } +} + +void sync_subgroup() { + subgroupMemoryBarrierShared(); + subgroupBarrier(); +} + +void main() { + const uint m_base = pc.m0 + gl_WorkGroupID.y * BM; + const uint n_base = gl_WorkGroupID.x * BN; + const uint sg = gl_SubgroupID; + const uint row0 = (sg % 4u) * 32u; + const uint col0 = (sg / 4u) * 16u * FN; + + ACC acc[2][4]; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) acc[r][c] = ACC(0.0); + } + + const uint steps = (pc.K + BK - 1u) / BK; + Regs regs; + load_step(m_base, n_base, 0u, regs); + store_step(0u, regs); + barrier(); + for (uint step = 0u; step < steps; ++step) { + const bool more = step + 1u < steps; + if (more) load_step(m_base, n_base, (step + 1u) * BK, regs); + const uint a_base = (step & 1u) * STAGE, b_base = a_base + A_TERMS * A_PIECE; + FRAG_A fa[2][3]; + FRAG_B fb[4][3]; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint t = 0u; t < 3u; ++t) { + if (t < A_TERMS) { + coopMatLoad(fa[r][t], smem, a_base + t * A_PIECE + (row0 + r * 16u) * 2u, 2u, gl_CooperativeMatrixLayoutRowMajor); + } + } + } + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + [[unroll]] for (uint t = 0u; t < 3u; ++t) { + if (c < FN && t < B_TERMS) { + coopMatLoad(fb[c][t], smem, b_base + t * B_PIECE + (col0 + c * 16u) * 2u, 2u, + gl_CooperativeMatrixLayoutColumnMajor); + } + } + } + // Smallest products first. + [[unroll]] for (int s = 4; s >= 0; --s) { + [[unroll]] for (uint i = 0u; i < 3u; ++i) { + [[unroll]] for (uint j = 0u; j < 3u; ++j) { + if (int(i + j) == s && i < A_TERMS && j < B_TERMS && i + j < TERMS) { + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + if (c < FN) acc[r][c] = coopMatMulAdd(fa[r][i], fb[c][j], acc[r][c]); + } + } + } + } + } + } + if (more) store_step((step + 1u) & 1u, regs); + barrier(); + } + + // Epilogue: each 16x16 fragment goes through the subgroup's staging, then lane l writes 8 outputs of row l / 2. + const uint staging = sg * STAGING; + const uint lane = gl_SubgroupInvocationID; + const uint frow = lane / 2u, fcol = (lane % 2u) * 8u; + [[unroll]] for (uint r = 0u; r < 2u; ++r) { + const uint m = m_base + row0 + r * 16u + frow; + [[unroll]] for (uint c = 0u; c < 4u; ++c) { + if (c >= FN) continue; + coopMatStore(acc[r][c], smem, staging, 4u, gl_CooperativeMatrixLayoutRowMajor); + sync_subgroup(); + const uint n = n_base + col0 + c * 16u + fcol; + if (m < pc.M && n < pc.N) { + float v[8]; + const vec4 lo = uintBitsToFloat(smem[staging + frow * 4u + fcol / 4u]); + const vec4 hi = uintBitsToFloat(smem[staging + frow * 4u + fcol / 4u + 1u]); + v[0] = lo.x; v[1] = lo.y; v[2] = lo.z; v[3] = lo.w; + v[4] = hi.x; v[5] = hi.y; v[6] = hi.z; v[7] = hi.w; + const uint64_t at = uint64_t(m) * pc.ldo + n; + [[unroll]] for (uint j = 0u; j < 8u; ++j) { + if (n + j < pc.N) { + precise float y = v[j]; + if (pc.bias != 0ul) y = y + Floats(pc.bias).d[n + j]; + if (pc.residual != 0ul) y = y + Floats(pc.residual + (at + j) * 4ul).d[0]; + v[j] = y; + } + } + if (n + 8u <= pc.N) { + Vec4s(pc.out_ + at * 4ul).d[0] = floatBitsToUint(vec4(v[0], v[1], v[2], v[3])); + Vec4s(pc.out_ + at * 4ul + 16ul).d[0] = floatBitsToUint(vec4(v[4], v[5], v[6], v[7])); + } else { + [[unroll]] for (uint j = 0u; j < 8u; ++j) { + if (n + j < pc.N) Floats(pc.out_ + (at + j) * 4ul).d[0] = v[j]; + } + } + } + sync_subgroup(); + } + } +} diff --git a/vulkan/shaders/vae_conv_narrow.comp b/vulkan/shaders/vae_conv_narrow.comp new file mode 100644 index 0000000..d9a71d5 --- /dev/null +++ b/vulkan/shaders/vae_conv_narrow.comp @@ -0,0 +1,113 @@ +// A convolution or GEMM with at most 4 output columns (conv_out's 3), without cooperative matrices or 16-bit types: +// out[m, n] = (sum_k A[m, k] * B[n, k] + bias[n]) + residual[m, n] (each + rounded separately), n < N <= 4 +// with fp32 accumulators, K ascending, and operands rounded to TF32 first with TF32. A and B are addressed as in +// vae_conv_plain.comp; K and C are multiples of 4. One invocation per output row: per tap it reads the input pixel's +// channels as vec4s and the 4 columns of B for the same k from shared memory, where B is staged transposed ([k] -> vec4 of +// the 4 columns) in chunks of up to 512 channels. Workgroup y takes rows m0 + 256 y .. +#version 460 +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const uint TAPS = 9; +layout(constant_id = 1) const bool UPSAMPLE = false; +layout(constant_id = 2) const bool TF32 = true; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { vec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint64_t bias; + uint64_t residual; + uint M; + uint N; + uint K; + uint lda; + uint ldb; + uint ldo; + uint H; + uint W; + uint C; + uint m0; +} pc; + +const uint KC = 512u; +shared vec4 Bs[KC]; + +float tf32(float v) { + const uint u = floatBitsToUint(v); + if ((u & 0x7f800000u) == 0x7f800000u) return v; + return uintBitsToFloat((u + 0x1000u) & 0xffffe000u); +} +vec4 round4(vec4 v) { return TF32 ? vec4(tf32(v.x), tf32(v.y), tf32(v.z), tf32(v.w)) : v; } + +// The address of A[m, tap, 0] (TAPS 0: the row m), or false in the zero padding. +bool row_base(uint m, uint tap, out uint64_t at) { + if (TAPS == 0u) { + at = pc.a + uint64_t(m) * pc.lda * 4ul; + return true; + } + if (TAPS == 1u) { + at = pc.a + uint64_t(m) * pc.C * 4ul; + return true; + } + const uint oy = m / pc.W, ox = m - oy * pc.W; + const int iy = int(oy) + int(tap / 3u) - 1, ix = int(ox) + int(tap % 3u) - 1; + if (iy < 0 || ix < 0 || iy >= int(pc.H) || ix >= int(pc.W)) return false; + if (UPSAMPLE) { + at = pc.a + (uint64_t(uint(iy) >> 1) * (pc.W >> 1) + (uint(ix) >> 1)) * pc.C * 4ul; + } else { + at = pc.a + (uint64_t(iy) * pc.W + uint(ix)) * pc.C * 4ul; + } + return true; +} + +void main() { + const uint t = gl_LocalInvocationIndex; + const uint m = pc.m0 + gl_WorkGroupID.y * 256u + t; + const bool row_ok = m < pc.M; + // K = taps * cols: a convolution's taps of C channels, or one "tap" of K columns for a GEMM. + const uint taps = TAPS == 0u ? 1u : TAPS, cols = TAPS == 0u ? pc.K : pc.C; + vec4 acc = vec4(0.0); + for (uint tap = 0u; tap < taps; ++tap) { + uint64_t base = 0ul; + const bool inside = row_ok && row_base(m, tap, base); + for (uint c0 = 0u; c0 < cols; c0 += KC) { + const uint kn = min(KC, cols - c0); + barrier(); // the previous chunk's reads are done + for (uint i = t; i < kn; i += 256u) { + const uint k = tap * cols + c0 + i; + vec4 col = vec4(0.0); + [[unroll]] for (uint n = 0u; n < 4u; ++n) { + if (n < pc.N) { + const float v = Floats(pc.b + (uint64_t(n) * pc.ldb + k) * 4ul).d[0]; + col[n] = TF32 ? tf32(v) : v; + } + } + Bs[i] = col; + } + barrier(); + if (inside) { + for (uint i = 0u; i < kn; i += 4u) { + const vec4 a = round4(Vec4s(base + uint64_t(c0 + i) * 4ul).d[0]); + acc += a.x * Bs[i]; + acc += a.y * Bs[i + 1u]; + acc += a.z * Bs[i + 2u]; + acc += a.w * Bs[i + 3u]; + } + } + } + } + if (!row_ok) return; + [[unroll]] for (uint n = 0u; n < 4u; ++n) { + if (n >= pc.N) continue; + const uint64_t at = uint64_t(m) * pc.ldo + n; + precise float y = acc[n]; + if (pc.bias != 0ul) y = y + Floats(pc.bias).d[n]; + if (pc.residual != 0ul) y = y + Floats(pc.residual + at * 4ul).d[0]; + Floats(pc.out_ + at * 4ul).d[0] = y; + } +} diff --git a/vulkan/shaders/vae_conv_plain.comp b/vulkan/shaders/vae_conv_plain.comp new file mode 100644 index 0000000..ec1919e --- /dev/null +++ b/vulkan/shaders/vae_conv_plain.comp @@ -0,0 +1,163 @@ +// The VAE's convolutions and fp32 GEMMs without cooperative matrices, 16-bit types or a particular subgroup size: +// out[m, n] = (sum_k A[m, k] * B[n, k] + bias[n]) + residual[m, n] (each + rounded separately, as the reference) +// with fp32 operands and accumulators, K ascending. TF32 rounds both operands to TF32 (half away from zero) first, as +// cuDNN's TF32 convolution does; their products are then exact in fp32, so only the accumulation order differs from it. +// A and B are addressed as in vae_conv_coop.comp: TAPS 9 / 1 is a 3x3 / 1x1 convolution over an NHWC input of C channels +// for an H x W output (UPSAMPLE: read through a nearest 2x upsample of an H/2 x W/2 input), TAPS 0 a row-major A with +// stride lda; B is row-major [N, K]. K, C, lda and ldb are multiples of 4. M and N are free. +// +// 256 invocations as 16 x 16; a workgroup computes a tile of 128 rows x 16 TN columns (TN 4 or 6), invocation (x, y) the +// rows 8 y .. +8 and columns TN x .. +TN, reading them from shared memory as vectors. K walks in steps of 16 through +// shared memory; the next step's global loads are issued before the current step's products. +#version 460 +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const uint TAPS = 9; +layout(constant_id = 1) const bool UPSAMPLE = false; +layout(constant_id = 2) const bool TF32 = true; +layout(constant_id = 3) const uint TN = 6; // 4 or 6 + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { vec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t a; + uint64_t b; + uint64_t out_; + uint64_t bias; + uint64_t residual; + uint M; + uint N; + uint K; + uint lda; + uint ldb; + uint ldo; + uint H; + uint W; + uint C; + uint m0; +} pc; + +const uint TM = 8u; +const uint BM = 128u; +const uint BN = 16u * TN; +const uint BK = 16u; +const uint LDA_S = BM + 4u; // floats per k row of the A tile: a multiple of 4, rows offset by 4 banks +const uint LDB_S = BN + 2u; // a multiple of 2 +shared vec4 As[BK * LDA_S / 4u]; +shared vec2 Bs[BK * LDB_S / 2u]; + +float tf32(float v) { + const uint u = floatBitsToUint(v); + if ((u & 0x7f800000u) == 0x7f800000u) return v; + return uintBitsToFloat((u + 0x1000u) & 0xffffe000u); +} +vec4 round4(vec4 v) { return TF32 ? vec4(tf32(v.x), tf32(v.y), tf32(v.z), tf32(v.w)) : v; } + +vec4 load_a(uint m, uint k) { + if (m >= pc.M || k >= pc.K) return vec4(0.0); + uint64_t at; + if (TAPS == 0u) { + at = uint64_t(m) * pc.lda + k; + } else { + const uint tap = k / pc.C, c = k - tap * pc.C; + if (TAPS == 1u) { + at = uint64_t(m) * pc.C + c; + } else { + const uint oy = m / pc.W, ox = m - oy * pc.W; + const int iy = int(oy) + int(tap / 3u) - 1, ix = int(ox) + int(tap % 3u) - 1; + if (iy < 0 || ix < 0 || iy >= int(pc.H) || ix >= int(pc.W)) return vec4(0.0); + if (UPSAMPLE) { + at = (uint64_t(uint(iy) >> 1) * (pc.W >> 1) + (uint(ix) >> 1)) * pc.C + c; + } else { + at = (uint64_t(iy) * pc.W + uint(ix)) * pc.C + c; + } + } + } + return round4(Vec4s(pc.a + at * 4ul).d[0]); +} + +vec4 load_b(uint n, uint k) { + if (n >= pc.N || k >= pc.K) return vec4(0.0); + return round4(Vec4s(pc.b + (uint64_t(n) * pc.ldb + k) * 4ul).d[0]); +} + +// Chunks of 4 k values of one row: BM * 4 = 512 of A, BN * 4 <= 384 of B; invocation t takes chunks t and t + 256. +const uint B_CHUNKS = BN * 4u; + +void store_a(uint chunk, vec4 v) { + const uint row = chunk / 4u, kk = (chunk % 4u) * 4u; + // As holds vec4s of 4 consecutive rows: scalar writes go through the vec4's components. + const uint at = row / 4u, lane = row % 4u; + [[unroll]] for (uint j = 0u; j < 4u; ++j) As[(kk + j) * (LDA_S / 4u) + at][lane] = v[j]; +} +void store_b(uint chunk, vec4 v) { + const uint row = chunk / 4u, kk = (chunk % 4u) * 4u; + const uint at = row / 2u, lane = row % 2u; + [[unroll]] for (uint j = 0u; j < 4u; ++j) Bs[(kk + j) * (LDB_S / 2u) + at][lane] = v[j]; +} + +void main() { + const uint t = gl_LocalInvocationIndex; + const uint tx = t % 16u, ty = t / 16u; + const uint m_base = pc.m0 + gl_WorkGroupID.y * BM; + const uint n_base = gl_WorkGroupID.x * BN; + + float acc[TM][6]; + [[unroll]] for (uint i = 0u; i < TM; ++i) { + [[unroll]] for (uint j = 0u; j < 6u; ++j) acc[i][j] = 0.0; + } + + vec4 ra0, ra1, rb0, rb1; + const uint steps = (pc.K + BK - 1u) / BK; + ra0 = load_a(m_base + t / 4u, (t % 4u) * 4u); + ra1 = load_a(m_base + (t + 256u) / 4u, (t % 4u) * 4u); + rb0 = load_b(n_base + t / 4u, (t % 4u) * 4u); + rb1 = t + 256u < B_CHUNKS ? load_b(n_base + (t + 256u) / 4u, (t % 4u) * 4u) : vec4(0.0); + for (uint step = 0u; step < steps; ++step) { + barrier(); // the previous step's reads of As and Bs are done + store_a(t, ra0); + store_a(t + 256u, ra1); + if (t < B_CHUNKS) store_b(t, rb0); + if (t + 256u < B_CHUNKS) store_b(t + 256u, rb1); + barrier(); + if (step + 1u < steps) { + const uint k = (step + 1u) * BK + (t % 4u) * 4u; + ra0 = load_a(m_base + t / 4u, k); + ra1 = load_a(m_base + (t + 256u) / 4u, k); + if (t < B_CHUNKS) rb0 = load_b(n_base + t / 4u, k); + if (t + 256u < B_CHUNKS) rb1 = load_b(n_base + (t + 256u) / 4u, k); + } + [[unroll]] for (uint k = 0u; k < BK; ++k) { + const vec4 a0 = As[k * (LDA_S / 4u) + ty * 2u], a1 = As[k * (LDA_S / 4u) + ty * 2u + 1u]; + const float av[8] = float[8](a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w); + float bv[6]; + [[unroll]] for (uint q = 0u; q < 3u; ++q) { + const vec2 b = q < TN / 2u ? Bs[k * (LDB_S / 2u) + tx * (TN / 2u) + q] : vec2(0.0); + bv[2u * q] = b.x; + bv[2u * q + 1u] = b.y; + } + [[unroll]] for (uint i = 0u; i < TM; ++i) { + [[unroll]] for (uint j = 0u; j < 6u; ++j) { + if (j < TN) acc[i][j] += av[i] * bv[j]; + } + } + } + } + + [[unroll]] for (uint i = 0u; i < TM; ++i) { + const uint m = m_base + ty * TM + i; + if (m >= pc.M) continue; + [[unroll]] for (uint j = 0u; j < 6u; ++j) { + const uint n = n_base + tx * TN + j; + if (j >= TN || n >= pc.N) continue; + const uint64_t at = uint64_t(m) * pc.ldo + n; + precise float y = acc[i][j]; + if (pc.bias != 0ul) y = y + Floats(pc.bias).d[n]; + if (pc.residual != 0ul) y = y + Floats(pc.residual + at * 4ul).d[0]; + Floats(pc.out_ + at * 4ul).d[0] = y; + } + } +} diff --git a/vulkan/shaders/vae_gn_apply.comp b/vulkan/shaders/vae_gn_apply.comp new file mode 100644 index 0000000..8f12009 --- /dev/null +++ b/vulkan/shaders/vae_gn_apply.comp @@ -0,0 +1,45 @@ +// GroupNorm apply on an NHWC fp32 tensor [P, C] (C a multiple of 32 and 4), in place or not: +// y = x * a + b with a = rstd * gamma[c], b = -a * mean + beta[c] (torch's fused GroupNorm parameters) +// then, with SILU, y / (1 + exp(-y)) as torch's silu. One invocation per 4 channels of one pixel, dispatched as +// x + y * x_count workgroups of 256. +#version 460 +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { vec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t y; + uint64_t stats; // [32] {mean, rstd} + uint64_t gamma; // [C] + uint64_t beta; // [C] + uint quads; // P * C / 4 + uint C; + uint silu; +} pc; + +void main() { + const uint i = (gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x) * 256u + gl_LocalInvocationIndex; + if (i >= pc.quads) return; + const uint c0 = (i % (pc.C / 4u)) * 4u, cpg = pc.C / 32u; + const vec4 v = Vec4s(pc.x + uint64_t(i) * 16ul).d[0]; + Floats st = Floats(pc.stats), gm = Floats(pc.gamma), bt = Floats(pc.beta); + vec4 r; + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint c = c0 + j, g = c / cpg; + precise const float a = st.d[2u * g + 1u] * gm.d[c]; + const float b = fma(-a, st.d[2u * g], bt.d[c]); + float y = fma(a, v[j], b); + if (pc.silu != 0u) { + precise const float denom = 1.0 + exp(-y); + precise const float s = y / denom; + y = s; + } + r[j] = y; + } + Vec4s(pc.y + uint64_t(i) * 16ul).d[0] = r; +} diff --git a/vulkan/shaders/vae_gn_finalize.comp b/vulkan/shaders/vae_gn_finalize.comp new file mode 100644 index 0000000..0d71fd7 --- /dev/null +++ b/vulkan/shaders/vae_gn_finalize.comp @@ -0,0 +1,50 @@ +// GroupNorm statistics, pass 2: merges the chunks' {count, mean, m2} of one group (Chan et al.'s pairwise update) and +// writes stats[g] = {mean, rstd} with rstd = 1 / sqrt(m2 / count + eps), the biased variance as GroupNorm uses. One +// workgroup per group: invocation t merges chunks t, t + 256, ... in order, then a halving tree. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t part; + uint64_t stats; + uint chunks; + float eps; +} pc; + +shared vec3 red[256]; + +vec3 merge(vec3 a, vec3 b) { + if (a.x == 0.0) return b; + if (b.x == 0.0) return a; + precise const float n = a.x + b.x; + precise const float delta = b.y - a.y; + precise const float mean = a.y + delta * (b.x / n); + precise const float m2 = a.z + b.z + delta * delta * (a.x * (b.x / n)); + return vec3(n, mean, m2); +} + +void main() { + const uint g = gl_WorkGroupID.x, t = gl_LocalInvocationIndex; + vec3 acc = vec3(0.0); + for (uint c = t; c < pc.chunks; c += 256u) { + Floats p = Floats(pc.part + (uint64_t(c) * 32u + g) * 16ul); + acc = merge(acc, vec3(p.d[0], p.d[1], p.d[2])); + } + red[t] = acc; + barrier(); + for (uint off = 128u; off > 0u; off >>= 1) { + if (t < off) red[t] = merge(red[t], red[t + off]); + barrier(); + } + if (t == 0u) { + const vec3 r = red[0]; + precise const float var = r.z / r.x; + precise const float var_eps = var + pc.eps; + Floats s = Floats(pc.stats + uint64_t(g) * 8ul); + s.d[0] = r.y; + s.d[1] = 1.0 / sqrt(var_eps); + } +} diff --git a/vulkan/shaders/vae_gn_stats.comp b/vulkan/shaders/vae_gn_stats.comp new file mode 100644 index 0000000..719d0a2 --- /dev/null +++ b/vulkan/shaders/vae_gn_stats.comp @@ -0,0 +1,85 @@ +// GroupNorm statistics, pass 1: per chunk of pixels of an NHWC fp32 tensor [P, C] and per group of C / 32 channels, the +// chunk's element count, mean and sum of squared deviations from that mean (two passes over the chunk, so there is no +// cancellation): part[(chunk * 32 + g) * 4 + {0, 1, 2}] = {count, mean, m2}. vae_gn_finalize.comp merges the chunks. +// C is a multiple of 4 (and of 32) up to 1024. One chunk of `ppc` pixels per workgroup, dispatched as x + y * x_count. +// With quads = C / 4 and rows = 256 / quads, invocation t < quads * rows owns channels 4 (t % quads) .. +4 of the pixels +// p0 + t / quads, + rows, ...; per-channel sums meet in shared memory and are summed per group in a fixed order. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { vec4 d[]; }; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t part; + uint pixels; + uint C; + uint ppc; // pixels per chunk + uint chunks; +} pc; + +shared float red[1024]; +shared float gmean[32]; + +void main() { + const uint chunk = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; + if (chunk >= pc.chunks) return; // uniform per workgroup + const uint t = gl_LocalInvocationIndex; + const uint quads = pc.C / 4u, rows = 256u / quads, cpg = pc.C / 32u; + const bool in_tile = t < quads * rows; + const uint quad = t % quads, row = t / quads; + const uint p0 = chunk * pc.ppc, p1 = min(p0 + pc.ppc, pc.pixels); + + vec4 s = vec4(0.0); + if (in_tile) { + for (uint p = p0 + row; p < p1; p += rows) s += Vec4s(pc.x + (uint64_t(p) * pc.C + quad * 4u) * 4ul).d[0]; + } + if (in_tile) { + red[row * pc.C + quad * 4u + 0u] = s.x; + red[row * pc.C + quad * 4u + 1u] = s.y; + red[row * pc.C + quad * 4u + 2u] = s.z; + red[row * pc.C + quad * 4u + 3u] = s.w; + } + barrier(); + const float count = float((p1 - p0) * cpg); + if (t < 32u) { + float sum = 0.0; + for (uint r = 0u; r < rows; ++r) { + for (uint c = t * cpg; c < (t + 1u) * cpg; ++c) sum += red[r * pc.C + c]; + } + gmean[t] = sum / count; + } + barrier(); + + vec4 q = vec4(0.0); + if (in_tile) { + const uint c0 = quad * 4u; + const vec4 mu = vec4(gmean[c0 / cpg], gmean[(c0 + 1u) / cpg], gmean[(c0 + 2u) / cpg], gmean[(c0 + 3u) / cpg]); + for (uint p = p0 + row; p < p1; p += rows) { + const vec4 d = Vec4s(pc.x + (uint64_t(p) * pc.C + c0) * 4ul).d[0] - mu; + q += d * d; + } + } + barrier(); // the group sums are done with red + if (in_tile) { + red[row * pc.C + quad * 4u + 0u] = q.x; + red[row * pc.C + quad * 4u + 1u] = q.y; + red[row * pc.C + quad * 4u + 2u] = q.z; + red[row * pc.C + quad * 4u + 3u] = q.w; + } + barrier(); + if (t < 32u) { + float m2 = 0.0; + for (uint r = 0u; r < rows; ++r) { + for (uint c = t * cpg; c < (t + 1u) * cpg; ++c) m2 += red[r * pc.C + c]; + } + Floats out_ = Floats(pc.part + (uint64_t(chunk) * 32u + t) * 16ul); + out_.d[0] = count; + out_.d[1] = gmean[t]; + out_.d[2] = m2; + out_.d[3] = 0.0; + } +} diff --git a/vulkan/shaders/vae_latent_unpack.comp b/vulkan/shaders/vae_latent_unpack.comp new file mode 100644 index 0000000..d50a794 --- /dev/null +++ b/vulkan/shaders/vae_latent_unpack.comp @@ -0,0 +1,30 @@ +// The pipeline's latent de-normalization and unpack (pipeline_moss_image.py:430-436): +// y[2i + p, 2j + q, c] = z[4c + 2p + q, i, j] * scale[4c + 2p + q] + mean[4c + 2p + q] +// for z [128, h, w] fp32 NCHW and y [2h, 2w, 32] fp32 NHWC, with scale = sqrt(var + eps) formed on the host (correctly +// rounded, as torch) and the * and + rounded separately. One invocation per output element. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t z; + uint64_t scale; + uint64_t mean; + uint64_t y; + uint h; + uint w; +} pc; + +void main() { + const uint o = (gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x) * 256u + gl_LocalInvocationIndex; + if (o >= pc.h * pc.w * 128u) return; + const uint c = o % 32u, px = o / 32u; + const uint ox = px % (2u * pc.w), oy = px / (2u * pc.w); + const uint ch = 4u * c + 2u * (oy & 1u) + (ox & 1u); + const float v = Floats(pc.z + ((uint64_t(ch) * pc.h + oy / 2u) * pc.w + ox / 2u) * 4ul).d[0]; + precise const float scaled = v * Floats(pc.scale).d[ch]; + precise const float shifted = scaled + Floats(pc.mean).d[ch]; + Floats(pc.y + uint64_t(o) * 4ul).d[0] = shifted; +} diff --git a/vulkan/shaders/vae_rgb8.comp b/vulkan/shaders/vae_rgb8.comp new file mode 100644 index 0000000..ad7e5de --- /dev/null +++ b/vulkan/shaders/vae_rgb8.comp @@ -0,0 +1,38 @@ +// The decoder output [n] fp32 (NHWC RGB) to RGB8 as the pipeline does: .clamp(-1, 1), then VaeImageProcessor.postprocess: +// (x * 0.5 + 0.5).clamp(0, 1), numpy's (x * 255).round() (half to even) and .astype(uint8); NaN becomes 0. Each +// invocation writes one 32-bit word of 4 bytes; bytes past n in the last word are written as 0. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(buffer_reference, std430, buffer_reference_align = 4) buffer Uints { uint d[]; }; + +layout(push_constant) uniform PC { + uint64_t x; + uint64_t rgb; + uint n; +} pc; + +uint to_u8(float v) { + if (isnan(v)) return 0u; + v = clamp(v, -1.0, 1.0); + precise const float half_ = v * 0.5; + precise const float shifted = half_ + 0.5; + v = clamp(shifted, 0.0, 1.0); + precise const float scaled = v * 255.0; + return uint(roundEven(scaled)); +} + +void main() { + const uint word = (gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x) * 256u + gl_LocalInvocationIndex; + const uint i0 = word * 4u; + if (i0 >= pc.n) return; + Floats X = Floats(pc.x); + uint packed = 0u; + for (uint j = 0u; j < 4u; ++j) { + if (i0 + j < pc.n) packed |= to_u8(X.d[i0 + j]) << (8u * j); + } + Uints(pc.rgb).d[word] = packed; +} diff --git a/vulkan/shaders/vae_softmax.comp b/vulkan/shaders/vae_softmax.comp new file mode 100644 index 0000000..83d6754 --- /dev/null +++ b/vulkan/shaders/vae_softmax.comp @@ -0,0 +1,58 @@ +// Row softmax of attention scores in place: s[r, :L] (row stride ld) -> exp(s * scale - max * scale) / sum, the +// probabilities SDPA multiplies V by. fp32; one workgroup per row, dispatched as x + y * x_count; three passes over the +// row (max, sum, write), each a strided per-invocation loop and a halving tree. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t s; + uint rows; + uint L; + uint ld; + float scale; +} pc; + +shared float red[256]; + +void main() { + const uint row = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; + if (row >= pc.rows) return; // uniform per workgroup + const uint t = gl_LocalInvocationIndex; + Floats S = Floats(pc.s + uint64_t(row) * pc.ld * 4ul); + + float m = -3.402823466e38; + for (uint j = t; j < pc.L; j += 256u) m = max(m, S.d[j]); + red[t] = m; + barrier(); + for (uint off = 128u; off > 0u; off >>= 1) { + if (t < off) red[t] = max(red[t], red[t + off]); + barrier(); + } + precise const float shift = red[0] * pc.scale; + barrier(); + + float sum = 0.0; + for (uint j = t; j < pc.L; j += 256u) { + precise const float z = S.d[j] * pc.scale - shift; + sum += exp(z); + } + red[t] = sum; + barrier(); + for (uint off = 128u; off > 0u; off >>= 1) { + if (t < off) { + precise const float v = red[t] + red[t + off]; + red[t] = v; + } + barrier(); + } + const float inv = 1.0 / red[0]; + + for (uint j = t; j < pc.L; j += 256u) { + precise const float z = S.d[j] * pc.scale - shift; + precise const float p = exp(z) * inv; + S.d[j] = p; + } +} diff --git a/vulkan/shaders/vae_transpose.comp b/vulkan/shaders/vae_transpose.comp new file mode 100644 index 0000000..00b0866 --- /dev/null +++ b/vulkan/shaders/vae_transpose.comp @@ -0,0 +1,32 @@ +// fp32 transpose: out[c, r] = in[r, c] for in [R, Cc] (row stride ld_in) and out [Cc, R] (row stride ld_out), through +// 32 x 33 shared tiles. Workgroup (x, y) moves the tile at columns 32 x, rows 32 y. +#version 460 +#extension GL_GOOGLE_include_directive : require +#include "common.glsl" + +layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in; + +layout(push_constant) uniform PC { + uint64_t in_; + uint64_t out_; + uint R; + uint Cc; + uint ld_in; + uint ld_out; +} pc; + +shared float tile[32 * 33]; + +void main() { + const uint c0 = gl_WorkGroupID.x * 32u, r0 = gl_WorkGroupID.y * 32u; + const uint tx = gl_LocalInvocationID.x, ty = gl_LocalInvocationID.y; + for (uint i = ty; i < 32u; i += 8u) { + const uint r = r0 + i, c = c0 + tx; + if (r < pc.R && c < pc.Cc) tile[i * 33u + tx] = Floats(pc.in_ + (uint64_t(r) * pc.ld_in + c) * 4ul).d[0]; + } + barrier(); + for (uint i = ty; i < 32u; i += 8u) { + const uint c = c0 + i, r = r0 + tx; + if (r < pc.R && c < pc.Cc) Floats(pc.out_ + (uint64_t(c) * pc.ld_out + r) * 4ul).d[0] = tile[tx * 33u + i]; + } +} diff --git a/vulkan/te_kernels.cpp b/vulkan/te_kernels.cpp new file mode 100644 index 0000000..b8c7ec2 --- /dev/null +++ b/vulkan/te_kernels.cpp @@ -0,0 +1,262 @@ +#include "te_kernels.h" + +#include +#include +#include +#include + +namespace mi::vk { +namespace { + +const std::uint32_t kNormSpirv[] = { +#include "te_rms_norm.spv.inc" +}; +const std::uint32_t kConvSpirv[] = { +#include "te_conv_silu.spv.inc" +}; +const std::uint32_t kGatesSpirv[] = { +#include "te_gates.spv.inc" +}; +const std::uint32_t kGdnQkSpirv[] = { +#include "te_gdn_qk.spv.inc" +}; +const std::uint32_t kGdnChunkSpirv[] = { +#include "te_gdn_chunk.spv.inc" +}; +const std::uint32_t kGdnSolveSpirv[] = { +#include "te_gdn_solve.spv.inc" +}; +const std::uint32_t kGdnStateSpirv[] = { +#include "te_gdn_state.spv.inc" +}; +const std::uint32_t kAttentionSpirv[] = { +#include "te_attention.spv.inc" +}; +const std::uint32_t kGateMulSpirv[] = { +#include "te_gate_mul.spv.inc" +}; + +struct NormPush { + std::uint64_t x, weight, gate, out, table; + std::uint32_t ldx, hsx, ldo, hso, ldg, hsg, table_ld, d, bw; + float factor, eps; +}; +struct ConvPush { + std::uint64_t x, w, y; + std::uint32_t ldx, ldy, channels; +}; +struct GatesPush { + std::uint64_t ba, a_log, dt_bias, g, beta; + std::uint32_t ldba, tokens; +}; +struct GdnQkPush { + std::uint64_t qkv, qn, kn; + std::uint32_t ldqkv, tokens; +}; +struct GdnChunkPush { + std::uint64_t qn, kn, g, beta, gcum, eg, el, a, attn; + std::uint32_t tokens, row0, len; +}; +struct GdnSolvePush { + std::uint64_t qkv, kn, beta, eg, a, u, w; + std::uint32_t ldqkv, tokens, row0, len; +}; +struct GdnStatePush { + std::uint64_t qn, kn, gcum, eg, el, attn, u, w, out; + std::uint32_t tokens, row0, len; +}; +struct AttentionPush { + std::uint64_t q, k, v, o; + std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_first, row_end, group; + float log2_scale; +}; +struct GateMulPush { + std::uint64_t a, b, out; + std::uint32_t lda, ldb, ldo, cols, block, block_stride; +}; + +constexpr std::uint32_t kMaxGroups = 65535; // the guaranteed maxComputeWorkGroupCount per dimension + +void require(bool condition, const std::string& message) { + if (!condition) throw std::invalid_argument("vulkan text encoder: " + message); +} + +template +std::unique_ptr pipeline(Device& device, const std::uint32_t (&spirv)[N], std::uint32_t push, int mode = -1) { + std::vector spec; + if (mode >= 0) spec.push_back(static_cast(mode)); + return std::make_unique(device, spirv, sizeof(spirv), push, 0, spec); +} + +std::uint32_t pow2_floor(std::uint64_t n) { + std::uint32_t p = 1; + while (static_cast(p) * 2 <= n && p < (1u << 30)) p *= 2; + return p; +} + +} // namespace + +std::uint32_t torch_reduce_threads(std::uint64_t rows, std::uint32_t d) { + require(d % 4 == 0 && d >= 128, "norm rows must be a multiple of 4 and at least 128 long"); + constexpr std::uint32_t kMaxThreads = 512; + const std::uint32_t dim0 = d / 4; + const std::uint32_t dim0_pow2 = dim0 < kMaxThreads ? pow2_floor(dim0) : kMaxThreads; + const std::uint32_t dim1_pow2 = rows < kMaxThreads ? pow2_floor(rows) : kMaxThreads; + std::uint32_t bw = std::min(dim0_pow2, 32u); + const std::uint32_t bh = std::min(dim1_pow2, kMaxThreads / bw); + bw = std::min(dim0_pow2, kMaxThreads / bh); + const std::uint32_t per_thread = (d + bw - 1) / bw; + require(per_thread < std::min(bh * 16, 256u), "row length " + std::to_string(d) + " needs torch's cross-warp split"); + return bw; +} + +std::uint64_t GdnWorkspace::bytes(std::uint32_t tokens) { + GdnWorkspace ws = at(0, tokens); + return ws.w + static_cast(kGdnValueHeads) * tokens * kGdnHeadDim * 4; +} + +GdnWorkspace GdnWorkspace::at(VkDeviceAddress base, std::uint32_t tokens) { + const auto align = [](std::uint64_t v) { return (v + 255) & ~std::uint64_t{255}; }; + const std::uint64_t t = tokens; + const std::uint64_t key = align(kGdnKeyHeads * t * kGdnHeadDim * 4), per_head = align(kGdnValueHeads * t * 4), + chunk = align(kGdnValueHeads * t * kGdnChunk * 4), + value = align(kGdnValueHeads * t * kGdnHeadDim * 4); + GdnWorkspace ws{}; + ws.qn = base; + ws.kn = ws.qn + key; + ws.gcum = ws.kn + key; + ws.eg = ws.gcum + per_head; + ws.el = ws.eg + per_head; + ws.a = ws.el + per_head; + ws.attn = ws.a + chunk; + ws.u = ws.attn + chunk; + ws.w = ws.u + value; + return ws; +} + +TeKernels::TeKernels(Device& device) { + norm_ = pipeline(device, kNormSpirv, sizeof(NormPush), 0); + norm_gated_ = pipeline(device, kNormSpirv, sizeof(NormPush), 1); + conv_ = pipeline(device, kConvSpirv, sizeof(ConvPush)); + gates_ = pipeline(device, kGatesSpirv, sizeof(GatesPush)); + gdn_qk_ = pipeline(device, kGdnQkSpirv, sizeof(GdnQkPush)); + gdn_chunk_ = pipeline(device, kGdnChunkSpirv, sizeof(GdnChunkPush)); + gdn_solve_ = pipeline(device, kGdnSolveSpirv, sizeof(GdnSolvePush)); + gdn_state_ = pipeline(device, kGdnStateSpirv, sizeof(GdnStatePush)); + attention_ = pipeline(device, kAttentionSpirv, sizeof(AttentionPush)); + silu_mul_ = pipeline(device, kGateMulSpirv, sizeof(GateMulPush), 0); + sigmoid_mul_ = pipeline(device, kGateMulSpirv, sizeof(GateMulPush), 1); +} + +TeKernels::~TeKernels() = default; + +void TeKernels::record_rms_norm(VkCommandBuffer cb, const RmsNormArgs& a) const { + require(a.d % 4 == 0 && a.ldx % 4 == 0 && a.hsx % 4 == 0 && a.ldo % 4 == 0 && a.hso % 4 == 0 && + (!a.gated || (a.ldg % 4 == 0 && a.hsg % 4 == 0)), + "norm: d and strides must be multiples of 4"); + require(a.rope == 0 || (!a.gated && a.d >= 64 && a.rope_ld % 4 == 0), "norm: RoPE needs the plain norm and d >= 64"); + require(a.heads >= 1 && a.heads <= kMaxGroups && a.tokens <= kMaxGroups, "norm: too many rows"); + if (a.tokens == 0) return; + const std::uint64_t rows = static_cast(a.tokens) * a.heads; + const float factor = static_cast(rows) / static_cast(rows * a.d); + const NormPush push{a.x, + a.weight, + a.gated ? a.gate : 0, + a.out, + a.gated ? 0 : a.rope, + a.ldx, + a.hsx, + a.ldo, + a.hso, + a.ldg, + a.hsg, + a.rope_ld, + a.d, + torch_reduce_threads(rows, a.d), + factor, + a.eps}; + (a.gated ? norm_gated_ : norm_)->dispatch(cb, &push, a.heads, a.tokens); + compute_barrier(cb); +} + +void TeKernels::record_conv_silu(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, + VkDeviceAddress y, std::uint32_t ldy, std::uint32_t tokens) const { + require(ldx % 4 == 0 && ldy % 4 == 0 && tokens <= kMaxGroups, "conv: strides must be multiples of 4"); + if (tokens == 0) return; + const ConvPush push{x, weight, y, ldx, ldy, kGdnConvDim}; + conv_->dispatch(cb, &push, kGdnConvDim / 4 / 64, tokens); + compute_barrier(cb); +} + +void TeKernels::record_gates(VkCommandBuffer cb, VkDeviceAddress ba, std::uint32_t ldba, VkDeviceAddress a_log, + VkDeviceAddress dt_bias, VkDeviceAddress g, VkDeviceAddress beta, + std::uint32_t tokens) const { + require(ldba % 2 == 0, "gates: the stride must be even"); + if (tokens == 0) return; + const GatesPush push{ba, a_log, dt_bias, g, beta, ldba, tokens}; + gates_->dispatch(cb, &push, (tokens * 16 + 63) / 64, 1); + compute_barrier(cb); +} + +void TeKernels::record_gated_delta_rule(VkCommandBuffer cb, const GdnArgs& a) const { + require(a.cu.size() >= 2 && a.cu.front() == 0, "gated delta rule: cu must start at 0 and hold a sequence"); + const auto tokens = static_cast(a.cu.back()); + require(tokens <= kMaxGroups && a.ldqkv % 4 == 0, "gated delta rule: too many tokens or a bad stride"); + if (tokens == 0) return; + const GdnQkPush qk{a.qkv, a.ws.qn, a.ws.kn, a.ldqkv, tokens}; + gdn_qk_->dispatch(cb, &qk, kGdnKeyHeads, tokens); + compute_barrier(cb); + const auto for_each = [&](auto&& f) { + for (std::size_t i = 0; i + 1 < a.cu.size(); ++i) { + require(a.cu[i + 1] >= a.cu[i], "gated delta rule: cu must not decrease"); + const auto row0 = static_cast(a.cu[i]), len = static_cast(a.cu[i + 1] - a.cu[i]); + if (len > 0) f(row0, len, (len + kGdnChunk - 1) / kGdnChunk); + } + }; + for_each([&](std::uint32_t row0, std::uint32_t len, std::uint32_t chunks) { + const GdnChunkPush push{a.ws.qn, a.ws.kn, a.g, a.beta, a.ws.gcum, a.ws.eg, + a.ws.el, a.ws.a, a.ws.attn, tokens, row0, len}; + gdn_chunk_->dispatch(cb, &push, chunks, kGdnValueHeads); + }); + compute_barrier(cb); + for_each([&](std::uint32_t row0, std::uint32_t len, std::uint32_t chunks) { + const GdnSolvePush push{a.qkv, a.ws.kn, a.beta, a.ws.eg, a.ws.a, a.ws.u, a.ws.w, a.ldqkv, tokens, row0, len}; + gdn_solve_->dispatch(cb, &push, chunks * 4, kGdnValueHeads); + }); + compute_barrier(cb); + for_each([&](std::uint32_t row0, std::uint32_t len, std::uint32_t) { + const GdnStatePush push{a.ws.qn, a.ws.kn, a.ws.gcum, a.ws.eg, a.ws.el, a.ws.attn, + a.ws.u, a.ws.w, a.out, tokens, row0, len}; + gdn_state_->dispatch(cb, &push, kGdnHeadDim / 16, kGdnValueHeads); + }); + compute_barrier(cb); +} + +void TeKernels::record_attention(VkCommandBuffer cb, const TeAttentionArgs& a) const { + require(a.ldq % 4 == 0 && a.ldk % 4 == 0 && a.ldv % 4 == 0 && a.ldo % 4 == 0 && a.hsq % 4 == 0 && a.hsk % 4 == 0 && + a.hsv % 4 == 0 && a.hso % 4 == 0, + "attention: strides must be multiples of 4"); + require(a.group >= 1 && a.heads % a.group == 0 && a.q_first % 16 == 0 && a.rows / 16 < kMaxGroups, + "attention: bad head grouping, q_first not a multiple of 16, or too many rows"); + if (a.rows == 0) return; + // In double then rounded, as FlashAttention (float * M_LOG2E). + const float log2_scale = static_cast(static_cast(a.scale) * 1.4426950408889634074); + const AttentionPush push{a.q, a.k, a.v, a.o, a.ldq, a.ldk, a.ldv, a.ldo, + a.hsq, a.hsk, a.hsv, a.hso, a.q_first, a.q_first + a.rows, a.group, log2_scale}; + attention_->dispatch(cb, &push, a.heads, (a.rows + 15) / 16); + compute_barrier(cb); +} + +void TeKernels::record_gate_mul(VkCommandBuffer cb, GateMul op, VkDeviceAddress a, std::uint32_t lda, VkDeviceAddress b, + std::uint32_t ldb, VkDeviceAddress out, std::uint32_t ldo, std::uint32_t rows, + std::uint32_t cols, std::uint32_t block, std::uint32_t block_stride) const { + require(lda % 4 == 0 && ldb % 4 == 0 && ldo % 4 == 0 && cols % 4 == 0 && block % 4 == 0 && block_stride % 4 == 0 && + block > 0 && rows <= kMaxGroups, + "gate_mul: strides, columns and blocks must be multiples of 4"); + if (rows == 0 || cols == 0) return; + const GateMulPush push{a, b, out, lda, ldb, ldo, cols, block, block_stride}; + (op == GateMul::kSilu ? silu_mul_ : sigmoid_mul_)->dispatch(cb, &push, (cols / 4 + 63) / 64, rows); + compute_barrier(cb); +} + +} // namespace mi::vk diff --git a/vulkan/te_kernels.h b/vulkan/te_kernels.h new file mode 100644 index 0000000..7e10916 --- /dev/null +++ b/vulkan/te_kernels.h @@ -0,0 +1,92 @@ +// The Qwen3.5 text encoder's kernels other than its GEMMs (those are BlockKernels'), in the engine's op order and +// rounding points (src/models/qwen35_text, src/kernels/gdn, src/kernels/basic). Plain compute: no cooperative matrices, +// 16-bit arithmetic or particular subgroup size, so both paths use them. Arguments are device addresses; strides are in +// elements. Each record_* adds its own barrier. +#pragma once + +#include +#include +#include + +#include "runtime.h" + +namespace mi::vk { + +inline constexpr int kGdnKeyHeads = 16; +inline constexpr int kGdnValueHeads = 32; +inline constexpr int kGdnHeadDim = 128; +inline constexpr int kGdnConvDim = 2 * kGdnKeyHeads * kGdnHeadDim + kGdnValueHeads * kGdnHeadDim; // [q | k | v], 8192 +inline constexpr int kGdnChunk = 64; + +// torch's threads per row for a last-dim fp32 reduction of `rows` rows of length d (setReduceConfig, 4-wide vector +// loads, 512 threads at most). Throws where torch would split a row across warps, which the norm does not reproduce. +std::uint32_t torch_reduce_threads(std::uint64_t rows, std::uint32_t d); + +// Rows are (token, head) pairs: row (t, h) of x starts at x + t * ldx + h * hsx, likewise out and gate. out may alias +// x. +struct RmsNormArgs { + VkDeviceAddress x = 0, weight = 0, gate = 0, out = 0; + std::uint32_t ldx = 0, hsx = 0, ldo = 0, hso = 0, ldg = 0, hsg = 0; + std::uint32_t tokens = 0, heads = 1, d = 0; + bool gated = false; // Qwen3_5RMSNormGated with `gate`; otherwise Qwen3_5RMSNorm + VkDeviceAddress rope = 0; // Qwen3_5RMSNorm only: [tokens, 64] bf16 table [cos(32) | sin(32)] applied after the norm + std::uint32_t rope_ld = 64; + float eps = 1e-6f; +}; + +// The gated delta rule's fp32 scratch for `tokens` packed tokens. +struct GdnWorkspace { + VkDeviceAddress qn, kn, gcum, eg, el, a, attn, u, w; + static std::uint64_t bytes(std::uint32_t tokens); + static GdnWorkspace at(VkDeviceAddress base, std::uint32_t tokens); +}; + +// Sequences packed along the rows: sequence i is rows cu[i] .. cu[i + 1]. +struct GdnArgs { + VkDeviceAddress qkv = 0; // conv output, bf16 rows [q | k | v] of stride ldqkv + std::uint32_t ldqkv = 0; + VkDeviceAddress g = 0; // f32 [tokens, 32] + VkDeviceAddress beta = 0; // bf16 [tokens, 32] + VkDeviceAddress out = 0; // bf16 [tokens, 32, 128] + GdnWorkspace ws{}; + std::vector cu; +}; + +// One sequence: q, k, v and o point at its first row. Query head h reads key/value head h / group. Rows q_first .. +// q_first + rows - 1 of the sequence are computed, each over keys 0 .. its own index; q_first is a multiple of 16. +struct TeAttentionArgs { + VkDeviceAddress q = 0, k = 0, v = 0, o = 0; + std::uint32_t ldq = 0, ldk = 0, ldv = 0, ldo = 0; + std::uint32_t hsq = 0, hsk = 0, hsv = 0, hso = 0; + std::uint32_t q_first = 0, rows = 0; + std::uint32_t heads = 0, group = 1; + float scale = 0.0625f; +}; + +class TeKernels { + public: + explicit TeKernels(Device& device); + ~TeKernels(); + + void record_rms_norm(VkCommandBuffer cb, const RmsNormArgs& args) const; + // One sequence of `tokens` rows: y = silu(causal_conv1d(x)), x and y bf16 [tokens, 8192] with strides ldx and ldy. + void record_conv_silu(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t ldx, VkDeviceAddress weight, + VkDeviceAddress y, std::uint32_t ldy, std::uint32_t tokens) const; + // ba: bf16 rows [b(32) | a(32)] of stride ldba; g f32 [tokens, 32]; beta bf16 [tokens, 32]. + void record_gates(VkCommandBuffer cb, VkDeviceAddress ba, std::uint32_t ldba, VkDeviceAddress a_log, + VkDeviceAddress dt_bias, VkDeviceAddress g, VkDeviceAddress beta, std::uint32_t tokens) const; + void record_gated_delta_rule(VkCommandBuffer cb, const GdnArgs& args) const; + void record_attention(VkCommandBuffer cb, const TeAttentionArgs& args) const; + // out = bf16(bf16(silu(a)) * b) (silu) or bf16(a * bf16(sigmoid(b))); column c of b is + // (c / block) * block_stride + c % block. + enum class GateMul { kSilu, kSigmoid }; + void record_gate_mul(VkCommandBuffer cb, GateMul op, VkDeviceAddress a, std::uint32_t lda, VkDeviceAddress b, + std::uint32_t ldb, VkDeviceAddress out, std::uint32_t ldo, std::uint32_t rows, + std::uint32_t cols, std::uint32_t block, std::uint32_t block_stride) const; + + private: + std::unique_ptr norm_, norm_gated_, conv_, gates_, gdn_qk_, gdn_chunk_, gdn_solve_, gdn_state_, attention_, + silu_mul_, sigmoid_mul_; +}; + +} // namespace mi::vk diff --git a/vulkan/text_encoder.cpp b/vulkan/text_encoder.cpp new file mode 100644 index 0000000..68dadec --- /dev/null +++ b/vulkan/text_encoder.cpp @@ -0,0 +1,638 @@ +#include "text_encoder.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "half.h" +#include "safetensors.h" +#include "te_kernels.h" + +namespace mi::vk { +namespace { + +using Clock = std::chrono::steady_clock; + +double seconds_since(Clock::time_point t0) { return std::chrono::duration(Clock::now() - t0).count(); } + +[[noreturn]] void fail(const std::string& what) { throw std::runtime_error("vulkan text encoder: " + what); } + +constexpr std::uint32_t kQkvDim = + 2 * kTeHeads * kTeHeadDim + 2 * kTeKvHeads * kTeHeadDim; // [q|gate per head] k v: 10240 +constexpr std::uint32_t kQgDim = 2 * kTeHeads * kTeHeadDim; // 8192 +constexpr std::uint32_t kKvDim = kTeKvHeads * kTeHeadDim; // 1024 +constexpr std::uint32_t kValueDim = kGdnValueHeads * kGdnHeadDim; // 4096 +constexpr std::uint32_t kQkvzDim = kGdnConvDim + kValueDim; // 12288 +constexpr std::uint32_t kGateUpDim = 2 * kTeIntermediate; // 24576 +constexpr std::uint32_t kRopeDim = 64; +constexpr std::uint32_t kMaxTokens = 65535; + +std::uint64_t align256(std::uint64_t v) { return (v + 255) & ~std::uint64_t{255}; } + +// Qwen3_5TextRotaryEmbedding inv_freq for rotary dim 64, theta 1e7, as torch's CPU computes it (kernels/basic/rope.h). +constexpr std::uint32_t kRopeInvFreqBits[32] = { + 0x3f800000, 0x3f1ab32b, 0x3ebaf81b, 0x3e61f835, 0x3e088d77, 0x3da50956, 0x3d47763f, 0x3cf11177, + 0x3c91ad39, 0x3c301052, 0x3bd4ca15, 0x3b80967d, 0x3b1b690d, 0x3abbd3ed, 0x3a6301e2, 0x3a092e02, + 0x39a5cb60, 0x394860c1, 0x38f22ce2, 0x3892587e, 0x3830df52, 0x37d5c441, 0x37812dab, 0x371c1fc4, + 0x36bcb0c1, 0x36640cc6, 0x3609cf4a, 0x35a68e4c, 0x35494c57, 0x34f3499d, 0x3493048e, 0x3431af44}; + +// Rows [cos(32) | sin(32)] of positions 0 .. len - 1: fp32 phase = position * inv_freq, cos and sin rounded to bf16. +std::vector rope_table(std::uint32_t len) { + std::vector out(static_cast(len) * kRopeDim); + for (std::uint32_t p = 0; p < len; ++p) { + for (std::uint32_t i = 0; i < 32; ++i) { + float f = 0; + std::memcpy(&f, &kRopeInvFreqBits[i], 4); + const float phase = static_cast(p) * f; + out[static_cast(p) * kRopeDim + i] = float_to_bf16(std::cos(phase)); + out[static_cast(p) * kRopeDim + 32 + i] = float_to_bf16(std::sin(phase)); + } + } + return out; +} + +struct Piece { + const st::File* file; + const st::TensorInfo* tensor; + std::uint64_t offset; // in the layer buffer +}; + +// Where a layer's tensors sit in the device buffer. Tensors read by one GEMM are adjacent, in this order: GatedDeltaNet +// in_proj_qkv | in_proj_z and in_proj_b | in_proj_a, attention q_proj | k_proj | v_proj, MLP gate_proj | up_proj. +struct LayerPlan { + bool full_attention = false; + std::vector pieces; + std::uint64_t bytes = 0; + std::uint64_t in_norm = 0, post_norm = 0, gate_up = 0, down = 0; + std::uint64_t qkvz = 0, ba = 0, conv = 0, a_log = 0, dt_bias = 0, gdn_norm = 0, gdn_out = 0; + std::uint64_t qkv = 0, o = 0, q_norm = 0, k_norm = 0; +}; + +// One scratch buffer, laid out per phase of a layer (offsets for `t` tokens). +struct GdnScratch { + std::uint64_t qkvz, ba, conv, g, beta, core, ws, end; + explicit GdnScratch(std::uint32_t t) { + qkvz = 0; + ba = align256(qkvz + std::uint64_t{t} * kQkvzDim * 2); + conv = align256(ba + std::uint64_t{t} * 64 * 2); + g = align256(conv + std::uint64_t{t} * kGdnConvDim * 2); + beta = align256(g + std::uint64_t{t} * kGdnValueHeads * 4); + core = align256(beta + std::uint64_t{t} * kGdnValueHeads * 2); + ws = align256(core + std::uint64_t{t} * kValueDim * 2); + end = ws + GdnWorkspace::bytes(t); + } +}; +struct AttnScratch { + std::uint64_t qkv, q, attn, end; + explicit AttnScratch(std::uint32_t t) { + qkv = 0; + q = align256(qkv + std::uint64_t{t} * kQkvDim * 2); + attn = align256(q + std::uint64_t{t} * kTeHidden * 2); + end = attn + std::uint64_t{t} * kTeHidden * 2; + } +}; +struct MlpScratch { + std::uint64_t gu, act, end; + explicit MlpScratch(std::uint32_t t) { + gu = 0; + act = align256(gu + std::uint64_t{t} * kGateUpDim * 2); + end = act + std::uint64_t{t} * kTeIntermediate * 2; + } +}; + +// Collects recorded work and submits it in pieces of at most `budget` estimated FLOP. +class Submitter { + public: + Submitter(Device& device, double budget, bool profile, TextEncoderStats& stats) + : device_(device), budget_(budget), profile_(profile), stats_(stats) {} + + // `name` labels the op in TextEncoderStats::op_seconds when profiling, which submits every op on its own. + void add(const char* name, double flop, std::function op) { + if (profile_) { + flush(); + ops_.push_back(std::move(op)); + const auto t0 = Clock::now(); + flush(); + stats_.op_seconds[name] += seconds_since(t0); + return; + } + if (!ops_.empty() && work_ + flop > budget_) flush(); + ops_.push_back(std::move(op)); + work_ += flop; + } + void flush() { + if (ops_.empty()) return; + const auto t0 = Clock::now(); + device_.run([&](VkCommandBuffer cb) { + for (const auto& op : ops_) op(cb); + }); + stats_.compute_s += seconds_since(t0); + ++stats_.submits; + ops_.clear(); + work_ = 0; + } + double budget() const { return budget_; } + + private: + Device& device_; + double budget_; + bool profile_; + TextEncoderStats& stats_; + std::vector> ops_; + double work_ = 0; +}; + +// The reader thread's handoff: slot s holds layer `filled[s]` once read; the main thread frees it after its copy. +struct Handoff { + std::mutex mu; + std::condition_variable cv; + int filled[2] = {-1, -1}; + bool busy[2] = {false, false}; + bool abort = false; + std::exception_ptr error; +}; + +} // namespace + +struct TextEncoder::Impl { + Device& device; + BlockPath path; + TextEncoderOptions options; + BlockKernels gemm; + TeKernels te; + std::map> files; + std::map> tensors; + const st::File* embed_file = nullptr; + const st::TensorInfo* embed = nullptr; + std::pair final_norm{}; + std::vector layers; + std::uint64_t layer_bytes = 0; + TextEncoderStats stats; + + Impl(Device& d, const std::filesystem::path& dir, BlockPath p, TextEncoderOptions o) + : device(d), path(p), options(o), gemm(d, p), te(d) { + const std::map index = st::read_index(dir / "model.safetensors.index.json"); + for (const auto& [name, shard] : index) { + if (name.rfind("language_model.", 0) != 0) continue; + auto& file = files[shard]; + if (!file) file = std::make_unique(dir / shard); + tensors[name] = {file.get(), &file->at(name)}; + } + const auto& e = lookup("language_model.embed_tokens.weight", {kTeVocab, kTeHidden}); + embed_file = e.first; + embed = e.second; + final_norm = lookup("language_model.norm.weight", {kTeHidden}); + for (int i = 0; i < kTeLayers; ++i) layers.push_back(plan_layer(i)); + for (const LayerPlan& l : layers) layer_bytes = std::max(layer_bytes, l.bytes); + } + + const std::pair& lookup(const std::string& name, + const std::vector& shape) const { + const auto it = tensors.find(name); + if (it == tensors.end()) fail("no tensor " + name); + if (it->second.second->dtype != "BF16" || it->second.second->shape != shape) + fail("tensor " + name + " has an unexpected dtype or shape"); + return it->second; + } + + LayerPlan plan_layer(int i) const { + LayerPlan l; + l.full_attention = i % 4 == 3; + const std::string p = "language_model.layers." + std::to_string(i) + "."; + // Tensors of one group are placed back to back: their sizes are multiples of 256 bytes, so alignment adds no gap. + const auto add = [&](const std::string& name, const std::vector& shape) { + const auto& [file, t] = lookup(p + name, shape); + if (t->bytes % 256 != 0 && shape.size() == 2) fail("tensor " + t->name + " is not a multiple of 256 bytes"); + l.bytes = align256(l.bytes); + const std::uint64_t at = l.bytes; + l.pieces.push_back({file, t, at}); + l.bytes += t->bytes; + return at; + }; + l.in_norm = add("input_layernorm.weight", {kTeHidden}); + l.post_norm = add("post_attention_layernorm.weight", {kTeHidden}); + l.gate_up = add("mlp.gate_proj.weight", {kTeIntermediate, kTeHidden}); + add("mlp.up_proj.weight", {kTeIntermediate, kTeHidden}); + l.down = add("mlp.down_proj.weight", {kTeHidden, kTeIntermediate}); + if (l.full_attention) { + l.qkv = add("self_attn.q_proj.weight", {kQgDim, kTeHidden}); + add("self_attn.k_proj.weight", {kKvDim, kTeHidden}); + add("self_attn.v_proj.weight", {kKvDim, kTeHidden}); + l.o = add("self_attn.o_proj.weight", {kTeHidden, kTeHeads * kTeHeadDim}); + l.q_norm = add("self_attn.q_norm.weight", {kTeHeadDim}); + l.k_norm = add("self_attn.k_norm.weight", {kTeHeadDim}); + } else { + l.qkvz = add("linear_attn.in_proj_qkv.weight", {kGdnConvDim, kTeHidden}); + add("linear_attn.in_proj_z.weight", {kValueDim, kTeHidden}); + l.ba = add("linear_attn.in_proj_b.weight", {kGdnValueHeads, kTeHidden}); + add("linear_attn.in_proj_a.weight", {kGdnValueHeads, kTeHidden}); + l.conv = add("linear_attn.conv1d.weight", {kGdnConvDim, 1, 4}); + l.a_log = add("linear_attn.A_log", {kGdnValueHeads}); + l.dt_bias = add("linear_attn.dt_bias", {kGdnValueHeads}); + l.gdn_norm = add("linear_attn.norm.weight", {kGdnHeadDim}); + l.gdn_out = add("linear_attn.out_proj.weight", {kTeHidden, kValueDim}); + } + l.bytes = align256(l.bytes); + return l; + } + + double budget() const { + const double gflop = options.submit_gflop > 0 ? options.submit_gflop : (path != BlockPath::kPlain ? 200.0 : 20.0); + return gflop * 1e9; + } + + std::vector> encode(const std::vector>& seqs, + const TeLayerHook& hook); + + // A layer's tensors into `dst` at their offsets, in chunks of 16 MiB that `reader_threads` threads take in turn: + // several reads in flight are what a cold NVMe drive needs to reach its speed. + void read_layer(const LayerPlan& l, std::byte* dst) const { + constexpr std::uint64_t kChunk = std::uint64_t{16} << 20; + struct Chunk { + const Piece* piece; + std::uint64_t rel, bytes; + }; + std::vector chunks; + for (const Piece& p : l.pieces) { + for (std::uint64_t rel = 0; rel < p.tensor->bytes; rel += kChunk) + chunks.push_back({&p, rel, std::min(kChunk, p.tensor->bytes - rel)}); + } + std::atomic next{0}; + std::vector errors(static_cast(std::max(1, options.reader_threads))); + const auto work = [&](std::size_t t) { + try { + for (std::size_t i = next++; i < chunks.size(); i = next++) { + const Chunk& c = chunks[i]; + c.piece->file->read(*c.piece->tensor, dst + c.piece->offset + c.rel, c.rel, c.bytes); + } + } catch (...) { + errors[t] = std::current_exception(); + next = chunks.size(); + } + }; + { + std::vector helpers; + for (std::size_t t = 1; t < errors.size(); ++t) helpers.emplace_back(work, t); + work(0); + } + for (const auto& e : errors) { + if (e) std::rethrow_exception(e); + } + } +}; + +TextEncoder::TextEncoder(Device& device, const std::filesystem::path& dir) + : TextEncoder(device, dir, default_block_path(device.info().caps)) {} + +TextEncoder::TextEncoder(Device& device, const std::filesystem::path& dir, BlockPath path, TextEncoderOptions options) + : impl_(std::make_unique(device, dir, path, options)) {} + +TextEncoder::~TextEncoder() = default; + +BlockPath TextEncoder::path() const { return impl_->path; } +const TextEncoderStats& TextEncoder::last_stats() const { return impl_->stats; } + +std::vector> TextEncoder::encode(const std::vector>& sequences, + const TeLayerHook& hook) { + return impl_->encode(sequences, hook); +} + +std::vector> TextEncoder::Impl::encode(const std::vector>& seqs, + const TeLayerHook& hook) { + const auto t_start = Clock::now(); + if (seqs.empty()) throw std::invalid_argument("vulkan text encoder: no sequences"); + std::vector cu = {0}; + std::uint32_t max_len = 0; + for (const auto& s : seqs) { + if (s.empty()) throw std::invalid_argument("vulkan text encoder: empty sequence"); + for (const std::int32_t id : s) { + if (id < 0 || id >= kTeVocab) + throw std::invalid_argument("vulkan text encoder: token id " + std::to_string(id) + " outside the vocabulary"); + } + if (static_cast(cu.back()) + s.size() > kMaxTokens) + throw std::invalid_argument("vulkan text encoder: more than 65535 tokens"); + cu.push_back(cu.back() + static_cast(s.size())); + max_len = std::max(max_len, static_cast(s.size())); + } + const auto T = static_cast(cu.back()); + stats = TextEncoderStats{}; + stats.tokens = T; + + const GdnScratch gs(T); + const AttnScratch as(T); + const MlpScratch ms(T); + const std::uint64_t scratch_bytes = std::max({gs.end, as.end, ms.end}); + const std::uint64_t act_bytes = std::uint64_t{T} * kTeHidden * 2; + + // Layer k runs from weights[k % 2] while layer k + 1 is copied into the other. + std::unique_ptr weights[2] = {std::make_unique(device, layer_bytes, Buffer::Kind::kDevice, true), + std::make_unique(device, layer_bytes, Buffer::Kind::kDevice, true)}; + Buffer x(device, act_bytes), h(device, act_bytes), scratch(device, scratch_bytes); + Buffer rope(device, std::uint64_t{max_len} * kRopeDim * 2), small(device, 3 * kTeHidden * 2); + stats.device_bytes = 2 * layer_bytes + x.bytes() + h.bytes() + scratch.bytes() + rope.bytes() + small.bytes(); + std::unique_ptr staging[2] = {std::make_unique(device, layer_bytes, Buffer::Kind::kStaging, true), + std::make_unique(device, layer_bytes, Buffer::Kind::kStaging, true)}; + stats.staging_bytes = 2 * layer_bytes; + + // Small constants: the final norm weight and a vector of ones (the gate that makes the gated-residual GEMM a plain + // residual add: bf16(x + bf16(ab) * 1) = bf16(x + bf16(ab))). + const VkDeviceAddress final_norm_w = small.address(), ones = small.address() + kTeHidden * 2; + { + std::vector host(2 * kTeHidden, 0x3f80); + final_norm.first->read(*final_norm.second, host.data()); + small.upload(host.data(), host.size() * 2); + const std::vector table = rope_table(max_len); + rope.upload(table.data(), table.size() * 2); + } + + Handoff handoff; + double read_s = 0; + std::thread reader([&] { + try { + for (int k = 0; k < kTeLayers; ++k) { + const int s = k % 2; + { + std::unique_lock lock(handoff.mu); + handoff.cv.wait(lock, [&] { return handoff.abort || (!handoff.busy[s] && handoff.filled[s] < 0); }); + if (handoff.abort) return; + handoff.busy[s] = true; + } + const auto t0 = Clock::now(); + read_layer(layers[static_cast(k)], static_cast(staging[s]->mapped())); + read_s += seconds_since(t0); + std::lock_guard lock(handoff.mu); + handoff.busy[s] = false; + handoff.filled[s] = k; + handoff.cv.notify_all(); + } + } catch (...) { + std::lock_guard lock(handoff.mu); + handoff.error = std::current_exception(); + handoff.abort = true; + handoff.cv.notify_all(); + } + }); + const auto stop_reader = [&] { + { + std::lock_guard lock(handoff.mu); + handoff.abort = true; + handoff.cv.notify_all(); + } + if (reader.joinable()) reader.join(); + }; + + std::vector> result; + std::vector in_flight; // copies into the buffers above, which must finish before the buffers go + try { + // Embedding rows of the ids, gathered on the host. + { + const auto t0 = Clock::now(); + std::vector rows(static_cast(T) * kTeHidden); + std::size_t r = 0; + for (const auto& s : seqs) { + for (const std::int32_t id : s) { + embed_file->read(*embed, rows.data() + r++ * kTeHidden, std::uint64_t(id) * kTeHidden * 2, kTeHidden * 2); + } + } + x.upload(rows.data(), rows.size() * 2); + stats.embed_s = seconds_since(t0); + } + + Submitter sub(device, budget(), options.profile, stats); + const VkDeviceAddress xa = x.address(), ha = h.address(), sa = scratch.address(); + const auto bf = [](VkDeviceAddress base, std::uint64_t elements) { return base + elements * 2; }; + + // GEMMs split along M so that each piece fits the submit budget; rows of the output do not depend on the split. + const auto rows_per_piece = [&](std::uint32_t n, std::uint32_t k) { + const double per_row = 2.0 * n * k; + auto rows = static_cast(std::max(1.0, std::floor(sub.budget() / per_row))); + if (rows >= 128) rows -= rows % 128; + return rows; + }; + const auto linear = [&](const char* name, VkDeviceAddress out, std::uint32_t ldo, VkDeviceAddress in, + std::uint32_t ldi, VkDeviceAddress w, std::uint32_t n, std::uint32_t k) { + const std::uint32_t step = rows_per_piece(n, k); + for (std::uint32_t m0 = 0; m0 < T; m0 += step) { + const std::uint32_t m = std::min(step, T - m0); + sub.add(name, 2.0 * m * n * k, [=, this](VkCommandBuffer cb) { + gemm.record_gemm(cb, bf(out, std::uint64_t{m0} * ldo), ldo, bf(in, std::uint64_t{m0} * ldi), ldi, w, k, m, n, + k); + }); + } + }; + const auto linear_residual = [&](const char* name, VkDeviceAddress in, std::uint32_t ldi, VkDeviceAddress w, + std::uint32_t k) { + const std::uint32_t step = rows_per_piece(kTeHidden, k); + for (std::uint32_t m0 = 0; m0 < T; m0 += step) { + const std::uint32_t m = std::min(step, T - m0); + sub.add(name, 2.0 * m * kTeHidden * k, [=, this](VkCommandBuffer cb) { + gemm.record_gemm_gated_residual(cb, bf(xa, std::uint64_t{m0} * kTeHidden), kTeHidden, + bf(in, std::uint64_t{m0} * ldi), ldi, w, k, ones, m, kTeHidden, k); + }); + } + }; + const auto for_each_sequence = [&](auto&& f) { + for (std::size_t i = 0; i + 1 < cu.size(); ++i) { + f(static_cast(cu[i]), static_cast(cu[i + 1] - cu[i])); + } + }; + const auto hidden_norm = [&](VkDeviceAddress weight, VkDeviceAddress out) { + for_each_sequence([&](std::uint32_t s0, std::uint32_t len) { + RmsNormArgs a; + a.x = bf(xa, std::uint64_t{s0} * kTeHidden); + a.out = bf(out, std::uint64_t{s0} * kTeHidden); + a.weight = weight; + a.ldx = a.ldo = kTeHidden; + a.tokens = len; + a.d = kTeHidden; + sub.add("rms_norm", 8.0 * len * kTeHidden, [=, this](VkCommandBuffer cb) { te.record_rms_norm(cb, a); }); + }); + }; + + const auto wait_staged = [&](int j) { + const auto t0 = Clock::now(); + std::unique_lock lock(handoff.mu); + handoff.cv.wait(lock, [&] { return handoff.error != nullptr || handoff.filled[j % 2] == j; }); + if (handoff.error) std::rethrow_exception(handoff.error); + stats.read_wait_s += seconds_since(t0); + }; + const auto release_staged = [&](int j) { + std::lock_guard lock(handoff.mu); + handoff.filled[j % 2] = -1; + handoff.cv.notify_all(); + }; + // Layer j's staging -> device copy, on the transfer queue where there is one. + const auto copy = [&](int j) { + in_flight.push_back( + device.copy_async(*staging[j % 2], 0, *weights[j % 2], 0, layers[static_cast(j)].bytes)); + return in_flight.back(); + }; + // Blocks until layer j's copy is done, then hands its staging slot back to the reader. + const auto copied = [&](int j, std::uint64_t ticket) { + const auto t0 = Clock::now(); + device.host_wait(ticket); + in_flight.erase(std::find(in_flight.begin(), in_flight.end(), ticket)); + stats.upload_s += seconds_since(t0); + release_staged(j); + }; + std::uint64_t tickets[kTeLayers] = {}; + wait_staged(0); + tickets[0] = copy(0); + copied(0, tickets[0]); + + for (int k = 0; k < kTeLayers; ++k) { + const LayerPlan& l = layers[static_cast(k)]; + device.gpu_wait(tickets[k]); // the layer's first submit waits for its weights + // The next layer's copy runs while this layer computes. + if (k + 1 < kTeLayers) { + wait_staged(k + 1); + tickets[k + 1] = copy(k + 1); + } + const VkDeviceAddress wa = weights[k % 2]->address(); + const auto W = [&](std::uint64_t offset) { return wa + offset; }; + + hidden_norm(W(l.in_norm), ha); + if (l.full_attention) { + const VkDeviceAddress qkv = sa + as.qkv, q = sa + as.q, attn = sa + as.attn; + linear("gemm q|k|v", qkv, kQkvDim, ha, kTeHidden, W(l.qkv), kQkvDim, kTeHidden); + for_each_sequence([&](std::uint32_t s0, std::uint32_t len) { + RmsNormArgs qn; + qn.x = bf(qkv, std::uint64_t{s0} * kQkvDim); + qn.ldx = kQkvDim; + qn.hsx = 2 * kTeHeadDim; + qn.out = bf(q, std::uint64_t{s0} * kTeHidden); + qn.ldo = kTeHidden; + qn.hso = kTeHeadDim; + qn.weight = W(l.q_norm); + qn.tokens = len; + qn.heads = kTeHeads; + qn.d = kTeHeadDim; + qn.rope = rope.address(); + RmsNormArgs kn = qn; + kn.x = kn.out = bf(qkv, std::uint64_t{s0} * kQkvDim + kQgDim); + kn.ldx = kn.ldo = kQkvDim; + kn.hsx = kn.hso = kTeHeadDim; + kn.weight = W(l.k_norm); + kn.heads = kTeKvHeads; + sub.add("q, k norm+rope", 8.0 * len * kQgDim, [=, this](VkCommandBuffer cb) { + te.record_rms_norm(cb, qn); + te.record_rms_norm(cb, kn); + }); + }); + for_each_sequence([&](std::uint32_t s0, std::uint32_t len) { + TeAttentionArgs a; + a.q = bf(q, std::uint64_t{s0} * kTeHidden); + a.k = bf(qkv, std::uint64_t{s0} * kQkvDim + kQgDim); + a.v = bf(qkv, std::uint64_t{s0} * kQkvDim + kQgDim + kKvDim); + a.o = bf(attn, std::uint64_t{s0} * kTeHidden); + a.ldq = a.ldo = kTeHidden; + a.ldk = a.ldv = kQkvDim; + a.hsq = a.hsk = a.hsv = a.hso = kTeHeadDim; + a.heads = kTeHeads; + a.group = kTeHeads / kTeKvHeads; + a.scale = 0.0625f; + // Query rows in pieces: row i costs about 4 * 256 * (i + 1) FLOP per head. + const double per_row = 4.0 * kTeHeadDim * kTeHeads * len; + const auto step = std::max(16u, static_cast(std::floor(sub.budget() / per_row)) / 16 * 16); + for (std::uint32_t r0 = 0; r0 < len; r0 += step) { + a.q_first = r0; + a.rows = std::min(step, len - r0); + sub.add("attention", per_row * a.rows, [=, this](VkCommandBuffer cb) { te.record_attention(cb, a); }); + } + }); + sub.add("sigmoid gate", 4.0 * T * kTeHidden, [=, this](VkCommandBuffer cb) { + te.record_gate_mul(cb, TeKernels::GateMul::kSigmoid, attn, kTeHidden, bf(qkv, kTeHeadDim), kQkvDim, attn, + kTeHidden, T, kTeHidden, kTeHeadDim, 2 * kTeHeadDim); + }); + linear_residual("gemm o+residual", attn, kTeHidden, W(l.o), kTeHidden); + } else { + const VkDeviceAddress qkvz = sa + gs.qkvz, ba = sa + gs.ba, conv = sa + gs.conv, g = sa + gs.g, + beta = sa + gs.beta, core = sa + gs.core; + linear("gemm qkv|z", qkvz, kQkvzDim, ha, kTeHidden, W(l.qkvz), kQkvzDim, kTeHidden); + linear("gemm b|a", ba, 64, ha, kTeHidden, W(l.ba), 64, kTeHidden); + for_each_sequence([&](std::uint32_t s0, std::uint32_t len) { + sub.add("conv1d+silu", 8.0 * len * kGdnConvDim, [=, this](VkCommandBuffer cb) { + te.record_conv_silu(cb, bf(qkvz, std::uint64_t{s0} * kQkvzDim), kQkvzDim, W(l.conv), + bf(conv, std::uint64_t{s0} * kGdnConvDim), kGdnConvDim, len); + }); + }); + GdnArgs ga; + ga.qkv = conv; + ga.ldqkv = kGdnConvDim; + ga.g = g; + ga.beta = beta; + ga.out = core; + ga.ws = GdnWorkspace::at(sa + gs.ws, T); + ga.cu = cu; + // A rough count per token and value head of the chunk products, the solve and the scan; only sizes the submits. + const double gdn_flop = 2.0 * T * kGdnValueHeads * (2.0 * 64 * 128 * 2 + 256.0 * 64 + 6.0 * 128 * 128); + sub.add("gates+gated delta rule", gdn_flop, [=, this](VkCommandBuffer cb) { + te.record_gates(cb, ba, 64, W(l.a_log), W(l.dt_bias), g, beta, T); + te.record_gated_delta_rule(cb, ga); + }); + for_each_sequence([&](std::uint32_t s0, std::uint32_t len) { + RmsNormArgs a; + a.x = a.out = bf(core, std::uint64_t{s0} * kValueDim); + a.ldx = a.ldo = kValueDim; + a.hsx = a.hso = kGdnHeadDim; + a.gate = bf(qkvz, std::uint64_t{s0} * kQkvzDim + kGdnConvDim); + a.ldg = kQkvzDim; + a.hsg = kGdnHeadDim; + a.weight = W(l.gdn_norm); + a.tokens = len; + a.heads = kGdnValueHeads; + a.d = kGdnHeadDim; + a.gated = true; + sub.add("gated rms_norm", 12.0 * len * kValueDim, + [=, this](VkCommandBuffer cb) { te.record_rms_norm(cb, a); }); + }); + linear_residual("gemm out+residual", core, kValueDim, W(l.gdn_out), kValueDim); + } + + hidden_norm(W(l.post_norm), ha); + const VkDeviceAddress gu = sa + ms.gu, act = sa + ms.act; + linear("gemm gate|up", gu, kGateUpDim, ha, kTeHidden, W(l.gate_up), kGateUpDim, kTeHidden); + sub.add("silu_mul", 8.0 * T * kTeIntermediate, [=, this](VkCommandBuffer cb) { + te.record_gate_mul(cb, TeKernels::GateMul::kSilu, gu, kGateUpDim, bf(gu, kTeIntermediate), kGateUpDim, act, + kTeIntermediate, T, kTeIntermediate, kTeIntermediate, 0); + }); + linear_residual("gemm down+residual", act, kTeIntermediate, W(l.down), kTeIntermediate); + sub.flush(); + if (k + 1 < kTeLayers) copied(k + 1, tickets[k + 1]); + if (hook) hook(k, x, T); + } + stop_reader(); + + hidden_norm(final_norm_w, ha); + sub.flush(); + std::vector out(static_cast(T) * kTeHidden); + h.download(out.data(), out.size() * 2); + for (std::size_t i = 0; i + 1 < cu.size(); ++i) { + result.emplace_back(out.begin() + std::ptrdiff_t{cu[i]} * kTeHidden, + out.begin() + std::ptrdiff_t{cu[i + 1]} * kTeHidden); + } + } catch (...) { + stop_reader(); + for (const std::uint64_t ticket : in_flight) { + try { + device.host_wait(ticket); + } catch (...) { // already failing; the device is likely lost + } + } + throw; + } + stats.read_s = read_s; + stats.total_s = seconds_since(t_start); + return result; +} + +} // namespace mi::vk diff --git a/vulkan/text_encoder.h b/vulkan/text_encoder.h new file mode 100644 index 0000000..e70a53f --- /dev/null +++ b/vulkan/text_encoder.h @@ -0,0 +1,88 @@ +// The Qwen3.5 text encoder as MOSS-Image uses it (src/models/qwen35_text/text_encoder.h): Qwen3_5Model(input_ids) +// .last_hidden_state for text-only input. 32 decoder layers, a GatedDeltaNet in three of every four and full gated +// attention in the fourth, each RMSNorm -> token mixer -> residual -> RMSNorm -> SiLU MLP -> residual, then the final +// norm. bf16 [S, 4096] per sequence. +// +// All sequences go through each layer together, packed along the rows: GEMMs and elementwise ops see every token at +// once (their outputs do not depend on the row count here), norms, the conv, the gated delta rule and attention run per +// sequence. Weights stream from the safetensors files: the embedding rows of the given ids are gathered on the host, a +// reader thread reads layer k + 2 into one of two host staging buffers while layer k runs, and layer k + 1 is copied +// from staging into one of two device buffers meanwhile, on the device's transfer queue where it has one +// (Device::copy_async). Work goes to the device in submits of a bounded size (TextEncoderOptions::submit_gflop), so no +// submit runs long enough to trip a display driver's watchdog. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "runtime.h" + +namespace mi::vk { + +inline constexpr int kTeHidden = 4096; +inline constexpr int kTeIntermediate = 12288; +inline constexpr int kTeLayers = 32; +inline constexpr int kTeHeads = 16; +inline constexpr int kTeKvHeads = 4; +inline constexpr int kTeHeadDim = 256; +inline constexpr int kTeVocab = 248320; + +struct TextEncoderOptions { + // Estimated work per submit, in GFLOP; 0 picks the path's default (20 plain, 200 cooperative matrix). A GEMM larger + // than this is split along its rows. + double submit_gflop = 0; + // Threads reading layer weights from the files. + int reader_threads = 4; + // Submit every op on its own and time it into TextEncoderStats::op_seconds (slower; for finding costs). + bool profile = false; +}; + +// Times of the last encode(), in seconds of host wall time. +struct TextEncoderStats { + std::uint32_t tokens = 0; + double total_s = 0; + double embed_s = 0; // embedding rows read and uploaded + double read_s = 0; // the reader thread reading layer weights from the files + double read_wait_s = 0; // waiting for the reader + double upload_s = 0; // waiting for staging -> device copies (the part not hidden behind compute) + double compute_s = 0; // the layers' and the final norm's submits + int submits = 0; + std::uint64_t device_bytes = 0; // device memory the encoder held at the peak of the call (its buffers) + std::uint64_t staging_bytes = 0; // host-visible staging it held + std::map op_seconds; // TextEncoderOptions::profile only: seconds per kind of op +}; + +// Raw output of decoder layer `layer` (before the final norm), all sequences packed: bf16 [tokens, 4096] in `hidden`. +// Called after the layer's work has completed; the hook may download from the buffer. +using TeLayerHook = std::function; + +class TextEncoder { + public: + // `dir`: the text_encoder directory, with model.safetensors.index.json and its shards. The default path is the + // cooperative-matrix one where the device has bf16 fragments. + TextEncoder(Device& device, const std::filesystem::path& dir); + TextEncoder(Device& device, const std::filesystem::path& dir, BlockPath path, TextEncoderOptions options = {}); + ~TextEncoder(); + TextEncoder(const TextEncoder&) = delete; + TextEncoder& operator=(const TextEncoder&) = delete; + + BlockPath path() const; + + // last_hidden_state of every sequence as bf16 bits [S_i * 4096]. Sequences are non-empty, ids below kTeVocab. + std::vector> encode(const std::vector>& sequences, + const TeLayerHook& hook = {}); + + const TextEncoderStats& last_stats() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace mi::vk diff --git a/vulkan/tools/moss_infer_vulkan.cpp b/vulkan/tools/moss_infer_vulkan.cpp new file mode 100644 index 0000000..6a1dc74 --- /dev/null +++ b/vulkan/tools/moss_infer_vulkan.cpp @@ -0,0 +1,344 @@ +// moss-infer-vulkan: text to image on the Vulkan backend, a standalone program for GPUs the CUDA engine does not support. +// It needs neither CUDA nor an NVIDIA driver, only a Vulkan driver. Its flags are those of `moss-infer generate` that the +// backend runs; the tokenizer, the seed's initial noise, the schedule and the PNG writer are the engine's own sources, +// so a seed gives the same inputs as on CUDA. Exit status: 0 done, 1 failed, 2 a usage error or an unsupported request. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "generate.h" +#include "mossinfer/mossinfer.h" +#include "npy.h" +#include "pipeline/png_writer.h" +#include "pipeline/samplers/initial_noise.h" +#include "pipeline/samplers/schedules.h" +#include "runtime.h" +#include "tokenizer/prompt_tokens.h" +#include "tokenizer/qwen_bpe.h" +#include "util/path.h" + +namespace { + +namespace fs = std::filesystem; +using namespace mi; + +constexpr const char* kUsage = + "usage: moss-infer-vulkan --model --prompt [-o ] [--latent-out ]\n" + " [--negative ] [--size x] [--width N] [--height N] [--steps N] [--cfg X]\n" + " [--shift X] [--seed N|-1] [--skip-layers \"22\"] [--max-seq N] [--tier reference]\n" + " [--device N] [--kernels coop|coop-f16|plain] [--no-metadata] [--quiet]\n" + " moss-infer-vulkan --list-devices\n" + "Text to image with MOSS-Image Base or SFT (the bundle's bf16 files) on a Vulkan GPU, with the official euler sampler\n" + "and schedule: the fallback for GPUs the CUDA engine (moss-infer) does not support. The flags mean what they mean for\n" + "`moss-infer generate`; the defaults are the same (1024x1024, 50 steps, cfg 3, shift 5, seed 42, skip layers 22).\n" + "At least one of -o / --latent-out is required; without -o the VAE decode is skipped. --device picks a Vulkan device\n" + "(default: the first usable one, a discrete GPU first); --kernels picks the cooperative-matrix kernels on bf16 or f16\n" + "fragments or the plain ones (default: the first of those the device has). The device and kernels are printed at the\n" + "start. Progress goes to stderr one line per event: `text encoder`, `denoise / `, `vae`, then\n" + "a summary.\n"; + +struct Args { + fs::path model, out_png, out_latent; + std::string prompt, negative; + bool have_prompt = false, metadata = true, quiet = false, list = false; + int width = 1024, height = 1024, steps = 50, max_seq = 4096, device = -1; + double cfg = 3.0, shift = 5.0; + std::int64_t seed = 42; + std::vector skip = {22}; + std::string kernels; +}; + +[[noreturn]] void usage(const std::string& what) { throw std::invalid_argument(what); } + +double number(const std::string& flag, const std::string& v) { + std::size_t used = 0; + double d = 0; + try { + d = std::stod(v, &used); + } catch (const std::exception&) { + usage(flag + " takes a number, got " + v); + } + if (used != v.size() || !std::isfinite(d)) usage(flag + " takes a finite number, got " + v); + return d; +} + +std::int64_t integer(const std::string& flag, const std::string& v) { + std::size_t used = 0; + long long i = 0; + try { + i = std::stoll(v, &used); + } catch (const std::exception&) { + usage(flag + " takes an integer, got " + v); + } + if (used != v.size()) usage(flag + " takes an integer, got " + v); + return i; +} + +int int32(const std::string& flag, const std::string& v) { + const std::int64_t i = integer(flag, v); + if (i < std::numeric_limits::min() || i > std::numeric_limits::max()) usage(flag + " is out of range"); + return static_cast(i); +} + +std::vector layers(const std::string& v) { + std::vector out; + std::size_t at = 0; + while (at < v.size()) { + const std::size_t comma = v.find(',', at); + std::string item = v.substr(at, comma == std::string::npos ? std::string::npos : comma - at); + item.erase(0, item.find_first_not_of(' ')); + item.erase(item.find_last_not_of(' ') + 1); + if (!item.empty()) { + const int l = int32("--skip-layers", item); + if (l < 0 || l >= vk::kDiTLayers) usage("--skip-layers takes block indices in [0, 41]"); + out.push_back(l); + } + if (comma == std::string::npos) break; + at = comma + 1; + } + return out; +} + +Args parse(const std::vector& argv) { + Args a; + for (std::size_t i = 1; i < argv.size(); ++i) { + const std::string& f = argv[i]; + if (f == "--no-metadata") { + a.metadata = false; + continue; + } + if (f == "--quiet") { + a.quiet = true; + continue; + } + if (f == "--list-devices") { + a.list = true; + continue; + } + if (f == "-h" || f == "--help") usage(""); + if (i + 1 >= argv.size()) usage(f + " needs a value"); + const std::string& v = argv[++i]; + if (f == "--model") a.model = util::path_from_utf8(v); + else if (f == "--prompt") a.prompt = v, a.have_prompt = true; + else if (f == "--negative") a.negative = v; + else if (f == "-o" || f == "--output") a.out_png = util::path_from_utf8(v); + else if (f == "--latent-out") a.out_latent = util::path_from_utf8(v); + else if (f == "--size") { + const std::size_t x = v.find('x'); + if (x == std::string::npos) usage("--size takes x"); + a.width = int32("--size", v.substr(0, x)); + a.height = int32("--size", v.substr(x + 1)); + } + else if (f == "--width") a.width = int32(f, v); + else if (f == "--height") a.height = int32(f, v); + else if (f == "--steps") a.steps = int32(f, v); + else if (f == "--cfg") a.cfg = number(f, v); + else if (f == "--shift") a.shift = number(f, v); + else if (f == "--seed") a.seed = integer(f, v); + else if (f == "--skip-layers") a.skip = layers(v); + else if (f == "--max-seq") a.max_seq = int32(f, v); + else if (f == "--device") a.device = int32(f, v); + else if (f == "--kernels") a.kernels = v; + else if (f == "--tier") { + // The backend runs the bf16 files exactly as the reference tier does; the lossy tiers need the CUDA kernels. + if (v != "reference") usage("moss-infer-vulkan runs --tier reference only"); + } + else usage("unknown or unsupported flag " + f); + } + if (a.list) return a; + if (a.model.empty() || !a.have_prompt) usage("--model and --prompt are required"); + if (a.out_png.empty() && a.out_latent.empty()) usage("at least one of -o / --latent-out is required"); + if (a.width < 16 || a.height < 16 || a.width % 16 != 0 || a.height % 16 != 0 || a.width > 4096 || a.height > 4096) { + usage("width and height must be multiples of 16 in [16, 4096]"); + } + if (a.steps < 1 || a.steps > 1000) usage("--steps must be in [1, 1000]"); + if (a.cfg < 0) usage("--cfg must be >= 0"); + if (a.shift <= 0) usage("--shift must be > 0"); + if (a.seed < -1 || a.seed >= (std::int64_t{1} << 53)) usage("--seed must be -1 (random) or in [0, 2^53)"); + if (a.max_seq < 1 || a.max_seq > static_cast(tok::kMossMaxSequenceLength)) usage("--max-seq must be in [1, 4096]"); + if (!a.kernels.empty() && a.kernels != "coop" && a.kernels != "coop-f16" && a.kernels != "plain") { + usage("--kernels takes coop, coop-f16 or plain"); + } + return a; +} + +int list_devices() { + const std::vector devices = vk::list_devices(); + if (devices.empty()) std::printf("no Vulkan devices\n"); + for (const vk::DeviceInfo& d : devices) { + std::printf("device %u: %s [%04x:%04x]%s\n", d.index, d.name.c_str(), d.vendor_id, d.device_id, + d.discrete ? " discrete" : ""); + if (!d.unusable_reason.empty()) { + std::printf(" unusable: %s\n", d.unusable_reason.c_str()); + } else { + const std::string no_coop = vk::no_coop_reason(d.caps); + std::printf(" kernels: %s%s\n", vk::block_path_name(vk::default_block_path(d.caps)), + no_coop.empty() ? "" : (" (no coop: " + no_coop + ")").c_str()); + } + } + return 0; +} + +int run(const Args& a) { + vk::Device device(a.device); + const vk::Capabilities& caps = device.info().caps; + const vk::BlockPath path = a.kernels.empty() ? vk::default_block_path(caps) : vk::block_path_from_name(a.kernels); + // An explicit choice the device cannot run is an unsupported request (exit 2), refused before anything loads. + if (path == vk::BlockPath::kCoop && !caps.coop_bf16) { + throw std::invalid_argument(device.info().name + " cannot run --kernels coop: " + caps.coop_bf16_reason); + } + if (path == vk::BlockPath::kCoopF16 && !caps.coop_f16) { + throw std::invalid_argument(device.info().name + " cannot run --kernels coop-f16: " + caps.coop_f16_reason); + } + // Always reported: the plain kernels are much slower, and which ones ran should be visible. + const std::string no_coop = vk::no_coop_reason(caps); + std::fprintf(stderr, "vulkan: device %u, %s; kernels: %s%s\n", device.info().index, device.info().name.c_str(), + vk::block_path_name(path), no_coop.empty() ? "" : (" (no coop: " + no_coop + ")").c_str()); + + const tok::Tokenizer tokenizer = tok::Tokenizer::load(a.model / "tokenizer" / "tokenizer.json"); + vk::GenerateRequest g; + g.model_dir = a.model; + tok::PromptTokenReport cond_report, uncond_report; + g.cond_ids = tok::encode_moss_prompt_report(tokenizer, a.prompt, static_cast(a.max_seq), cond_report); + const bool cfg = a.cfg != 1.0; + if (cfg) { + g.uncond_ids = tok::encode_moss_prompt_report(tokenizer, a.negative, static_cast(a.max_seq), uncond_report); + } + for (const auto* t : {&cond_report, &uncond_report}) { + if (t->truncated) { + std::fprintf(stderr, "warning: %s truncated to %zu of %zu tokens (%zu dropped; --max-seq %zu)\n", + t == &cond_report ? "prompt" : "negative prompt", t->kept_tokens, t->total_tokens, t->dropped_tokens(), + t->max_length); + } + } + + std::uint64_t seed = 0; + if (a.seed >= 0) { + seed = static_cast(a.seed); + } else { // as the engine: a random 32-bit seed, reported + std::random_device rd; + seed = std::uniform_int_distribution(0, (std::uint64_t{1} << 32) - 1)(rd); + } + g.width = static_cast(a.width); + g.height = static_cast(a.height); + g.noise = pipeline::initial_noise(pipeline::NoiseLayout::kOfficial, seed, vk::kDiTLatentChannels, a.height / 16, a.width / 16); + const pipeline::Schedule schedule = pipeline::moss_schedule(a.steps, a.shift); + g.sigmas = schedule.sigmas; + g.timesteps = schedule.timesteps; + g.cfg = static_cast(a.cfg); + for (const int l : a.skip) g.uncond_skip.set(static_cast(l)); + g.decode = !a.out_png.empty(); + if (!a.out_latent.empty()) { + g.on_latents = [&](const std::vector& latents) { + vk::npy::save(a.out_latent, "(std::chrono::steady_clock::now().time_since_epoch()).count(); + if (std::string_view(stage) == "denoise" && done > 0) { + std::fprintf(stderr, "denoise %d/%d %.2f s\n", done, total, now - last); + } else if (done == 0) { + std::fprintf(stderr, "%s\n", stage); + } + last = now; + }; + } + + vk::GenerateResult r; + try { + r = vk::generate(device, path, g); + } catch (const std::invalid_argument& e) { + throw std::runtime_error(e.what()); // an internal check, not a usage error: the request was checked above + } + if (!a.out_png.empty()) { + std::vector text; + if (a.metadata) { + pipeline::GenerationMetadata m; + m.engine_version = std::to_string(MI_VERSION_MAJOR) + "." + std::to_string(MI_VERSION_MINOR) + "." + + std::to_string(MI_VERSION_PATCH); + m.tier = "reference"; + m.models.push_back({"bundle", util::path_utf8(a.model.filename()), ""}); + m.prompt = a.prompt; + m.negative_prompt = a.negative; + m.sampler = "euler"; + m.scheduler = pipeline::scheduler_name(pipeline::SchedulerKind::kMoss); + m.steps = a.steps; + m.cfg = a.cfg; + m.shift = a.shift; + m.seed = seed; + m.width = a.width; + m.height = a.height; + m.extra.emplace_back("backend", std::string("vulkan ") + vk::block_path_name(path)); + text = pipeline::generation_text_chunks(m); + } + pipeline::write_png_rgb8(a.out_png, r.rgb, static_cast(r.width), static_cast(r.height), text); + } + if (!a.quiet) { + std::fprintf(stderr, "done: seed %llu; text encoder %.1f s, DiT load %.1f s (%d of 42 blocks on the device), denoise %.1f s " + "(%.2f s/step), VAE %.1f s\n", + static_cast(seed), r.text_encoder_s, r.dit_load_s, r.dit_resident_blocks, r.denoise_s, + r.denoise_s / static_cast(g.timesteps.size()), r.vae_s); + } + return 0; +} + +int run_main(const std::vector& argv) { + Args a; + try { + a = parse(argv); + } catch (const std::invalid_argument& e) { + if (e.what()[0] != '\0') std::fprintf(stderr, "moss-infer-vulkan: %s\n", e.what()); + std::fputs(kUsage, stderr); + return 2; + } + try { + return a.list ? list_devices() : run(a); + } catch (const std::invalid_argument& e) { + std::fprintf(stderr, "moss-infer-vulkan: %s\n", e.what()); + return 2; + } catch (const std::exception& e) { + std::fprintf(stderr, "moss-infer-vulkan: %s\n", e.what()); + return 1; + } +} + +} // namespace + +#ifdef _WIN32 +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include + +// char** argv is in the ANSI code page on Windows and loses characters outside it: take the UTF-16 command line. +int wmain(int argc, wchar_t** argv) { + std::vector utf8; + for (int i = 0; i < argc; ++i) { + const int n = WideCharToMultiByte(CP_UTF8, 0, argv[i], -1, nullptr, 0, nullptr, nullptr); + std::string s(n > 0 ? static_cast(n) : 1, '\0'); + if (n > 0) WideCharToMultiByte(CP_UTF8, 0, argv[i], -1, s.data(), n, nullptr, nullptr); + s.pop_back(); + utf8.push_back(std::move(s)); + } + return run_main(utf8); +} +#else +int main(int argc, char** argv) { return run_main(std::vector(argv, argv + argc)); } +#endif diff --git a/vulkan/tools/probe.py b/vulkan/tools/probe.py new file mode 100644 index 0000000..a44c3ee --- /dev/null +++ b/vulkan/tools/probe.py @@ -0,0 +1,123 @@ +"""Lists a Vulkan device's cooperative-matrix shapes and memory heaps, using only the Vulkan loader. + +The C++ tools need the Vulkan headers to build; this script needs nothing but Python (ctypes), so it works on a +machine with just the GPU driver installed, e.g. to see whether a device has the bf16 shape the coop kernels use. + + py -3 probe.py +""" +import ctypes as C +import sys + +VK_SUCCESS = 0 +VK_STRUCTURE_TYPE_APPLICATION_INFO = 0 +VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO = 1 +VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR = 1000506001 + +COMPONENT = { + 0: "f16", 1: "f32", 2: "f64", 3: "s8", 4: "s16", 5: "s32", 6: "s64", + 7: "u8", 8: "u16", 9: "u32", 10: "u64", 1000141000: "bf16", +} +SCOPE = {1: "device", 2: "workgroup", 3: "subgroup"} + + +class VkApplicationInfo(C.Structure): + _fields_ = [ + ("sType", C.c_int32), ("pNext", C.c_void_p), ("pApplicationName", C.c_char_p), + ("applicationVersion", C.c_uint32), ("pEngineName", C.c_char_p), ("engineVersion", C.c_uint32), + ("apiVersion", C.c_uint32), + ] + + +class VkInstanceCreateInfo(C.Structure): + _fields_ = [ + ("sType", C.c_int32), ("pNext", C.c_void_p), ("flags", C.c_uint32), + ("pApplicationInfo", C.POINTER(VkApplicationInfo)), ("enabledLayerCount", C.c_uint32), + ("ppEnabledLayerNames", C.c_void_p), ("enabledExtensionCount", C.c_uint32), + ("ppEnabledExtensionNames", C.c_void_p), + ] + + +class VkCooperativeMatrixProperties(C.Structure): + _fields_ = [ + ("sType", C.c_int32), ("pNext", C.c_void_p), + ("MSize", C.c_uint32), ("NSize", C.c_uint32), ("KSize", C.c_uint32), + ("AType", C.c_int32), ("BType", C.c_int32), ("CType", C.c_int32), ("ResultType", C.c_int32), + ("saturatingAccumulation", C.c_uint32), ("scope", C.c_int32), + ] + + +class VkMemoryType(C.Structure): + _fields_ = [("propertyFlags", C.c_uint32), ("heapIndex", C.c_uint32)] + + +class VkMemoryHeap(C.Structure): + _fields_ = [("size", C.c_uint64), ("flags", C.c_uint32)] + + +class VkPhysicalDeviceMemoryProperties(C.Structure): + _fields_ = [ + ("memoryTypeCount", C.c_uint32), ("memoryTypes", VkMemoryType * 32), + ("memoryHeapCount", C.c_uint32), ("memoryHeaps", VkMemoryHeap * 16), + ] + + +def main(): + vk = C.WinDLL("vulkan-1.dll") if sys.platform == "win32" else C.CDLL("libvulkan.so.1") + vk.vkGetInstanceProcAddr.restype = C.c_void_p + vk.vkGetInstanceProcAddr.argtypes = [C.c_void_p, C.c_char_p] + gipa = vk.vkGetInstanceProcAddr + + def proc(instance, name, restype, argtypes): + address = gipa(instance, name.encode()) + return C.CFUNCTYPE(restype, *argtypes)(address) if address else None + + app = VkApplicationInfo(VK_STRUCTURE_TYPE_APPLICATION_INFO, None, b"mossinfer-probe", 1, b"probe", 1, (1 << 22) | (3 << 12)) + info = VkInstanceCreateInfo(VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO, None, 0, C.pointer(app), 0, None, 0, None) + instance = C.c_void_p() + if proc(None, "vkCreateInstance", C.c_int32, + [C.POINTER(VkInstanceCreateInfo), C.c_void_p, C.POINTER(C.c_void_p)])(C.byref(info), None, + C.byref(instance)) != VK_SUCCESS: + print("vkCreateInstance failed") + return 1 + + enumerate_devices = proc(instance, "vkEnumeratePhysicalDevices", C.c_int32, + [C.c_void_p, C.POINTER(C.c_uint32), C.POINTER(C.c_void_p)]) + count = C.c_uint32(0) + enumerate_devices(instance, C.byref(count), None) + devices = (C.c_void_p * count.value)() + enumerate_devices(instance, C.byref(count), devices) + get_memory = proc(instance, "vkGetPhysicalDeviceMemoryProperties", None, + [C.c_void_p, C.POINTER(VkPhysicalDeviceMemoryProperties)]) + print(f"physical devices: {count.value}") + + for index in range(count.value): + device = devices[index] + print(f"\ndevice {index}:") + query = proc(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR", C.c_int32, + [C.c_void_p, C.POINTER(C.c_uint32), C.POINTER(VkCooperativeMatrixProperties)]) + n = C.c_uint32(0) + if query is None or query(device, C.byref(n), None) != VK_SUCCESS or n.value == 0: + print(" no VK_KHR_cooperative_matrix shapes") + else: + shapes = (VkCooperativeMatrixProperties * n.value)() + for shape in shapes: + shape.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR + if query(device, C.byref(n), shapes) == VK_SUCCESS: + print(f" cooperative matrix: {n.value} shapes") + for shape in shapes[:n.value]: + print(f" {shape.MSize}x{shape.NSize}x{shape.KSize} " + f"A={COMPONENT.get(shape.AType, shape.AType)} B={COMPONENT.get(shape.BType, shape.BType)} " + f"C={COMPONENT.get(shape.CType, shape.CType)} R={COMPONENT.get(shape.ResultType, shape.ResultType)} " + f"sat={shape.saturatingAccumulation} scope={SCOPE.get(shape.scope, shape.scope)}") + if get_memory is not None: + memory = VkPhysicalDeviceMemoryProperties() + get_memory(device, C.byref(memory)) + for heap in range(memory.memoryHeapCount): + info = memory.memoryHeaps[heap] + kind = " (device local)" if info.flags & 1 else "" + print(f" heap {heap}: {info.size / (1024 * 1024):.0f} MiB{kind}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vulkan/tools/vulkan_block_check.cpp b/vulkan/tools/vulkan_block_check.cpp new file mode 100644 index 0000000..e6d7cba --- /dev/null +++ b/vulkan/tools/vulkan_block_check.cpp @@ -0,0 +1,493 @@ +// Checks the Vulkan DiT block's kernels against host replays of their op order, with tails on every axis, then runs and +// times one block on synthetic activations. The replays and limits come from the HIP prototype's known-answer check: `same` +// is the share of outputs with the replay's bits, and ULPs are bf16 ULPs of max(|replay|, floor). Exit status: 0 ran, +// 1 failed. Usage: vulkan_block_check [--device N] [--rows L1,L2,...] [--iters N] [--no-block] +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "half.h" +#include "runtime.h" + +namespace { + +using namespace mi::vk; +using bits = std::uint16_t; + +float from_bf16(bits b) { return bf16_to_float(b); } +bits to_bf16(float v) { + const auto u = std::bit_cast(v); + if ((u & 0x7fffffffu) > 0x7f800000u) return static_cast((u >> 16) | 0x0040u); + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} +float host_rbf(float v) { return from_bf16(to_bf16(v)); } + +// splitmix64-hashed bf16 values +-m * 2^(scale_exp - 7 - e), m in [128, 255], e in 0..3. +std::vector synth_bf16(std::uint64_t seed, std::size_t n, int scale_exp) { + std::vector out(n); + for (std::size_t i = 0; i < n; ++i) { + std::uint64_t z = i + seed * 0x9E3779B97F4A7C15ull; + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBull; + z ^= z >> 31; + const float m = static_cast(128 + (z & 0x7F)); + const auto e = static_cast((z >> 7) & 3u); + out[i] = to_bf16(std::ldexp((z >> 10) & 1 ? -m : m, scale_exp - 7 - e)); + } + return out; +} + +std::vector synth_rope(std::uint64_t seed, std::size_t tokens) { + const std::vector angle = synth_bf16(seed, tokens * 64, 1); + std::vector out(tokens * 128); + for (std::size_t t = 0; t < tokens; ++t) { + for (std::size_t i = 0; i < 64; ++i) { + const float a = from_bf16(angle[t * 64 + i]); + out[t * 128 + i] = to_bf16(std::sin(a)); + out[t * 128 + 64 + i] = to_bf16(std::cos(a)); + } + } + return out; +} + +float host_rsqrt_of_mean(const bits* row, int d, float factor) { + std::array v{}; + for (int t = 0; t < 32; ++t) { + std::array acc{}; + for (int idx = t; idx * 4 + 3 < d; idx += 32) { + for (int i = 0; i < 4; ++i) { + const float e = from_bf16(row[idx * 4 + i]); + const float sq = e * e; + acc[static_cast(i)] += sq; + } + } + v[static_cast(t)] = acc[0] + acc[1] + acc[2] + acc[3]; + } + for (int off = 16; off > 0; off >>= 1) { + for (int t = 0; t < off; ++t) v[static_cast(t)] += v[static_cast(t + off)]; + } + const float mean = v[0] * factor; + return 1.0f / std::sqrt(mean + 1e-6f); +} + +bits host_norm_modulate(bits x, float r, bits weight, bits scale, bits shift) { + const float xn = host_rbf(from_bf16(x) * r); + const float y = host_rbf(xn * from_bf16(weight)); + const float one_plus = host_rbf(1.0f + from_bf16(scale)); + const float scaled = host_rbf(y * one_plus); + return to_bf16(scaled + from_bf16(shift)); +} + +void host_qk_norm_rope_head(bits* row, const bits* weight, const bits* table, float factor) { + const float r = host_rsqrt_of_mean(row, kDiTHeadDim, factor); + std::array y{}; + for (int d = 0; d < kDiTHeadDim; ++d) { + y[static_cast(d)] = host_rbf(host_rbf(from_bf16(row[d]) * r) * from_bf16(weight[d])); + } + for (int d = 0; d < kDiTHeadDim / 2; ++d) { + const float lo = y[static_cast(d)], hi = y[static_cast(d + 64)]; + const float si = from_bf16(table[d]), co = from_bf16(table[64 + d]); + row[d] = to_bf16(host_rbf(lo * co) + host_rbf(-hi * si)); + row[d + 64] = to_bf16(host_rbf(hi * co) + host_rbf(lo * si)); + } +} + +float host_situ(float h, float g) { + constexpr float kInv25 = 1.0f / 25.0f; + const float g4 = host_rbf(std::tanh(host_rbf(g * 0.25f))); + const float sigmoid = 1.0f / (1.0f + std::exp(-g)); + const float left = host_rbf(host_rbf(4.0f * g4) * host_rbf(sigmoid)); + const float right = host_rbf(25.0f * host_rbf(std::tanh(host_rbf(h * kInv25)))); + return left * right; +} + +constexpr int kStep = 32; + +std::array replay_attention_row(const bits* q, const bits* k, std::int64_t ldk, const bits* v, + std::int64_t ldv, int seqlen_k, float scale) { + const float log2_scale = static_cast(static_cast(scale) * 1.4426950408889634074); + std::array acc{}; + std::array s{}, p{}; + float row_max = 0, row_sum = 0; + const int steps = (seqlen_k + kStep - 1) / kStep; + for (int step = steps - 1; step >= 0; --step) { + const bool first = step == steps - 1; + const int n0 = step * kStep; + const int cols = std::min(kStep, seqlen_k - n0); + for (int c = 0; c < cols; ++c) { + const bits* key = k + static_cast(n0 + c) * ldk; + float sum = 0; + for (int d = 0; d < kDiTHeadDim; ++d) { + const float product = from_bf16(q[d]) * from_bf16(key[d]); + sum += product; + } + s[static_cast(c)] = sum; + } + const float prev = row_max; + float cur = first ? s[0] : std::fmax(prev, s[0]); + for (int c = 1; c < cols; ++c) cur = std::fmax(cur, s[static_cast(c)]); + row_max = cur; + float sum = 0; + if (!first) { + const float diff = prev - cur; + const float rescale = std::exp2(diff * log2_scale); + sum = row_sum * rescale; + for (float& e : acc) e *= rescale; + } + const float max_scaled = std::isinf(cur) && cur < 0 ? 0.0f : cur * log2_scale; + for (int c = 0; c < cols; ++c) { + const float w = std::exp2(std::fma(s[static_cast(c)], log2_scale, -max_scaled)); + sum += w; + p[static_cast(c)] = host_rbf(w); + } + row_sum = sum; + for (int c = 0; c < cols; ++c) { + const bits* val = v + static_cast(n0 + c) * ldv; + for (int d = 0; d < kDiTHeadDim; ++d) { + const float product = p[static_cast(c)] * from_bf16(val[d]); + acc[static_cast(d)] += product; + } + } + } + const float inv_sum = row_sum == 0.0f || row_sum != row_sum ? 1.0f : 1.0f / row_sum; + for (float& e : acc) e *= inv_sum; + return acc; +} + +std::unique_ptr device_buffer(Device& device, const std::vector& values) { + auto b = std::make_unique(device, std::max(values.size(), 1) * 2); + b->upload(values.data(), values.size() * 2); + return b; +} + +std::vector read(Buffer& b, std::size_t count) { + std::vector out(count); + b.download(out.data(), count * 2); + return out; +} + +double ulps_off(float got, double ref, double floor_mag) { + int e = 0; + (void)std::frexp(std::max(std::abs(ref), floor_mag), &e); + return std::abs(static_cast(got) - ref) / std::ldexp(1.0, e - 8); +} + +double rms_of(const std::vector& v) { + double s = 0.0; + for (const double x : v) s += x * x; + return std::sqrt(s / static_cast(std::max(v.size(), 1))); +} + +constexpr double kMinSamePercent = 99.0; + +struct Limit { + double ulps; + double floor_of_rms; +}; + +// row_len > 0: the floor below which errors count as at that magnitude is taken per row of row_len values instead of over +// the whole output, for outputs whose rows differ in scale by many orders of magnitude. +bool report(const char* name, const std::vector& got, const std::vector& ref, Limit limit, + std::size_t row_len = 0) { + const std::size_t len = row_len > 0 ? row_len : std::max(ref.size(), 1); + double worst = 0.0; + std::size_t nans = 0, same = 0; + for (std::size_t row = 0; row * len < got.size(); ++row) { + const std::size_t first = row * len, last = std::min(got.size(), first + len); + const double floor_mag = rms_of(std::vector(ref.begin() + static_cast(first), + ref.begin() + static_cast(last))) * + limit.floor_of_rms; + for (std::size_t i = first; i < last; ++i) { + const float g = from_bf16(got[i]); + if (std::isnan(g)) ++nans; + same += g == from_bf16(to_bf16(static_cast(ref[i]))); // by value: -0 and +0 are the same + worst = std::max(worst, ulps_off(g, ref[i], floor_mag)); + } + } + const double same_pct = 100.0 * static_cast(same) / static_cast(got.size()); + const bool ok = nans == 0 && worst <= limit.ulps && same_pct >= kMinSamePercent; + std::printf(" %-22s %s (%zu values, %.3f %% same, limit %.0f %%; worst %.2f ULPs, limit %.1f; %zu NaN)\n", name, + ok ? "ok" : "FAILED", got.size(), same_pct, kMinSamePercent, worst, limit.ulps, nans); + return ok; +} + +double host_dot(const std::vector& x, int ldx, const std::vector& w, int ldw, int r, int n, int k) { + double s = 0.0; + for (int i = 0; i < k; ++i) { + s += static_cast(from_bf16(x[static_cast(r) * ldx + i])) * from_bf16(w[static_cast(n) * ldw + i]); + } + return s; +} + +bool check_norms(Device& device, const BlockKernels& kernels) { + constexpr int rows = 37, d = kDiTDim; + const auto x_h = synth_bf16(11, static_cast(rows) * d, 0), w_h = synth_bf16(12, d, 0), + scale_h = synth_bf16(13, d, -2), shift_h = synth_bf16(14, d, -2); + auto x = device_buffer(device, x_h), w = device_buffer(device, w_h), scale = device_buffer(device, scale_h), + shift = device_buffer(device, shift_h), out = device_buffer(device, std::vector(x_h.size())); + const float factor = static_cast(rows) / static_cast(rows * d); + std::vector ref(x_h.size()); + for (int r = 0; r < rows; ++r) { + const bits* row = x_h.data() + static_cast(r) * d; + const float rs = host_rsqrt_of_mean(row, d, factor); + for (int c = 0; c < d; ++c) { + ref[static_cast(r) * d + c] = from_bf16(host_norm_modulate(row[c], rs, w_h[static_cast(c)], + scale_h[static_cast(c)], shift_h[static_cast(c)])); + } + } + device.run([&](VkCommandBuffer cb) { + kernels.record_norm_modulate(cb, x->address(), d, w->address(), scale->address(), shift->address(), out->address(), d, rows, d); + }); + bool ok = report("rms_norm_modulate", read(*out, x_h.size()), ref, {3.0, 1.0}); + for (int r = 0; r < rows; ++r) { + const bits* row = x_h.data() + static_cast(r) * d; + const float rs = host_rsqrt_of_mean(row, d, factor); + for (int c = 0; c < d; ++c) { + ref[static_cast(r) * d + c] = host_rbf(host_rbf(from_bf16(row[c]) * rs) * from_bf16(w_h[static_cast(c)])); + } + } + device.run([&](VkCommandBuffer cb) { kernels.record_norm(cb, x->address(), d, w->address(), out->address(), d, rows, d); }); + ok = report("rms_norm", read(*out, x_h.size()), ref, {3.0, 1.0}) && ok; + + { + constexpr std::uint32_t n = 4 * 1001; + const auto a_h = synth_bf16(19, n, 2), b_h = synth_bf16(20, n, 3); + auto a = device_buffer(device, a_h), b = device_buffer(device, b_h), o = device_buffer(device, std::vector(n)); + std::vector silu_ref(n), situ_ref(n); + for (std::uint32_t i = 0; i < n; ++i) { + const float v = from_bf16(a_h[i]); + silu_ref[i] = v / (1.0f + std::exp(-v)); + situ_ref[i] = host_situ(v, from_bf16(b_h[i])); + } + device.run([&](VkCommandBuffer cb) { kernels.record_silu(cb, a->address(), o->address(), n); }); + ok = report("silu", read(*o, n), silu_ref, {1.0, 1.0}) && ok; + device.run([&](VkCommandBuffer cb) { kernels.record_situ(cb, a->address(), b->address(), o->address(), n); }); + ok = report("situ", read(*o, n), situ_ref, {1.5, 1.0}) && ok; + } + + constexpr int tokens = 37, ld = 3 * kDiTDim; + const auto qkv_h = synth_bf16(15, static_cast(tokens) * ld, 0), q_h = synth_bf16(16, kDiTHeadDim, 0), + k_h = synth_bf16(17, kDiTHeadDim, 0), table_h = synth_rope(18, tokens); + auto qkv = device_buffer(device, qkv_h), qw = device_buffer(device, q_h), kw = device_buffer(device, k_h), table = device_buffer(device, table_h); + std::vector want = qkv_h; + const std::int64_t head_rows = std::int64_t{tokens} * kDiTHeads; + const float qk_factor = static_cast(head_rows) / static_cast(head_rows * kDiTHeadDim); + for (int t = 0; t < tokens; ++t) { + for (int h = 0; h < 2 * kDiTHeads; ++h) { + host_qk_norm_rope_head(want.data() + static_cast(t) * ld + static_cast(h) * kDiTHeadDim, + (h < kDiTHeads ? q_h : k_h).data(), table_h.data() + static_cast(t) * kDiTHeadDim, qk_factor); + } + } + device.run([&](VkCommandBuffer cb) { + kernels.record_qk_norm_rope(cb, qkv->address(), ld, tokens, kDiTHeads, qw->address(), kw->address(), table->address(), kDiTHeadDim); + }); + std::vector rope_ref(want.size()); + for (std::size_t i = 0; i < want.size(); ++i) rope_ref[i] = from_bf16(want[i]); + return report("qk_norm_rope", read(*qkv, want.size()), rope_ref, {3.0, 1.0}) && ok; +} + +// m, n and k: tails on every axis of both paths' tiles; the larger shape also crosses tiles and bands of tiles. +// wide_rows: row r of x is scaled by 2^e, e cycling through values far outside f16's range, which the coop-f16 path must +// bring in range row by row (and which the other paths take as they are). +bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int k, int f, bool wide_rows = false) { + const int ld = 3 * k; + auto x_h = synth_bf16(21, static_cast(m) * ld, 0); + const auto w_h = synth_bf16(22, static_cast(n) * k, -2), gate_h = synth_bf16(23, n, -2), + res_h = synth_bf16(24, static_cast(m) * n, 0); + if (wide_rows) { + constexpr int kExponents[] = {0, 24, -40, 60, -90, 17, 100, -20}; + for (int r = 0; r < m; ++r) { + for (int c = 0; c < ld; ++c) { + bits& v = x_h[static_cast(r) * ld + c]; + v = to_bf16(std::ldexp(from_bf16(v), kExponents[r % 8])); + } + } + } + auto x = device_buffer(device, x_h), w = device_buffer(device, w_h), gate = device_buffer(device, gate_h), + out = device_buffer(device, std::vector(static_cast(m) * n)); + std::vector plain(static_cast(m) * n), gated(plain.size()), situ(static_cast(m) * f); + for (int r = 0; r < m; ++r) { + for (int c = 0; c < n; ++c) { + const double ab = host_dot(x_h, ld, w_h, k, r, c, k); + plain[static_cast(r) * n + c] = ab; + const float delta = host_rbf(host_rbf(static_cast(ab)) * from_bf16(gate_h[static_cast(c)])); + gated[static_cast(r) * n + c] = from_bf16(res_h[static_cast(r) * n + c]) + delta; + } + for (int c = 0; c < f; ++c) { + const float h = host_rbf(static_cast(host_dot(x_h, ld, w_h, k, r, c, k))); + const float g = host_rbf(static_cast(host_dot(x_h, ld, w_h, k, r, f + c, k))); + situ[static_cast(r) * f + c] = host_situ(h, g); + } + } + const std::size_t rows_n = wide_rows ? static_cast(n) : 0, rows_f = wide_rows ? static_cast(f) : 0; + std::printf(" GEMMs at M %d, N %d, K %d, SiTU F %d%s\n", m, n, k, f, wide_rows ? ", rows scaled by 2^-90 .. 2^100" : ""); + device.run([&](VkCommandBuffer cb) { kernels.record_gemm(cb, out->address(), n, x->address(), ld, w->address(), k, m, n, k); }); + bool ok = report("gemm", read(*out, plain.size()), plain, {1.0, 1.0 / 64.0}, rows_n); + out->upload(res_h.data(), res_h.size() * 2); + device.run([&](VkCommandBuffer cb) { + kernels.record_gemm_gated_residual(cb, out->address(), n, x->address(), ld, w->address(), k, gate->address(), m, n, k); + }); + ok = report("gemm_gated_residual", read(*out, plain.size()), gated, {1.0, 1.0}, rows_n) && ok; + std::vector biased(plain.size()), biased_res(plain.size()); + for (int r = 0; r < m; ++r) { + for (int c = 0; c < n; ++c) { + const std::size_t i = static_cast(r) * n + c; + const float bias = from_bf16(gate_h[static_cast(c)]); + biased[i] = plain[i] + bias; + biased_res[i] = from_bf16(res_h[i]) + host_rbf(static_cast(plain[i]) + bias); + } + } + device.run([&](VkCommandBuffer cb) { + kernels.record_gemm_bias(cb, out->address(), n, x->address(), ld, w->address(), k, gate->address(), m, n, k); + }); + ok = report("gemm_bias", read(*out, plain.size()), biased, {1.0, 1.0 / 64.0}, rows_n) && ok; + out->upload(res_h.data(), res_h.size() * 2); + device.run([&](VkCommandBuffer cb) { + kernels.record_gemm_bias_residual(cb, out->address(), n, x->address(), ld, w->address(), k, gate->address(), m, n, k); + }); + // Two roundings: bf16(ab + bias) can land one ULP of its own off when ab differs in the last bits, then the sum rounds. + ok = report("gemm_bias_residual", read(*out, plain.size()), biased_res, {1.5, 1.0}, rows_n) && ok; + auto ffn = device_buffer(device, std::vector(static_cast(m) * f)); + device.run([&](VkCommandBuffer cb) { kernels.record_gemm_situ_glu(cb, ffn->address(), f, x->address(), ld, w->address(), k, m, f, k); }); + return report("gemm_situ_glu", read(*ffn, situ.size()), situ, {1.5, 1.0}, rows_f) && ok; +} + +bool check_attention(Device& device, const BlockKernels& kernels) { + constexpr int heads = 2, ld = heads * kDiTHeadDim; + const std::vector lengths = {37, 64, 5, 33, 200}; + std::vector cu = {0}; + for (const int l : lengths) cu.push_back(cu.back() + l); + const auto total = static_cast(cu.back()); + const float scale = static_cast(1.0 / std::sqrt(static_cast(kDiTHeadDim))); + const auto q_h = synth_bf16(31, total * ld, 0), k_h = synth_bf16(32, total * ld, 0), v_h = synth_bf16(33, total * ld, 0); + auto q = device_buffer(device, q_h), k = device_buffer(device, k_h), v = device_buffer(device, v_h), + o = device_buffer(device, std::vector(total * ld)); + std::vector ref(total * ld); + for (std::size_t b = 0; b < lengths.size(); ++b) { + const std::int64_t r0 = cu[b]; + for (int h = 0; h < heads; ++h) { + for (std::int64_t r = r0; r < cu[b + 1]; ++r) { + const auto row = replay_attention_row(q_h.data() + r * ld + h * kDiTHeadDim, k_h.data() + r0 * ld + h * kDiTHeadDim, ld, + v_h.data() + r0 * ld + h * kDiTHeadDim, ld, lengths[b], scale); + std::copy(row.begin(), row.end(), ref.begin() + r * ld + h * kDiTHeadDim); + } + } + } + const AttentionArgs args{q->address(), k->address(), v->address(), o->address(), ld, ld, ld, ld, kDiTHeadDim, kDiTHeadDim, + kDiTHeadDim, kDiTHeadDim, cu, cu, heads, scale}; + device.run([&](VkCommandBuffer cb) { kernels.record_attention(cb, args); }); + return report("attention", read(*o, total * ld), ref, {1.0, 1.0}); +} + +double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +bool run_block(Device& device, const BlockKernels& kernels, const std::vector& rows, int iters) { + std::vector cu = {0}; + for (const int l : rows) cu.push_back(cu.back() + l); + const auto n = static_cast(cu.back()); + const std::size_t dim = kDiTDim, ffn_dim = kDiTFfnDim; + auto norm1 = device_buffer(device, synth_bf16(1, dim, 0)), norm2 = device_buffer(device, synth_bf16(2, dim, 0)), + q_norm = device_buffer(device, synth_bf16(3, kDiTHeadDim, 0)), k_norm = device_buffer(device, synth_bf16(4, kDiTHeadDim, 0)), + qkv_w = device_buffer(device, synth_bf16(5, 3 * dim * dim, -5)), o_w = device_buffer(device, synth_bf16(6, dim * dim, -5)), + up = device_buffer(device, synth_bf16(7, 2 * ffn_dim * dim, -5)), down = device_buffer(device, synth_bf16(8, dim * ffn_dim, -6)); + const auto x_init = synth_bf16(9, n * dim, 0); + auto x = device_buffer(device, x_init), rope = device_buffer(device, synth_rope(10, n)), mod = device_buffer(device, synth_bf16(11, 6 * dim, -2)); + Buffer y(device, n * dim * 2), qkv(device, n * 3 * dim * 2), attn(device, n * dim * 2), ffn(device, n * ffn_dim * 2); + const BlockWeights w{norm1->address(), norm2->address(), q_norm->address(), k_norm->address(), qkv_w->address(), o_w->address(), up->address(), down->address()}; + const BlockBuffers b{x->address(), y.address(), qkv.address(), attn.address(), ffn.address(), rope->address(), mod->address(), cu, static_cast(n)}; + + // One submit per op, so a slow device's watchdog sees short work and the times are per op. + std::vector> op_ms(kBlockOps); + std::vector block_ms; + for (int it = 0; it < iters; ++it) { + x->upload(x_init.data(), x_init.size() * 2); + double total = 0.0; + for (int op = 0; op < kBlockOps; ++op) { + const double t0 = now_ms(); + device.run([&](VkCommandBuffer cb) { record_block_op(kernels, cb, op, w, b); }); + const double ms = now_ms() - t0; + op_ms[static_cast(op)].push_back(ms); + total += ms; + } + block_ms.push_back(total); + } + const auto out = read(*x, x_init.size()); + const auto nans = static_cast(std::count_if(out.begin(), out.end(), [](bits v) { return (v & 0x7fffu) > 0x7f80u; })); + std::printf(" %-22s %s (%zu of %zu output values are NaN)\n", "whole block", nans == 0 ? "ok" : "FAILED", nans, out.size()); + std::printf("\n%zu tokens in %zu rows, %d iterations\n", n, rows.size(), iters); + const auto row = [](const char* name, std::vector ms) { + std::sort(ms.begin(), ms.end()); + std::printf("| %-20s | %10.2f | %10.2f |\n", name, ms.front(), ms[ms.size() / 2]); + }; + std::printf("| %-20s | %10s | %10s |\n", "op", "min ms", "median ms"); + for (int op = 0; op < kBlockOps; ++op) row(block_op_name(op), op_ms[static_cast(op)]); + row("whole block", block_ms); + return nans == 0; +} + +} // namespace + +int main(int argc, char** argv) { + int device_index = -1, iters = 3; + bool block = true; + std::string path_name; + std::vector rows = {4378, 4106}; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--device" && i + 1 < argc) { + device_index = std::atoi(argv[++i]); + } else if (arg == "--iters" && i + 1 < argc) { + iters = std::atoi(argv[++i]); + } else if (arg == "--path" && i + 1 < argc) { + path_name = argv[++i]; + } else if (arg == "--no-block") { + block = false; + } else if (arg == "--rows" && i + 1 < argc) { + rows.clear(); + std::stringstream ss(argv[++i]); + std::string item; + while (std::getline(ss, item, ',')) rows.push_back(std::atoi(item.c_str())); + } else { + std::printf("usage: %s [--device N] [--path coop|coop-f16|plain] [--rows L1,L2,...] [--iters N] [--no-block]\n", argv[0]); + return 2; + } + } + try { + Device device(device_index); + std::printf("device %u: %s\nknown-answer check against host replays:\n", device.info().index, device.info().name.c_str()); + if (!path_name.empty() && path_name != "coop" && path_name != "coop-f16" && path_name != "plain") { + std::printf("--path takes coop, coop-f16 or plain\n"); + return 2; + } + const BlockKernels kernels = path_name.empty() ? BlockKernels(device) + : BlockKernels(device, block_path_from_name(path_name)); + const std::string no_coop = no_coop_reason(device.info().caps); + std::printf("path: %s%s\n", block_path_name(kernels.path()), no_coop.empty() ? "" : (" (no coop: " + no_coop + ")").c_str()); + bool ok = check_norms(device, kernels); + ok = check_gemms(device, kernels, 37, 104, 40, 24) && ok; + ok = check_gemms(device, kernels, 300, 1048, 264, 520) && ok; + ok = check_gemms(device, kernels, 300, 1048, 264, 520, true) && ok; + ok = check_attention(device, kernels) && ok; + if (!ok) { + std::printf("known-answer check FAILED: not timing.\n"); + return 1; + } + if (block) ok = run_block(device, kernels, rows, iters); + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/tools/vulkan_block_real.cpp b/vulkan/tools/vulkan_block_real.cpp new file mode 100644 index 0000000..015e22c --- /dev/null +++ b/vulkan/tools/vulkan_block_real.cpp @@ -0,0 +1,152 @@ +// Runs one DiT block's kernels on real weights and activations, and compares each output with the engine's +// own, captured from its forward at the 1024x1024 golden case. Reads files only, so it needs no CUDA. +// Usage: vulkan_block_real DIR [--device N] +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "half.h" +#include "runtime.h" + +namespace { + +using namespace mi::vk; +using bits = std::uint16_t; +namespace fs = std::filesystem; + +std::vector load(const fs::path& path) { + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) throw std::runtime_error("cannot read " + path.string()); + const auto bytes = static_cast(in.tellg()); + std::vector out(bytes / 2); + in.seekg(0); + in.read(reinterpret_cast(out.data()), static_cast(bytes)); + return out; +} + +std::unique_ptr device_buffer(Device& device, const std::vector& values) { + auto b = std::make_unique(device, values.size() * 2); + b->upload(values.data(), values.size() * 2); + return b; +} + +// ULPs of max(|engine|, rms): near-zero outputs of cancelling sums are not relatively accurate in any accumulation order. +double scaled_ulps(float got, float engine, double rms) { + int e = 0; + (void)std::frexp(std::max(std::abs(engine), rms), &e); + return std::abs(static_cast(got) - static_cast(engine)) / std::ldexp(1.0, e - 8); +} + +bool compare(const std::vector& got, const std::vector& engine) { + double squares = 0.0; + for (const bits b : engine) squares += static_cast(bf16_to_float(b)) * bf16_to_float(b); + const double rms = std::sqrt(squares / static_cast(engine.size())); + std::vector errors(got.size()); + std::size_t same = 0, nans = 0, engine_nans = 0; + double dot = 0.0, norm_got = 0.0, norm_engine = 0.0; + for (std::size_t i = 0; i < got.size(); ++i) { + const float g = bf16_to_float(got[i]), e = bf16_to_float(engine[i]); + if (std::isnan(g)) ++nans; + if (std::isnan(e)) ++engine_nans; + same += got[i] == engine[i]; + errors[i] = scaled_ulps(g, e, rms); + dot += static_cast(g) * e; + norm_got += static_cast(g) * g; + norm_engine += static_cast(e) * e; + } + std::sort(errors.begin(), errors.end()); + const auto at = [&](double q) { return errors[static_cast(q * static_cast(errors.size() - 1))]; }; + const double cosine = dot / std::sqrt(norm_got * norm_engine); + std::printf(" %zu values: %.2f %% same bits, scaled ULPs p50 %.2f p99 %.2f p99.9 %.2f max %.2f, cosine %.9f, NaN %zu (engine %zu)\n", + got.size(), 100.0 * static_cast(same) / static_cast(got.size()), at(0.5), at(0.99), at(0.999), + errors.back(), cosine, nans, engine_nans); + return nans == engine_nans && cosine >= 0.999; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { + std::printf("usage: %s DIR [--device N] [--path coop|coop-f16|plain]\n", argv[0]); + return 2; + } + const fs::path dir = argv[1]; + int device_index = -1; + std::string path_name; + for (int i = 2; i + 1 < argc; i += 2) { + const std::string arg = argv[i]; + if (arg == "--device") device_index = std::atoi(argv[i + 1]); + else if (arg == "--path") path_name = argv[i + 1]; + } + try { + Device device(device_index); + const BlockKernels kernels = path_name.empty() ? BlockKernels(device) + : BlockKernels(device, block_path_from_name(path_name)); + std::printf("device %u: %s, path %s\n", device.info().index, device.info().name.c_str(), block_path_name(kernels.path())); + + std::vector lengths; + struct Case { + int block, tokens, skip_uncond; + }; + std::vector cases; + std::ifstream manifest(dir / "manifest.txt"); + std::string line; + while (std::getline(manifest, line)) { + std::istringstream in(line); + std::string kind, unused; + in >> kind; + if (kind == "lengths") { + std::int32_t a = 0, b = 0; + in >> a >> b; + lengths = {a, b}; + } else if (kind == "block") { + Case c{}; + in >> c.block >> unused >> c.tokens >> unused >> c.skip_uncond; + cases.push_back(c); + } + } + if (lengths.empty() || cases.empty()) throw std::runtime_error("no manifest in " + dir.string()); + + auto rope = device_buffer(device, load(dir / "rope.bin")); + auto mod = device_buffer(device, load(dir / "mod.bin")); + bool ok = true; + for (const Case& c : cases) { + const std::string p = "block" + std::to_string(c.block) + "."; + std::printf("block %d: %d tokens%s\n", c.block, c.tokens, c.skip_uncond ? " (the uncond row skips it)" : ""); + auto norm1 = device_buffer(device, load(dir / (p + "norm1.bin"))), norm2 = device_buffer(device, load(dir / (p + "norm2.bin"))), + q_norm = device_buffer(device, load(dir / (p + "q_norm.bin"))), k_norm = device_buffer(device, load(dir / (p + "k_norm.bin"))), + qkv_w = device_buffer(device, load(dir / (p + "qkv.bin"))), o_w = device_buffer(device, load(dir / (p + "o.bin"))), + up = device_buffer(device, load(dir / (p + "up_gate.bin"))), down = device_buffer(device, load(dir / (p + "down.bin"))); + auto x = device_buffer(device, load(dir / (p + "x_in.bin"))); + const auto engine = load(dir / (p + "x_out.bin")); + const auto n = static_cast(c.tokens); + Buffer y(device, n * kDiTDim * 2), qkv(device, n * 3 * kDiTDim * 2), attn(device, n * kDiTDim * 2), ffn(device, n * kDiTFfnDim * 2); + std::vector cu = {0, lengths[0]}; + if (!c.skip_uncond) cu.push_back(lengths[0] + lengths[1]); + const BlockWeights w{norm1->address(), norm2->address(), q_norm->address(), k_norm->address(), qkv_w->address(), o_w->address(), up->address(), down->address()}; + const BlockBuffers b{x->address(), y.address(), qkv.address(), attn.address(), ffn.address(), rope->address(), mod->address(), cu, static_cast(n)}; + const auto start = std::chrono::steady_clock::now(); + for (int op = 0; op < kBlockOps; ++op) device.run([&](VkCommandBuffer cb) { record_block_op(kernels, cb, op, w, b); }); + std::printf(" %.0f ms\n", std::chrono::duration(std::chrono::steady_clock::now() - start).count()); + std::vector got(engine.size()); + x->download(got.data(), got.size() * 2); + ok = compare(got, engine) && ok; + } + std::printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/tools/vulkan_check.cpp b/vulkan/tools/vulkan_check.cpp new file mode 100644 index 0000000..bd5fcc4 --- /dev/null +++ b/vulkan/tools/vulkan_check.cpp @@ -0,0 +1,204 @@ +// Bring-up check for the Vulkan backend: lists the devices, checks each GEMM path against an fp64 CPU reference, then +// times the engine's shapes. Exits non-zero when a check fails. +// +// Usage: vulkan_check [--list] [--device N] [--path auto|coop|plain] [--no-bench] [--bench M N K iterations]... +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "gemm.h" +#include "half.h" +#include "runtime.h" + +namespace { + +using mi::vk::Buffer; +using mi::vk::Device; +using mi::vk::Gemm; +using mi::vk::GemmArgs; +using mi::vk::GemmPath; +using mi::vk::InputType; + +struct Shape { + std::uint32_t m, n, k; + int iterations; +}; + +std::vector random_values(std::size_t count, std::uint32_t seed, InputType input) { + std::mt19937 rng(seed); + std::uniform_real_distribution values(-1.0f, 1.0f); + std::vector out(count); + for (auto& v : out) { + const float f = values(rng); + v = input == InputType::kBF16 ? mi::vk::float_to_bf16(f) : mi::vk::float_to_half(f); + } + return out; +} + +float decode(std::uint16_t v, InputType input) { + return input == InputType::kBF16 ? mi::vk::bf16_to_float(v) : mi::vk::half_to_float(v); +} + +// Strides larger than the rows and a poisoned output, so a write past a row or past M fails the check. +bool check_gemm(Device& device, const Gemm& gemm, const Shape& s, InputType input) { + const std::uint32_t lda = s.k + 16, ldb = s.k + 32, ldc = s.n + 8; + const std::size_t a_count = static_cast(s.m) * lda, b_count = static_cast(s.n) * ldb; + const std::size_t c_count = (static_cast(s.m) + 1) * ldc; + const auto a = random_values(a_count, 1234, input), b = random_values(b_count, 5678, input); + const std::vector poison(c_count, -123.0f); + Buffer da(device, a_count * 2), db(device, b_count * 2), dc(device, c_count * 4); + da.upload(a.data(), a_count * 2); + db.upload(b.data(), b_count * 2); + dc.upload(poison.data(), c_count * 4); + const GemmArgs args{&da, &db, &dc, s.m, s.n, s.k, lda, ldb, ldc}; + device.run([&](VkCommandBuffer cb) { gemm.record(cb, args, input); }); + std::vector c(c_count); + dc.download(c.data(), c_count * 4); + + double max_error = 0.0, dot = 0.0, norm_got = 0.0, norm_want = 0.0, squared = 0.0, toward = 0.0; + std::size_t overwritten = 0, wrong_nan = 0; + for (std::size_t i = 0; i < c_count; ++i) { + const std::size_t row = i / ldc, column = i % ldc; + if (row >= s.m || column >= s.n) { + if (c[i] != -123.0f) ++overwritten; + continue; + } + double want = 0.0; + for (std::uint32_t k = 0; k < s.k; ++k) { + want += static_cast(decode(a[row * lda + k], input)) * static_cast(decode(b[column * ldb + k], input)); + } + const double got = c[i]; + if (std::isnan(got)) ++wrong_nan; + max_error = std::max(max_error, std::fabs(got - want)); + squared += (got - want) * (got - want); + toward += (got - want) * (want >= 0.0 ? 1.0 : -1.0); + dot += got * want; + norm_got += got * got; + norm_want += want * want; + } + const double cosine = dot / std::sqrt(norm_got * norm_want); + // Tensor cores round their fp32 accumulator with a bias toward zero, so the error grows with K. Measured on an RTX 4090 + // Laptop: 4.8e-5 at K = 512 and 1.3e-4 at K = 1040, matching cuBLASLt on the same data, and 8.1e-4 at K = 4096, where + // cuBLASLt reaches 1.3e-4. The plain path stays well inside this. Far below one bf16 step at these magnitudes. + const double tolerance = 2.5e-7 * s.k + 1e-5; + const bool ok = overwritten == 0 && wrong_nan == 0 && max_error <= tolerance && cosine >= 0.9999; + const double cells = static_cast(s.m) * s.n; + std::printf("%-5s %-4s %5ux%5ux%5u: max error %.3g (limit %.3g), rms %.3g, mean toward |value| %.3g, cosine %.9f, " + "outside writes %zu %s\n", + mi::vk::gemm_path_name(gemm.path()), input == InputType::kBF16 ? "bf16" : "f16", s.m, s.n, s.k, max_error, + tolerance, std::sqrt(squared / cells), toward / cells, cosine, overwritten, ok ? "ok" : "FAILED"); + return ok; +} + +void bench_gemm(Device& device, const Gemm& gemm, const Shape& s) { + const std::size_t a_count = static_cast(s.m) * s.k, b_count = static_cast(s.n) * s.k; + const auto a = random_values(a_count, 1234, InputType::kF16), b = random_values(b_count, 5678, InputType::kF16); + Buffer da(device, a_count * 2), db(device, b_count * 2); + Buffer dc(device, static_cast(s.m) * s.n * 4); + da.upload(a.data(), a_count * 2); + db.upload(b.data(), b_count * 2); + const GemmArgs args{&da, &db, &dc, s.m, s.n, s.k, s.k, s.k, s.n}; + const auto record = [&](VkCommandBuffer cb) { + for (int i = 0; i < s.iterations; ++i) gemm.record(cb, args); + }; + device.run(record); // warm-up + const auto start = std::chrono::steady_clock::now(); + device.run(record); + const double seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count() / s.iterations; + std::printf("%-5s %5ux%5ux%5u: %9.3f ms %7.2f TFLOPS\n", mi::vk::gemm_path_name(gemm.path()), s.m, s.n, s.k, + seconds * 1e3, 2.0 * s.m * s.n * s.k / seconds / 1e12); +} + +void print_devices() { + for (const auto& d : mi::vk::list_devices()) { + std::printf("device %u: %s [%04x:%04x]%s\n", d.index, d.name.c_str(), d.vendor_id, d.device_id, + d.discrete ? " discrete" : ""); + if (!d.unusable_reason.empty()) { + std::printf(" unusable: %s\n", d.unusable_reason.c_str()); + continue; + } + std::printf(" shared memory %u bytes, subgroup size %u-%u\n", d.caps.shared_memory_bytes, d.caps.min_subgroup, + d.caps.max_subgroup); + std::printf(" GEMM: %s\n", d.caps.coop ? "cooperative matrix" : ("plain compute (" + d.caps.coop_reason + ")").c_str()); + const std::string no_coop = mi::vk::no_coop_reason(d.caps); + std::printf(" model kernels: %s%s\n", mi::vk::block_path_name(mi::vk::default_block_path(d.caps)), + no_coop.empty() ? "" : (" (" + no_coop + ")").c_str()); + if (d.caps.coop_f16 && !d.caps.coop_bf16) std::printf(" no bf16 fragments: %s\n", d.caps.coop_bf16_reason.c_str()); + } +} + +} // namespace + +int main(int argc, char** argv) { + int device_index = -1; + std::string path = "auto"; + bool bench = true, list_only = false; + std::vector shapes; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--list") { + list_only = true; + } else if (arg == "--device" && i + 1 < argc) { + device_index = std::atoi(argv[++i]); + } else if (arg == "--path" && i + 1 < argc) { + path = argv[++i]; + } else if (arg == "--no-bench") { + bench = false; + } else if (arg == "--bench" && i + 4 < argc) { + Shape s{}; + s.m = static_cast(std::atoi(argv[++i])); + s.n = static_cast(std::atoi(argv[++i])); + s.k = static_cast(std::atoi(argv[++i])); + s.iterations = std::atoi(argv[++i]); + shapes.push_back(s); + } else { + std::printf("usage: %s [--list] [--device N] [--path auto|coop|plain] [--no-bench] [--bench M N K iterations]...\n", + argv[0]); + return 2; + } + } + if (path != "auto" && path != "coop" && path != "plain") { + std::printf("--path must be auto, coop or plain\n"); + return 2; + } + try { + print_devices(); + if (list_only) return 0; + Device device(device_index); + std::printf("using device %u: %s\n", device.info().index, device.info().name.c_str()); + const Gemm gemm = path == "auto" ? Gemm(device) + : path == "coop" ? Gemm(device, GemmPath::kCoop) + : Gemm(device, GemmPath::kPlain); + + // One tile; tails in M, N and K; several bands with a partial last one; the largest row-count tile boundaries. + const std::vector checks = {{1, 4, 16, 1}, {3, 8, 16, 1}, {64, 64, 64, 1}, {65, 68, 48, 1}, + {67, 136, 1040, 1}, {129, 256, 4096, 1}, {128, 640, 64, 1}, {94, 644, 64, 1}, + {1030, 9216, 64, 1}, {1025, 136, 1040, 1}, {512, 512, 512, 1}}; + bool ok = true; + for (const Shape& s : checks) ok = check_gemm(device, gemm, s, InputType::kF16) && ok; + if (gemm.path() == GemmPath::kPlain) { + for (const Shape& s : checks) ok = check_gemm(device, gemm, s, InputType::kBF16) && ok; + } + if (bench) { + // The engine's shapes at 1024x1024 with CFG: the DiT's QKV, projection, up/gate and down, then the text encoder's. + if (shapes.empty()) { + shapes = {{8222, 9216, 3072, 5}, {8222, 3072, 3072, 5}, {8222, 16384, 3072, 3}, + {8222, 3072, 8192, 3}, {282, 12288, 4096, 10}, {282, 4096, 12288, 10}}; + } + for (const Shape& s : shapes) bench_gemm(device, gemm, s); + } + std::printf("%s\n", ok ? "all checks passed" : "CHECKS FAILED"); + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/tools/vulkan_dit_run.cpp b/vulkan/tools/vulkan_dit_run.cpp new file mode 100644 index 0000000..75011b7 --- /dev/null +++ b/vulkan/tools/vulkan_dit_run.cpp @@ -0,0 +1,173 @@ +// Runs the DiT's sampling loop on Vulkan from a golden case's text-encoder outputs and initial noise (the official PyTorch +// pipeline's, tests/golden), with CFG and the official Euler step on the host in fp32, and compares step 0's predictions +// and the final latent with the golden ones. Writes the final latent as .npy for the VAE. +// Usage: vulkan_dit_run TRANSFORMER_DIR GOLDEN_DIR CASE [--steps N] [--out FILE.npy] [--path coop|coop-f16|plain] [--device N] +// [--dit-budget GiB] (device memory for the weights; the blocks past it stream) +// [--submit-gflop X] (work per submit; 0: the path's default) +// GOLDEN_DIR is tests/golden, CASE e.g. g_fridge_1024; the case's cfg and uncond skip layers are 3.0 and {22}. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "dit.h" +#include "npy.h" +#include "runtime.h" + +namespace { + +using namespace mi::vk; +namespace fs = std::filesystem; + +std::uint16_t to_bf16(float v) { + const auto u = std::bit_cast(v); + if ((u & 0x7fffffffu) > 0x7f800000u) return static_cast((u >> 16) | 0x0040u); + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +std::vector load_f32(const fs::path& p) { return npy::load(p).to_vector(); } + +struct Compare { + double cosine, max_abs, rms_ref, rms_diff; + std::size_t nans; +}; + +Compare compare(const std::vector& got, const std::vector& ref) { + double dot = 0, gg = 0, rr = 0, dd = 0, max_abs = 0; + std::size_t nans = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + if (std::isnan(got[i])) ++nans; + const double g = got[i], r = ref[i], d = g - r; + dot += g * r; + gg += g * g; + rr += r * r; + dd += d * d; + max_abs = std::max(max_abs, std::abs(d)); + } + const double n = static_cast(got.size()); + return {dot / std::sqrt(gg * rr), max_abs, std::sqrt(rr / n), std::sqrt(dd / n), nans}; +} + +void print(const char* name, const Compare& c) { + std::printf(" %-22s cosine %.9f, rms diff %.3e (rms of golden %.3e), max abs diff %.3e, NaN %zu\n", name, c.cosine, + c.rms_diff, c.rms_ref, c.max_abs, c.nans); +} + +double now_s() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 4) { + std::printf("usage: %s TRANSFORMER_DIR GOLDEN_DIR CASE [--steps N] [--out FILE.npy] [--path coop|coop-f16|plain] [--device N]\n", + argv[0]); + return 2; + } + const fs::path transformer = argv[1], golden = argv[2]; + const std::string name = argv[3]; + int steps = -1, device_index = -1; + double budget_gib = 0, submit_gflop = 0; + fs::path out_path; + std::string path_name; + for (int i = 4; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--steps" && i + 1 < argc) steps = std::atoi(argv[++i]); + else if (arg == "--out" && i + 1 < argc) out_path = argv[++i]; + else if (arg == "--path" && i + 1 < argc) path_name = argv[++i]; + else if (arg == "--device" && i + 1 < argc) device_index = std::atoi(argv[++i]); + else if (arg == "--dit-budget" && i + 1 < argc) budget_gib = std::atof(argv[++i]); + else if (arg == "--submit-gflop" && i + 1 < argc) submit_gflop = std::atof(argv[++i]); + else { + std::printf("unknown argument %s\n", arg.c_str()); + return 2; + } + } + try { + const fs::path dir = golden / name; + const npy::Array noise = npy::load(dir / "initial_noise.npy"); + if (noise.shape.size() != 4 || noise.shape[1] != kDiTLatentChannels) throw std::runtime_error("unexpected noise shape"); + const auto height = static_cast(noise.shape[2]), width = static_cast(noise.shape[3]); + std::vector latents = noise.to_vector(); + const std::vector sigmas = load_f32(dir / "sigmas.npy"), timesteps = load_f32(dir / "timesteps.npy"); + const int total_steps = static_cast(timesteps.size()); + if (steps < 0 || steps > total_steps) steps = total_steps; + constexpr float kCfg = 3.0f; + SkipSet uncond_skip; + uncond_skip.set(22); + + Device device(device_index); + const BlockKernels kernels = path_name.empty() ? BlockKernels(device) + : BlockKernels(device, block_path_from_name(path_name)); + std::printf("device %u: %s, path %s\n", device.info().index, device.info().name.c_str(), block_path_name(kernels.path())); + double t0 = now_s(); + DitOptions options; + options.device_budget = static_cast(budget_gib * (1u << 30)); + options.submit_gflop = submit_gflop; + Dit dit(device, kernels, transformer, options); + std::printf("weights: %.2f GiB on the device, %.2f GiB in host memory (%d of %d blocks resident), loaded in %.1f s\n", + static_cast(dit.weight_bytes()) / (1u << 30), static_cast(dit.host_weight_bytes()) / (1u << 30), + dit.resident_blocks(), kDiTLayers, now_s() - t0); + + std::vector> texts; + std::vector rows; + for (const char* which : {"cond", "uncond"}) { + const npy::Array te = npy::load(golden / (name + "." + which + ".te_last_hidden.npy")); + const std::vector values = te.to_vector(); + std::vector bits(values.size()); + for (std::size_t i = 0; i < values.size(); ++i) bits[i] = to_bf16(values[i]); // exact: the encoder's bf16 output + const auto length = static_cast(te.shape[0]); + t0 = now_s(); + texts.push_back(dit.encode_text(bits, length)); + std::printf("text refiner, %s: %u tokens in %.3f s\n", which, length, now_s() - t0); + rows.push_back({texts.back().get(), length, std::string(which) == "uncond" ? uncond_skip : SkipSet{}}); + } + + const std::size_t n = latents.size(); + for (int s = 0; s < steps; ++s) { + t0 = now_s(); + const std::vector> pred = dit.forward(latents, height, width, timesteps[static_cast(s)], rows); + const double dt_s = now_s() - t0; + if (s == 0) { + std::printf("step 0 against the golden predictions:\n"); + print("cond", compare(pred[0], load_f32(dir / "dit_step0_cond.npy"))); + print("uncond (skips 22)", compare(pred[1], load_f32(dir / "dit_step0_uncond.npy"))); + } + // CFG then the official Euler step, the engine's op order (sampler_kernels.cu): u + g (c - u); x + h v. + const float h = static_cast(static_cast(sigmas[static_cast(s) + 1]) - + static_cast(sigmas[static_cast(s)])); + for (std::size_t i = 0; i < n; ++i) { + const float diff = pred[0][i] - pred[1][i]; + const float scaled = kCfg * diff; + const float v = pred[1][i] + scaled; + const float step = h * v; + latents[i] = latents[i] + step; + } + const std::vector& ms = dit.last_op_ms(); + std::printf("step %2d: %.2f s (", s, dt_s); + for (int op = 0; op < kBlockOps; ++op) std::printf("%s %.0f, ", block_op_name(op), ms[static_cast(op)]); + std::printf("rest %.0f ms, of which block copies %.0f ms)\n", ms[kBlockOps], dit.last_copy_ms()); + std::fflush(stdout); + } + if (steps == total_steps) { + std::printf("final latent against the golden one:\n"); + print("final latent", compare(latents, load_f32(dir / "final_latent.npy"))); + } + if (!out_path.empty()) { + npy::save(out_path, " +#include +#include +#include +#include +#include +#include +#include + +#include "half.h" +#include "npy.h" +#include "runtime.h" +#include "text_encoder.h" + +namespace { + +using namespace mi::vk; + +struct Metrics { + double cosine = 0, min_row_cosine = 1, max_row_rel = 0, mean_row_rel = 0, same_pct = 0, max_ulps = 0; + std::size_t nans = 0; +}; + +Metrics compare(const std::uint16_t* got, const std::vector& ref) { + Metrics m; + const std::size_t rows = ref.size() / kTeHidden; + double rr_all = 0; + for (const float v : ref) rr_all += static_cast(v) * v; + const double floor_mag = std::sqrt(rr_all / static_cast(ref.size())) / 64.0; + double dot = 0, gg = 0, rr = 0; + std::size_t same = 0; + for (std::size_t r = 0; r < rows; ++r) { + double rdot = 0, rg = 0, rref = 0, rdiff = 0; + for (std::size_t c = 0; c < kTeHidden; ++c) { + const std::size_t i = r * kTeHidden + c; + const double g = bf16_to_float(got[i]), w = ref[i]; + if (std::isnan(g)) ++m.nans; + same += got[i] == float_to_bf16(ref[i]); + rdot += g * w; + rg += g * g; + rref += w * w; + rdiff += (g - w) * (g - w); + int e = 0; + (void)std::frexp(std::max(std::abs(w), floor_mag), &e); + m.max_ulps = std::max(m.max_ulps, std::abs(g - w) / std::ldexp(1.0, e - 8)); + } + dot += rdot; + gg += rg; + rr += rref; + m.min_row_cosine = std::min(m.min_row_cosine, rdot / std::sqrt(rg * rref)); + m.max_row_rel = std::max(m.max_row_rel, std::sqrt(rdiff / rref)); + m.mean_row_rel += std::sqrt(rdiff / rref) / static_cast(rows); + } + m.cosine = dot / std::sqrt(gg * rr); + m.same_pct = 100.0 * static_cast(same) / static_cast(ref.size()); + return m; +} + +void print(const std::string& what, const Metrics& m) { + std::printf( + " %-16s cos %.7f min row cos %.6f row relL2 max %.3e mean %.3e same %6.2f %% max %.1f ULPs %zu NaN\n", + what.c_str(), m.cosine, m.min_row_cosine, m.max_row_rel, m.mean_row_rel, m.same_pct, m.max_ulps, m.nans); +} + +std::vector golden(const std::string& dir, const std::string& key, const std::string& what, std::size_t rows) { + const npy::Array a = npy::load(dir + "/" + key + "." + what + ".npy"); + if (a.shape != std::vector{static_cast(rows), kTeHidden}) { + throw std::runtime_error("unexpected golden shape: " + key + "." + what); + } + return a.to_vector(); +} + +} // namespace + +int main(int argc, char** argv) { + int device_index = -1, iters = 1; + TextEncoderOptions options; + std::string path_name, weights = "/workspace/Models/MOSS-Image/base/text_encoder", golden_dir = "tests/golden"; + std::vector cases; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--device" && i + 1 < argc) { + device_index = std::atoi(argv[++i]); + } else if (arg == "--path" && i + 1 < argc) { + path_name = argv[++i]; + } else if (arg == "--weights" && i + 1 < argc) { + weights = argv[++i]; + } else if (arg == "--golden" && i + 1 < argc) { + golden_dir = argv[++i]; + } else if (arg == "--case" && i + 1 < argc) { + cases.emplace_back(argv[++i]); + } else if (arg == "--profile") { + options.profile = true; + } else if (arg == "--iters" && i + 1 < argc) { + iters = std::max(1, std::atoi(argv[++i])); + } else { + std::printf( + "usage: %s [--device N] [--path coop|coop-f16|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] " + "[--profile]\n", + argv[0]); + return 2; + } + } + if (cases.empty()) cases.emplace_back("g_fridge_1024"); + if (!path_name.empty() && path_name != "coop" && path_name != "coop-f16" && path_name != "plain") { + std::printf("--path takes coop, coop-f16 or plain\n"); + return 2; + } + try { + Device device(device_index); + const BlockPath path = + path_name.empty() ? default_block_path(device.info().caps) : block_path_from_name(path_name); + TextEncoder encoder(device, weights, path, options); + std::printf("device %u: %s, path %s\n", device.info().index, device.info().name.c_str(), + block_path_name(encoder.path())); + bool ok = true; + for (const std::string& id : cases) { + const std::vector keys = {id + ".cond", id + ".uncond"}; + std::vector> seqs; + for (const std::string& key : keys) { + const auto ids = npy::load(golden_dir + "/" + key + ".token_ids.npy").to_vector(); + seqs.emplace_back(ids.begin(), ids.end()); + } + std::map> layers; + const TeLayerHook hook = [&](int layer, Buffer& hidden, std::uint32_t tokens) { + if (layer != 0 && layer != 3 && layer != 31) return; + auto& dst = layers[layer]; + dst.resize(static_cast(tokens) * kTeHidden); + hidden.download(dst.data(), dst.size() * 2); + }; + std::vector> out; + for (int it = 0; it < iters; ++it) { + out = encoder.encode(seqs, it == 0 ? hook : TeLayerHook{}); + const TextEncoderStats& s = encoder.last_stats(); + std::printf( + "%s: %u tokens, encode %.3f s (embed %.3f, file reads %.3f in the reader thread, waiting for them %.3f, " + "uploads %.3f, compute %.3f, %d submits), device %.1f MiB, host staging %.1f MiB\n", + id.c_str(), s.tokens, s.total_s, s.embed_s, s.read_s, s.read_wait_s, s.upload_s, s.compute_s, s.submits, + static_cast(s.device_bytes) / (1 << 20), static_cast(s.staging_bytes) / (1 << 20)); + for (const auto& [op, seconds] : s.op_seconds) std::printf(" %-24s %8.3f s\n", op.c_str(), seconds); + } + std::size_t row0 = 0; + for (std::size_t i = 0; i < keys.size(); ++i) { + const std::size_t rows = seqs[i].size(); + std::printf("%s (%zu tokens)\n", keys[i].c_str(), rows); + for (const int layer : {0, 3, 31}) { + char name[16]; + std::snprintf(name, sizeof(name), "te_layer%02d", layer); + const Metrics m = compare(layers[layer].data() + row0 * kTeHidden, golden(golden_dir, keys[i], name, rows)); + print(name, m); + ok = ok && m.nans == 0; + } + const Metrics m = compare(out[i].data(), golden(golden_dir, keys[i], "te_last_hidden", rows)); + print("te_last_hidden", m); + ok = ok && m.nans == 0; + row0 += rows; + } + } + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/tools/vulkan_te_kernel_check.cpp b/vulkan/tools/vulkan_te_kernel_check.cpp new file mode 100644 index 0000000..bd47e3f --- /dev/null +++ b/vulkan/tools/vulkan_te_kernel_check.cpp @@ -0,0 +1,756 @@ +// Checks the text encoder's kernels (te_kernels.h) against host replays of their op order, on small shapes with tails +// and several sequences, and the gated delta rule also against an fp64 token-by-token recurrence. `same` is the share +// of outputs with the replay's bits; ULPs are bf16 ULPs of max(|replay|, floor), fp32 ULPs for fp32 outputs. The +// replays use the host's exp, sqrt and division, which differ from the device's in the last bits of fp32. Exit status: +// 0 passed, 1 failed. Usage: vulkan_te_kernel_check [--device N] +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "half.h" +#include "runtime.h" +#include "te_kernels.h" + +namespace { + +using namespace mi::vk; +using bits = std::uint16_t; + +float from_bf16(bits b) { return bf16_to_float(b); } +bits to_bf16(float v) { return float_to_bf16(v); } +float rbf(float v) { return from_bf16(to_bf16(v)); } + +std::uint64_t mix(std::uint64_t z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBull; + return z ^ (z >> 31); +} + +// bf16 values +-m * 2^(scale_exp - 7 - e), m in [128, 255], e in 0..3. +std::vector synth_bf16(std::uint64_t seed, std::size_t n, int scale_exp) { + std::vector out(n); + for (std::size_t i = 0; i < n; ++i) { + const std::uint64_t z = mix(i + seed * 0x9E3779B97F4A7C15ull); + const float m = static_cast(128 + (z & 0x7F)); + const auto e = static_cast((z >> 7) & 3u); + out[i] = to_bf16(std::ldexp((z >> 10) & 1 ? -m : m, scale_exp - 7 - e)); + } + return out; +} + +template +std::unique_ptr device_buffer(Device& device, const std::vector& values) { + auto b = std::make_unique(device, std::max(values.size(), 1) * sizeof(T)); + if (!values.empty()) b->upload(values.data(), values.size() * sizeof(T)); + return b; +} + +template +std::vector read(Buffer& b, std::size_t count) { + std::vector out(count); + b.download(out.data(), count * sizeof(T)); + return out; +} + +double ulps_off(double got, double ref, double floor_mag, int mantissa_bits) { + int e = 0; + (void)std::frexp(std::max(std::abs(ref), floor_mag), &e); + return std::abs(got - ref) / std::ldexp(1.0, e - mantissa_bits); +} + +double rms_of(const std::vector& v) { + double s = 0.0; + for (const double x : v) s += x * x; + return std::sqrt(s / static_cast(std::max(v.size(), 1))); +} + +struct Limit { + double ulps; + double floor_of_rms; + double min_same_percent = 99.0; +}; + +// bf16 outputs against replayed values (rounded to bf16 for `same`). +bool report(const char* name, const std::vector& got, const std::vector& ref, Limit limit) { + const double floor_mag = rms_of(ref) * limit.floor_of_rms; + double worst = 0.0; + std::size_t nans = 0, same = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + const float g = from_bf16(got[i]); + if (std::isnan(g)) ++nans; + same += got[i] == to_bf16(static_cast(ref[i])); + worst = std::max(worst, ulps_off(g, ref[i], floor_mag, 8)); + } + const double same_pct = 100.0 * static_cast(same) / static_cast(got.size()); + const bool ok = nans == 0 && worst <= limit.ulps && same_pct >= limit.min_same_percent; + std::printf(" %-26s %s (%zu values, %.3f %% same, limit %.1f %%; worst %.2f ULPs, limit %.1f; %zu NaN)\n", name, + ok ? "ok" : "FAILED", got.size(), same_pct, limit.min_same_percent, worst, limit.ulps, nans); + return ok; +} + +// fp32 outputs. +bool report_f32(const char* name, const std::vector& got, const std::vector& ref, double max_ulps) { + double worst = 0.0; + std::size_t nans = 0, same = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + if (std::isnan(got[i])) ++nans; + same += std::bit_cast(got[i]) == std::bit_cast(ref[i]); + worst = std::max(worst, ulps_off(got[i], ref[i], 0.0, 24)); + } + const double same_pct = 100.0 * static_cast(same) / static_cast(got.size()); + const bool ok = nans == 0 && worst <= max_ulps; + std::printf(" %-26s %s (%zu values, %.3f %% same; worst %.2f fp32 ULPs, limit %.1f; %zu NaN)\n", name, + ok ? "ok" : "FAILED", got.size(), same_pct, worst, max_ulps, nans); + return ok; +} + +std::vector widen(const std::vector& v) { + std::vector out(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) out[i] = from_bf16(v[i]); + return out; +} + +float silu(float x) { return x / (1.0f + std::exp(-x)); } +float sigmoid(float x) { return 1.0f / (1.0f + std::exp(-x)); } + +// ---- norms ---- + +float host_inv_rms(const bits* row, int d, int bw, float factor, float eps) { + std::vector v(static_cast(bw)); + for (int t = 0; t < bw; ++t) { + std::array acc{}; + for (int idx = t; idx * 4 + 3 < d; idx += bw) { + for (int i = 0; i < 4; ++i) { + const float e = from_bf16(row[idx * 4 + i]); + const float sq = e * e; + acc[static_cast(i)] += sq; + } + } + v[static_cast(t)] = ((acc[0] + acc[1]) + acc[2]) + acc[3]; + } + for (int off = bw / 2; off >= 32; off >>= 1) { + for (int t = 0; t < off; ++t) v[static_cast(t)] += v[static_cast(t + off)]; + } + for (int off = 16; off > 0; off >>= 1) { + for (int t = 0; t < off; ++t) v[static_cast(t)] += v[static_cast(t + off)]; + } + const float mean = v[0] * factor; + return 1.0f / std::sqrt(mean + eps); +} + +// x and out [tokens][heads] rows at (t * ld + h * hs); the replay of te_rms_norm. +std::vector replay_norm(const std::vector& x, int ldx, int hsx, const std::vector& w, + const std::vector* gate, int ldg, int hsg, const std::vector* table, + std::vector out, int ldo, int hso, int tokens, int heads, int d) { + const std::uint64_t rows = static_cast(tokens) * heads; + const int bw = static_cast(torch_reduce_threads(rows, static_cast(d))); + const float factor = static_cast(rows) / static_cast(rows * d); + for (int t = 0; t < tokens; ++t) { + for (int h = 0; h < heads; ++h) { + const bits* xr = x.data() + static_cast(t) * ldx + static_cast(h) * hsx; + const float r = host_inv_rms(xr, d, bw, factor, 1e-6f); + std::vector y(static_cast(d)); + for (int c = 0; c < d; ++c) { + const float xn = from_bf16(xr[c]) * r; + if (gate == nullptr) { + const float one_plus = 1.0f + from_bf16(w[static_cast(c)]); + y[static_cast(c)] = xn * one_plus; + } else { + const float g = from_bf16((*gate)[static_cast(t) * ldg + static_cast(h) * hsg + c]); + const float weighted = rbf(from_bf16(w[static_cast(c)]) * rbf(xn)); + y[static_cast(c)] = weighted * silu(g); + } + } + bits* orow = out.data() + static_cast(t) * ldo + static_cast(h) * hso; + for (int c = 0; c < d; ++c) orow[c] = to_bf16(y[static_cast(c)]); + if (table != nullptr) { + const bits* tab = table->data() + static_cast(t) * 64; + for (int c = 0; c < 32; ++c) { + const float lo = rbf(y[static_cast(c)]), hi = rbf(y[static_cast(c + 32)]); + const float co = from_bf16(tab[c]), si = from_bf16(tab[32 + c]); + orow[c] = to_bf16(rbf(lo * co) + rbf(-hi * si)); + orow[c + 32] = to_bf16(rbf(hi * co) + rbf(lo * si)); + } + } + } + } + return out; +} + +bool check_norms(Device& device, const TeKernels& k) { + bool ok = true; + // The hidden-state norm at row counts that give torch 512, 128, 64 and 32 threads per row. + for (const int tokens : {1, 5, 9, 37}) { + constexpr int d = 4096, ld = 4100; + const auto x_h = synth_bf16(100 + tokens, static_cast(tokens) * ld, 2), w_h = synth_bf16(7, d, -3); + auto x = device_buffer(device, x_h), w = device_buffer(device, w_h), + out = device_buffer(device, std::vector(static_cast(tokens) * ld, 0x7fc0)); + RmsNormArgs a; + a.x = x->address(); + a.ldx = ld; + a.weight = w->address(); + a.out = out->address(); + a.ldo = ld; + a.tokens = static_cast(tokens); + a.d = d; + device.run([&](VkCommandBuffer cb) { k.record_rms_norm(cb, a); }); + std::vector want = + replay_norm(x_h, ld, 0, w_h, nullptr, 0, 0, nullptr, + std::vector(static_cast(tokens) * ld, 0x7fc0), ld, 0, tokens, 1, d); + std::vector got = read(*out, want.size()); + // Compare the rows only, not the stride's padding (which must stay untouched: NaN both sides). + std::vector g2, w2; + for (int t = 0; t < tokens; ++t) { + for (int c = 0; c < d; ++c) { + g2.push_back(got[static_cast(t) * ld + c]); + w2.push_back(want[static_cast(t) * ld + c]); + } + for (int c = d; c < ld; ++c) ok = ok && got[static_cast(t) * ld + c] == 0x7fc0; + } + const std::string name = "rms_norm rows " + std::to_string(tokens) + " (bw " + + std::to_string(torch_reduce_threads(static_cast(tokens), d)) + ")"; + ok = report(name.c_str(), g2, widen(w2), {1.0, 0.0}) && ok; + } + + // q norm + RoPE out of place from the [q | gate] per-head layout, and k norm + RoPE in place, as the attention layer. + for (const int tokens : {2, 19}) { + constexpr int heads = 3, kv_heads = 2, d = 256, ld = 2 * heads * d + kv_heads * d + 8; + const auto qkv_h = synth_bf16(200 + tokens, static_cast(tokens) * ld, 1), qw = synth_bf16(8, d, -3), + kw = synth_bf16(9, d, -3); + std::vector table_h(static_cast(tokens) * 64); + for (int t = 0; t < tokens; ++t) { + for (int i = 0; i < 32; ++i) { + const float phase = static_cast(t) * std::pow(1e7f, -static_cast(i) / 32.0f); + table_h[static_cast(t) * 64 + i] = to_bf16(std::cos(phase)); + table_h[static_cast(t) * 64 + 32 + i] = to_bf16(std::sin(phase)); + } + } + auto qkv = device_buffer(device, qkv_h), qwb = device_buffer(device, qw), kwb = device_buffer(device, kw), + table = device_buffer(device, table_h), + q = device_buffer(device, std::vector(static_cast(tokens) * heads * d)); + RmsNormArgs qa; + qa.x = qkv->address(); + qa.ldx = ld; + qa.hsx = 2 * d; + qa.out = q->address(); + qa.ldo = heads * d; + qa.hso = d; + qa.weight = qwb->address(); + qa.tokens = static_cast(tokens); + qa.heads = heads; + qa.d = d; + qa.rope = table->address(); + RmsNormArgs ka = qa; + ka.x = ka.out = qkv->address() + 2ull * heads * d * 2; + ka.ldx = ka.ldo = ld; + ka.hsx = ka.hso = d; + ka.weight = kwb->address(); + ka.heads = kv_heads; + device.run([&](VkCommandBuffer cb) { + k.record_rms_norm(cb, qa); + k.record_rms_norm(cb, ka); + }); + const auto q_want = + replay_norm(qkv_h, ld, 2 * d, qw, nullptr, 0, 0, &table_h, + std::vector(static_cast(tokens) * heads * d), heads * d, d, tokens, heads, d); + std::vector k_in(qkv_h.begin() + 2 * heads * d, qkv_h.end()); + k_in.resize(qkv_h.size()); + const auto k_want = replay_norm(k_in, ld, d, kw, nullptr, 0, 0, &table_h, k_in, ld, d, tokens, kv_heads, d); + const auto qkv_got = read(*qkv, qkv_h.size()); + std::vector k_got, k_ref; + for (int t = 0; t < tokens; ++t) { + for (int c = 0; c < kv_heads * d; ++c) { + k_got.push_back(qkv_got[static_cast(t) * ld + 2 * heads * d + c]); + k_ref.push_back(k_want[static_cast(t) * ld + c]); + } + for (int c = 0; c < 2 * heads * d; ++c) + ok = ok && qkv_got[static_cast(t) * ld + c] == qkv_h[static_cast(t) * ld + c]; + } + const std::string qn = "q norm+rope, " + std::to_string(tokens) + " tokens", + kn = "k norm+rope in place, " + std::to_string(tokens); + ok = report(qn.c_str(), read(*q, q_want.size()), widen(q_want), {1.0, 0.0}) && ok; + ok = report(kn.c_str(), k_got, widen(k_ref), {1.0, 0.0}) && ok; + } + + // The gated norm in place over 128-wide heads, the gate in another buffer's rows. + { + constexpr int tokens = 7, heads = 5, d = 128, ld = heads * d, ldg = 3 * ld; + const auto x_h = synth_bf16(300, static_cast(tokens) * ld, 0), + g_h = synth_bf16(301, static_cast(tokens) * ldg, 1), w_h = synth_bf16(302, d, 0); + auto x = device_buffer(device, x_h), g = device_buffer(device, g_h), w = device_buffer(device, w_h); + RmsNormArgs a; + a.x = a.out = x->address(); + a.ldx = a.ldo = ld; + a.hsx = a.hso = d; + a.gate = g->address() + 2ull * ld * 2; + a.ldg = ldg; + a.hsg = d; + a.weight = w->address(); + a.tokens = tokens; + a.heads = heads; + a.d = d; + a.gated = true; + device.run([&](VkCommandBuffer cb) { k.record_rms_norm(cb, a); }); + std::vector g_shift(g_h.begin() + 2 * ld, g_h.end()); + g_shift.resize(g_h.size()); + const auto want = replay_norm(x_h, ld, d, w_h, &g_shift, ldg, d, nullptr, x_h, ld, d, tokens, heads, d); + ok = report("gated rms_norm in place", read(*x, want.size()), widen(want), {1.0, 0.0}) && ok; + } + return ok; +} + +// ---- conv, gates, gate_mul ---- + +bool check_conv_gates(Device& device, const TeKernels& k) { + bool ok = true; + const std::vector lengths = {5, 70, 1}; + constexpr int ch = kGdnConvDim, ldx = ch + 4096, ldy = ch; + int total = 0; + for (const int l : lengths) total += l; + const auto x_h = synth_bf16(400, static_cast(total) * ldx, 1), + w_h = synth_bf16(401, static_cast(ch) * 4, -1); + auto x = device_buffer(device, x_h), w = device_buffer(device, w_h), + y = device_buffer(device, std::vector(static_cast(total) * ldy)); + std::vector want(static_cast(total) * ldy); + int s0 = 0; + for (const int len : lengths) { + const int start = s0; + device.run([&](VkCommandBuffer cb) { + k.record_conv_silu(cb, x->address() + 2ull * start * ldx, ldx, w->address(), y->address() + 2ull * start * ldy, + ldy, static_cast(len)); + }); + for (int t = 0; t < len; ++t) { + for (int c = 0; c < ch; ++c) { + float acc = 0.0f; + for (int kk = 0; kk < 4; ++kk) { + const int src = t - 3 + kk; + if (src < 0) continue; + acc = std::fma(from_bf16(w_h[static_cast(c) * 4 + kk]), + from_bf16(x_h[static_cast(start + src) * ldx + c]), acc); + } + want[static_cast(start + t) * ldy + c] = from_bf16(to_bf16(silu(rbf(acc)))); + } + } + s0 += len; + } + ok = report("conv1d+silu, 3 sequences", read(*y, want.size()), want, {1.0, 0.0}) && ok; + + // Gates over [b | a] rows of stride 72. + { + constexpr int tokens = 45, ld = 72; + const auto ba_h = synth_bf16(410, static_cast(tokens) * ld, 3), alog_h = synth_bf16(411, 32, 1), + dt_h = synth_bf16(412, 32, 2); + auto ba = device_buffer(device, ba_h), alog = device_buffer(device, alog_h), dt = device_buffer(device, dt_h); + Buffer g(device, tokens * 32 * 4), beta(device, tokens * 32 * 2); + device.run([&](VkCommandBuffer cb) { + k.record_gates(cb, ba->address(), ld, alog->address(), dt->address(), g.address(), beta.address(), tokens); + }); + std::vector g_want(tokens * 32); + std::vector beta_want(tokens * 32); + for (int t = 0; t < tokens; ++t) { + for (int h = 0; h < 32; ++h) { + const float b = from_bf16(ba_h[static_cast(t) * ld + h]), + a = from_bf16(ba_h[static_cast(t) * ld + 32 + h]); + beta_want[static_cast(t) * 32 + h] = rbf(sigmoid(b)); + const float xx = a + from_bf16(dt_h[static_cast(h)]); + const float sp = xx > 20.0f ? xx : std::log1p(std::exp(xx)); + g_want[static_cast(t) * 32 + h] = -std::exp(from_bf16(alog_h[static_cast(h)])) * sp; + } + } + ok = report("gates beta", read(beta, beta_want.size()), beta_want, {1.0, 0.0}) && ok; + // exp_f and log1p_f are fma polynomials, and Vulkan lets fma round twice (SPIR-V's Fma inherits the precision of a + // multiply then an add): lavapipe on LLVM 15 does, and lands up to 17 fp32 ULPs off, about 1e-6 relative. + ok = report_f32("gates g", read(g, g_want.size()), g_want, 32.0) && ok; + } + + // silu(gate) * up from one [gate | up] row, and the attention gate over 256-wide heads of a [q | gate] layout. + { + constexpr int rows = 9, cols = 520, ld = 2 * cols + 8; + const auto gu_h = synth_bf16(420, static_cast(rows) * ld, 1); + auto gu = device_buffer(device, gu_h), + out = device_buffer(device, std::vector(static_cast(rows) * cols)); + device.run([&](VkCommandBuffer cb) { + k.record_gate_mul(cb, TeKernels::GateMul::kSilu, gu->address(), ld, gu->address() + 2ull * cols, ld, + out->address(), cols, rows, cols, cols, 0); + }); + std::vector want(static_cast(rows) * cols); + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + const float g = from_bf16(gu_h[static_cast(r) * ld + c]), + u = from_bf16(gu_h[static_cast(r) * ld + cols + c]); + want[static_cast(r) * cols + c] = rbf(rbf(silu(g)) * u); + } + } + ok = report("silu_mul", read(*out, want.size()), want, {1.0, 0.0}) && ok; + + constexpr int heads = 3, hd = 256, ldq = 2 * heads * hd + 16, lda = heads * hd; + const auto qg_h = synth_bf16(421, static_cast(rows) * ldq, 2), + a_h = synth_bf16(422, static_cast(rows) * lda, 0); + auto qg = device_buffer(device, qg_h), at = device_buffer(device, a_h); + device.run([&](VkCommandBuffer cb) { + k.record_gate_mul(cb, TeKernels::GateMul::kSigmoid, at->address(), lda, qg->address() + 2ull * hd, ldq, + at->address(), lda, rows, lda, hd, 2 * hd); + }); + std::vector want2(static_cast(rows) * lda); + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < lda; ++c) { + const float g = from_bf16(qg_h[static_cast(r) * ldq + (c / hd) * 2 * hd + hd + c % hd]); + want2[static_cast(r) * lda + c] = + rbf(from_bf16(a_h[static_cast(r) * lda + c]) * rbf(sigmoid(g))); + } + } + ok = report("sigmoid_mul in place", read(*at, want2.size()), want2, {1.0, 0.0}) && ok; + } + return ok; +} + +// ---- the gated delta rule ---- + +constexpr int kKeyDim = kGdnKeyHeads * kGdnHeadDim; + +// The engine's chunked rule op by op (chunk_rule.cu): bf16 [T, 32, 128] out. +std::vector replay_gdn(const std::vector& qkv, int ld, const std::vector& g, + const std::vector& beta, const std::vector& cu) { + const int T = cu.back(); + constexpr int D = kGdnHeadDim, C = kGdnChunk; + std::vector qn(static_cast(kGdnKeyHeads) * T * D), kn(qn.size()); + for (int t = 0; t < T; ++t) { + for (int kh = 0; kh < kGdnKeyHeads; ++kh) { + for (int which = 0; which < 2; ++which) { + const bits* x = qkv.data() + static_cast(t) * ld + which * kKeyDim + kh * D; + std::array lanes{}; + for (int l = 0; l < 32; ++l) { + float s = 0.0f; + for (int e = 0; e < 4; ++e) s = std::fma(from_bf16(x[4 * l + e]), from_bf16(x[4 * l + e]), s); + lanes[static_cast(l)] = s; + } + for (int off = 16; off > 0; off >>= 1) { + for (int l = 0; l < off; ++l) lanes[static_cast(l)] += lanes[static_cast(l + off)]; + } + const float inv = 1.0f / std::sqrt(lanes[0] + 1e-6f); + float* dst = (which == 0 ? qn : kn).data() + (static_cast(kh) * T + t) * D; + for (int d = 0; d < D; ++d) + dst[d] = which == 0 ? (from_bf16(x[d]) * inv) * 0.08838834764831845f : from_bf16(x[d]) * inv; + } + } + } + std::vector out(static_cast(T) * 32 * D); + for (std::size_t s = 0; s + 1 < cu.size(); ++s) { + const int start = cu[s], len = cu[s + 1] - cu[s]; + for (int h = 0; h < kGdnValueHeads; ++h) { + const int kh = h / 2; + const auto Q = [&](int t, int d) { return qn[(static_cast(kh) * T + t) * D + d]; }; + const auto K = [&](int t, int d) { return kn[(static_cast(kh) * T + t) * D + d]; }; + std::vector S(static_cast(D) * D, 0.0f); // [k][v] + for (int c0 = 0; c0 < len; c0 += C) { + const int t0 = start + c0, n = std::min(C, len - c0); + std::array gc{}, bt{}, eg{}, el{}; + float acc = 0.0f; + for (int i = 0; i < C; ++i) { + acc += i < n ? g[static_cast(t0 + i) * 32 + h] : 0.0f; + gc[static_cast(i)] = acc; + bt[static_cast(i)] = i < n ? from_bf16(beta[static_cast(t0 + i) * 32 + h]) : 0.0f; + } + for (int i = 0; i < C; ++i) { + eg[static_cast(i)] = std::exp(gc[static_cast(i)]); + el[static_cast(i)] = + std::exp(gc[static_cast(n - 1)] - gc[static_cast(i)]); + } + std::vector a(C * C, 0.0f), attn(C * C, 0.0f); + for (int i = 0; i < n; ++i) { + for (int j = 0; j <= i; ++j) { + float kk = 0.0f, qk = 0.0f; + for (int d = 0; d < D; ++d) { + kk = std::fma(K(t0 + i, d), K(t0 + j, d), kk); + qk = std::fma(Q(t0 + i, d), K(t0 + j, d), qk); + } + const float decay = std::exp(gc[static_cast(i)] - gc[static_cast(j)]); + if (j < i) a[static_cast(i) * C + j] = (kk * bt[static_cast(i)]) * decay; + attn[static_cast(i) * C + j] = qk * decay; + } + } + std::vector u(static_cast(C) * D, 0.0f), w(u.size(), 0.0f); + for (int col = 0; col < D; ++col) { + for (int i = 0; i < n; ++i) { + const float b = bt[static_cast(i)]; + float xu = from_bf16(qkv[static_cast(t0 + i) * ld + 2 * kKeyDim + h * D + col]) * b; + float xw = (K(t0 + i, col) * b) * eg[static_cast(i)]; + for (int j = 0; j < i; ++j) { + xu = std::fma(-a[static_cast(i) * C + j], u[static_cast(j) * D + col], xu); + xw = std::fma(-a[static_cast(i) * C + j], w[static_cast(j) * D + col], xw); + } + u[static_cast(i) * D + col] = xu; + w[static_cast(i) * D + col] = xw; + } + } + std::vector vn(static_cast(C) * D, 0.0f); + for (int i = 0; i < n; ++i) { + for (int c = 0; c < D; ++c) { + float pw = 0.0f, pq = 0.0f; + for (int d = 0; d < D; ++d) { + const float sv = S[static_cast(d) * D + c]; + pw = std::fma(w[static_cast(i) * D + d], sv, pw); + pq = std::fma(Q(t0 + i, d) * eg[static_cast(i)], sv, pq); + } + vn[static_cast(i) * D + c] = u[static_cast(i) * D + c] - pw; + u[static_cast(i) * D + c] = pq; // reused: q S + } + } + for (int i = 0; i < n; ++i) { + for (int c = 0; c < D; ++c) { + float o = u[static_cast(i) * D + c]; + for (int j = 0; j <= i; ++j) + o = std::fma(attn[static_cast(i) * C + j], vn[static_cast(j) * D + c], o); + out[(static_cast(t0 + i) * 32 + h) * D + c] = o; + } + } + const float cd = std::exp(gc[static_cast(n - 1)]); + for (int d = 0; d < D; ++d) { + for (int c = 0; c < D; ++c) { + float upd = 0.0f; + for (int j = 0; j < n; ++j) + upd = std::fma(K(t0 + j, d) * el[static_cast(j)], vn[static_cast(j) * D + c], + upd); + S[static_cast(d) * D + c] = std::fma(S[static_cast(d) * D + c], cd, upd); + } + } + } + } + } + return out; +} + +// torch_recurrent_gated_delta_rule in fp64: the algorithm the chunked form rearranges. +std::vector recurrent_gdn(const std::vector& qkv, int ld, const std::vector& g, + const std::vector& beta, const std::vector& cu) { + constexpr int D = kGdnHeadDim; + std::vector out(static_cast(cu.back()) * 32 * D); + for (std::size_t s = 0; s + 1 < cu.size(); ++s) { + for (int h = 0; h < kGdnValueHeads; ++h) { + const int kh = h / 2; + std::vector S(static_cast(D) * D, 0.0); + for (int t = cu[s]; t < cu[s + 1]; ++t) { + const bits* row = qkv.data() + static_cast(t) * ld; + std::array q{}, k{}, v{}; + double sq = 0, sk = 0; + for (int d = 0; d < D; ++d) { + q[static_cast(d)] = from_bf16(row[kh * D + d]); + k[static_cast(d)] = from_bf16(row[kKeyDim + kh * D + d]); + v[static_cast(d)] = from_bf16(row[2 * kKeyDim + h * D + d]); + sq += q[static_cast(d)] * q[static_cast(d)]; + sk += k[static_cast(d)] * k[static_cast(d)]; + } + for (int d = 0; d < D; ++d) { + q[static_cast(d)] /= std::sqrt(sq + 1e-6) * std::sqrt(static_cast(D)); + k[static_cast(d)] /= std::sqrt(sk + 1e-6); + } + const double decay = std::exp(static_cast(g[static_cast(t) * 32 + h])); + const double b = from_bf16(beta[static_cast(t) * 32 + h]); + for (double& e : S) e *= decay; + for (int c = 0; c < D; ++c) { + double kv = 0; + for (int d = 0; d < D; ++d) kv += S[static_cast(d) * D + c] * k[static_cast(d)]; + const double delta = (v[static_cast(c)] - kv) * b; + double o = 0; + for (int d = 0; d < D; ++d) { + S[static_cast(d) * D + c] += k[static_cast(d)] * delta; + o += S[static_cast(d) * D + c] * q[static_cast(d)]; + } + out[(static_cast(t) * 32 + h) * D + c] = o; + } + } + } + } + return out; +} + +bool check_gdn(Device& device, const TeKernels& k) { + const std::vector lengths = {1, 5, 64, 130}; + std::vector cu = {0}; + for (const int l : lengths) cu.push_back(cu.back() + l); + const int T = cu.back(); + constexpr int ld = kGdnConvDim + 8; + // Conv outputs are SiLU outputs: mostly small positive values with some negative ones. + std::vector qkv_h = synth_bf16(500, static_cast(T) * ld, 0); + for (bits& b : qkv_h) b = to_bf16(silu(2.0f * from_bf16(b))); + std::vector g_h(static_cast(T) * 32); + std::vector beta_h(g_h.size()); + const auto gs = synth_bf16(501, g_h.size(), 1), bs = synth_bf16(502, g_h.size(), 1); + for (std::size_t i = 0; i < g_h.size(); ++i) { + g_h[i] = -0.6f * std::abs(from_bf16(gs[i])); // decays exp(g) in (0.3, 1] + beta_h[i] = to_bf16(sigmoid(from_bf16(bs[i]))); + } + auto qkv = device_buffer(device, qkv_h), g = device_buffer(device, g_h), beta = device_buffer(device, beta_h); + Buffer out(device, static_cast(T) * 32 * kGdnHeadDim * 2), + ws(device, GdnWorkspace::bytes(static_cast(T))); + GdnArgs a; + a.qkv = qkv->address(); + a.ldqkv = ld; + a.g = g->address(); + a.beta = beta->address(); + a.out = out.address(); + a.ws = GdnWorkspace::at(ws.address(), static_cast(T)); + a.cu.assign(cu.begin(), cu.end()); + device.run([&](VkCommandBuffer cb) { k.record_gated_delta_rule(cb, a); }); + const auto got = read(out, static_cast(T) * 32 * kGdnHeadDim); + const auto want = replay_gdn(qkv_h, ld, g_h, beta_h, cu); + bool ok = report("gated delta rule, 4 seqs", got, want, {2.0, 1.0 / 64.0, 98.0}); + + // The algorithm: per (token, head) row, relative L2 against the fp64 recurrence. + const auto exact = recurrent_gdn(qkv_h, ld, g_h, beta_h, cu); + double worst = 0, worst_replay = 0; + for (std::size_t r = 0; r < exact.size() / kGdnHeadDim; ++r) { + double diff = 0, diff_replay = 0, ref = 0; + for (int c = 0; c < kGdnHeadDim; ++c) { + const std::size_t i = r * kGdnHeadDim + c; + diff += std::pow(from_bf16(got[i]) - exact[i], 2); + diff_replay += std::pow(want[i] - exact[i], 2); + ref += exact[i] * exact[i]; + } + worst = std::max(worst, std::sqrt(diff / ref)); + worst_replay = std::max(worst_replay, std::sqrt(diff_replay / ref)); + } + const bool alg_ok = worst < 1e-2; + std::printf( + " %-26s %s (worst row rel-L2 %.3e vs fp64 recurrence, limit 1e-2; fp32 replay before bf16 rounding %.3e)\n", + "gated delta rule algorithm", alg_ok ? "ok" : "FAILED", worst, worst_replay); + return ok && alg_ok; +} + +// ---- attention ---- + +std::vector replay_attention_row(const bits* q, const bits* k, int ldk, const bits* v, int ldv, int seqlen_k, + float scale) { + constexpr int D = 256, kStep = 16; + const float log2_scale = static_cast(static_cast(scale) * 1.4426950408889634074); + std::vector acc(D, 0.0f); + std::array s{}, p{}; + float row_max = 0, row_sum = 0; + const int steps = (seqlen_k + kStep - 1) / kStep; + for (int step = steps - 1; step >= 0; --step) { + const bool first = step == steps - 1; + const int n0 = step * kStep, cols = std::min(kStep, seqlen_k - n0); + for (int c = 0; c < cols; ++c) { + const bits* key = k + static_cast(n0 + c) * ldk; + float sum = 0; + for (int d = 0; d < D; ++d) { + const float product = from_bf16(q[d]) * from_bf16(key[d]); + sum += product; + } + s[static_cast(c)] = sum; + } + const float prev = row_max; + float cur = first ? s[0] : std::fmax(prev, s[0]); + for (int c = 1; c < cols; ++c) cur = std::fmax(cur, s[static_cast(c)]); + row_max = cur; + float sum = 0; + if (!first) { + const float rescale = std::exp2((prev - cur) * log2_scale); + sum = row_sum * rescale; + for (float& e : acc) e *= rescale; + } + const float max_scaled = cur * log2_scale; + for (int c = 0; c < cols; ++c) { + const float w = std::exp2(std::fma(s[static_cast(c)], log2_scale, -max_scaled)); + sum += w; + p[static_cast(c)] = rbf(w); + } + row_sum = sum; + for (int c = 0; c < cols; ++c) { + const bits* val = v + static_cast(n0 + c) * ldv; + for (int d = 0; d < D; ++d) { + const float product = p[static_cast(c)] * from_bf16(val[d]); + acc[static_cast(d)] += product; + } + } + } + const float inv_sum = row_sum == 0.0f || row_sum != row_sum ? 1.0f : 1.0f / row_sum; + for (float& e : acc) e *= inv_sum; + return acc; +} + +bool check_attention(Device& device, const TeKernels& k) { + constexpr int heads = 4, kv_heads = 2, D = 256, ldq = heads * D + 8, ldkv = 2 * kv_heads * D + 16, ldo = heads * D; + const std::vector lengths = {37, 70, 1}; + std::vector cu = {0}; + for (const int l : lengths) cu.push_back(cu.back() + l); + const auto T = static_cast(cu.back()); + const auto q_h = synth_bf16(600, T * ldq, 1), kv_h = synth_bf16(601, T * ldkv, 1); + auto q = device_buffer(device, q_h), kv = device_buffer(device, kv_h), + o = device_buffer(device, std::vector(T * ldo)); + std::vector want(T * ldo); + for (std::size_t s = 0; s < lengths.size(); ++s) { + const int r0 = cu[s], len = lengths[s]; + TeAttentionArgs a; + a.q = q->address() + 2ull * r0 * ldq; + a.k = kv->address() + 2ull * r0 * ldkv; + a.v = a.k + 2ull * kv_heads * D; + a.o = o->address() + 2ull * r0 * ldo; + a.ldq = ldq; + a.ldk = a.ldv = ldkv; + a.ldo = ldo; + a.hsq = a.hsk = a.hsv = a.hso = D; + a.heads = heads; + a.group = heads / kv_heads; + // Two pieces of rows, as the encoder splits long sequences. + const std::uint32_t split = static_cast(len) / 2 / 16 * 16; + device.run([&](VkCommandBuffer cb) { + if (split > 0) { + a.q_first = 0; + a.rows = split; + k.record_attention(cb, a); + } + a.q_first = split; + a.rows = static_cast(len) - split; + k.record_attention(cb, a); + }); + for (int h = 0; h < heads; ++h) { + const int kvh = h / (heads / kv_heads); + for (int r = 0; r < len; ++r) { + const auto row = replay_attention_row(q_h.data() + static_cast(r0 + r) * ldq + h * D, + kv_h.data() + static_cast(r0) * ldkv + kvh * D, ldkv, + kv_h.data() + static_cast(r0) * ldkv + (kv_heads + kvh) * D, + ldkv, r + 1, 0.0625f); + std::copy(row.begin(), row.end(), want.begin() + static_cast((r0 + r) * ldo + h * D)); + } + } + } + return report("causal GQA attention", read(*o, want.size()), want, {1.0, 1.0}); +} + +} // namespace + +int main(int argc, char** argv) { + int device_index = -1; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--device" && i + 1 < argc) { + device_index = std::atoi(argv[++i]); + } else { + std::printf("usage: %s [--device N]\n", argv[0]); + return 2; + } + } + try { + Device device(device_index); + std::printf("device %u: %s\ntext encoder kernels against host replays:\n", device.info().index, + device.info().name.c_str()); + const TeKernels kernels(device); + bool ok = check_norms(device, kernels); + ok = check_conv_gates(device, kernels) && ok; + ok = check_attention(device, kernels) && ok; + ok = check_gdn(device, kernels) && ok; + std::printf("%s\n", ok ? "all passed" : "FAILED"); + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/tools/vulkan_vae_check.cpp b/vulkan/tools/vulkan_vae_check.cpp new file mode 100644 index 0000000..984d3a9 --- /dev/null +++ b/vulkan/tools/vulkan_vae_check.cpp @@ -0,0 +1,806 @@ +// Checks the Vulkan VAE decoder: each kernel against a host replay of its op order on small shapes with tails, then the +// real decoder on a final latent against the reference pipeline's PNG (PSNR, max |diff|, identical bytes), with timing. +// Errors are scaled: |got - replay| / (|replay| + floor), where a convolution's floor is the sum of its |products| (the +// size an fp32 sum can be wrong by, whatever its order) and an elementwise op's is a fraction of the replay's rms. +// Exit status: 0 everything passed, 1 a check failed, 2 bad usage. +// Usage: vulkan_vae_check [--device N] [--path coop|plain] [--kat-only | --no-kat | --bench] [--vae FILE] [--latent NPY] +// [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile] [--calibrate] +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "block.h" +#include "runtime.h" +#include "vae.h" + +namespace { + +using namespace mi::vk; + +// ---- Files ------------------------------------------------------------------------------------------------------------ + +std::vector read_file(const std::string& path) { + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) throw std::runtime_error("cannot read " + path); + std::vector data(static_cast(in.tellg())); + in.seekg(0); + in.read(data.data(), static_cast(data.size())); + return data; +} + +// A little-endian fp32 C-order .npy (versions 1-3). +std::vector read_npy_f32(const std::string& path, std::vector& shape) { + const std::vector d = read_file(path); + if (d.size() < 10 || std::memcmp(d.data(), "\x93NUMPY", 6) != 0) throw std::runtime_error(path + ": not a .npy file"); + const int major = static_cast(d[6]); + std::size_t header_len = 0, at = 0; + if (major == 1) { + header_len = static_cast(d[8]) | static_cast(static_cast(d[9])) << 8; + at = 10; + } else { + if (d.size() < 12) throw std::runtime_error(path + ": truncated"); + for (int i = 3; i >= 0; --i) header_len = header_len << 8 | static_cast(d[8 + static_cast(i)]); + at = 12; + } + if (at + header_len > d.size()) throw std::runtime_error(path + ": truncated header"); + const std::string h(d.data() + at, header_len); + if (h.find("'descr': '(shape.back()); + } + at += header_len; + if (d.size() - at != count * 4) throw std::runtime_error(path + ": data size does not match the shape"); + std::vector out(count); + std::memcpy(out.data(), d.data() + at, count * 4); + return out; +} + +std::uint32_t be32(const unsigned char* p) { + return static_cast(p[0]) << 24 | static_cast(p[1]) << 16 | + static_cast(p[2]) << 8 | p[3]; +} + +// 8-bit RGB or RGBA, non-interlaced PNG -> RGB. +std::vector read_png_rgb(const std::string& path, int& width, int& height) { + const std::vector raw = read_file(path); + const auto* d = reinterpret_cast(raw.data()); + if (raw.size() < 8 || std::memcmp(d, "\x89PNG\r\n\x1a\n", 8) != 0) throw std::runtime_error(path + ": not a PNG"); + std::size_t at = 8; + int channels = 0; + std::vector idat; + while (at + 12 <= raw.size()) { + const std::uint32_t len = be32(d + at); + const std::string type(reinterpret_cast(d + at + 4), 4); + if (at + 12 + len > raw.size()) throw std::runtime_error(path + ": truncated chunk"); + const unsigned char* body = d + at + 8; + if (type == "IHDR") { + width = static_cast(be32(body)); + height = static_cast(be32(body + 4)); + if (body[8] != 8 || (body[9] != 2 && body[9] != 6) || body[12] != 0) { + throw std::runtime_error(path + ": only 8-bit RGB/RGBA non-interlaced PNGs are read"); + } + channels = body[9] == 2 ? 3 : 4; + } else if (type == "IDAT") { + idat.insert(idat.end(), body, body + len); + } else if (type == "IEND") { + break; + } + at += 12 + len; + } + if (channels == 0) throw std::runtime_error(path + ": no IHDR"); + const std::size_t stride = static_cast(width) * channels; + std::vector filtered((stride + 1) * static_cast(height)); + uLongf out_len = static_cast(filtered.size()); + if (uncompress(filtered.data(), &out_len, idat.data(), static_cast(idat.size())) != Z_OK || out_len != filtered.size()) { + throw std::runtime_error(path + ": cannot inflate the image data"); + } + std::vector img(stride * static_cast(height)); + for (int y = 0; y < height; ++y) { + const unsigned char* src = filtered.data() + static_cast(y) * (stride + 1); + unsigned char* row = img.data() + static_cast(y) * stride; + const unsigned char* prev = y > 0 ? row - stride : nullptr; + const int filter = src[0]; + for (std::size_t x = 0; x < stride; ++x) { + const int a = x >= static_cast(channels) ? row[x - channels] : 0; + const int b = prev ? prev[x] : 0; + const int c = prev && x >= static_cast(channels) ? prev[x - channels] : 0; + int pred = 0; + switch (filter) { + case 0: pred = 0; break; + case 1: pred = a; break; + case 2: pred = b; break; + case 3: pred = (a + b) / 2; break; + case 4: { + const int p = a + b - c, pa = std::abs(p - a), pb = std::abs(p - b), pc = std::abs(p - c); + pred = pa <= pb && pa <= pc ? a : pb <= pc ? b : c; + break; + } + default: throw std::runtime_error(path + ": bad filter type"); + } + row[x] = static_cast(src[x + 1] + pred); + } + } + if (channels == 3) return img; + std::vector rgb(static_cast(width) * height * 3); + for (std::size_t i = 0; i < static_cast(width) * height; ++i) { + for (int ch = 0; ch < 3; ++ch) rgb[i * 3 + static_cast(ch)] = img[i * 4 + static_cast(ch)]; + } + return rgb; +} + +void put32(std::vector& v, std::uint32_t x) { + for (int s = 24; s >= 0; s -= 8) v.push_back(static_cast(x >> s)); +} +void chunk(std::vector& out, const char* type, const std::vector& body) { + put32(out, static_cast(body.size())); + const std::size_t start = out.size(); + out.insert(out.end(), type, type + 4); + out.insert(out.end(), body.begin(), body.end()); + put32(out, static_cast(crc32(0, out.data() + start, static_cast(out.size() - start)))); +} + +void write_png_rgb(const std::string& path, const std::vector& rgb, int width, int height) { + const std::size_t stride = static_cast(width) * 3; + std::vector raw((stride + 1) * static_cast(height)); + for (int y = 0; y < height; ++y) { + raw[static_cast(y) * (stride + 1)] = 0; + std::memcpy(raw.data() + static_cast(y) * (stride + 1) + 1, rgb.data() + static_cast(y) * stride, stride); + } + uLongf len = compressBound(static_cast(raw.size())); + std::vector z(len); + if (compress2(z.data(), &len, raw.data(), static_cast(raw.size()), 6) != Z_OK) throw std::runtime_error("deflate failed"); + z.resize(len); + std::vector png = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'}; + std::vector ihdr; + put32(ihdr, static_cast(width)); + put32(ihdr, static_cast(height)); + ihdr.insert(ihdr.end(), {8, 2, 0, 0, 0}); + chunk(png, "IHDR", ihdr); + chunk(png, "IDAT", z); + chunk(png, "IEND", {}); + std::ofstream out(path, std::ios::binary); + out.write(reinterpret_cast(png.data()), static_cast(png.size())); + if (!out) throw std::runtime_error("cannot write " + path); +} + +// ---- Host replays ----------------------------------------------------------------------------------------------------- + +float host_tf32(float v) { + std::uint32_t u = std::bit_cast(v); + if ((u & 0x7f800000u) == 0x7f800000u) return v; + u = (u + 0x1000u) & 0xffffe000u; + return std::bit_cast(u); +} +float host_bf16(float v) { + std::uint32_t u = std::bit_cast(v); + u = (u + 0x7fffu + ((u >> 16) & 1u)) & 0xffff0000u; + return std::bit_cast(u); +} + +std::vector random_f32(std::uint64_t seed, std::size_t n, float lo, float hi, bool bf16_values = false) { + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(lo, hi); + std::vector v(n); + for (float& x : v) x = bf16_values ? host_bf16(dist(rng)) : dist(rng); + return v; +} + +std::unique_ptr upload(Device& device, const std::vector& v) { + auto b = std::make_unique(device, std::max(v.size(), 4) * 4); + b->upload(v.data(), v.size() * 4); + return b; +} +std::vector download(Buffer& b, std::size_t n) { + std::vector v(n); + b.download(v.data(), n * 4); + return v; +} + +struct Limit { + double max_err; + const char* why; +}; + +// err_i = |got - ref| / (|ref| + floor_i). +bool report(const char* name, const std::vector& got, const std::vector& ref, const std::vector& floor, + Limit limit) { + double worst = 0.0; + std::size_t same = 0, nans = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + if (std::isnan(got[i])) { + ++nans; + continue; + } + same += got[i] == static_cast(ref[i]); + worst = std::max(worst, std::abs(static_cast(got[i]) - ref[i]) / (std::abs(ref[i]) + floor[i])); + } + const bool ok = nans == 0 && worst <= limit.max_err; + std::printf(" %-40s %s %zu values, %.2f %% same as the fp32-rounded replay, worst scaled error %.3g (limit %.3g: %s), %zu NaN\n", + name, ok ? "ok " : "FAILED", got.size(), 100.0 * static_cast(same) / static_cast(got.size()), + worst, limit.max_err, limit.why, nans); + return ok; +} + +struct ConvCase { + const char* name; + std::uint32_t taps; // 0 GEMM, 1, 9 + bool upsample; + std::uint32_t h, w, c; // conv output grid and input channels + std::uint32_t m, k, lda; // GEMM + std::uint32_t n, ldb, ldo; + bool bias, residual, b_bf16; + VaePrecision precision; + std::uint32_t band; // rows per dispatch (0: one) +}; + +// Operand as the kernel sees it: TF32-rounded or exact. +double operand(float v, VaePrecision p) { return p == VaePrecision::kTf32 ? host_tf32(v) : v; } + +bool check_conv(Device& device, const VaeKernels& k, const ConvCase& cc) { + const bool conv = cc.taps != 0; + const std::uint32_t m = conv ? cc.h * cc.w : cc.m, kk = conv ? cc.taps * cc.c : cc.k, lda = conv ? 0 : cc.lda; + const std::uint32_t in_h = cc.upsample ? cc.h / 2 : cc.h, in_w = cc.upsample ? cc.w / 2 : cc.w; + const std::size_t a_count = conv ? static_cast(in_h) * in_w * cc.c : static_cast(m) * lda; + const auto a = random_f32(1, a_count, -2.0f, 2.0f); + const auto b = random_f32(2, static_cast(cc.n) * cc.ldb, -0.5f, 0.5f, cc.b_bf16); + const auto bias = random_f32(3, cc.n, -1.0f, 1.0f); + const auto res = random_f32(4, static_cast(m) * cc.ldo, -4.0f, 4.0f); + auto da = upload(device, a), db = upload(device, b), dbias = upload(device, bias); + auto dout = upload(device, res); // poisoned with the residual; written in place when residual is on + const double nan = std::nan(""); + std::vector ref(static_cast(m) * cc.ldo, nan), floor(ref.size(), 0.0); + for (std::uint32_t r = 0; r < m; ++r) { + for (std::uint32_t n = 0; n < cc.n; ++n) { + double sum = 0.0, mag = 0.0; + for (std::uint32_t q = 0; q < kk; ++q) { + float av = 0.0f; + if (!conv) { + av = a[static_cast(r) * lda + q]; + } else { + const std::uint32_t tap = q / cc.c, ch = q % cc.c, oy = r / cc.w, ox = r % cc.w; + int iy = static_cast(oy), ix = static_cast(ox); + if (cc.taps == 9) { + iy += static_cast(tap / 3) - 1; + ix += static_cast(tap % 3) - 1; + } + if (iy < 0 || ix < 0 || iy >= static_cast(cc.h) || ix >= static_cast(cc.w)) continue; + if (cc.upsample) { + iy /= 2; + ix /= 2; + } + av = a[(static_cast(iy) * in_w + static_cast(ix)) * cc.c + ch]; + } + const double p = operand(av, cc.precision) * operand(b[static_cast(n) * cc.ldb + q], cc.precision); + sum += p; + mag += std::abs(p); + } + float y = static_cast(sum); + if (cc.bias) { + y = y + bias[n]; + mag += std::abs(bias[n]); + } + if (cc.residual) { + y = y + res[static_cast(r) * cc.ldo + n]; + mag += std::abs(res[static_cast(r) * cc.ldo + n]); + } + ref[static_cast(r) * cc.ldo + n] = y; + floor[static_cast(r) * cc.ldo + n] = mag; + } + } + ConvArgs args; + args.a = da->address(); + args.b = db->address(); + args.out = dout->address(); + args.bias = cc.bias ? dbias->address() : 0; + args.residual = cc.residual ? dout->address() : 0; + args.m = m; + args.n = cc.n; + args.k = kk; + args.lda = lda; + args.ldb = cc.ldb; + args.ldo = cc.ldo; + args.taps = cc.taps; + args.upsample = cc.upsample; + args.h = cc.h; + args.w = cc.w; + args.c = cc.c; + args.b_bf16 = cc.b_bf16; + device.run([&](VkCommandBuffer cb) { + const std::uint32_t band = cc.band == 0 ? m : cc.band; + for (std::uint32_t m0 = 0; m0 < m; m0 += band) { + ConvArgs part = args; + part.m0 = m0; + part.m = std::min(m, m0 + band); + k.record_conv(cb, part, cc.precision); + } + }); + const std::vector got_all = download(*dout, ref.size()); + // Columns past n keep the poison (the residual) untouched: checked bit for bit. + std::vector got; + std::vector want, fl; + std::size_t clobbered = 0; + for (std::size_t i = 0; i < ref.size(); ++i) { + if (std::isnan(ref[i])) { + clobbered += got_all[i] != res[i]; + continue; + } + got.push_back(got_all[i]); + want.push_back(ref[i]); + fl.push_back(floor[i]); + } + // Products are exact (TF32, or a bf16-valued B with A in 3 pieces) or carry 2^-17 relative error per operand (2 pieces); + // the rest is fp32 accumulation, at most about 2^-24 per add on |products|. + const bool exact_products = cc.precision == VaePrecision::kTf32 || (cc.precision == VaePrecision::kFp32 && cc.b_bf16) || + k.path() == VaePath::kPlain; + const Limit limit = exact_products || cc.n < 64 ? Limit{2e-6, "fp32 accumulation"} + : Limit{4e-5, "16-bit operand pieces + fp32 accumulation"}; + bool ok = report(cc.name, got, want, fl, limit); + if (clobbered != 0) { + std::printf(" %-40s FAILED %zu values outside the output columns were written\n", cc.name, clobbered); + ok = false; + } + return ok; +} + +bool check_convs(Device& device, const VaeKernels& k) { + const ConvCase cases[] = { + // name, taps, up, h, w, c, m, k, lda, n, ldb, ldo, bias, res, b_bf16, precision, band + {"conv3x3 tf32 bias+residual, banded", 9, false, 13, 21, 32, 0, 0, 0, 96, 288, 100, true, true, true, VaePrecision::kTf32, 128}, + {"conv3x3 tf32 upsample, fp32 B", 9, true, 14, 18, 48, 0, 0, 0, 136, 432, 140, true, false, false, VaePrecision::kTf32, 0}, + {"conv1x1 tf32", 1, false, 11, 13, 40, 0, 0, 0, 72, 40, 76, true, false, true, VaePrecision::kTf32, 0}, + {"conv3x3 tf32 narrow (n = 3)", 9, false, 9, 30, 24, 0, 0, 0, 3, 216, 4, true, false, true, VaePrecision::kTf32, 0}, + {"conv3x3 tf32 narrow, c = 96 (n = 3)", 9, false, 5, 7, 96, 0, 0, 0, 3, 864, 3, true, true, true, VaePrecision::kTf32, 0}, + {"gemm fp32 narrow, k > 512 (n = 4)", 0, false, 0, 0, 0, 300, 1100, 1104, 4, 1100, 4, true, false, false, VaePrecision::kFp32, 0}, + {"conv1x1 tf32 narrow (n = 32)", 1, false, 10, 15, 32, 0, 0, 0, 32, 32, 32, true, false, true, VaePrecision::kTf32, 0}, + {"gemm fp32 bf16 B bias+residual", 0, false, 0, 0, 0, 150, 100, 104, 136, 100, 140, true, true, true, VaePrecision::kFp32, 0}, + {"gemm fp32, fp32 B", 0, false, 0, 0, 0, 77, 64, 64, 384, 68, 388, false, false, false, VaePrecision::kFp32, 0}, + {"gemm split16", 0, false, 0, 0, 0, 200, 68, 72, 260, 68, 264, false, false, false, VaePrecision::kSplit16, 0}, + }; + bool ok = true; + for (const ConvCase& c : cases) ok = check_conv(device, k, c) && ok; + return ok; +} + +bool check_group_norm(Device& device, const VaeKernels& k, std::uint32_t pixels, std::uint32_t c, bool silu, bool in_place, + float offset) { + std::vector x = random_f32(10 + c, static_cast(pixels) * c, -3.0f, 3.0f); + for (std::size_t i = 0; i < x.size(); ++i) x[i] += offset * static_cast(1 + (i % c) / (c / 32)); + const auto gamma = random_f32(11, c, 0.5f, 1.5f), beta = random_f32(12, c, -0.5f, 0.5f); + auto dx = upload(device, x), dg = upload(device, gamma), db = upload(device, beta); + auto dy = upload(device, std::vector(x.size(), 0.0f)); + Buffer part(device, VaeKernels::gn_part_bytes(pixels, c)), stats(device, 256); + const VkDeviceAddress y = in_place ? dx->address() : dy->address(); + device.run([&](VkCommandBuffer cb) { + k.record_gn_stats(cb, dx->address(), pixels, c, part.address(), stats.address(), 1e-6f); + k.record_gn_apply(cb, dx->address(), y, pixels, c, stats.address(), dg->address(), db->address(), silu); + }); + const std::uint32_t cpg = c / 32; + std::vector ref(x.size()), floor(x.size()); + for (std::uint32_t g = 0; g < 32; ++g) { + double s = 0.0; + for (std::uint32_t p = 0; p < pixels; ++p) { + for (std::uint32_t j = 0; j < cpg; ++j) s += x[static_cast(p) * c + g * cpg + j]; + } + const double n = static_cast(pixels) * cpg, mean = s / n; + double q = 0.0; + for (std::uint32_t p = 0; p < pixels; ++p) { + for (std::uint32_t j = 0; j < cpg; ++j) { + const double d = x[static_cast(p) * c + g * cpg + j] - mean; + q += d * d; + } + } + const auto mean_f = static_cast(mean), rstd = static_cast(1.0 / std::sqrt(q / n + 1e-6)); + for (std::uint32_t j = 0; j < cpg; ++j) { + const std::uint32_t ch = g * cpg + j; + const float a = rstd * gamma[ch], b = std::fma(-a, mean_f, beta[ch]); + for (std::uint32_t p = 0; p < pixels; ++p) { + const std::size_t i = static_cast(p) * c + ch; + float v = std::fma(a, x[i], b); + if (silu) v = v / (1.0f + std::exp(-v)); + ref[i] = v; + // The statistics carry fp32 rounding: up to a few ulps of mean and rstd, i.e. of |a x| and |a mean|. + floor[i] = std::abs(a * x[i]) + std::abs(a * mean_f); + } + } + } + const auto got = download(in_place ? *dx : *dy, x.size()); + const std::string name = "group_norm c=" + std::to_string(c) + (silu ? " silu" : "") + (in_place ? " in place" : "") + + " offset " + std::to_string(static_cast(offset)); + return report(name.c_str(), got, ref, floor, {2e-6, "fp32 statistics, exp"}); +} + +bool check_softmax(Device& device, const VaeKernels& k) { + constexpr std::uint32_t rows = 37, l = 300, ld = 304; + const float scale = 0.05103104f; + const auto s = random_f32(20, static_cast(rows) * ld, -60.0f, 60.0f); + auto ds = upload(device, s); + device.run([&](VkCommandBuffer cb) { k.record_softmax(cb, ds->address(), rows, l, ld, scale); }); + const auto all = download(*ds, s.size()); + std::vector got; + std::vector ref, floor; + std::size_t clobbered = 0; + for (std::uint32_t r = 0; r < rows; ++r) { + const float* row = s.data() + static_cast(r) * ld; + float mx = row[0]; + for (std::uint32_t j = 1; j < l; ++j) mx = std::max(mx, row[j]); + const float shift = mx * scale; + double sum = 0.0; + for (std::uint32_t j = 0; j < l; ++j) sum += std::exp(static_cast(row[j] * scale - shift)); + for (std::uint32_t j = 0; j < ld; ++j) { + const std::size_t i = static_cast(r) * ld + j; + if (j >= l) { + clobbered += all[i] != s[i]; + continue; + } + got.push_back(all[i]); + ref.push_back(std::exp(static_cast(row[j] * scale - shift)) / sum); + floor.push_back(1e-30); + } + } + // exp is accurate to a few ulps of its result in Vulkan (3 + 2 |x| ulps), here |x| up to ~6. + bool ok = report("softmax", got, ref, floor, {3e-6, "Vulkan exp precision"}); + if (clobbered != 0) { + std::printf(" softmax: %zu values past l were written\n", clobbered); + ok = false; + } + return ok; +} + +// Bit-exact ops: the transpose, the latent unpack and RGB8. +bool check_exact(Device& device, const VaeKernels& k) { + bool ok = true; + { + constexpr std::uint32_t r = 70, c = 45, ld_in = 48, ld_out = 72; + const auto in = random_f32(30, static_cast(r) * ld_in, -1.0f, 1.0f); + auto din = upload(device, in), dout = upload(device, std::vector(static_cast(c) * ld_out, -7.0f)); + device.run([&](VkCommandBuffer cb) { k.record_transpose(cb, din->address(), dout->address(), r, c, ld_in, ld_out); }); + const auto got = download(*dout, static_cast(c) * ld_out); + std::size_t bad = 0; + for (std::uint32_t i = 0; i < c; ++i) { + for (std::uint32_t j = 0; j < ld_out; ++j) { + const float want = j < r ? in[static_cast(j) * ld_in + i] : -7.0f; + bad += std::bit_cast(got[static_cast(i) * ld_out + j]) != std::bit_cast(want); + } + } + std::printf(" %-40s %s %zu of %zu values differ\n", "transpose", bad == 0 ? "ok " : "FAILED", bad, got.size()); + ok = ok && bad == 0; + } + { + constexpr std::uint32_t h = 3, w = 5; + const auto z = random_f32(31, 128ull * h * w, -4.0f, 4.0f); + const auto var = random_f32(32, 128, 0.1f, 9.0f), mean = random_f32(33, 128, -1.0f, 1.0f); + std::vector scale(128); + for (int i = 0; i < 128; ++i) scale[static_cast(i)] = std::sqrt(var[static_cast(i)] + 1e-4f); + auto dz = upload(device, z), ds = upload(device, scale), dm = upload(device, mean); + Buffer y(device, 4ull * h * w * 32 * 4); + device.run([&](VkCommandBuffer cb) { k.record_latent_unpack(cb, dz->address(), ds->address(), dm->address(), y.address(), h, w); }); + const auto got = download(y, 4ull * h * w * 32); + std::size_t bad = 0; + for (std::uint32_t oy = 0; oy < 2 * h; ++oy) { + for (std::uint32_t ox = 0; ox < 2 * w; ++ox) { + for (std::uint32_t c = 0; c < 32; ++c) { + const std::uint32_t ch = 4 * c + 2 * (oy & 1) + (ox & 1); + const float v = z[(static_cast(ch) * h + oy / 2) * w + ox / 2]; + const float scaled = v * scale[ch]; + const float want = scaled + mean[ch]; + bad += std::bit_cast(got[(static_cast(oy) * 2 * w + ox) * 32 + c]) != std::bit_cast(want); + } + } + } + std::printf(" %-40s %s %zu of %zu values differ\n", "latent unpack", bad == 0 ? "ok " : "FAILED", bad, got.size()); + ok = ok && bad == 0; + } + { + std::vector x = random_f32(34, 1001, -1.3f, 1.3f); + const float special[] = {std::nanf(""), -1.0f, 1.0f, 0.0f, -0.0f, 1.0f / 255.0f, 0.5f / 255.0f * 2.0f - 1.0f, + 1.5f / 255.0f * 2.0f - 1.0f, 2.5f / 255.0f * 2.0f - 1.0f, INFINITY, -INFINITY}; + std::copy(std::begin(special), std::end(special), x.begin()); + auto dx = upload(device, x); + Buffer rgb(device, 1004); + device.run([&](VkCommandBuffer cb) { k.record_rgb8(cb, dx->address(), rgb.address(), static_cast(x.size())); }); + std::vector got(1004); + rgb.download(got.data(), got.size()); + std::size_t bad = 0; + for (std::size_t i = 0; i < got.size(); ++i) { + std::uint8_t want = 0; + if (i < x.size() && !std::isnan(x[i])) { + float v = std::min(std::max(x[i], -1.0f), 1.0f); + const float half = v * 0.5f; + v = std::min(std::max(half + 0.5f, 0.0f), 1.0f); + want = static_cast(std::nearbyint(v * 255.0f)); + } + bad += got[i] != want; + } + std::printf(" %-40s %s %zu of %zu bytes differ\n", "rgb8", bad == 0 ? "ok " : "FAILED", bad, got.size()); + ok = ok && bad == 0; + } + return ok; +} + +bool run_kat(Device& device, const VaeKernels& k) { + std::printf("known-answer checks against host replays (path %s):\n", vae_path_name(k.path())); + bool ok = check_convs(device, k); + ok = check_group_norm(device, k, 1000, 96, true, false, 0.0f) && ok; + ok = check_group_norm(device, k, 777, 384, false, true, 20.0f) && ok; + ok = check_group_norm(device, k, 2050, 192, true, true, 5.0f) && ok; + ok = check_softmax(device, k) && ok; + ok = check_exact(device, k) && ok; + return ok; +} + +// ---- Kernel timing ---------------------------------------------------------------------------------------------------- + +double now_s() { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); } + +// The decoder's heaviest shapes at 1024x1024, one submit each, best of 3. +void run_bench(Device& device, const VaeKernels& k) { + struct Shape { + const char* name; + std::uint32_t h, w, c, n, taps; + bool upsample; + }; + const Shape shapes[] = {{"up3 conv 96->96 3x3 @1024", 1024, 1024, 96, 96, 9, false}, + {"up2 upsampler 192 3x3 @1024", 1024, 1024, 192, 192, 9, true}, + {"up2 conv 384->192 3x3 @512", 512, 512, 384, 192, 9, false}, + {"up1 upsampler 384 3x3 @512", 512, 512, 384, 384, 9, true}, + {"mid conv 384 3x3 @128", 128, 128, 384, 384, 9, false}, + {"conv_out 96->3 3x3 @1024", 1024, 1024, 96, 3, 9, false}}; + std::printf("\nkernel timing (path %s), best of 3 single submits:\n", vae_path_name(k.path())); + for (const Shape& s : shapes) { + const std::uint64_t in_px = static_cast(s.h) * s.w / (s.upsample ? 4 : 1); + Buffer x(device, in_px * s.c * 4), wt(device, 9ull * s.c * s.n * 4), y(device, static_cast(s.h) * s.w * s.n * 4); + device.run([&](VkCommandBuffer cb) { + vkCmdFillBuffer(cb, x.handle(), 0, VK_WHOLE_SIZE, 0x3c000000u); + vkCmdFillBuffer(cb, wt.handle(), 0, VK_WHOLE_SIZE, 0x3c000000u); + }); + ConvArgs a; + a.a = x.address(); + a.b = wt.address(); + a.out = y.address(); + a.m = s.h * s.w; + a.n = s.n; + a.k = s.taps * s.c; + a.ldb = a.k; + a.ldo = s.n; + a.taps = s.taps; + a.upsample = s.upsample; + a.h = s.h; + a.w = s.w; + a.c = s.c; + a.b_bf16 = true; + double best = 1e30; + for (int it = 0; it < 3; ++it) { + const double t0 = now_s(); + device.run([&](VkCommandBuffer cb) { + for (std::uint32_t m0 = 0; m0 < a.m; m0 += 65535u * 64u) { + ConvArgs band = a; + band.m0 = m0; + band.m = std::min(a.m, m0 + 65535u * 64u); + k.record_conv(cb, band, VaePrecision::kTf32); + } + }); + best = std::min(best, now_s() - t0); + } + const double macs = static_cast(a.m) * a.n * a.k; + std::printf(" %-30s %9.2f ms %7.2f TFLOPS (useful)\n", s.name, best * 1e3, 2.0 * macs / best / 1e12); + } +} + +// The DiT block's bf16 GEMM at 8222 x 3072 x 3072 (README: 33.1 TFLOPS on the coop path, 1.62 on the plain path, on an +// idle RTX 4090 Laptop), timed next to a decode: when another process shares the GPU, its rate shows how much. +double calibration_tflops(Device& device) { + const BlockKernels block(device); + constexpr std::uint32_t m = 8222, n = 3072, k = 3072; + Buffer a(device, 2ull * m * k), b(device, 2ull * n * k), c(device, 2ull * m * n); + device.run([&](VkCommandBuffer cb) { + vkCmdFillBuffer(cb, a.handle(), 0, VK_WHOLE_SIZE, 0x3c003c00u); + vkCmdFillBuffer(cb, b.handle(), 0, VK_WHOLE_SIZE, 0x3c003c00u); + }); + double best = 1e30; + for (int it = 0; it < 3; ++it) { + const double t0 = now_s(); + device.run([&](VkCommandBuffer cb) { block.record_gemm(cb, c.address(), n, a.address(), k, b.address(), k, m, n, k); }); + best = std::min(best, now_s() - t0); + } + return 2.0 * m * n * k / best / 1e12; +} + +// ---- Real decode ------------------------------------------------------------------------------------------------------ + +struct Options { + int device = -1; + std::string path; + bool kat = true, decode = true, bench = false, profile = false, calibrate = false; + std::string vae = "/workspace/Models/MOSS-Image/base/vae/diffusion_pytorch_model.safetensors"; + std::string latent = "/workspace/Projects/MossInfer/tests/golden/g_fridge_1024/final_latent.npy"; + std::string golden = "/workspace/Projects/MossInfer/tests/golden/g_fridge_1024/image.png"; + std::string out; + std::uint32_t crop_h = 0, crop_w = 0; + int iters = 2; + std::uint64_t max_macs = 0; +}; + +bool run_decode(Device& device, const Options& o, VaePath path) { + std::vector shape; + std::vector latent = read_npy_f32(o.latent, shape); + if (shape.size() == 4 && shape[0] == 1) shape.erase(shape.begin()); + if (shape.size() != 3 || shape[0] != 128) throw std::runtime_error(o.latent + ": expected [1, 128, h, w] or [128, h, w]"); + auto h = static_cast(shape[1]), w = static_cast(shape[2]); + const bool cropped = o.crop_h != 0; + if (cropped) { + // The top-left crop_h x crop_w latents: no reference image exists for a crop; the decode runs and is written only. + if (o.crop_h > h || o.crop_w > w) throw std::runtime_error("--crop exceeds the latent grid"); + std::vector c(128ull * o.crop_h * o.crop_w); + for (std::uint32_t ch = 0; ch < 128; ++ch) { + for (std::uint32_t y = 0; y < o.crop_h; ++y) { + std::memcpy(&c[(static_cast(ch) * o.crop_h + y) * o.crop_w], + &latent[(static_cast(ch) * h + y) * w], o.crop_w * 4); + } + } + latent.swap(c); + h = o.crop_h; + w = o.crop_w; + } + VaeDecoder decoder(device, o.vae, path); + if (o.max_macs != 0) decoder.set_max_macs_per_submit(o.max_macs); + decoder.set_profile(o.profile); + std::printf("\ndecode %ux%u (latents %ux%u, decoder c = %d, path %s): weights %.1f MB, decode allocates %.3f GiB\n", 16 * w, + 16 * h, h, w, decoder.base_channels(), vae_path_name(path), static_cast(decoder.weight_bytes()) / 1e6, + static_cast(decoder.workspace_bytes(h, w)) / (1024.0 * 1024.0 * 1024.0)); + if (o.calibrate) { + std::printf(" calibration before: the block's GEMM 8222x3072x3072 at %.2f TFLOPS (%s path)\n", calibration_tflops(device), + device.info().caps.coop_bf16 ? "coop" : "plain"); + } + std::vector rgb; + std::vector image; + std::vector times; + VaeDecodeStats st; + for (int i = 0; i < std::max(1, o.iters); ++i) { + st = decoder.decode(latent.data(), h, w, rgb, i == 0 ? &image : nullptr); + times.push_back(st.seconds); + std::printf(" run %d: %.3f s (%d submits)%s\n", i, st.seconds, st.submits, i == 0 ? ", pipelines created" : ""); + } + if (o.calibrate) std::printf(" calibration after: %.2f TFLOPS\n", calibration_tflops(device)); + if (o.profile) { + double total = 0.0; + for (const auto& [name, s] : st.profile) total += s; + std::printf(" profile of the last run (one submit per op, %.3f s in submits):\n", total); + for (const auto& [name, s] : st.profile) std::printf(" %-48s %9.2f ms %5.1f %%\n", name.c_str(), s * 1e3, 100.0 * s / total); + } + std::size_t nans = 0; + for (const float v : image) nans += std::isnan(v); + std::printf(" decoder output: %zu NaN of %zu\n", nans, image.size()); + if (!o.out.empty()) { + write_png_rgb(o.out, rgb, static_cast(16 * w), static_cast(16 * h)); + std::printf(" wrote %s\n", o.out.c_str()); + } + if (cropped || o.golden.empty()) return nans == 0; + int gw = 0, gh = 0; + const std::vector golden = read_png_rgb(o.golden, gw, gh); + if (static_cast(gw) != 16 * w || static_cast(gh) != 16 * h) { + throw std::runtime_error(o.golden + " is " + std::to_string(gw) + "x" + std::to_string(gh) + ", not the decode's size"); + } + double se = 0.0, sum_abs = 0.0; + int max_abs = 0; + std::size_t same = 0; + std::vector hist(4, 0); + for (std::size_t i = 0; i < rgb.size(); ++i) { + const int d = std::abs(static_cast(rgb[i]) - static_cast(golden[i])); + se += static_cast(d) * d; + sum_abs += d; + max_abs = std::max(max_abs, d); + same += d == 0; + ++hist[static_cast(std::min(d, 3))]; + } + const double n = static_cast(rgb.size()), mse = se / n; + const double psnr = mse == 0.0 ? INFINITY : 10.0 * std::log10(255.0 * 255.0 / mse); + std::printf(" against %s:\n PSNR %.2f dB, max |diff| %d, mean |diff| %.5f, identical bytes %.4f %%\n", o.golden.c_str(), + psnr, max_abs, sum_abs / n, 100.0 * static_cast(same) / n); + std::printf(" |diff| 0: %zu, 1: %zu, 2: %zu, >= 3: %zu bytes\n", hist[0], hist[1], hist[2], hist[3]); + return nans == 0; +} + +} // namespace + +int main(int argc, char** argv) { + Options o; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + const auto next = [&]() -> std::string { + if (i + 1 >= argc) throw std::invalid_argument(arg + " needs a value"); + return argv[++i]; + }; + try { + if (arg == "--device") { + o.device = std::stoi(next()); + } else if (arg == "--path") { + o.path = next(); + } else if (arg == "--kat-only") { + o.decode = false; + } else if (arg == "--bench") { + o.bench = true; + o.decode = false; + } else if (arg == "--calibrate") { + o.calibrate = true; + } else if (arg == "--profile") { + o.profile = true; + } else if (arg == "--no-kat") { + o.kat = false; + } else if (arg == "--vae") { + o.vae = next(); + } else if (arg == "--latent") { + o.latent = next(); + } else if (arg == "--golden") { + o.golden = next(); + } else if (arg == "--out") { + o.out = next(); + } else if (arg == "--crop") { + const std::string v = next(); + const std::size_t comma = v.find(','); + o.crop_h = static_cast(std::stoul(v.substr(0, comma))); + o.crop_w = static_cast(std::stoul(v.substr(comma + 1))); + } else if (arg == "--iters") { + o.iters = std::stoi(next()); + } else if (arg == "--max-macs") { + o.max_macs = std::stoull(next()); + } else { + throw std::invalid_argument("unknown argument " + arg); + } + } catch (const std::exception& e) { + std::printf("%s\nusage: %s [--device N] [--path coop|plain] [--kat-only | --no-kat | --bench] [--vae FILE] [--latent NPY]\n" + " [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile] [--calibrate]\n", + e.what(), argv[0]); + return 2; + } + } + if (!o.path.empty() && o.path != "coop" && o.path != "plain") { + std::printf("--path takes coop or plain\n"); + return 2; + } + try { + Device device(o.device); + const VaePath path = o.path.empty() ? (device.info().caps.coop_bf16 ? VaePath::kCoop : VaePath::kPlain) + : (o.path == "coop" ? VaePath::kCoop : VaePath::kPlain); + std::printf("device %u: %s, path %s%s\n", device.info().index, device.info().name.c_str(), vae_path_name(path), + device.info().caps.coop_bf16 ? "" : (" (no coop: " + device.info().caps.coop_bf16_reason + ")").c_str()); + bool ok = true; + if (o.kat) { + const VaeKernels kernels(device, path); + ok = run_kat(device, kernels); + if (!ok) std::printf("known-answer check FAILED\n"); + } + if (o.bench) { + const VaeKernels kernels(device, path); + run_bench(device, kernels); + } + if (o.decode) ok = run_decode(device, o, path) && ok; + return ok ? 0 : 1; + } catch (const std::exception& e) { + std::printf("error: %s\n", e.what()); + return 1; + } +} diff --git a/vulkan/vae.cpp b/vulkan/vae.cpp new file mode 100644 index 0000000..08cdd3d --- /dev/null +++ b/vulkan/vae.cpp @@ -0,0 +1,769 @@ +#include "vae.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "safetensors.h" + +namespace mi::vk { +namespace { + +const std::uint32_t kConvCoopSpirv[] = { +#include "vae_conv_coop.spv.inc" +}; +const std::uint32_t kConvPlainSpirv[] = { +#include "vae_conv_plain.spv.inc" +}; +const std::uint32_t kConvNarrowSpirv[] = { +#include "vae_conv_narrow.spv.inc" +}; +const std::uint32_t kGnStatsSpirv[] = { +#include "vae_gn_stats.spv.inc" +}; +const std::uint32_t kGnFinalizeSpirv[] = { +#include "vae_gn_finalize.spv.inc" +}; +const std::uint32_t kGnApplySpirv[] = { +#include "vae_gn_apply.spv.inc" +}; +const std::uint32_t kSoftmaxSpirv[] = { +#include "vae_softmax.spv.inc" +}; +const std::uint32_t kTransposeSpirv[] = { +#include "vae_transpose.spv.inc" +}; +const std::uint32_t kUnpackSpirv[] = { +#include "vae_latent_unpack.spv.inc" +}; +const std::uint32_t kRgb8Spirv[] = { +#include "vae_rgb8.spv.inc" +}; + +struct ConvPush { + std::uint64_t a, b, out, bias, residual; + std::uint32_t m, n, k, lda, ldb, ldo, h, w, c, m0; +}; +struct GnStatsPush { + std::uint64_t x, part; + std::uint32_t pixels, c, ppc, chunks; +}; +struct GnFinalizePush { + std::uint64_t part, stats; + std::uint32_t chunks; + float eps; +}; +struct GnApplyPush { + std::uint64_t x, y, stats, gamma, beta; + std::uint32_t quads, c, silu; +}; +struct SoftmaxPush { + std::uint64_t s; + std::uint32_t rows, l, ld; + float scale; +}; +struct TransposePush { + std::uint64_t in, out; + std::uint32_t r, c, ld_in, ld_out; +}; +struct UnpackPush { + std::uint64_t z, scale, mean, y; + std::uint32_t h, w; +}; +struct Rgb8Push { + std::uint64_t x, rgb; + std::uint32_t n; +}; + +constexpr std::uint32_t kCoopSubgroup = 32; +constexpr std::uint32_t kGroups = 32; +constexpr float kResnetEps = 1e-6f; +constexpr float kBatchNormEps = 1e-4f; // vae.config.batch_norm_eps +constexpr std::uint32_t kGnChunkElems = 32768; + +void require(bool condition, const std::string& message) { + if (!condition) throw std::invalid_argument("vulkan vae: " + message); +} + +std::uint32_t div_up(std::uint64_t a, std::uint64_t b) { return static_cast((a + b - 1) / b); } + +// A 1-D count of workgroups as x + y * x_count, within the guaranteed 65535 per dimension. +void dispatch_1d(const Pipeline& p, VkCommandBuffer cb, const void* push, std::uint32_t groups) { + if (groups == 0) return; + const std::uint32_t x = std::min(groups, 65535); + p.dispatch(cb, push, x, div_up(groups, x)); +} + +std::uint32_t gn_ppc(std::uint32_t c) { return std::max(1, kGnChunkElems / c); } + +} // namespace + +const char* vae_path_name(VaePath path) { return path == VaePath::kCoop ? "coop" : "plain"; } + +VaeKernels::VaeKernels(Device& device) + : VaeKernels(device, device.info().caps.coop_bf16 ? VaePath::kCoop : VaePath::kPlain) {} + +VaeKernels::VaeKernels(Device& device, VaePath path) : device_(device), path_(path) { + if (path == VaePath::kCoop && !device.info().caps.coop_bf16) { + throw std::runtime_error(device.info().name + " cannot run the cooperative-matrix VAE: " + + device.info().caps.coop_bf16_reason); + } + gn_stats_ = std::make_unique(device, kGnStatsSpirv, sizeof(kGnStatsSpirv), sizeof(GnStatsPush)); + gn_finalize_ = std::make_unique(device, kGnFinalizeSpirv, sizeof(kGnFinalizeSpirv), sizeof(GnFinalizePush)); + gn_apply_ = std::make_unique(device, kGnApplySpirv, sizeof(kGnApplySpirv), sizeof(GnApplyPush)); + softmax_ = std::make_unique(device, kSoftmaxSpirv, sizeof(kSoftmaxSpirv), sizeof(SoftmaxPush)); + transpose_ = std::make_unique(device, kTransposeSpirv, sizeof(kTransposeSpirv), sizeof(TransposePush)); + unpack_ = std::make_unique(device, kUnpackSpirv, sizeof(kUnpackSpirv), sizeof(UnpackPush)); + rgb8_ = std::make_unique(device, kRgb8Spirv, sizeof(kRgb8Spirv), sizeof(Rgb8Push)); +} + +VaeKernels::~VaeKernels() = default; + +// Coop: outputs at least 64 wide run on bf16 fragments; narrower ones (post_quant_conv's 32, conv_out's 3) on the plain +// kernels, whose operands then get the same rounding. Outputs at most 4 wide take the narrow kernel on either path. +const Pipeline& VaeKernels::conv_pipeline(const ConvArgs& a, VaePrecision precision, std::uint32_t& tile_rows, + std::uint32_t& tile_cols) const { + const std::uint32_t taps = a.taps, up = a.upsample ? 1 : 0; + if (path_ == VaePath::kCoop && a.n >= 64) { + const std::uint32_t fn = a.n % 128 == 0 ? 4 : a.n % 96 == 0 ? 3 : a.n % 64 == 0 ? 2 : 4; + std::uint32_t a_terms = 2, b_terms = a.b_bf16 ? 1 : 2, tf32 = 0; + if (precision == VaePrecision::kTf32) tf32 = 1; + if (precision == VaePrecision::kFp32) a_terms = 3; + const std::uint32_t shared = 2 * (a_terms * 128 + b_terms * 32 * fn) * 2 * 16; + require(shared <= device_.info().caps.shared_memory_bytes, + "the cooperative-matrix convolution needs " + std::to_string(shared) + " bytes of shared memory"); + tile_rows = 128; + tile_cols = 32 * fn; + const std::vector key = {0, taps, up, a_terms, b_terms, tf32, fn}; + auto& p = conv_[key]; + if (!p) { + p = std::make_unique(device_, kConvCoopSpirv, sizeof(kConvCoopSpirv), sizeof(ConvPush), kCoopSubgroup, + std::vector{taps, up, a_terms, b_terms, tf32, fn}); + } + return *p; + } + const std::uint32_t tf32 = precision == VaePrecision::kTf32 ? 1 : 0; + if (a.n <= 4) { + tile_rows = 256; + tile_cols = 4; + const std::vector key = {2, taps, up, tf32}; + auto& p = conv_[key]; + if (!p) { + p = std::make_unique(device_, kConvNarrowSpirv, sizeof(kConvNarrowSpirv), sizeof(ConvPush), 0, + std::vector{taps, up, tf32}); + } + return *p; + } + // The decoder's widths are multiples of 96: tiles 96 wide, others 64. + const std::uint32_t tn = a.n % 96 == 0 ? 6 : 4; + tile_rows = 128; + tile_cols = 16 * tn; + const std::vector key = {1, taps, up, tf32, tn}; + auto& p = conv_[key]; + if (!p) { + p = std::make_unique(device_, kConvPlainSpirv, sizeof(kConvPlainSpirv), sizeof(ConvPush), 0, + std::vector{taps, up, tf32, tn}); + } + return *p; +} + +void VaeKernels::record_conv(VkCommandBuffer cb, const ConvArgs& a, VaePrecision precision) const { + require(a.taps == 0 || a.taps == 1 || a.taps == 9, "conv: taps must be 0, 1 or 9"); + require(a.n > 0 && a.k > 0 && a.k % 4 == 0 && a.m0 <= a.m, "conv: bad shape (k must be a multiple of 4)"); + require(a.ldb % 4 == 0 && a.ldb >= a.k && a.ldo >= a.n, "conv: ldb must be a multiple of 4; ldb, ldo must cover k, n"); + // The cooperative-matrix kernel stores 4 outputs at a time; the plain one (and narrow outputs) one. + require(a.ldo % 4 == 0 || path_ == VaePath::kPlain || a.n < 64, "conv: ldo must be a multiple of 4"); + require(a.a % 16 == 0 && a.b % 16 == 0 && a.out % 16 == 0 && a.residual % 16 == 0, "conv: addresses must be 16-byte aligned"); + if (a.taps == 0) { + require(a.lda % 4 == 0 && a.lda >= a.k, "conv: lda must be a multiple of 4 covering k"); + } else { + require(a.c % 8 == 0 && a.k == a.taps * a.c && a.h > 0 && a.w > 0 && a.m <= a.h * a.w, + "conv: c must be a multiple of 8, k = taps * c and rows within h * w"); + require(!a.upsample || (a.taps == 9 && a.h % 2 == 0 && a.w % 2 == 0), "conv: upsample needs a 3x3 conv and an even output"); + } + if (a.m == a.m0) return; + std::uint32_t tile_rows = 0, tile_cols = 0; + const Pipeline& p = conv_pipeline(a, precision, tile_rows, tile_cols); + const ConvPush push{a.a, a.b, a.out, a.bias, a.residual, a.m, a.n, a.k, a.lda, a.ldb, a.ldo, a.h, a.w, a.c, a.m0}; + const std::uint32_t rows_groups = div_up(a.m - a.m0, tile_rows); + require(rows_groups <= 65535, "conv: too many rows for one dispatch; split it into bands"); + p.dispatch(cb, &push, div_up(a.n, tile_cols), rows_groups); + compute_barrier(cb); +} + +std::uint64_t VaeKernels::gn_part_bytes(std::uint32_t pixels, std::uint32_t c) { + return static_cast(div_up(pixels, gn_ppc(c))) * kGroups * 16; +} + +void VaeKernels::record_gn_stats(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t pixels, std::uint32_t c, + VkDeviceAddress part, VkDeviceAddress stats, float eps) const { + require(c % kGroups == 0 && c % 4 == 0 && c <= 1024 && pixels > 0, "group norm: c must be a multiple of 32 up to 1024"); + require(x % 16 == 0 && part % 16 == 0 && stats % 8 == 0, "group norm: x and part must be 16-byte aligned"); + const std::uint32_t ppc = gn_ppc(c), chunks = div_up(pixels, ppc); + const GnStatsPush sp{x, part, pixels, c, ppc, chunks}; + dispatch_1d(*gn_stats_, cb, &sp, chunks); + compute_barrier(cb); + const GnFinalizePush fp{part, stats, chunks, eps}; + gn_finalize_->dispatch(cb, &fp, kGroups, 1); + compute_barrier(cb); +} + +void VaeKernels::record_gn_apply(VkCommandBuffer cb, VkDeviceAddress x, VkDeviceAddress y, std::uint32_t pixels, + std::uint32_t c, VkDeviceAddress stats, VkDeviceAddress gamma, VkDeviceAddress beta, + bool silu) const { + require(c % kGroups == 0 && c % 4 == 0, "group norm: c must be a multiple of 32"); + require(x % 16 == 0 && y % 16 == 0, "group norm: x and y must be 16-byte aligned"); + const std::uint64_t quads = static_cast(pixels) * c / 4; + require(quads < (1ull << 32), "group norm: tensor too large"); + const GnApplyPush push{x, y, stats, gamma, beta, static_cast(quads), c, silu ? 1u : 0u}; + dispatch_1d(*gn_apply_, cb, &push, div_up(quads, 256)); + compute_barrier(cb); +} + +void VaeKernels::record_softmax(VkCommandBuffer cb, VkDeviceAddress s, std::uint32_t rows, std::uint32_t l, + std::uint32_t ld, float scale) const { + require(ld >= l, "softmax: ld must cover l"); + const SoftmaxPush push{s, rows, l, ld, scale}; + dispatch_1d(*softmax_, cb, &push, rows); + compute_barrier(cb); +} + +void VaeKernels::record_transpose(VkCommandBuffer cb, VkDeviceAddress in, VkDeviceAddress out, std::uint32_t r_count, + std::uint32_t c_count, std::uint32_t ld_in, std::uint32_t ld_out) const { + require(ld_in >= c_count && ld_out >= r_count, "transpose: strides must cover the rows"); + require(div_up(r_count, 32) <= 65535 && div_up(c_count, 32) <= 65535, "transpose: too large"); + const TransposePush push{in, out, r_count, c_count, ld_in, ld_out}; + transpose_->dispatch(cb, &push, div_up(c_count, 32), div_up(r_count, 32)); + compute_barrier(cb); +} + +void VaeKernels::record_latent_unpack(VkCommandBuffer cb, VkDeviceAddress z, VkDeviceAddress scale, VkDeviceAddress mean, + VkDeviceAddress y, std::uint32_t h, std::uint32_t w) const { + const std::uint64_t n = static_cast(h) * w * 128; + require(n < (1ull << 32), "latent unpack: too large"); + const UnpackPush push{z, scale, mean, y, h, w}; + dispatch_1d(*unpack_, cb, &push, div_up(n, 256)); + compute_barrier(cb); +} + +void VaeKernels::record_rgb8(VkCommandBuffer cb, VkDeviceAddress x, VkDeviceAddress rgb, std::uint32_t n) const { + require(rgb % 4 == 0, "rgb8: rgb must be 4-byte aligned"); + const Rgb8Push push{x, rgb, n}; + dispatch_1d(*rgb8_, cb, &push, div_up(div_up(n, 4), 256)); + compute_barrier(cb); +} + +// ---- Decoder ---------------------------------------------------------------------------------------------------------- + +namespace { + +struct Conv { + VkDeviceAddress w = 0, b = 0; // fp32 [k, r, r, c] (KRSC), [k] + std::uint32_t k = 0, c = 0, r = 0; + bool bf16 = false; // every weight is a bf16 value +}; +struct Norm { + VkDeviceAddress gamma = 0, beta = 0; + std::uint32_t c = 0; +}; +struct Resnet { + Norm norm1, norm2; + Conv conv1, conv2, shortcut; // shortcut.k == 0 when in == out + std::uint32_t in = 0, out = 0; +}; + +constexpr int kUpBlocks = 4, kResnetsPerBlock = 3; +constexpr std::uint64_t kAlign = 256; +std::uint64_t align_up(std::uint64_t v, std::uint64_t a = kAlign) { return (v + a - 1) / a * a; } + +bool all_bf16(const std::vector& v) { + return std::all_of(v.begin(), v.end(), [](float f) { + std::uint32_t u = 0; + std::memcpy(&u, &f, 4); + return (u & 0xffffu) == 0; + }); +} + +} // namespace + +struct VaeDecoder::Weights { + std::unique_ptr buffer; + VkDeviceAddress bn_scale = 0, bn_mean = 0; + Conv post_quant, conv_in, conv_out, qkv, attn_out; + Norm attn_norm, norm_out; + Resnet mid[2]; + Resnet up[kUpBlocks][kResnetsPerBlock]; + Conv upsampler[kUpBlocks - 1]; +}; + +namespace { + +// Packs tensors into one host array at 64-byte offsets; addresses are offsets until relocated by the buffer's address. +class Packer { + public: + explicit Packer(st::File& file) : file_(file) { data_.resize(16); } // offset 0 stays "none" + + VkDeviceAddress add(const std::vector& v) { + const std::size_t at = data_.size(); + data_.insert(data_.end(), v.begin(), v.end()); + data_.resize(align_up(data_.size(), 16)); + return at * 4; + } + std::vector read(const std::string& name, std::initializer_list shape) { + const st::TensorInfo* t = file_.find(name); + if (t == nullptr) throw std::runtime_error("vulkan vae: weight '" + name + "' missing"); + if (t->shape != std::vector(shape)) { + throw std::runtime_error("vulkan vae: weight '" + name + "' has an unexpected shape"); + } + return file_.read_f32(name); + } + Conv conv(const std::string& p, std::uint32_t k, std::uint32_t c, std::uint32_t r) { + const std::vector kcrs = read(p + ".weight", {k, c, r, r}); + std::vector krsc(kcrs.size()); + for (std::uint32_t ki = 0; ki < k; ++ki) { + for (std::uint32_t ci = 0; ci < c; ++ci) { + for (std::uint32_t t = 0; t < r * r; ++t) krsc[(ki * r * r + t) * c + ci] = kcrs[(ki * c + ci) * r * r + t]; + } + } + Conv cv; + cv.k = k; + cv.c = c; + cv.r = r; + cv.bf16 = all_bf16(krsc); + cv.w = add(krsc); + cv.b = add(read(p + ".bias", {k})); + return cv; + } + Norm norm(const std::string& p, std::uint32_t c) { + Norm n; + n.c = c; + n.gamma = add(read(p + ".weight", {c})); + n.beta = add(read(p + ".bias", {c})); + return n; + } + Resnet resnet(const std::string& p, std::uint32_t in, std::uint32_t out) { + Resnet r; + r.in = in; + r.out = out; + r.norm1 = norm(p + ".norm1", in); + r.conv1 = conv(p + ".conv1", out, in, 3); + r.norm2 = norm(p + ".norm2", out); + r.conv2 = conv(p + ".conv2", out, out, 3); + if (in != out) r.shortcut = conv(p + ".conv_shortcut", out, in, 1); + return r; + } + const std::vector& data() const { return data_; } + + private: + st::File& file_; + std::vector data_; +}; + +void relocate(VkDeviceAddress& a, VkDeviceAddress base) { + if (a != 0) a += base; +} +void relocate(Conv& c, VkDeviceAddress base) { + relocate(c.w, base); + relocate(c.b, base); +} +void relocate(Norm& n, VkDeviceAddress base) { + relocate(n.gamma, base); + relocate(n.beta, base); +} +void relocate(Resnet& r, VkDeviceAddress base) { + relocate(r.norm1, base); + relocate(r.norm2, base); + relocate(r.conv1, base); + relocate(r.conv2, base); + relocate(r.shortcut, base); +} + +} // namespace + +VaeDecoder::VaeDecoder(Device& device, const std::filesystem::path& path) + : VaeDecoder(device, path, device.info().caps.coop_bf16 ? VaePath::kCoop : VaePath::kPlain) {} + +VaeDecoder::VaeDecoder(Device& device, const std::filesystem::path& path, VaePath vae_path) + : device_(device), kernels_(std::make_unique(device, vae_path)), weights_(std::make_unique()) { + st::File file(path); + const st::TensorInfo* conv_in = file.find("decoder.conv_in.weight"); + if (conv_in == nullptr || conv_in->shape.size() != 4 || conv_in->shape[0] % 4 != 0 || conv_in->shape[1] != 32) { + throw std::runtime_error("vulkan vae: " + file.path() + " has no FLUX.2 decoder.conv_in.weight [4c, 32, 3, 3]"); + } + base_ = static_cast(conv_in->shape[0] / 4); + require(base_ % 32 == 0, "decoder widths must be multiples of 32"); + const auto c = static_cast(base_), mid = 4 * c; + Weights& w = *weights_; + Packer pk(file); + + std::vector mean = pk.read("bn.running_mean", {128}), var = pk.read("bn.running_var", {128}); + std::vector scale(128); + for (int i = 0; i < 128; ++i) scale[static_cast(i)] = std::sqrt(var[static_cast(i)] + kBatchNormEps); + w.bn_scale = pk.add(scale); + w.bn_mean = pk.add(mean); + w.post_quant = pk.conv("post_quant_conv", 32, 32, 1); + w.conv_in = pk.conv("decoder.conv_in", mid, 32, 3); + for (int i = 0; i < 2; ++i) w.mid[i] = pk.resnet("decoder.mid_block.resnets." + std::to_string(i), mid, mid); + const std::string a = "decoder.mid_block.attentions.0."; + w.attn_norm = pk.norm(a + "group_norm", mid); + { + // to_q, to_k, to_v as one [3 mid, mid] linear: the same products and sums per output as three. + std::vector wqkv, bqkv; + for (const char* p : {"to_q", "to_k", "to_v"}) { + const std::vector wt = pk.read(a + p + ".weight", {mid, mid}), bt = pk.read(a + p + ".bias", {mid}); + wqkv.insert(wqkv.end(), wt.begin(), wt.end()); + bqkv.insert(bqkv.end(), bt.begin(), bt.end()); + } + w.qkv = Conv{pk.add(wqkv), pk.add(bqkv), 3 * mid, mid, 1, all_bf16(wqkv)}; + const std::vector wo = pk.read(a + "to_out.0.weight", {mid, mid}); + w.attn_out = Conv{pk.add(wo), pk.add(pk.read(a + "to_out.0.bias", {mid})), mid, mid, 1, all_bf16(wo)}; + } + std::uint32_t prev = mid; + for (int b = 0; b < kUpBlocks; ++b) { + const std::uint32_t out = c * (b < 2 ? 4 : b == 2 ? 2 : 1); + const std::string p = "decoder.up_blocks." + std::to_string(b); + for (int r = 0; r < kResnetsPerBlock; ++r) { + w.up[b][r] = pk.resnet(p + ".resnets." + std::to_string(r), r == 0 ? prev : out, out); + } + if (b + 1 < kUpBlocks) w.upsampler[b] = pk.conv(p + ".upsamplers.0.conv", out, out, 3); + prev = out; + } + w.norm_out = pk.norm("decoder.conv_norm_out", c); + w.conv_out = pk.conv("decoder.conv_out", 3, c, 3); + + const std::vector& data = pk.data(); + w.buffer = std::make_unique(device_, data.size() * 4); + w.buffer->upload(data.data(), data.size() * 4); + const VkDeviceAddress base = w.buffer->address(); + relocate(w.bn_scale, base); + relocate(w.bn_mean, base); + for (Conv* cv : {&w.post_quant, &w.conv_in, &w.conv_out, &w.qkv, &w.attn_out}) relocate(*cv, base); + for (Conv& cv : w.upsampler) relocate(cv, base); + relocate(w.attn_norm, base); + relocate(w.norm_out, base); + for (Resnet& r : w.mid) relocate(r, base); + for (auto& block : w.up) { + for (Resnet& r : block) relocate(r, base); + } +} + +VaeDecoder::~VaeDecoder() = default; + +std::uint64_t VaeDecoder::weight_bytes() const { return weights_->buffer->bytes(); } + +namespace { + +// The device memory of one decode. +struct Layout { + std::uint64_t region = 0; // bytes per ping-pong region; each holds two halves + std::uint64_t part = 0; // GroupNorm partials + std::uint32_t slab = 0; // query rows per attention slab + std::uint64_t latents = 0, rgb = 0; + std::uint64_t total() const { return 2 * region + align_up(part) + kAlign + latents + rgb; } +}; + +Layout layout(std::uint32_t c, std::uint32_t h, std::uint32_t w, std::uint64_t max_macs) { + require(h > 0 && w > 0, "latent grid must be at least 1x1"); + const std::uint64_t full = 256ull * h * w, latent_px = 4ull * h * w, mid = 4ull * c; + Layout l; + // Attention scratch: xn [L, mid], qkv [L, 3 mid], V^T [mid, L], then query-row slabs of scores [rows, L]. + const std::uint64_t l_tokens = latent_px; + const std::uint64_t budget_rows = std::max(VaeKernels::kRowAlign, max_macs / (l_tokens * mid) / + VaeKernels::kRowAlign * + VaeKernels::kRowAlign); + l.slab = static_cast(std::min(budget_rows, align_up(l_tokens, VaeKernels::kRowAlign))); + const std::uint64_t attention = 3 * align_up(l_tokens * mid * 4) + align_up(l_tokens * 3 * mid * 4) + + align_up(static_cast(l.slab) * l_tokens * 4); + const std::uint64_t tail = 2ull * c * full * 4; // up_blocks[2]'s upsampler output: 2c channels at full resolution + l.region = align_up(std::max(tail, attention), 2 * kAlign); + // Every GroupNorm input: {pixels, channels}. + const std::uint64_t norms[][2] = {{latent_px, mid}, {4 * latent_px, mid}, {16 * latent_px, mid}, {16 * latent_px, 2ull * c}, + {full, 2ull * c}, {full, c}}; + for (const auto& g : norms) { + l.part = std::max(l.part, VaeKernels::gn_part_bytes(static_cast(g[0]), static_cast(g[1]))); + } + l.latents = align_up(128ull * h * w * 4); + l.rgb = align_up(full * 3); + return l; +} + +// Groups recorded ops into submits of at most `budget` multiply-adds (elementwise ops count kElementCost per element). +// With a profile, every op is its own submit and its wall time is added to the profile under the op's name. +class Submitter { + public: + static constexpr std::uint64_t kElementCost = 32; + Submitter(Device& device, std::uint64_t budget, std::vector>* profile) + : device_(device), budget_(budget), profile_(profile) {} + void add(const std::string& name, std::uint64_t cost, std::function op) { + if (profile_ != nullptr) { + const auto t0 = std::chrono::steady_clock::now(); + device_.run(op); + ++submits_; + const double s = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + auto it = std::find_if(profile_->begin(), profile_->end(), [&](const auto& e) { return e.first == name; }); + if (it == profile_->end()) { + profile_->emplace_back(name, s); + } else { + it->second += s; + } + return; + } + if (!ops_.empty() && cost_ + cost > budget_) flush(); + ops_.push_back(std::move(op)); + cost_ += cost; + } + void flush() { + if (ops_.empty()) return; + device_.run([&](VkCommandBuffer cb) { + for (const auto& op : ops_) op(cb); + }); + ops_.clear(); + cost_ = 0; + ++submits_; + } + int submits() const { return submits_; } + + private: + Device& device_; + std::uint64_t budget_; + std::vector>* profile_; + std::vector> ops_; + std::uint64_t cost_ = 0; + int submits_ = 0; +}; + +std::string dims(std::uint32_t h, std::uint32_t w) { return " @" + std::to_string(h) + "x" + std::to_string(w); } + +class Run { + public: + Run(const VaeKernels& k, Submitter& s, std::uint64_t max_macs, VkDeviceAddress part, VkDeviceAddress stats) + : k_(k), s_(s), max_macs_(max_macs), part_(part), stats_(stats) {} + + // A convolution or GEMM in bands of output rows, each at most max_macs multiply-adds (at least one row tile). + void gemm(const std::string& name, ConvArgs a, VaePrecision precision) { + const std::uint64_t per_row = static_cast(a.n) * a.k; + // 65535 row tiles of 128 is what one dispatch addresses. + const std::uint64_t rows = std::min( + 65535ull * VaeKernels::kRowAlign, + std::max(VaeKernels::kRowAlign, max_macs_ / per_row / VaeKernels::kRowAlign * VaeKernels::kRowAlign)); + const std::uint32_t end = a.m; + for (std::uint64_t m0 = a.m0; m0 < end; m0 += rows) { + a.m0 = static_cast(m0); + a.m = static_cast(std::min(end, m0 + rows)); + s_.add(name, (a.m - a.m0) * per_row, [this, a, precision](VkCommandBuffer cb) { k_.record_conv(cb, a, precision); }); + } + } + void conv(const Conv& cv, VkDeviceAddress x, std::uint32_t h, std::uint32_t w, VkDeviceAddress y, + VkDeviceAddress residual = 0, bool upsample = false) { + ConvArgs a; + a.a = x; + a.b = cv.w; + a.out = y; + a.bias = cv.b; + a.residual = residual; + a.m = h * w; + a.n = cv.k; + a.k = cv.r * cv.r * cv.c; + a.ldb = a.k; + a.ldo = cv.k; + a.taps = cv.r * cv.r; + a.upsample = upsample; + a.h = h; + a.w = w; + a.c = cv.c; + a.b_bf16 = cv.bf16; + const std::string name = "conv " + std::to_string(cv.r) + "x" + std::to_string(cv.r) + " " + std::to_string(cv.c) + "->" + + std::to_string(cv.k) + (upsample ? " upsampled" : "") + dims(h, w); + gemm(name, a, VaePrecision::kTf32); + } + void norm(const Norm& n, VkDeviceAddress x, std::uint32_t pixels, VkDeviceAddress y, bool silu) { + const std::uint64_t cost = static_cast(pixels) * n.c * Submitter::kElementCost; + const std::string shape = " " + std::to_string(n.c) + " ch, " + std::to_string(pixels) + " px"; + s_.add("group_norm stats" + shape, cost, [this, x, pixels, c = n.c](VkCommandBuffer cb) { k_.record_gn_stats(cb, x, pixels, c, part_, stats_, kResnetEps); }); + s_.add(std::string(silu ? "group_norm+silu apply" : "group_norm apply") + shape, cost, [this, n, x, y, pixels, silu](VkCommandBuffer cb) { + k_.record_gn_apply(cb, x, y, pixels, n.c, stats_, n.gamma, n.beta, silu); + }); + } + // ResnetBlock2D (as Flux2VaeDecoder::resnet): returns the buffer holding the output, x or a. + VkDeviceAddress resnet(const Resnet& r, VkDeviceAddress x, VkDeviceAddress a, VkDeviceAddress b, std::uint32_t h, + std::uint32_t w) { + const std::uint32_t px = h * w; + if (r.shortcut.k != 0) { + conv(r.shortcut, x, h, w, a); + norm(r.norm1, x, px, x, true); + conv(r.conv1, x, h, w, b); + norm(r.norm2, b, px, b, true); + conv(r.conv2, b, h, w, a, a); + return a; + } + norm(r.norm1, x, px, a, true); + conv(r.conv1, a, h, w, b); + norm(r.norm2, b, px, b, true); + conv(r.conv2, b, h, w, x, x); + return x; + } + // x [L, mid] += to_out(SDPA(to_q/k/v(GN(x)))), scratch in ws (a region). + void attention(const VaeDecoder::Weights& w, VkDeviceAddress x, VkDeviceAddress ws, std::uint32_t l_tokens, + std::uint32_t mid, std::uint32_t slab) { + const VkDeviceAddress xn = ws, qkv = xn + align_up(static_cast(l_tokens) * mid * 4), + vt = qkv + align_up(static_cast(l_tokens) * 3 * mid * 4), + scores = vt + align_up(static_cast(l_tokens) * mid * 4); + norm(w.attn_norm, x, l_tokens, xn, false); + ConvArgs p; + p.a = xn; + p.b = w.qkv.w; + p.bias = w.qkv.b; + p.out = qkv; + p.m = l_tokens; + p.n = 3 * mid; + p.k = mid; + p.lda = mid; + p.ldb = mid; + p.ldo = 3 * mid; + p.b_bf16 = w.qkv.bf16; + gemm("attention qkv", p, VaePrecision::kFp32); + s_.add("attention transpose V", static_cast(l_tokens) * mid * Submitter::kElementCost, [this, qkv, vt, l_tokens, mid](VkCommandBuffer cb) { + k_.record_transpose(cb, qkv + 2ull * mid * 4, vt, l_tokens, mid, 3 * mid, l_tokens); + }); + // SDPA's scale as the engine forms it: 1 / sqrt(E) in fp32. + const float scale = 1.0f / std::sqrt(static_cast(mid)); + const VkDeviceAddress o = xn; // GN output is dead after the projections + for (std::uint32_t r0 = 0; r0 < l_tokens; r0 += slab) { + const std::uint32_t rows = std::min(slab, l_tokens - r0); + ConvArgs s; + s.a = qkv + static_cast(r0) * 3 * mid * 4; + s.b = qkv + static_cast(mid) * 4; + s.out = scores; + s.m = rows; + s.n = l_tokens; + s.k = mid; + s.lda = 3 * mid; + s.ldb = 3 * mid; + s.ldo = l_tokens; + gemm("attention scores", s, VaePrecision::kSplit16); + s_.add("attention softmax", static_cast(rows) * l_tokens * Submitter::kElementCost, + [this, scores, rows, l_tokens, scale](VkCommandBuffer cb) { k_.record_softmax(cb, scores, rows, l_tokens, l_tokens, scale); }); + ConvArgs pv; + pv.a = scores; + pv.b = vt; + pv.out = o + static_cast(r0) * mid * 4; + pv.m = rows; + pv.n = mid; + pv.k = l_tokens; + pv.lda = l_tokens; + pv.ldb = l_tokens; + pv.ldo = mid; + gemm("attention P V", pv, VaePrecision::kSplit16); + } + // Reference order: round(to_out(o) + bias), then + residual. + ConvArgs out; + out.a = o; + out.b = w.attn_out.w; + out.bias = w.attn_out.b; + out.residual = x; + out.out = x; + out.m = l_tokens; + out.n = mid; + out.k = mid; + out.lda = mid; + out.ldb = mid; + out.ldo = mid; + out.b_bf16 = w.attn_out.bf16; + gemm("attention out", out, VaePrecision::kFp32); + } + + private: + const VaeKernels& k_; + Submitter& s_; + std::uint64_t max_macs_; + VkDeviceAddress part_, stats_; +}; + +} // namespace + +std::uint64_t VaeDecoder::workspace_bytes(std::uint32_t h, std::uint32_t w) const { + return layout(static_cast(base_), h, w, max_macs_).total(); +} + +std::uint64_t VaeDecoder::largest_buffer_bytes(std::uint32_t h, std::uint32_t w) const { + return layout(static_cast(base_), h, w, max_macs_).region; +} + +VaeDecodeStats VaeDecoder::decode(const float* latents, std::uint32_t h, std::uint32_t w, std::vector& rgb, + std::vector* image) { + const auto t0 = std::chrono::steady_clock::now(); + const auto c = static_cast(base_), mid = 4 * c; + const Layout l = layout(c, h, w, max_macs_); + const std::uint32_t hh = 2 * h, ww = 2 * w, H = 16 * h, W = 16 * w; + require(static_cast(H) * W * 3 < (1ull << 32), "image too large"); + Buffer r0(device_, l.region), r1(device_, l.region), aux(device_, align_up(l.part) + kAlign), + z(device_, l.latents), out(device_, l.rgb); + z.upload(latents, 128ull * h * w * 4); + const VkDeviceAddress part = aux.address(), stats = aux.address() + align_up(l.part); + const std::uint64_t half = l.region / 2; + const Weights& wt = *weights_; + + VaeDecodeStats st; + Submitter sub(device_, max_macs_, profile_ ? &st.profile : nullptr); + Run run(*kernels_, sub, max_macs_, part, stats); + VkDeviceAddress p = r0.address(), q = r1.address(); + + // Head at H/8: unpack, post_quant_conv, conv_in, mid resnet 0, mid attention. + const VkDeviceAddress unpacked = q, pq = q + half; + sub.add("latent unpack", 128ull * h * w * Submitter::kElementCost, [&, unpacked](VkCommandBuffer cb) { + kernels_->record_latent_unpack(cb, z.address(), wt.bn_scale, wt.bn_mean, unpacked, h, w); + }); + run.conv(wt.post_quant, unpacked, hh, ww, pq); + run.conv(wt.conv_in, pq, hh, ww, p); + run.resnet(wt.mid[0], p, q, q + half, hh, ww); + run.attention(wt, p, q, hh * ww, mid, l.slab); + + // Tail: mid resnet 1, the up blocks, norm_out + SiLU, conv_out, RGB8. + std::uint32_t ch = hh, cw = ww; + run.resnet(wt.mid[1], p, q, q + half, ch, cw); + for (int b = 0; b < kUpBlocks; ++b) { + VkDeviceAddress x = p; + for (int r = 0; r < kResnetsPerBlock; ++r) { + const VkDeviceAddress other = x == p ? q : p; + x = run.resnet(wt.up[b][r], x, other, other + half, ch, cw); + } + if (b + 1 < kUpBlocks) { + const VkDeviceAddress y = x == p ? q : p; + ch *= 2; + cw *= 2; + run.conv(wt.upsampler[b], x, ch, cw, y, 0, true); + x = y; + } + p = x; + q = x == r0.address() ? r1.address() : r0.address(); + } + run.norm(wt.norm_out, p, ch * cw, p, true); + run.conv(wt.conv_out, p, ch, cw, q); + const std::uint32_t n = H * W * 3; + sub.add("rgb8", static_cast(n) * Submitter::kElementCost, + [&, q, n](VkCommandBuffer cb) { kernels_->record_rgb8(cb, q, out.address(), n); }); + sub.flush(); + + rgb.resize(n); + { + std::vector words(l.rgb); + out.download(words.data(), l.rgb); + std::memcpy(rgb.data(), words.data(), n); + } + if (image != nullptr) { + // conv_out's output starts its region. + image->resize(n); + (q == r0.address() ? r0 : r1).download(image->data(), static_cast(n) * 4); + } + st.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + st.device_bytes = l.total(); + st.submits = sub.submits(); + return st; +} + +} // namespace mi::vk diff --git a/vulkan/vae.h b/vulkan/vae.h new file mode 100644 index 0000000..eebf469 --- /dev/null +++ b/vulkan/vae.h @@ -0,0 +1,152 @@ +// The FLUX.2 VAE decoder as the MOSS-Image pipeline runs it (src/models/flux2_vae, pipeline_moss_image.py:430-438): the +// DiT's final latents, fp32 NCHW [128, h, w], to RGB8 [16h, 16w, 3]. Activations are fp32 NHWC throughout, as in the +// engine's reference precision (DecodeMode::reference()): +// latent de-normalization + unpack -> post_quant_conv 1x1 -> conv_in 3x3 -> mid resnet 0 -> mid attention +// -> mid resnet 1 -> 4 up blocks (3 resnets each, nearest 2x upsample + 3x3 conv after the first three) +// -> GroupNorm + SiLU -> conv_out 3x3 -> clamp and the image processor's rounding to uint8. +// Convolutions round both operands to TF32 like cuDNN's (torch's allow_tf32 convolutions) and accumulate in fp32; the +// nearest upsample is read inside the following convolution. The attention's linears are fp32 GEMMs (allow_tf32 off in the +// reference), its scores and P V are fp32-accurate to 16 significant bits per operand on the cooperative-matrix path and +// plain fp32 on the other. GroupNorm statistics are fp32 per chunk of pixels (mean and sum of squared deviations), merged +// pairwise. Accumulation orders differ from cuDNN's and torch's, so outputs are close to the engine's, not bit-identical. +// +// Two paths (generate() runs kCoop with BlockPath::kCoop, kPlain with the others): kCoop runs the convolutions and GEMMs on bf16 cooperative matrices +// (Capabilities::coop_bf16), splitting each fp32 operand into bf16 pieces whose products are exact (TF32 = 2 pieces); +// kPlain needs no cooperative matrices, 16-bit types or particular subgroup size. +// +// Memory (untiled): two regions of 2c x 16h x 16w fp32 each (c = 96 for the small decoder), as the engine's ping-pong +// schedule, plus GroupNorm partials, the input latents and the RGB output; the mid attention's scratch lives in the second +// region. At 1024x1024 that is 1.5 GiB plus about 112 MB of fp32 weights. +// Submits: every op is recorded into submits of at most `max_macs_per_submit` multiply-adds (a convolution larger than +// that runs in bands of output rows, each its own submit), so no single submit runs long on a slow GPU. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "runtime.h" + +namespace mi::vk { + +enum class VaePath { kCoop, kPlain }; +const char* vae_path_name(VaePath path); + +// How a convolution or GEMM rounds its operands (see vae_conv_coop.comp). +enum class VaePrecision { + kTf32, // both operands rounded to TF32, products exact: torch's TF32 convolution up to accumulation order + kFp32, // fp32 GEMM: exact products on the coop path when B is bf16-valued (A in 3 bf16 pieces), else 16-bit B pieces + kSplit16, // coop: 2 bf16 pieces per operand (16 significant bits each); plain: fp32 +}; + +// out[m, n] = (sum_k A[m, k] B[n, k] + bias[n]) + residual[m, n], fp32, for rows m in [m0, m), with B row-major [n, k]. +// taps 0: A is row-major with stride lda. taps 1 / 9: a 1x1 / 3x3 convolution (zero padding 1) over an NHWC input of c +// channels for an h x w output, whose pixel index is the row; upsample: the input is h/2 x w/2, read through a nearest 2x +// upsample. k = taps * c for a convolution. Requirements: c a multiple of 8, k, lda and ldb multiples of 4, ldo a multiple +// of 4 on the coop path for n >= 64, 16-byte aligned addresses. residual may equal out; out must not overlap A or B. +struct ConvArgs { + VkDeviceAddress a = 0, b = 0, out = 0, bias = 0, residual = 0; + std::uint32_t m = 0, n = 0, k = 0, lda = 0, ldb = 0, ldo = 0; + std::uint32_t taps = 0; + bool upsample = false; + std::uint32_t h = 0, w = 0, c = 0; + std::uint32_t m0 = 0; + bool b_bf16 = false; // every B value is a bf16 value (fewer pieces on the coop path) +}; + +class VaeKernels { + public: + // kCoop where the device has bf16 cooperative matrices, unless `path` names one. Throws if the device cannot run it. + explicit VaeKernels(Device& device); + VaeKernels(Device& device, VaePath path); + ~VaeKernels(); + VaeKernels(const VaeKernels&) = delete; + VaeKernels& operator=(const VaeKernels&) = delete; + + VaePath path() const { return path_; } + + // Every record_* records its dispatches and a barrier after them. + void record_conv(VkCommandBuffer cb, const ConvArgs& args, VaePrecision precision) const; + // Rows of one band of a convolution or GEMM: a multiple of the row tile every kernel uses. + static constexpr std::uint32_t kRowAlign = 128; + + // GroupNorm with 32 groups over x [pixels, c] (fp32 NHWC, c a multiple of 32): statistics into stats ({mean, rstd} per + // group, 256 bytes) through part (gn_part_bytes), then y = GN(x) * gamma + beta, SiLU optionally; y may be x. + static std::uint64_t gn_part_bytes(std::uint32_t pixels, std::uint32_t c); + void record_gn_stats(VkCommandBuffer cb, VkDeviceAddress x, std::uint32_t pixels, std::uint32_t c, VkDeviceAddress part, + VkDeviceAddress stats, float eps) const; + void record_gn_apply(VkCommandBuffer cb, VkDeviceAddress x, VkDeviceAddress y, std::uint32_t pixels, std::uint32_t c, + VkDeviceAddress stats, VkDeviceAddress gamma, VkDeviceAddress beta, bool silu) const; + // s[r, :l] (row stride ld) = softmax(s[r, :l] * scale), in place, rows r < rows. + void record_softmax(VkCommandBuffer cb, VkDeviceAddress s, std::uint32_t rows, std::uint32_t l, std::uint32_t ld, + float scale) const; + // out[c, r] = in[r, c]: in [r_count, c_count] with row stride ld_in, out with row stride ld_out. + void record_transpose(VkCommandBuffer cb, VkDeviceAddress in, VkDeviceAddress out, std::uint32_t r_count, + std::uint32_t c_count, std::uint32_t ld_in, std::uint32_t ld_out) const; + // z [128, h, w] NCHW -> y [2h, 2w, 32] NHWC: z * scale + mean per latent channel (scale = sqrt(var + eps), host-made). + void record_latent_unpack(VkCommandBuffer cb, VkDeviceAddress z, VkDeviceAddress scale, VkDeviceAddress mean, + VkDeviceAddress y, std::uint32_t h, std::uint32_t w) const; + // x [n] fp32 -> rgb, n bytes packed into 32-bit words (the last word's bytes past n are 0). + void record_rgb8(VkCommandBuffer cb, VkDeviceAddress x, VkDeviceAddress rgb, std::uint32_t n) const; + + private: + const Pipeline& conv_pipeline(const ConvArgs& args, VaePrecision precision, std::uint32_t& tile_rows, + std::uint32_t& tile_cols) const; + + Device& device_; + VaePath path_; + mutable std::map, std::unique_ptr> conv_; + std::unique_ptr gn_stats_, gn_finalize_, gn_apply_, softmax_, transpose_, unpack_, rgb8_; +}; + +struct VaeDecodeStats { + double seconds = 0; // decode() wall time, uploads and downloads included + std::uint64_t device_bytes = 0; // device memory decode() allocated (weights not included) + int submits = 0; + // With set_profile(true): wall seconds per op kind (one submit per op), in order of first use. + std::vector> profile; +}; + +class VaeDecoder { + public: + // Loads the decoder ("decoder.*", "post_quant_conv.*") and the latent statistics ("bn.running_mean", "bn.running_var") + // from a FLUX.2 VAE safetensors file in diffusers' naming, fp32 or bf16. The widths come from the weights' shapes + // (decoder_block_out_channels c, 2c, 4c, 4c). + VaeDecoder(Device& device, const std::filesystem::path& safetensors_path); + VaeDecoder(Device& device, const std::filesystem::path& safetensors_path, VaePath path); + ~VaeDecoder(); + VaeDecoder(const VaeDecoder&) = delete; + VaeDecoder& operator=(const VaeDecoder&) = delete; + + const VaeKernels& kernels() const { return *kernels_; } + int base_channels() const { return base_; } + std::uint64_t weight_bytes() const; + // Device bytes decode() allocates for an h x w latent grid. + std::uint64_t workspace_bytes(std::uint32_t h, std::uint32_t w) const; + // The largest single buffer of that workspace (each ping-pong region). + std::uint64_t largest_buffer_bytes(std::uint32_t h, std::uint32_t w) const; + void set_max_macs_per_submit(std::uint64_t macs) { max_macs_ = macs; } + // Every op in its own submit, timed (VaeDecodeStats::profile): slower, for finding where time goes. + void set_profile(bool on) { profile_ = on; } + + // latents: fp32 NCHW [128, h, w], the pipeline's final latents. rgb: resized to 16h x 16w x 3. image: when not null, the + // decoder's fp32 output [16h, 16w, 3] before the clamp. + VaeDecodeStats decode(const float* latents, std::uint32_t h, std::uint32_t w, std::vector& rgb, + std::vector* image = nullptr); + + struct Weights; + + private: + Device& device_; + std::unique_ptr kernels_; + std::unique_ptr weights_; + int base_ = 0; + std::uint64_t max_macs_ = 8ull << 30; + bool profile_ = false; +}; + +} // namespace mi::vk diff --git a/vulkan/vk_api.cpp b/vulkan/vk_api.cpp new file mode 100644 index 0000000..5ba76b4 --- /dev/null +++ b/vulkan/vk_api.cpp @@ -0,0 +1,68 @@ +#include "vk_api.h" + +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + +#define MI_VK_DEFINE(name) PFN_##name name = nullptr; +MI_VK_FUNCTIONS(MI_VK_DEFINE) +#undef MI_VK_DEFINE + +namespace mi::vk { +namespace { + +void* open_loader() { +#if defined(_WIN32) + // From the system directory only, as the Vulkan loader's documentation asks. + return reinterpret_cast(LoadLibraryExW(L"vulkan-1.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32)); +#elif defined(__APPLE__) + void* h = dlopen("libvulkan.1.dylib", RTLD_NOW | RTLD_LOCAL); + return h != nullptr ? h : dlopen("libMoltenVK.dylib", RTLD_NOW | RTLD_LOCAL); +#else + void* h = dlopen("libvulkan.so.1", RTLD_NOW | RTLD_LOCAL); + return h != nullptr ? h : dlopen("libvulkan.so", RTLD_NOW | RTLD_LOCAL); +#endif +} + +void* symbol(void* library, const char* name) { +#if defined(_WIN32) + return reinterpret_cast(GetProcAddress(static_cast(library), name)); +#else + return dlsym(library, name); +#endif +} + +} // namespace + +bool load_vulkan(std::string* error) { + static std::once_flag once; + static std::string failure; + std::call_once(once, [] { + void* library = open_loader(); + if (library == nullptr) { + failure = "no Vulkan loader found (libvulkan.so.1 / vulkan-1.dll): install the GPU's Vulkan driver"; + return; + } + std::string missing; +#define MI_VK_LOAD(name) \ + name = reinterpret_cast(symbol(library, #name)); \ + if (name == nullptr) missing += std::string(missing.empty() ? "" : ", ") + #name; + MI_VK_FUNCTIONS(MI_VK_LOAD) +#undef MI_VK_LOAD + if (!missing.empty()) failure = "the Vulkan loader lacks " + missing + " (it is older than Vulkan 1.2)"; + }); + if (!failure.empty() && error != nullptr) *error = failure; + return failure.empty(); +} + +} // namespace mi::vk diff --git a/vulkan/vk_api.h b/vulkan/vk_api.h new file mode 100644 index 0000000..a7d761c --- /dev/null +++ b/vulkan/vk_api.h @@ -0,0 +1,74 @@ +// The Vulkan API as the backend uses it: the headers without prototypes, and the core functions as pointers that +// load_vulkan() resolves from the system's Vulkan loader (libvulkan.so.1, vulkan-1.dll, libvulkan.1.dylib) at run time. +// So a program with the backend built in starts on a machine without Vulkan and fails only when it asks for a device. +// Extension functions come from vkGetInstanceProcAddr / vkGetDeviceProcAddr where they are used. +#pragma once + +#ifndef VK_NO_PROTOTYPES +#define VK_NO_PROTOTYPES +#endif +#include + +#include + +#define MI_VK_FUNCTIONS(X) \ + X(vkGetInstanceProcAddr) \ + X(vkGetDeviceProcAddr) \ + X(vkCreateInstance) \ + X(vkDestroyInstance) \ + X(vkEnumeratePhysicalDevices) \ + X(vkEnumerateDeviceExtensionProperties) \ + X(vkGetPhysicalDeviceProperties) \ + X(vkGetPhysicalDeviceProperties2) \ + X(vkGetPhysicalDeviceFeatures2) \ + X(vkGetPhysicalDeviceMemoryProperties) \ + X(vkGetPhysicalDeviceMemoryProperties2) \ + X(vkGetPhysicalDeviceQueueFamilyProperties) \ + X(vkCreateDevice) \ + X(vkDestroyDevice) \ + X(vkGetDeviceQueue) \ + X(vkQueueSubmit) \ + X(vkQueueWaitIdle) \ + X(vkCreateCommandPool) \ + X(vkDestroyCommandPool) \ + X(vkAllocateCommandBuffers) \ + X(vkFreeCommandBuffers) \ + X(vkBeginCommandBuffer) \ + X(vkEndCommandBuffer) \ + X(vkCreateBuffer) \ + X(vkDestroyBuffer) \ + X(vkGetBufferMemoryRequirements) \ + X(vkGetBufferDeviceAddress) \ + X(vkAllocateMemory) \ + X(vkFreeMemory) \ + X(vkBindBufferMemory) \ + X(vkMapMemory) \ + X(vkCreateShaderModule) \ + X(vkDestroyShaderModule) \ + X(vkCreatePipelineLayout) \ + X(vkDestroyPipelineLayout) \ + X(vkCreateComputePipelines) \ + X(vkDestroyPipeline) \ + X(vkCmdBindPipeline) \ + X(vkCmdPushConstants) \ + X(vkCmdDispatch) \ + X(vkCmdPipelineBarrier) \ + X(vkCmdCopyBuffer) \ + X(vkCmdFillBuffer) \ + X(vkCreateFence) \ + X(vkDestroyFence) \ + X(vkWaitForFences) \ + X(vkCreateSemaphore) \ + X(vkDestroySemaphore) + +#define MI_VK_DECLARE(name) extern PFN_##name name; +MI_VK_FUNCTIONS(MI_VK_DECLARE) +#undef MI_VK_DECLARE + +namespace mi::vk { + +// Opens the Vulkan loader once and resolves every function above. Returns false with a reason when the loader is +// missing or lacks one of them (a loader older than Vulkan 1.2). Thread-safe; later calls return the first result. +bool load_vulkan(std::string* error); + +} // namespace mi::vk