Repository navigation
moss-infer-vulkan: the Vulkan fallback program, with Linux and Windows packages - #8
Merged
Merged
Conversation
Two documents for a possible Vulkan compute backend: - docs/vulkan-backend-research.md: device audit on the RTX 2080 (Vulkan 1.4.357, cooperative-matrix shapes measured with a probe: f16 16x16x16 with f32 accumulate, int8 16x16x32, bf16 reported but emulated on Turing; 48 KiB shared memory, memory heaps), prior art (ggml-vulkan, stable-diffusion.cpp, ncnn, IREE), performance model against the measured CUDA numbers, numerics, risks. - docs/vulkan-backend-design.md: two-phase plan (standalone runner first, then a runtime/kernel abstraction), module layout and build integration, runtime design (memory, queues, dispatch), kernel plan for the DiT/text encoder/VAE, GEMM and attention tiling for 48 KiB, quantized formats, planner/CLI/server integration, test strategy and milestones. Research only: no engine code changes.
- docs/vulkan-backend-{research,design}.md: device audit on the RTX 2080 (Vulkan 1.4.357,
cooperative-matrix shapes, 48 KiB shared memory, heaps), prior art, performance model, and
a two-phase design for a Vulkan backend.
- tools/vulkan_gemm: a standalone tool (not part of the engine build) that runs an f16
cooperative-matrix GEMM at the engine's shapes, checks it against a CPU reference and times
it. Measured on the RTX 2080: 1e-5 max error at 64^3, ~8-9 TFLOPS at 3072^3 (memory-bound;
shared-memory staging is the next step).
- tools/vulkan_gemm/probe.py: lists the device's cooperative-matrix shapes and memory heaps
through vulkan-1.dll, no SDK needed.
No engine code changes.
Tool: requires full 32-wide subgroups and checks the device (shape, features, subgroup size) instead of assuming them; picks the first usable device; barriers between upload, dispatches and download; enables the features the shader declares, including vulkanMemoryModel, which was missing; runs free M through padded tiles, so the engine's real shapes run; checks several tiles, band tails and row tails, and exits non-zero on a failed check; correct fp16 conversions (exhaustively compared with the compiler's); probe.py works on Linux. Notes: Vulkan is the fallback for hardware the engine does not otherwise support. Measured against cuBLASLt on one GPU. Corrected DiT and text encoder shapes, attention tiles that fit 48 KiB, and stale facts.
…e GEMM Standalone vulkan/ replaces tools/vulkan_gemm. It stays outside src/ because the engine's build globs that directory, so default builds do not change. The coopmat GEMM stages A and B in shared memory and writes through it, so M and N are free and the down GEMM nearly doubles. Devices without the fast path get a plain compute GEMM that needs only 64-bit integers and buffer device addresses, and takes bf16 as well as f16. Checks pass on an RTX 4090 Laptop on both paths and on llvmpipe, with core validation clean.
…eplays RMSNorm with modulation, QK norm and RoPE, varlen attention with online softmax, and the GEMM's plain, gated residual and SiTU-GLU epilogues, all bf16 with the engine's rounding points. They need no cooperative matrices, 16-bit support or particular subgroup size. Known-answer checks pass on an RTX 4090 Laptop and on llvmpipe, with core validation clean. Shaders other than the cooperative-matrix one now compile for Vulkan 1.2: Vulkan 1.2 devices reject the newer SPIR-V the build was producing for them.
vulkan_block_real runs a block's plain compute kernels on real weights and activations and compares each output with the engine's, from files captured from its forward, so it needs no CUDA. At the 1024x1024 golden case the cosine is 0.999997 or better on blocks 1, 22 and 41, with no NaN.
…ues and flash attention Devices with VK_KHR_shader_bfloat16 fragments run the block's GEMMs as 128x128 tiles with their bf16, gated residual and SiTU-GLU epilogues, and attention as a flash attention over 32-key steps that rescales the output only when a row's maximum grew. One block of 8484 tokens on an RTX 4090 Laptop: 1.39 s on the plain path, 169 ms here. Against the engine on real blocks 1, 22 and 41: 96 to 99 % same bits, cosine 0.999999919 or better.
…op from a golden case dit.h runs MOSSImageTransformer on the block kernels with every weight resident: the time embedding and shared modulation, the text refiner, the joint sequences of several rows with per-row skips, and the head. New kernels: GEMMs with a bias alone or before a residual add, RMSNorm without modulation, SiLU and SiTU over separate tensors. vulkan_dit_run feeds a golden case's text-encoder outputs and noise through the 50-step CFG loop. Step 0 of g_fridge_1024 against the PyTorch pipeline on the coop path: cosine 0.999976 (cond) and 0.999983 (uncond), where the engine's own test asks for 0.99988.
…e statistics on request The plain path's GEMM stages fp32 tiles of 128 rows by 8 k in two shared-memory stages (17 KiB) and gives each of 256 invocations an 8x8 block of outputs with fused multiply-adds; it still needs no subgroups and no 16-bit arithmetic. Measured back to back with the old one on a throttled RTX 4090 Laptop (35 W cap, 660-780 MHz): 5 to 8 times faster on the DiT's shapes. MI_VK_PIPELINE_STATS=1 prints each pipeline's registers and spills where the driver reports them.
Each workgroup of 128 invocations takes 32 query rows of one head through 32-key steps staged in shared memory (30 KiB, so llvmpipe's 32 KiB still fits), with S and acc_o register-blocked 2 rows by 4 keys and 2 rows by 16 columns, and the softmax of a row split over four invocations that meet through shared memory. No subgroup operations or 16-bit arithmetic. Same steps as FlashAttention; passes the replay check on an RTX 4090 Laptop and on llvmpipe. One plain block back to back with the previous kernels on the throttled GPU: 14.4 s before, 2.9 s now (attention 3.5 s to 0.5 s).
…perative matrices and plain compute Decodes the golden 1024x1024 latent to within 1 of the reference PNG on 0.75 % of the bytes on both paths (PSNR 69.4 dB); each kernel is checked against a host replay.
cmake/modules/vulkan.cmake builds vulkan/ into the engine with -DMOSSINFER_VULKAN=ON (off by default). The CLI takes --backend cuda|vulkan, --vulkan-device and --vulkan-kernels; with vulkan it refuses every CUDA-only flag and every request the backend cannot run (other samplers and schedules, LoRAs, img2img, hires, control, accelerators), reports the device and kernels it runs, and computes the token ids, the seed's initial noise and the schedule with the engine's host code before calling vulkan/generate.h. generate.cpp joins the library with the text encoder and the VAE.
DitOptions::device_budget keeps the first blocks resident and the rest in mapped host memory, each copied into one device slot right before it runs; generate() sizes the budget from VK_EXT_memory_budget (else the device-local heap) minus the run's activations and a 512 MiB reserve. With a 4 GiB budget on g_fridge_1024 (15 of 42 blocks resident, 5.7 GiB in host memory) step 0 has the resident run's bits.
TextEncoder runs Qwen3_5Model(input_ids).last_hidden_state for several sequences packed along the rows: GatedDeltaNet layers (conv1d+SiLU, gates, the engine's chunked gated delta rule in fp32 split into four kernels that fit 32 KiB of shared memory), full gated attention (q/k norm with RoPE, causal GQA attention at head_dim 256), the MLP, with the GEMMs on BlockKernels and residual adds as its gated-residual epilogue with a gate of ones. Weights stream from the safetensors files: a reader thread fills one of two host staging buffers per layer, the layer is copied to one device buffer, and work is submitted in pieces of bounded size. The new kernels need no cooperative matrices, 16-bit arithmetic or subgroup size, so both paths share them. Against the official pipeline's goldens (g_fridge_1024) on an RTX 4090 Laptop, cooperative-matrix path: last_hidden cosine 0.9998972 (cond) and 0.9998012 (uncond), against 0.9998949 and 0.9998333 for torch itself with one kernel swapped. Known-answer checks of each kernel pass on that GPU and on llvmpipe, with core validation clean.
DiT block ops and the text refiner's linears run in token ranges of at most DitOptions::submit_gflop (200 GFLOP on the cooperative-matrix path, 20 on the plain one), attention in query ranges, so a slow device's submits stay under the Windows display watchdog; g_fridge_1024's step 0 keeps its bits at 5 GFLOP per submit. vk_api.h compiles against the headers without prototypes and resolves the core functions from libvulkan.so.1 / vulkan-1.dll on first use, so a binary with the backend starts on a machine without Vulkan and says so only when asked for a device.
…onv_out, per-op profile The plain tile reads its operands from shared memory as vectors and prefetches the next step during the current one; outputs of at most 4 channels take one invocation per pixel. --bench times the decoder's heaviest convolutions and --profile times every op in its own submit.
Device gains asynchronous copies on a transfer-only queue family (the copy engine) where the device has one, else on the compute queue: copy_async returns a ticket, gpu_wait makes the next run() wait for it through a semaphore, host_wait blocks until it completed. Buffers they touch are created with `transfer` set and shared concurrently by both families. The text encoder copies layer k + 1 into a second device buffer this way while layer k computes, and reads each layer from the files in 16 MiB chunks on four threads. On an RTX 4090 Laptop that other processes kept busy, a 420 MiB staging copy took 35-125 ms on the transfer queue against 410-780 ms on the compute queue; the g_fridge_1024 encode went from 29 s to 7.2 s, now bound by reading 16 GB from a mostly cold drive at 2.5 GB/s. Outputs are bit-identical; core validation stays clean.
Resolved: one safetensors reader (the text encoder's File and read_index, with Index for whole model directories built on them), the text encoder's npy reader with a writer added, both runtime.h additions, and the text encoder's sources in the library ahead of the tools switch. The loader table gains the fence and semaphore functions its async copies use.
Resolved: the VAE reads through the shared safetensors File (read_f32 added), and its sources join the one library with generate.cpp, so the whole text-to-image path builds.
…ious block computes Two device slots alternate; each streamed block's copy is issued while the block before it runs and waited for on the device by its first op. g_fridge_1024 with a 4 GiB budget (14 of 42 blocks resident): the bits of the resident run, and the copies' share of a step down from about 9.3 s to 0.1-0.2 s on the throttled laptop GPU.
…it runs The CLI requires --tier; the Vulkan backend runs the bf16 model files as the reference tier does and refuses high and turbo. First image: 512x512, 20 steps, seed 7 on an RTX 4090 Laptop, 45.97 dB PSNR against the CUDA backend's image of the same command (mean pixel difference 0.8).
… decode; bands within one dispatch's rows
One workgroup of 64 now takes 16 query rows of a head and walks key tiles of 16 held in shared memory with the Q tile, rows padded against bank conflicts, 25 KiB in all; each invocation keeps 4 output dims of all 16 rows. The online softmax is unchanged except for its step, 16 keys instead of 32. On an RTX 4090 Laptop, attention over the 4096-token golden prompt took 0.71 s for the 8 layers instead of 2.97 s, and the encode 8.4 s instead of 11.5 s. The README gains the text encoder's section: results against the goldens on both paths next to torch's own deviation, and times.
- DiT streaming: with an odd number of streamed blocks, the prefetch of the first one wrapped into the slot the last one was about to run from, so the last block ran on another block's weights. A prefetch into the running block's slot now waits until that block's runs are done, and the slot logic runs once per block around all of its runs. With 27 blocks streamed (a 4.25 GiB budget), g_fridge_1024's step 0 has the resident run's bits. - Elementwise ops dispatch a 2-D grid past 65535 workgroups (the text refiner's SiTU over a 2048+ token prompt). - Buffers past the device's maxBufferSize / maxMemoryAllocationSize are refused with a message; generate() checks the VAE's workspace against them and the free memory before the text encoder runs, and skips the VAE without a decode. - The final latents reach the caller before the decode, so --latent-out survives a failed decode. - Model files open by their wide name on Windows; paths reach the readers as std::filesystem::path. - mossinfer_vulkan compiles with hidden visibility: its entry-point-named pointers are not exported. - The coop path also requires compute subgroup vote and shuffle; attention_coop syncs the subgroup before the rescale overwrites S's staging. - Device::run frees its command buffer when recording throws; Dit's destructor does not throw; --vulkan-device rejects indices below -1 and names a missing one; internal checks no longer exit as usage errors.
…gine's build is unchanged again The frontend chooses the backend, so the fallback is its own program: vulkan/ builds moss-infer-vulkan without CUDA, with the engine's host sources it shares with moss-infer (tokenizer, initial noise, schedule, PNG writer) compiled in, and packages it with CPack (mossinfer-VERSION-PLATFORM-vulkan-ARCH). Its flags are moss-infer generate's that the backend runs; anything else is refused by name with exit status 2; --list-devices says which kernels each device gets. The --backend vulkan integration in moss-infer and cmake/modules/vulkan.cmake are removed: src/ and the engine's CMake match the base again. 512x512, 20 steps, seed 7: 1.96 s/step on the RTX 4090 Laptop.
ci.yml and release.yml call build-vulkan.yml next to the CUDA flavours. The release's package step exempts the Vulkan archives from the -system-cuda pairing and requires both of them; the publish globs already take them.
…d, README, vulkan/README.md)
docs/vulkan.md and vulkan/README.md describe what exists; the notes stay readable in PR mossimage-cpp#8's history.
xycld
changed the title
- CMake finds the Vulkan headers and glslc only (the SDK's Linux archive has no libvulkan.so, and FindVulkan required it; nothing links the loader) and checks that the headers define VK_KHR_shader_bfloat16. - te_attention.comp: NEG_INF is a macro; glslc 2026.3 rejects a global const initialized by uintBitsToFloat. Built with that SDK's glslc and headers, and its validation layer clean on lavapipe for all three checks.
…runs every check Vulkan lets fma round twice, and lavapipe on LLVM 15 (Ubuntu 22.04) does: the fma polynomials of exp and log1p then land up to 17 fp32 ULPs off (about 1e-6 relative) against the 8 the check allowed. The lavapipe step now runs all three checks before failing.
LunarG's server stalled one Windows download for 34 minutes before breaking it off; curl retries, and the SDK is cached per runner image and version.
… devices without bf16 ones Mesa's RADV exposes bf16 fragments only on RDNA4, so RDNA3 and other devices with f16 matrix cores ran the plain kernels, about 5x slower. The coop GEMM and attention shaders now also build with F16 (gemm_f16_coop, attention_f16_coop); BlockPath::kCoopF16 runs them, `--kernels coop-f16` picks it, and the default order is coop, coop-f16, plain. The operands stay bf16 in memory and convert 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 finite value in [2^15, 2^16), where every bf16 value is an f16 value; the GEMM scales rows by it on load and the accumulators back (one combined exponent, two exact halves). Every product is then the bf16 one. Attention converts q, k and v unscaled and feeds P as bf16(p) * 2^15, undone with 1 / row_sum. On the RTX 4090 Laptop: real DiT blocks 98.14 / 96.16 / 98.56 % same bits against the engine, the text encoder's golden cosines and step 0 (0.999976 / 0.999985) equal coop's; 1024x1024 step 8.71 s vs 8.70 s, 512x512 step 2.11 s vs 1.97 s. The VAE runs plain on it (its coop kernels need bf16 fragments). Scaling only A rows was not enough: about 0.1-0.4 % of the weights are below 2^-14, which the driver flushed. Also: - tanh_safe: below |x| = 1/16 the series, as 1 - exp(-2|x|) cancelled to 0 under ~2^-25 (wrong bits for tiny SiTU inputs on every path). - vulkan_block_check: a GEMM case whose rows span 2^-90 to 2^100, judged with per-row floors; "same" compares values. - moss-infer-vulkan: an explicit --kernels the device cannot run is refused before anything loads (exit 2); --list-devices and vulkan_check --list name the kernels a device gets.
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
moss-infer-vulkan: MOSS-Image text to image on Vulkan, a standalone program for GPUs the CUDA engine does not support (AMD, Intel, older NVIDIA cards). It started from CloverIris's cooperative-matrix GEMM bring-up (the first two commits) and now runs the whole pipeline. The frontend chooses which program to run; the CUDA engine (src/,include/, its CMake and presets) is unchanged.vulkan/is its own CMake project, built without CUDA. It compiles in the engine's tokenizer, seed noise, schedule and PNG writer, so a seed gives the same inputs as on CUDA. CPack makesmossinfer-VERSION-linux-vulkan-x86_64.tar.gzandmossinfer-VERSION-windows-vulkan-x86_64.zip. It needs a Vulkan driver, not CUDA or an NVIDIA driver; the loader is opened at run time.moss-infer generatethat the backend runs, with the same defaults, at--tier reference. Anything else (LoRAs, img2img, other samplers, lossy tiers, ...) is refused by name with exit status 2.--list-devicessays which kernels each GPU gets; progress goes to stderr one line per event.--kernelspicks one):coop: bf16VK_KHR_cooperative_matrixGEMMs with the block's epilogues, flash attention, TF32-exact VAE convolutions.coop-f16: the same kernels on f16 fragments, for devices with cooperative matrices but no bf16 ones (RADV gives bf16 only on RDNA4, so RDNA3 gets this path). Each GEMM row of both operands is scaled by a power of two before the exact bf16-to-f16 conversion and scaled back in the epilogue; on the 4090 its blocks, text encoder and DiT step 0 matchcoopexactly. The VAE runsplainon it.plain: register-blocked compute kernels with no cooperative matrices, 16-bit arithmetic or particular subgroup size; it also passes on llvmpipe.build-vulkan.ymlinstalls the Vulkan SDK 1.4.357.0 (pinned, SHA-256 checked), builds both packages, and on Linux runs the backend's kernel checks on lavapipe under the Khronos validation layer.ci.ymlandrelease.ymlcall it; a release requires both Vulkan archives.docs/vulkan.mdfor users,vulkan/README.mdfor the checks and kernels, one sentence in the README. The original research and design notes are dropped from the tree (still in this PR's history).Accuracy
Against the official PyTorch pipeline's golden data (1024x1024,
tests/golden/g_fridge_1024),cooppath:plain: 69-84 %), cosine >= 0.9999999The whole 50-step image keeps the composition and the handwritten Chinese text; its final latent is at cosine 0.974 to the reference (CUDA engine: 0.992). That is the sampler magnifying any difference, not an error: the same Vulkan steps fed torch's own text embeddings instead of the Vulkan encoder's (0.9998 apart) end 0.971 apart from each other.
Not run
plain), socoop-f16has not run on the GPUs it is for.Known limits
plainkernels: thecoopkernels need a subgroup size of 32.