From b0afeca7aee95134101dab11b53a6aea553b4577 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Mon, 2 Nov 2026 08:52:28 +0800 Subject: [PATCH 01/41] docs: Vulkan backend research and design notes 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-design.md | 222 ++++++++++++++++++++++++++++++++ docs/vulkan-backend-research.md | 199 ++++++++++++++++++++++++++++ 2 files changed, 421 insertions(+) create mode 100644 docs/vulkan-backend-design.md create mode 100644 docs/vulkan-backend-research.md diff --git a/docs/vulkan-backend-design.md b/docs/vulkan-backend-design.md new file mode 100644 index 0000000..95c09b9 --- /dev/null +++ b/docs/vulkan-backend-design.md @@ -0,0 +1,222 @@ +# Vulkan backend: design proposal + +Companion to [`vulkan-backend-research.md`](vulkan-backend-research.md) (device audit, prior art, performance model). +Status: proposal; nothing in this document is implemented yet. + +## Scope and goal + +Add a **Vulkan compute backend** that can run the whole MossInfer pipeline on a non-CUDA GPU (or on NVIDIA without the +CUDA toolkit), at a usable fraction of the CUDA speed, without regressing the CUDA path. + +Requirements: + +- Same model files (official bf16 bundle, ComfyUI INT8/INT4 ConvRot files) and the same public API/CLI/server surface. +- CPU side identical: tokenizer, samplers, PNG, planner, reports. +- The backend is selected at run time (or at build time) and the probe reports which one ran. +- Correctness first: every kernel validated against the CUDA implementation on the same weights; end-to-end images + compared statistically (Vulkan numerics differ in accumulation order, like sm_75 does today). + +Non-goals for the first iteration: macOS/MoltenVK, mobile, multi-GPU, Vulkan ray tracing/video, FP8 (not on Turing), +NVFP4 (Blackwell only), and training. + +## Where the CUDA coupling is + +The engine is not CUDA-abstracted today: + +- `Tensor` carries `Device::kCUDA` and a raw `void*`; every kernel takes `cudaStream_t` and uses `cudaMalloc`-style + arenas (`runtime/DeviceArena`). +- `Planner`/`pipeline` call `cudaMemGetInfo`-based free-memory helpers, `cudaEvent`s and cuBLASLt handles, and the + model code calls kernel functions directly (`kernels::Gemm::linear`, `moss_dit::MossDit`, ...). +- ~120 CUDA translation units implement the kernels (GEMM, attention, norm, RoPE, GDN, VAE, quantized formats). + +A Vulkan backend therefore needs two layers: a **device runtime** (buffers, queues, commands, events, memory budget) and +a **kernel set** behind the same call shapes the model code uses. + +## Options + +| Option | What it means | Cost | Risk | +|---|---|---|---| +| **A. Backend abstraction inside the engine** | Introduce `mi::runtime::Backend` (device/buffer/stream/event) and `mi::kernels::Backend` (GEMM, attention, norm, ...) interfaces; CUDA and Vulkan implement them; model/pipeline code talks to the interfaces. | High upfront refactor (~120 TUs touched), but one pipeline for both. | Refactor stalls the CUDA path; hard to land incrementally. | +| **B. Separate Vulkan runner first** | A standalone `moss-vulkan` executable (or tool) that loads the same bundle and runs the pipeline with Vulkan kernels; the engine stays untouched. Merge later through A. | Medium; no risk to CUDA; proves kernels and performance with real weights. | Duplicated pipeline logic until merged; second code path to maintain. | +| **C. CUDA shim over Vulkan** | Implement the CUDA runtime API subset (`cudaMalloc`, `cudaStream_t`, `cuBLASLt` calls) on top of Vulkan so the engine compiles unchanged. | Tempting but the semantics (streams/events, cuBLASLt descriptors, error model) do not map cleanly; CUTLASS kernels are SASS, not SPIR-V, so "compile the same kernels" is impossible. | Dead end for the kernel layer. | + +**Recommendation: B first, then A.** Build a standalone Vulkan runner that exercises the real weights and produces +images; once the kernels and the memory model are proven, introduce the runtime/kernel interfaces (A) and move the +pipeline behind them, with CUDA as the default and the Vulkan runner's kernels reused as the backend implementation. +This keeps `main` shippable at every step and gives measurable checkpoints before the refactor. + +## Architecture (target state) + +``` + pipeline / planner / models (unchanged, backend-agnostic) + | + mi::runtime (Device, Buffer, Stream, Event, Budget) + / \ + runtime/cuda runtime/vulkan + (DeviceArena, StreamSet, ...) (VkDevice, queues, allocator, ...) + | | + kernels (CUDA) kernels (Vulkan) + gemm / attention / norm / rope / gdn / vae same interfaces, SPIR-V shaders +``` + +- `mi::runtime::Device`: allocation (`Buffer` with size/usage/alignment), upload/download, `Stream` (queue + command + buffer), `Event` (timeline semaphore value), `free_budget_bytes()`, `properties()` (name, vendor, subgroup, coopmat + shapes). +- `mi::kernels::*`: free functions in the existing headers, with the CUDA implementation in `src/kernels/cuda/...` and the + Vulkan one in `src/kernels/vulkan/...`; both built into their own targets and linked according to the selected backend. +- The planner's memory model is backend-neutral: it consumes `free_budget_bytes()` and the transfer-rate profile it + already measures. + +## Module layout and build + +``` +cmake/modules/vulkan.cmake option MOSSINFER_VULKAN, find Vulkan (loader + headers), shader compilation +src/runtime/vulkan/ device, queues, allocator, staging ring, pipeline cache, descriptor/BDA helpers +src/kernels/vulkan/ GEMM, attention, norm, rope, gdn, vae, quantized formats (C++ dispatch) +shaders/vulkan/ *.comp GLSL sources; compiled to SPIR-V at build time by glslc +tests/vulkan/ kernel parity tests, end-to-end image tests +tools/vulkan_check.cpp bring-up tool (like tools/sm75_check.cu): device report, GEMM/attention/VAE checks +``` + +- **Shaders**: GLSL 4.60 + `GL_KHR_cooperative_matrix` + `GL_EXT_shader_16bit_storage`/`8bit_storage` + + `GL_EXT_shader_explicit_arithmetic_types_*` + `GL_KHR_shader_subgroup_*`, compiled by `glslc` to SPIR-V at build time + and embedded as `const uint32_t[]` (a generated header), so the binary has no external files. A pipeline cache file is + written next to the model cache and reused. +- **Build**: `-DMOSSINFER_VULKAN=ON` adds the runtime and kernels, and the CLI/server gain `--backend vulkan` (or + `MOSSINFER_BACKEND=vulkan`). Without the option nothing changes, and the package stays as today. A `vulkan` CPack suffix + like the sm75 one can follow once the backend ships. +- **Dependencies**: Vulkan loader and headers only (`vulkan-1.dll`/`libvulkan.so.1` at run time); no Vulkan SDK at run + time. Shader compilation needs `glslc` at build time (vendored or found). + +## Runtime design + +### Device and queues + +- One `VkDevice`, one compute queue (family 0 with `COMPUTE`), one transfer queue if the family differs. +- Single command buffer stream for compute; a second for uploads, synchronized with timeline semaphores. The engine's + CUDA path already separates compute and copy streams (`runtime::StreamSet`), so the abstraction matches. +- `Stream` = a command buffer in recording; `submit(stream)` batches recorded commands; `Event` = timeline semaphore + value; `wait(event)` maps to `vkWaitSemaphores`. + +### Memory + +- A size-class suballocator over a few large `VkDeviceMemory` blocks (256 B alignment, dedicated allocations above + 64 MiB), with `BDA` (buffer device address) so kernels take plain 64-bit device pointers in push constants. +- Heaps: device-local (weights, activations), host-visible coherent (staging ring, 8–16 MiB, mapped), and a + host-cached path for streamed blocks (readback-free upload through a mapped ring). +- Budget: `VK_EXT_memory_budget` (`heap_budget`, `heap_usage`) plus the allocator's own accounting; report + `free_budget_bytes = heap_budget - heap_usage - reserve` to the planner, the same contract `cudaMemGetInfo` provides. +- Weight streaming: the engine's DiT block streaming (5.16 GB INT8 at 256², 9.6 GB bf16) becomes: host-side ring buffer → + transfer queue → device-local block buffers, with timeline semaphores per block and a per-step upload budget. The + planner's existing `HiddenFraction` model (copy hidden behind compute) applies unchanged. + +### Dispatch and command recording + +- Each kernel is a `VkPipeline` (specialization constants for tile sizes/strides, push constants for pointers/offsets). +- The Vulkan runner records the whole block's kernel sequence into one command buffer (norm → modulation → QKV GEMM → + attention → out GEMM → FFN up/gate GEMM → SiTU → down GEMM → residual), so the CPU cost is one submit per block, not + 20–40. `vkCmdPipelineBarrier2` with compute-to-compute barriers is inserted by the runtime (or explicit per kernel). +- Secondary command buffers (or one pre-recorded buffer per shape) make the per-step recording cost negligible; a + block resident in device memory can replay a pre-recorded buffer, comparable to the CUDA `fused_graphs` path. + +## Kernel plan + +Legend: CM = cooperative matrix, SIMT = plain compute shader. + +| Stage | Operation | Shapes (DiT 3072, FFN 12288, VAE 96–384) | Vulkan kernel | +|---|---|---|---| +| DiT block | modulation/norm | [tokens, 3072] | SIMT, vectorized f16 | +| DiT block | QKV / out / FFN up / FFN down / final | 3072×3072, 3072×12288, 12288×3072 | CM f16 16×16×16, f32 acc; epilogues in shader (bias, gated residual) | +| DiT block | joint attention | tokens 4096 (1024²), heads 24, head_dim 128, non-causal varlen | tiled FA: CM QK^T, f32 softmax, SIMT PV or CM PV | +| DiT block | RoPE, SiTU-GLU, residual | elementwise | SIMT | +| Text encoder | RMSNorm, linear, RoPE, causal attention, GDN recurrence | hidden 4096, 32 layers | GEMM CM; attention CM/SIMT; GDN SIMT scan | +| VAE decode | conv2d (NHWC/KRSC), groupnorm, mid attention, upsample | 96/192/384, 64×64→1024² | SIMT f32 conv (like the sm_75 kernels), attention SIMT | +| INT8 path | ConvRot GEMM + dequant | same GEMM shapes | CM s8 16×16×32, s32 acc, dequant epilogue | +| INT4 path | unpack to INT8 then GEMM | same | SIMT unpack + INT8 CM | + +### GEMM tiling (48 KiB shared memory, f16, f32 accumulate) + +- Workgroup 64×64 output tile, 4 warps (2×2), each warp computes 32×32 via four 16×16×16 cooperative-matrix ops. +- K step 32: A tile 64×32 f16 = 4 KiB, B tile 32×64 f16 = 4 KiB; 3-stage pipeline = 24 KiB, leaves room for epilogue + staging and the bias/gate vectors. +- Registers: 32×32 f32 accumulator per warp = 1024 floats = 32 per thread, plus A/B fragments; well within a 255-register + budget at 256 threads/workgroup. +- Split-K for skinny GEMMs (e.g., final layer, modulation) or when M is small (prompt tokens); the engine's CUDA GEMM + already chunks rows, so the Vulkan dispatcher can reuse that policy. +- Epilogues are fused in the shader: bias add, `ffn_down` residual add, gated residual (multiply by gate), SiTU-GLU + (tanh/exp on f32), and the INT8 dequant (`acc * x_scale[row] * w_scale[col] + bias`). + +### Attention kernel + +- One workgroup per (block of 64 queries, head, sequence chunk): 4 warps, BM=64, BN=64, head_dim 128 (DiT) or 256 + (Qwen). +- QK^T on cooperative matrix (f16 inputs, f32 accumulate), softmax in SIMT f32 with the same online algorithm as the + CUDA sm_75 kernel (`src/kernels/sm75/attention_sm75.cu`), PV on cooperative matrix or SIMT. +- Varlen/GQA: per-sequence offsets and head mapping in push constants; the DiT joint attention is non-causal over + [text + image] tokens, Qwen is causal. +- 4096-token attention at 1024² needs 64×64 tiles × 64 query blocks × 24 heads; the same tiling as the CUDA kernel. + +### Quantization formats + +- INT8 ConvRot weights are stored as int8 with per-row (activation) and per-column (weight) f32 scales; the engine's + `Int8Operands` contract carries them. The Vulkan INT8 GEMM uses CM s8 16×16×32 (exact int32 accumulation) and the same + dequant epilogue, so results should be bit-identical to CUTLASS. +- INT4 ConvRot: reuse the CUDA engine's INT4→INT8 unpack (a SIMT kernel) and the INT8 GEMM, as on Hopper/Blackwell and + in the sm_75 port. +- ConvRot (Hadamard rotation) is already folded into the stored weights, so the GEMM kernel is rotation-agnostic. + +## Integration with the engine + +- **Probe**: `src/backend/probe.*` gains a Vulkan path when built with `MOSSINFER_VULKAN`: enumerate physical devices, + report name/vendor/heap/coopmat shapes, recommend `vulkan` when no CUDA device is usable (or when `--backend vulkan`). +- **Planner**: the cost model's rates (`bf16_tflops`, PCIe rate) are learned per backend and per device; the tier rules + are backend-neutral (reference/high/turbo), except that on Vulkan the 8-bit attention modes may not exist. +- **Tier/attention**: Vulkan on Turing-class hardware resolves an unset attention to `exact` (same rule as sm_75); on + devices with f16 cooperative matrix a `q8_fp16` port can come later. +- **CLI/server**: `--backend cuda|vulkan|auto`; the report gains `quality.backend = "vulkan"` and the device string; + `moss-server` passes it through `load`'s new optional `backend` parameter (protocol addition). +- **Packaging**: a `-vulkan` package flavour (loader + shaders embedded); the CUDA package stays the default. + +## Testing and validation + +1. **Kernel parity** (`tests/vulkan/`): for each kernel, run the same inputs through CUDA and Vulkan, compare exact + (INT8), cosine ≥ 0.99999 / max|diff| (f16), and report. Reuses the vectors `tests/data` already stores for the CUDA + kernels. +2. **Block parity**: one DiT block (same weights, same input) CUDA vs Vulkan, cosine on the block output. +3. **End-to-end**: generate the same seed/prompt/size on CUDA and Vulkan; compare PSNR/SSIM and the report; expected + differences are accumulation-order only. +4. **Performance gates**: step time at 256²/512²/1024², per-tier, on the 2080 (CUDA baseline in `docs/sm75.md`) and on + at least one AMD device (RX 6000/7000) and one Intel Arc if available. +5. **Stability**: validation layers (`VK_LAYER_KHRONOS_validation`) in debug builds, pipeline cache round-trip, memory + budget accounting, 100 consecutive generations without leaks. + +## Milestones + +| Milestone | Deliverable | Acceptance | Estimate | +|---|---|---|---| +| **V0** | Probe tool + microbenchmarks (GEMM, bandwidth, dispatch overhead, attention) | Numbers for all five unknowns in the research doc; go/no-go | 2–3 days | +| **V1** | Vulkan runtime (device, allocator, commands, pipeline cache) + f16 GEMM validated vs CPU/CUDA | GEMM cosine ≥ 0.9999 at engine shapes; ≥ 40 % of cuBLAS | 1–2 weeks | +| **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights | Block output matches CUDA block (cosine ≥ 0.999) | 1–2 weeks | +| **V3** | Text encoder (GEMM, RMSNorm, RoPE, causal attention, GDN) and VAE decoder | Encode/decode parity; first full Vulkan image | 2–3 weeks | +| **V4** | Pipeline integration (runtime abstraction, planner, streaming, CLI/server) | `moss-infer --backend vulkan` and `moss-server` produce images on the 2080 | 2–4 weeks | +| **V5** | Quantized paths (INT8/INT4), packaging, CI, docs | `high`/`turbo` tiers on Vulkan; `-vulkan` package; tests in CI | 1–2 weeks | + +The estimates assume one experienced engineer and the existing CUDA code as the specification; the GDN recurrence and +the VAE convolutions are the least predictable parts. + +## Risks and mitigations + +- **Scope explosion.** B-first keeps the engine shippable; V4 (refactor) is gated on V2/V3 parity. +- **Vendor differences.** Cooperative matrix shapes and shared-memory limits vary; the kernel set is selected from a + capability table (NVIDIA first, AMD/Intel fallbacks second). +- **Performance shortfall.** If V0 shows < 40 % of cuBLAS on the 2080, the Vulkan backend is positioned as a + portability backend (AMD/Intel), not a replacement on NVIDIA. +- **Maintenance.** Two kernel sets are expensive; the shared parity tests and the runtime abstraction are the + mitigation. If the backend cannot be maintained, V1–V3 remain a standalone tool. + +## Decisions requested + +1. Confirm the target: portability first (AMD/Intel) or NVIDIA parity first? +2. Approve the two-phase plan (standalone runner, then runtime abstraction) or prefer a full abstraction up front? +3. Choose the first device matrix for validation (2080 only, or also an AMD/Intel card). +4. Decide whether the Vulkan path ships as a separate package flavour or inside the default package (loader dependency). diff --git a/docs/vulkan-backend-research.md b/docs/vulkan-backend-research.md new file mode 100644 index 0000000..6742b22 --- /dev/null +++ b/docs/vulkan-backend-research.md @@ -0,0 +1,199 @@ +# Vulkan backend: research notes + +Status: research only; no Vulkan code is in the engine yet. This document and +[`vulkan-backend-design.md`](vulkan-backend-design.md) are the basis for a decision on whether, and how, MossInfer gets a +Vulkan compute backend. + +## Why a Vulkan backend + +MossInfer is CUDA-only: the pipeline calls cuBLASLt, CUTLASS kernels and CUDA kernels directly, and the probe refuses any +GPU whose compute capability is below 8.0 (7.5 with `MOSSINFER_SM75`). That confines the engine to NVIDIA on Windows and +Linux, with the CUDA toolkit's driver and memory-management model. + +A Vulkan compute backend would give: + +- **Other GPUs.** AMD (RDNA2/3/4), Intel (Arc/Xe), and NVIDIA on systems where the CUDA toolkit is not desirable. The + engine's model code is portable C++/CUDA; only the kernels and the device runtime are CUDA-specific. +- **One packaging story** for Windows and Linux without the CUDA runtime dependency, and a path to Steam Deck (RADV) + and other low-power devices. +- **An escape hatch from CUDA's memory model.** Vulkan exposes `VK_EXT_memory_budget`, explicit heaps and suballocation, + which fits the engine's planner/streaming design; CUDA's `cudaMemGetInfo` is coarser. + +The cost is a second implementation of every tensor-core kernel, plus a device/memory/command abstraction the engine +does not have today (see the design document). + +## What this machine exposes (primary evidence) + +Probed on the target machine: Windows 11, RTX 2080 8 GiB, NVIDIA driver 621.01 (Vulkan driver 621.1.0.0). The Vulkan +loader in `C:\Windows\System32\vulkan-1.dll` reports **instance 1.4.357**; the NVIDIA ICD reports **device API 1.4.357** +and one physical device. + +### Limits that shape kernel design + +| Limit | Value | Consequence | +|---|---|---| +| `maxComputeSharedMemorySize` | 48 KiB | 64 KiB is a CUDA opt-in only; Vulkan tiles must fit 48 KiB (Turing's hardware limit is 64 KiB but Vulkan caps at 48 KiB here). | +| `maxComputeWorkGroupInvocations` | 1024 | Up to 32 warps per workgroup. | +| `subgroupSize` | 32 | Cooperative-matrix `gl_ScopeSubgroup` = one warp. | +| `maxStorageBufferRange` | 2^32-1 bytes | 4 GiB per buffer; weights are 75 MB at most per tensor, so a small number of big buffers is fine. | +| Device-local heap | 7989 MiB | The engine's planner sees the same ~8 GiB. A second device-local heap of 214 MiB exists; host-visible heap is 16 GiB. | + +### Extensions a compute backend needs (all present) + +`VK_KHR_cooperative_matrix` (revision 2), `VK_KHR_16bit_storage`, `VK_KHR_8bit_storage`, `VK_KHR_shader_float16_int8`, +`VK_KHR_shader_integer_dot_product`, `VK_KHR_buffer_device_address`, `VK_EXT_memory_budget`, `VK_KHR_timeline_semaphore`, +`VK_KHR_synchronization2`, `VK_EXT_descriptor_indexing`, `VK_KHR_maintenance4`, `VK_KHR_maintenance5`, +`VK_EXT_subgroup_size_control`, `VK_KHR_pipeline_library`, `VK_EXT_shader_module_identifier`, +`VK_KHR_shader_float_controls` (and `2`), `VK_KHR_vulkan_memory_model`, `VK_KHR_shader_clock`, +`VK_KHR_performance_query`, `VK_NV_cooperative_matrix`, `VK_NV_cooperative_matrix2`, `VK_NV_cooperative_vector`. + +Shader features: `storageBuffer16BitAccess`, `uniformAndStorageBuffer16BitAccess`, `storageBuffer8BitAccess`, +`uniformAndStorageBuffer8BitAccess`, `shaderFloat16`, `shaderInt8`, `shaderBFloat16Type`, +`shaderBFloat16CooperativeMatrix`, `shaderSubgroupClock`, `shaderDeviceClock`, `shaderFmaFloat16`. + +### Cooperative-matrix shapes (queried with `vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR`) + +17 combinations, all `gl_ScopeSubgroup`, non-saturating: + +| Shape (M×N×K) | A | B | C | Result | +|---|---|---|---|---| +| 16×16×16, 16×8×16, 16×8×8 | f16 | f16 | f32 | f32 | +| 16×16×16, 16×8×16, 16×8×8 | f16 | f16 | f16 | f16 | +| 16×16×32, 16×8×32 | s8 | s8 | s32 | s32 | +| 16×16×32, 16×8×32 | s8 | u8 / u8×s8 / u8×u8 | s32 / u32 | s32 / u32 | +| 16×16×16 | bf16 | bf16 | f32 | f32 | + +Notes: + +- **f16 with f32 accumulate is available at 16×16×16 and 16×8×16**, which is the same HMMA shape the CUDA sm_75 port + uses (`m16n8k16`; the 16×8×16 cooperative matrix maps to it). A Vulkan GEMM can therefore reach the same tensor-core + instruction on the same hardware. +- **INT8 at 16×16×32 / 16×8×32 with s32 accumulate** covers the engine's INT8 ConvRot GEMMs (CUTLASS `m8n8k16` on + sm_75 is a different instruction, but the arithmetic is an exact int32 dot product either way). +- **bf16 cooperative matrix is reported by this Turing driver**, but Turing has no bf16 tensor core: the driver must be + emulating it (likely f16 down-conversion or split-f16 math). The engine already converts bf16 weights to fp16 for + sm_75, so the Vulkan backend should keep fp16 as its compute type and treat bf16 as a storage conversion, not rely on + the bf16 cooperative-matrix path. +- `VK_NV_cooperative_matrix` is still exposed (the pre-1.4 NV API); `vkGetPhysicalDeviceCooperativeMatrixPropertiesNV` + returned a table that our probe could not decode reliably (the struct's enums are narrower than the KHR ones). The + KHR query is the authoritative one on this driver; the backend should prefer it and fall back to the NV query only on + older drivers. + +### Memory heaps + +| Heap | Size | Flags | Use | +|---|---|---|---| +| 0 | 7989 MiB | device-local | Model weights, activations | +| 1 | 16347 MiB | host-visible (0x6 = visible + coherent; 0xe adds cached) | Staging ring, streamed blocks, host-cached weights | +| 2 | 214 MiB | device-local | Small warm allocations | + +That matches the CUDA picture the planner already models (8 GiB device + ~16 GiB host), so the existing planner cost +model (bytes per block, PCIe rate, hidden-copy fraction) transfers directly; the Vulkan backend only needs to report the +free device budget from `VK_EXT_memory_budget` instead of `cudaMemGetInfo`. + +## Prior art and what to take from it + +### ggml-vulkan (llama.cpp, whisper.cpp, stable-diffusion.cpp) + +The most complete Vulkan inference backend in the open-source ecosystem. Design points worth copying: + +- **One backend file plus generated shaders.** `ggml-vulkan.cpp` (~20k lines) houses device setup, memory arenas, + pipelines and per-op dispatch; the compute shaders are GLSL templates compiled offline to SPIR-V and embedded (the + build generates a C++ header with the SPIR-V arrays). +- **Specialization constants and push constants.** Workgroup sizes, tile sizes and strides are specialization + constants, so one shader covers many shapes; per-dispatch arguments (pointers via BDA, offsets, strides) go in push + constants, avoiding descriptor-set churn. +- **Cooperative matrix paths.** `coopmat` (KHR 16×16×16) and `coopmat2` (VK_NV_cooperative_matrix2) variants for matrix + multiply and flash attention; a scalar fallback for devices without cooperative matrix. +- **Pipeline cache.** The driver's pipeline cache is serialized to disk and reused; shader compilation is a large part + of a cold start otherwise. +- **Memory arenas.** Buffers are suballocated from a few large `VkDeviceMemory` blocks with size classes and a + buffer-device-address allocator; the engine's streaming design can use the same idea for its block buffers. +- **Flash attention shaders.** Tiled FA with online softmax, GQA support (`coopmat2` shader in llama.cpp #12559) and + separate paths for quantized KV; attention is the piece most likely to be slower than CUDA. + +### stable-diffusion.cpp + +Runs SD1.5/SDXL/Flux families through ggml, with the same Vulkan backend. It demonstrates diffusion inference on Vulkan +end-to-end, including VAE decode and img2img, on consumer GPUs. Its known weaknesses map to ours: performance depends +strongly on the GPU vendor's Vulkan driver, AMD RDNA3/4 is often close to CUDA, small Intel Arc parts are slower, and +attention remains the bottleneck. It is the closest prior art to what MossInfer needs. + +### Others + +- **ncnn Vulkan**: mobile/desktop CNN inference; mature Vulkan kernels (convolution, pooling), good evidence that + Vulkan convolution is practical (relevant for the VAE decoder). +- **Kompute / vuh**: C++ Vulkan compute wrappers; convenient for prototyping but no tensor-core help. +- **IREE HAL / ONNX Runtime**: IREE has a Vulkan HAL; ONNX Runtime's GPU EP is DirectML on Windows, not Vulkan. Neither + offers a drop-in path for a custom engine, but their Vulkan memory/scheduling designs are useful references. +- **wgpu / WebGPU**: a higher-level portability layer with a compute language (WGSL) and no cooperative matrix in the + standard; shader-f16 and subgroup support are optional. It would be easier to port but cannot reach tensor-core + throughput on Turing; not a candidate for the "fast" backend, only for a "works everywhere, slowly" fallback. + +## Performance model + +Measured CUDA numbers on this machine (RTX 2080 8 GiB, SFT bundle), from `docs/sm75.md`: + +| Tier | 256×256 step | Compute per step (approx.) | +|---|---|---| +| reference (bf16→fp16 GEMMs) | 0.80 s | 42 blocks × ~0.015 s + text encoder ~3 s per prompt | +| high (INT8 DiT) | 0.19 s | 42 blocks × ~0.004 s | +| turbo (INT4 via INT8) | 0.22 s | unpack overhead + INT8 | + +Turing's fp16 tensor-core peak on the 2080 is roughly 50–60 TFLOPS (NVIDIA's figures for the card are in that range, +depending on fp16 vs fp32 accumulate); cuBLASLt reaches a large fraction of it on the GEMM shapes of this model +(3072×12288×3072 etc.). A Vulkan cooperative-matrix GEMM with 16×16×16 f16×f16→f32 tiles can use +the same instruction; the realistic outcome is **50–80 % of cuBLAS** for large GEMMs, dropping for small ones (the DiT's +per-block GEMMs are 3072-wide, large enough; the modulation/norm kernels are bandwidth-bound and vendor-independent). + +What that means per step at 256² (INT8 high, 0.19 s): a Vulkan backend at 60 % of the CUDA GEMM throughput would be +~0.30 s/step; the planner's cost model would then choose the Vulkan tier/streaming plan accordingly. The VAE decode is +fp32/conv-heavy and would be the second-largest cost; Vulkan convolution on Turing has no tensor-core help for fp32, so +it stays bandwidth-bound (same as the CUDA fp32 VAE path). + +The critical unknowns to measure in a prototype (M0): + +1. Cooperative-matrix GEMM throughput at the engine's shapes (f16, f32 accumulate, 48 KiB tiles). +2. Dispatch overhead per kernel (engine runs ~20–40 dispatches per block; command-buffer batching must keep the CPU + cost per step below the GPU cost). +3. Attention throughput (fp16 QK^T + softmax + PV, 4096-token joint attention at 1024²). +4. PCIe upload rate for streamed DiT blocks (the CUDA path measures 2.9–4.6 GB/s from file/page cache). +5. fp16 vs bf16 numerics: the engine's reference tier is bf16 storage with fp32 accumulate; converting to fp16 for + Vulkan changes rounding in the same way the sm_75 CUDA path does (documented in `docs/sm75.md`). + +## Numerics and formats + +- **Storage:** the engine's official bundle is bf16; the sm_75 port converts it to fp16 in place, and the Vulkan backend + would do the same (bf16→fp16 is exact inside fp16's range, and the reverse conversion recovers the bf16 bits). + Community INT8/INT4 ConvRot files are already quantized; no conversion is needed for weights. +- **Compute:** f16 inputs with f32 accumulate via cooperative matrix (matching cuBLASLt's `CUBLAS_COMPUTE_32F` use in the + engine's fp16 path); INT8 weights use s32 accumulate (exact, same result as CUTLASS on CUDA); the VAE decoder stays + fp32 (SIMT), since Vulkan exposes no fp32 tensor core. +- **No FP8 on Turing** (`VK_KHR_shader_fp8` is not exposed), so the engine's `q8_fp8` attention path cannot be ported + to this GPU; `q8_fp16` would be the 8-bit attention variant on Vulkan, and `exact` remains the default on sm_75-class + hardware. +- **Denormals:** f16 cooperative matrix on NVIDIA flushes subnormal inputs/outputs? The CUDA sm_75 port relies on fp16 + exactly matching bf16 inside the normal range; the prototype must compare bit-for-bit or cosine against the CUDA path + and document the deviation. + +## Risks + +| Risk | Impact | Mitigation | +|---|---|---| +| Cooperative matrix on AMD/Intel differs (16×16×16 only, no int8) | Kernels must specialize per vendor; Vulkan backend becomes NVIDIA-first | Query properties at startup and select kernels; scalar/DP4A fallbacks | +| Shader compilation time and driver bugs | Cold start and hangs | Precompile SPIR-V at build time, ship pipeline cache, validation layers in CI | +| 48 KiB shared memory | Smaller tiles than CUDA, lower GEMM efficiency | 3-stage 64×64×32 f16 tiles; split-K for skinny shapes | +| Dispatch overhead (CPU-bound at small sizes) | Step time worse than CUDA | Batch dispatches into one command buffer; reuse command buffers; push constants + BDA | +| `VK_EXT_memory_budget` imprecision | Planner over/under-estimates free VRAM | Use budget + own allocation accounting; keep the CUDA planner model | +| bf16 coopmat is emulated on Turing | Wrong performance/numerics assumptions | Compute in f16, store bf16 as f16, never use bf16 coopmat | +| Windows WDDM TDR (2 s) on long kernels | Device reset on big attention/GEMM | Keep dispatches short; split work (the engine already tiles) | +| Two backends to test | Maintenance cost | Shared test vectors and CI matrix; backend-agnostic golden images | + +## Open questions for the design phase + +1. Is the target "NVIDIA parity plus AMD/Intel support", or "works on anything"? That decides how many fallbacks the + kernel layer needs. +2. Should the Vulkan path reuse the engine's existing planner/streaming code through a device abstraction, or start as a + separate `moss-vulkan` runner and merge later? (The design document recommends the latter first.) +3. Which attention variant is acceptable on Vulkan: `exact` only (like sm_75), or a `q8_fp16` port? +4. Is a WebGPU/wgpu fallback worth a third backend, or is Vulkan enough for the non-NVIDIA target? From 957a95f935d9e73261cc9e62f4f44632fdf01dfe Mon Sep 17 00:00:00 2001 From: CloverIris Date: Mon, 2 Nov 2026 11:10:01 +0800 Subject: [PATCH 02/41] vulkan: research notes and a cooperative-matrix GEMM bring-up tool - 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. --- tools/vulkan_gemm/CMakeLists.txt | 44 +++ tools/vulkan_gemm/README.md | 115 ++++++++ tools/vulkan_gemm/gemm_f16.comp | 84 ++++++ tools/vulkan_gemm/main.cpp | 491 +++++++++++++++++++++++++++++++ tools/vulkan_gemm/probe.py | 142 +++++++++ 5 files changed, 876 insertions(+) create mode 100644 tools/vulkan_gemm/CMakeLists.txt create mode 100644 tools/vulkan_gemm/README.md create mode 100644 tools/vulkan_gemm/gemm_f16.comp create mode 100644 tools/vulkan_gemm/main.cpp create mode 100644 tools/vulkan_gemm/probe.py diff --git a/tools/vulkan_gemm/CMakeLists.txt b/tools/vulkan_gemm/CMakeLists.txt new file mode 100644 index 0000000..7c82eb1 --- /dev/null +++ b/tools/vulkan_gemm/CMakeLists.txt @@ -0,0 +1,44 @@ +# Standalone build for the Vulkan cooperative-matrix GEMM bring-up tool. It is deliberately separate from the engine's +# CMake project: no engine code is linked, and nothing here affects a normal build. See README.md. +cmake_minimum_required(VERSION 3.20) +project(mossinfer_vulkan_gemm LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# Prefer the Vulkan SDK's CMake package; fall back to explicit paths for machines without the SDK (the loader's import +# library can be generated from vulkan-1.dll, see README.md). +find_package(Vulkan QUIET) +set(VULKAN_INCLUDE_DIR "" CACHE PATH "Vulkan include directory when no SDK is installed") +set(VULKAN_LIBRARY "" CACHE FILEPATH "Vulkan loader import library when no SDK is installed") + +if(VULKAN_INCLUDE_DIR AND VULKAN_LIBRARY) + set(VULKAN_INCLUDE ${VULKAN_INCLUDE_DIR}) + set(VULKAN_TARGET ${VULKAN_LIBRARY}) +elseif(Vulkan_FOUND) + set(VULKAN_INCLUDE ${Vulkan_INCLUDE_DIRS}) + set(VULKAN_TARGET Vulkan::Vulkan) +else() + message(FATAL_ERROR + "Vulkan not found: install the Vulkan SDK, or pass -DVULKAN_INCLUDE_DIR= -DVULKAN_LIBRARY=") +endif() + +find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN_SDK}/bin") +if(NOT GLSL_COMPILER) + message(FATAL_ERROR "glslc not found: it ships with the Vulkan SDK") +endif() + +set(GERM_SPV ${CMAKE_CURRENT_BINARY_DIR}/gemm_f16.spv) +add_custom_command( + OUTPUT ${GERM_SPV} + COMMAND ${GLSL_COMPILER} --target-env=vulkan1.3 -fshader-stage=comp -O ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp + -o ${GERM_SPV} + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp + COMMENT "glslc gemm_f16.comp") +add_custom_target(gemm_spv ALL DEPENDS ${GERM_SPV}) + +add_executable(vulkan_gemm main.cpp) +target_include_directories(vulkan_gemm PRIVATE ${VULKAN_INCLUDE}) +target_link_libraries(vulkan_gemm PRIVATE ${VULKAN_TARGET}) +target_compile_definitions(vulkan_gemm PRIVATE "DEFAULT_SPV_PATH=\"${GERM_SPV}\"") +add_dependencies(vulkan_gemm gemm_spv) diff --git a/tools/vulkan_gemm/README.md b/tools/vulkan_gemm/README.md new file mode 100644 index 0000000..e1a0098 --- /dev/null +++ b/tools/vulkan_gemm/README.md @@ -0,0 +1,115 @@ +# Vulkan GEMM bring-up + +The first artefact of the Vulkan backend work ([research](../docs/vulkan-backend-research.md), +[design](../docs/vulkan-backend-design.md)): a standalone tool that runs one f16 GEMM on a Vulkan device with +cooperative matrices, checks it against a CPU reference and times it at the engine's matrix shapes. + +It links no engine code and is not part of the engine build. The point is to answer the two questions the design +document leaves open before any integration work starts: + +1. Does `VK_KHR_cooperative_matrix` give a correct f16 GEMM with f32 accumulation on this device? +2. How fast is it, compared with the cuBLASLt path the engine uses? + +## What the kernel does + +`gemm_f16.comp` computes + +``` +C[M,N] = A[M,K] * B[K,N] +``` + +with `A` row-major (stride `K`) and `B` column-major (stride `K`), i.e. `B[k,n]` lives at `n*K + k`. That is exactly the +engine's weight layout (`[N,K]` row-major), so the GEMM consumes weights as stored, with no transpose. + +- One workgroup per 64x64 output tile, 128 threads (four 32-wide subgroups). +- Each subgroup owns a 32x32 quarter, computed as four `16x16x16` cooperative-matrix operations (`f16` inputs, `f32` + accumulate) per K step, 16 elements at a time. +- Output tiles are walked in bands of 8 N-tiles (the same idea as the engine's `BandedSwizzle` for INT8 GEMMs): a band + of weight tiles stays in L2 while every M tile streams past it. +- Arguments (buffer device addresses, sizes, strides) go through push constants; the buffers are addressed with + `GL_EXT_buffer_reference`, so there are no descriptor sets. + +## Build + +With the Vulkan SDK installed: + +```sh +cmake -S tools/vulkan_gemm -B build-vulkan-gemm -DCMAKE_BUILD_TYPE=Release +cmake --build build-vulkan-gemm +``` + +The SDK's `glslc` compiles the shader; the build fails with a clear message when either the SDK or `glslc` is missing. + +Without the SDK (for example a CI machine or a user who only has the driver), pass the two paths explicitly: + +```sh +cmake -S tools/vulkan_gemm -B build-vulkan-gemm \ + -DVULKAN_INCLUDE_DIR=/include \ + -DVULKAN_LIBRARY= +``` + +Two notes from a machine without administrator rights (where the SDK installer cannot run): + +- Headers: the [Vulkan-Headers](https://github.com/KhronosGroup/Vulkan-Headers) repository's `include/vulkan` is enough. + Its `vulkan_core.h` includes `vk_video/*`, which is only in the SDK; a local copy either adds those headers or + defines the handful of `StdVideo*` types the core header names (the tool never touches video). +- Import library: `dumpbin /exports C:\Windows\System32\vulkan-1.dll`, keep the `vk*` names in a `.def` file, and + `lib /def:vulkan-1.def /machine:x64 /out:vulkan-1.lib`. +- `glslc` can be extracted from the SDK installer with 7-Zip (`7z x VulkanSDK.exe -o`); it lands in `Bin/`. + +## Run + +```sh +build-vulkan-gemm/vulkan_gemm # check 64^3, then the engine shapes +build-vulkan-gemm/vulkan_gemm --bench 3072 3072 3072 10 +build-vulkan-gemm/vulkan_gemm --bench 3072 12288 3072 5 --no-check +``` + +`M` and `N` must be multiples of 64 and `K` a multiple of 16 (the kernel has no bounds checks yet). `--spv` selects a +different SPIR-V file. The CPU reference is `O(M*N*K)`, so it only runs for shapes up to 2^26 multiply-adds (64³, not +512³ or the engine shapes); `--no-check` disables it. + +Output on an RTX 2080 8 GiB, NVIDIA driver 621.01, Vulkan 1.4.357: + +``` +device: NVIDIA GeForce RTX 2080 (Vulkan 1.4.357) +features: shaderFloat16=1 storage16=1 storage8=1 bda=1 coopmat=1 +cooperative matrix shapes (17): + 16x16x16 A=f16 B=f16 C=f16 R=f16 sat=0 + 16x8x16 A=f16 B=f16 C=f16 R=f16 sat=0 + ... + 16x16x16 A=f16 B=f16 C=f32 R=f32 sat=0 <- the shape the kernel uses + ... + 16x16x32 A=s8 B=s8 C=s32 R=s32 sat=0 <- the INT8 path, later + 16x16x16 A=bf16 B=bf16 C=f32 R=f32 sat=0 <- driver-emulated on Turing; not used +check 64x64x64: max|C - CPU| = 1.01342e-05 +gemm 64x64x64: 0.102 ms 0.0 TFLOPS +gemm 512x512x512: 0.038 ms 7.0 TFLOPS +gemm 3072x3072x3072: 6.349 ms 9.1 TFLOPS +gemm 3072x12288x3072: 24.712 ms 9.4 TFLOPS +gemm 3072x3072x12288: 49.082 ms 4.7 TFLOPS +``` + +The small shapes mostly measure dispatch overhead (one dispatch is ~0.1 ms here); the large ones are what matters. + +## Status and limitations + +- **Correctness**: the 64³ case matches a double-precision CPU reference to `1e-5` (fp16 rounding of the inputs and the + f32 accumulator), which is the expected error for this data. The same kernel runs the engine's shapes. +- **Performance**: ~8-9 TFLOPS at 3072³, roughly 10-15 % of the card's fp16 tensor peak. The kernel is memory-bound: + the cooperative-matrix fragments load straight from global memory, so every subgroup issues strided per-lane loads and + each A/B element is fetched by several workgroups. Making the tiles larger (64x128) did not help, which confirmed the + diagnosis. +- **Next step**: stage A and B tiles in shared memory with vectorised coalesced copies and double buffering, then load + the fragments from shared memory. The 48 KiB shared-memory limit and the 64x64x32 tile staging are worked out in the + design document; the expected result is 20-30 TFLOPS, which would put a Vulkan DiT step within 2x of the CUDA path. +- No INT8/INT4 path yet, no bias/epilogue fusion, no multi-queue overlap, and one queue only. + +## Files + +| File | Purpose | +|---|---| +| `gemm_f16.comp` | the cooperative-matrix GEMM (GLSL 460, compiled by `glslc`) | +| `main.cpp` | device setup, buffers, pipeline, CPU check, timing | +| `probe.py` | optional: lists the device's cooperative-matrix shapes and memory heaps using only `vulkan-1.dll` (no SDK) | +| `CMakeLists.txt` | standalone build | diff --git a/tools/vulkan_gemm/gemm_f16.comp b/tools/vulkan_gemm/gemm_f16.comp new file mode 100644 index 0000000..5385d32 --- /dev/null +++ b/tools/vulkan_gemm/gemm_f16.comp @@ -0,0 +1,84 @@ +// C[M,N] = A[M,K] * B[K,N] with A row-major (stride lda) and B column-major (stride ldb == K), i.e. B[k,n] lives at +// n*ldb + k: exactly the engine's [N,K] weight layout, so the GEMM consumes weights as stored (no transpose). +// +// One workgroup per 64x64 output tile, 128 threads (4 subgroups of 32). Each subgroup owns a 32x32 quarter, computed as +// four 16x16x16 cooperative-matrix operations (f16 inputs, f32 accumulate) per K step. Output tiles are walked in bands +// of BAND N-tiles (the same idea as the engine's BandedSwizzle for INT8 GEMMs): a band of weight tiles stays in L2 +// while every M tile streams past it. +// +// The fragments load straight from global memory in this version. That is the correctness baseline; the shared-memory +// staged version (vectorised coalesced copies, double buffering) is the next step and is expected to be 3-4x faster. +#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 = 2) buffer F16Buf { float16_t d[]; }; +layout(buffer_reference, std430, buffer_reference_align = 4) buffer F32Buf { 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 BAND = 8u; + +void main() { + F16Buf A = F16Buf(pc.a); + F16Buf B = F16Buf(pc.b); + F32Buf C = F32Buf(pc.c); + + 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_tile = pos / width; + const uint n_tile = band * BAND + (pos - m_tile * width); + + const uint warp = gl_SubgroupID; + const uint m0 = m_tile * BM + (warp >> 1) * 32u; + const uint n0 = n_tile * BN + (warp & 1u) * 32u; + + coopmat fa0, fa1; + coopmat fb0, fb1; + coopmat acc00, acc01, acc10, acc11; + const float zero = 0.0; + acc00 = coopmat(zero); + acc01 = acc00; + acc10 = acc00; + acc11 = acc00; + + for (uint k = 0u; k < pc.K; k += 16u) { + coopMatLoad(fa0, A.d, m0 * pc.lda + k, pc.lda, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(fa1, A.d, (m0 + 16u) * pc.lda + k, pc.lda, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(fb0, B.d, n0 * pc.ldb + k, pc.ldb, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(fb1, B.d, (n0 + 16u) * pc.ldb + k, pc.ldb, gl_CooperativeMatrixLayoutColumnMajor); + acc00 = coopMatMulAdd(fa0, fb0, acc00); + acc01 = coopMatMulAdd(fa0, fb1, acc01); + acc10 = coopMatMulAdd(fa1, fb0, acc10); + acc11 = coopMatMulAdd(fa1, fb1, acc11); + } + + coopMatStore(acc00, C.d, m0 * pc.ldc + n0, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc01, C.d, m0 * pc.ldc + n0 + 16u, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc10, C.d, (m0 + 16u) * pc.ldc + n0, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); + coopMatStore(acc11, C.d, (m0 + 16u) * pc.ldc + n0 + 16u, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); +} diff --git a/tools/vulkan_gemm/main.cpp b/tools/vulkan_gemm/main.cpp new file mode 100644 index 0000000..02b4005 --- /dev/null +++ b/tools/vulkan_gemm/main.cpp @@ -0,0 +1,491 @@ +// Vulkan cooperative-matrix GEMM bring-up: one f16 GEMM kernel, checked against a CPU reference and timed at the +// engine's matrix shapes. This is the first artefact of the Vulkan backend work in docs/vulkan-backend-*.md; it is a +// standalone tool (no engine code is linked), so it can run before any backend integration exists. +// +// Kernel contract: C[M,N] = A[M,K] * B[K,N] with A row-major (stride K) and B column-major (stride K), i.e. B[k,n] lives +// at n*K + k, which is exactly the engine's [N,K] weight layout. That lets the GEMM consume weights as stored, with no +// transposes. M and N must be multiples of 64 (the tool checks). +// +// Usage: vulkan_gemm [--spv ] [--bench M N K iters]... [--no-check] +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef DEFAULT_SPV_PATH +#define DEFAULT_SPV_PATH "gemm_f16.spv" +#endif + +namespace { + +constexpr uint32_t kTile = 64; // BM == BN, one workgroup per tile + +[[noreturn]] void Fail(const char* what, VkResult r) { + std::printf("%s failed: VkResult %d\n", what, static_cast(r)); + std::exit(1); +} + +void Check(VkResult r, const char* what) { + if (r != VK_SUCCESS) Fail(what, r); +} + +float HalfToFloat(uint16_t h) { + const uint32_t sign = static_cast(h >> 15) << 31; + const uint32_t exp = (h >> 10) & 0x1f; + const uint32_t mant = h & 0x3ff; + uint32_t bits = 0; + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + uint32_t e = 127 - 15 + 1; + uint32_t m = mant; + while ((m & 0x400) == 0) { + m <<= 1; + --e; + } + bits = sign | ((e - 1) << 23) | ((m & 0x3ff) << 13); + } + } else if (exp == 31) { + bits = sign | 0x7f800000u | (mant << 13); + } else { + bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); + } + float f = 0.0f; + std::memcpy(&f, &bits, 4); + return f; +} + +uint16_t FloatToHalf(float f) { + uint32_t x = 0; + std::memcpy(&x, &f, 4); + const uint32_t sign = (x >> 16) & 0x8000u; + const int32_t exp = static_cast((x >> 23) & 0xff) - 127 + 15; + uint32_t mant = x & 0x7fffffu; + if (exp <= 0) { + if (exp < -10) return static_cast(sign); + mant = (mant | 0x800000u) >> (1 - exp); + return static_cast(sign | ((mant + 0x1000u) >> 13)); + } + if (exp >= 31) return static_cast(sign | 0x7c00u); + return static_cast(sign | (static_cast(exp) << 10) | ((mant + 0x1000u) >> 13)); +} + +struct Buffer { + VkBuffer buffer = VK_NULL_HANDLE; + VkDeviceMemory memory = VK_NULL_HANDLE; + VkDeviceSize size = 0; + VkDeviceAddress address = 0; + void* mapped = nullptr; +}; + +// Everything the tool needs from one device. Not a general runtime: the backend's runtime (buffers, arenas, queues) +// replaces this when the engine integration starts (docs/vulkan-backend-design.md). +struct Device { + VkInstance instance = VK_NULL_HANDLE; + VkPhysicalDevice physical = VK_NULL_HANDLE; + VkDevice device = VK_NULL_HANDLE; + VkQueue queue = VK_NULL_HANDLE; + uint32_t queue_family = 0; + VkCommandPool pool = VK_NULL_HANDLE; + VkPhysicalDeviceMemoryProperties memory{}; + VkPhysicalDeviceProperties properties{}; +}; + +uint32_t FindMemoryType(const Device& d, uint32_t type_bits, VkMemoryPropertyFlags flags) { + for (uint32_t i = 0; i < d.memory.memoryTypeCount; ++i) { + if ((type_bits & (1u << i)) != 0 && (d.memory.memoryTypes[i].propertyFlags & flags) == flags) return i; + } + std::printf("no memory type with flags 0x%x\n", static_cast(flags)); + std::exit(1); +} + +Buffer CreateBuffer(Device& d, VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags flags, + bool device_address) { + Buffer b; + b.size = size; + VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; + info.size = size; + info.usage = usage; + if (device_address) info.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; + info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; + Check(vkCreateBuffer(d.device, &info, nullptr, &b.buffer), "vkCreateBuffer"); + + VkMemoryRequirements requirements{}; + vkGetBufferMemoryRequirements(d.device, b.buffer, &requirements); + VkMemoryAllocateFlagsInfo allocate_flags{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO}; + allocate_flags.flags = device_address ? VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT : 0; + VkMemoryAllocateInfo allocate{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; + allocate.pNext = &allocate_flags; + allocate.allocationSize = requirements.size; + allocate.memoryTypeIndex = FindMemoryType(d, requirements.memoryTypeBits, flags); + Check(vkAllocateMemory(d.device, &allocate, nullptr, &b.memory), "vkAllocateMemory"); + Check(vkBindBufferMemory(d.device, b.buffer, b.memory, 0), "vkBindBufferMemory"); + if ((flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0) { + Check(vkMapMemory(d.device, b.memory, 0, VK_WHOLE_SIZE, 0, &b.mapped), "vkMapMemory"); + } + if (device_address) { + VkBufferDeviceAddressInfo info{VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO}; + info.buffer = b.buffer; + b.address = vkGetBufferDeviceAddress(d.device, &info); + } + return b; +} + +void DestroyBuffer(Device& d, Buffer& b) { + if (b.buffer != VK_NULL_HANDLE) vkDestroyBuffer(d.device, b.buffer, nullptr); + if (b.memory != VK_NULL_HANDLE) vkFreeMemory(d.device, b.memory, nullptr); + b = Buffer{}; +} + +void Submit(Device& d, const std::function& record) { + VkCommandBufferAllocateInfo allocate{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; + allocate.commandPool = d.pool; + allocate.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + allocate.commandBufferCount = 1; + VkCommandBuffer command_buffer = VK_NULL_HANDLE; + Check(vkAllocateCommandBuffers(d.device, &allocate, &command_buffer), "vkAllocateCommandBuffers"); + VkCommandBufferBeginInfo begin{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + begin.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; + Check(vkBeginCommandBuffer(command_buffer, &begin), "vkBeginCommandBuffer"); + record(command_buffer); + Check(vkEndCommandBuffer(command_buffer), "vkEndCommandBuffer"); + VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO}; + submit.commandBufferCount = 1; + submit.pCommandBuffers = &command_buffer; + Check(vkQueueSubmit(d.queue, 1, &submit, VK_NULL_HANDLE), "vkQueueSubmit"); + Check(vkQueueWaitIdle(d.queue), "vkQueueWaitIdle"); + vkFreeCommandBuffers(d.device, d.pool, 1, &command_buffer); +} + +void Upload(Device& d, Buffer& dst, const void* data, VkDeviceSize size) { + Buffer staging = CreateBuffer(d, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, false); + std::memcpy(staging.mapped, data, static_cast(size)); + Submit(d, [&](VkCommandBuffer cb) { + VkBufferCopy copy{0, 0, size}; + vkCmdCopyBuffer(cb, staging.buffer, dst.buffer, 1, ©); + }); + DestroyBuffer(d, staging); +} + +void Download(Device& d, Buffer& src, void* data, VkDeviceSize size) { + Buffer staging = CreateBuffer(d, size, VK_BUFFER_USAGE_TRANSFER_DST_BIT, + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, false); + Submit(d, [&](VkCommandBuffer cb) { + VkBufferCopy copy{0, 0, size}; + vkCmdCopyBuffer(cb, src.buffer, staging.buffer, 1, ©); + }); + std::memcpy(data, staging.mapped, static_cast(size)); + DestroyBuffer(d, staging); +} + +std::vector ReadFile(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + std::printf("cannot open %s\n", path.c_str()); + std::exit(1); + } + const auto size = static_cast(file.tellg()); + if (size == 0 || size % 4 != 0) { + std::printf("%s is not a SPIR-V blob\n", path.c_str()); + std::exit(1); + } + std::vector words(size / 4); + file.seekg(0); + file.read(reinterpret_cast(words.data()), static_cast(size)); + return words; +} + +const char* ComponentName(int type) { + switch (type) { + case 0: return "f16"; + case 1: return "f32"; + case 2: return "f64"; + case 3: return "s8"; + case 4: return "s16"; + case 5: return "s32"; + case 6: return "s64"; + case 7: return "u8"; + case 8: return "u16"; + case 9: return "u32"; + case 10: return "u64"; + case 1000141000: return "bf16"; + default: return "?"; + } +} +void PrintCooperativeMatrixShapes(Device& d) { + auto query = reinterpret_cast( + vkGetInstanceProcAddr(d.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); + if (query == nullptr) return; + uint32_t count = 0; + if (query(d.physical, &count, nullptr) != VK_SUCCESS || count == 0) return; + std::vector shapes(count); + for (auto& s : shapes) s.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; + if (query(d.physical, &count, shapes.data()) != VK_SUCCESS) return; + std::printf("cooperative matrix shapes (%u):\n", count); + for (const auto& s : shapes) { + std::printf(" %ux%ux%u A=%s B=%s C=%s R=%s sat=%u\n", s.MSize, s.NSize, s.KSize, + ComponentName(static_cast(s.AType)), ComponentName(static_cast(s.BType)), + ComponentName(static_cast(s.CType)), ComponentName(static_cast(s.ResultType)), + s.saturatingAccumulation); + } +} + +Device OpenDevice() { + Device d; + VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; + app.pApplicationName = "mossinfer-vulkan-gemm"; + app.apiVersion = VK_API_VERSION_1_3; + VkInstanceCreateInfo instance_info{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; + instance_info.pApplicationInfo = &app; + Check(vkCreateInstance(&instance_info, nullptr, &d.instance), "vkCreateInstance"); + + uint32_t count = 0; + Check(vkEnumeratePhysicalDevices(d.instance, &count, nullptr), "vkEnumeratePhysicalDevices"); + if (count == 0) { + std::printf("no Vulkan physical device\n"); + std::exit(1); + } + std::vector devices(count); + Check(vkEnumeratePhysicalDevices(d.instance, &count, devices.data()), "vkEnumeratePhysicalDevices"); + d.physical = devices[0]; + vkGetPhysicalDeviceProperties(d.physical, &d.properties); + vkGetPhysicalDeviceMemoryProperties(d.physical, &d.memory); + + uint32_t families = 0; + vkGetPhysicalDeviceQueueFamilyProperties(d.physical, &families, nullptr); + std::vector queue_families(families); + vkGetPhysicalDeviceQueueFamilyProperties(d.physical, &families, queue_families.data()); + bool found = false; + for (uint32_t i = 0; i < families; ++i) { + if ((queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) != 0) { + d.queue_family = i; + found = true; + break; + } + } + if (!found) { + std::printf("no compute queue\n"); + std::exit(1); + } + + VkPhysicalDeviceShaderFloat16Int8Features float16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES}; + VkPhysicalDevice16BitStorageFeatures storage16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES}; + VkPhysicalDevice8BitStorageFeatures storage8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_8BIT_STORAGE_FEATURES}; + VkPhysicalDeviceBufferDeviceAddressFeatures address{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR coop{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; + features.pNext = &float16; + float16.pNext = &storage16; + storage16.pNext = &storage8; + storage8.pNext = &address; + address.pNext = &coop; + vkGetPhysicalDeviceFeatures2(d.physical, &features); + + std::printf("device: %s (Vulkan %u.%u.%u)\n", d.properties.deviceName, + VK_API_VERSION_MAJOR(d.properties.apiVersion), VK_API_VERSION_MINOR(d.properties.apiVersion), + VK_API_VERSION_PATCH(d.properties.apiVersion)); + std::printf("features: shaderFloat16=%u storage16=%u storage8=%u bda=%u coopmat=%u\n", float16.shaderFloat16, + storage16.storageBuffer16BitAccess, storage8.storageBuffer8BitAccess, address.bufferDeviceAddress, + coop.cooperativeMatrix); + if (float16.shaderFloat16 == 0 || storage16.storageBuffer16BitAccess == 0 || address.bufferDeviceAddress == 0 || + coop.cooperativeMatrix == 0) { + std::printf("the device lacks a feature the kernel needs\n"); + std::exit(1); + } + float16.shaderFloat16 = VK_TRUE; + float16.shaderInt8 = VK_TRUE; + storage16.storageBuffer16BitAccess = VK_TRUE; + storage16.uniformAndStorageBuffer16BitAccess = VK_TRUE; + storage8.storageBuffer8BitAccess = VK_TRUE; + address.bufferDeviceAddress = VK_TRUE; + coop.cooperativeMatrix = VK_TRUE; + + const char* extensions[] = {"VK_KHR_cooperative_matrix"}; + const float priority = 1.0f; + VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; + queue_info.queueFamilyIndex = d.queue_family; + queue_info.queueCount = 1; + queue_info.pQueuePriorities = &priority; + VkDeviceCreateInfo device_info{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO}; + device_info.pNext = &features; + device_info.queueCreateInfoCount = 1; + device_info.pQueueCreateInfos = &queue_info; + device_info.enabledExtensionCount = 1; + device_info.ppEnabledExtensionNames = extensions; + Check(vkCreateDevice(d.physical, &device_info, nullptr, &d.device), "vkCreateDevice"); + vkGetDeviceQueue(d.device, d.queue_family, 0, &d.queue); + + VkCommandPoolCreateInfo pool_info{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; + pool_info.queueFamilyIndex = d.queue_family; + Check(vkCreateCommandPool(d.device, &pool_info, nullptr, &d.pool), "vkCreateCommandPool"); + return d; +} + +// One GEMM run: fills A and B with deterministic data, optionally checks a 64x64 case against a CPU reference, +// then times `iterations` dispatches. +void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M, uint32_t N, uint32_t K, + int iterations, bool check) { + const VkDeviceSize a_bytes = static_cast(M) * K * 2; + const VkDeviceSize b_bytes = static_cast(N) * K * 2; + const VkDeviceSize c_bytes = static_cast(M) * N * 4; + + std::vector a(static_cast(M) * K); + std::vector b(static_cast(N) * K); + std::mt19937 rng(1234); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (auto& v : a) v = FloatToHalf(dist(rng)); + for (auto& v : b) v = FloatToHalf(dist(rng)); + + Buffer buffer_a = CreateBuffer(d, a_bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); + Buffer buffer_b = CreateBuffer(d, b_bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); + Buffer buffer_c = CreateBuffer(d, c_bytes, + VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | + VK_BUFFER_USAGE_TRANSFER_SRC_BIT, + VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); + Upload(d, buffer_a, a.data(), a_bytes); + Upload(d, buffer_b, b.data(), b_bytes); + + struct Push { + uint64_t a; + uint64_t b; + uint64_t c; + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t lda; + uint32_t ldb; + uint32_t ldc; + } push{buffer_a.address, buffer_b.address, buffer_c.address, M, N, K, K, K, N}; + + const auto record = [&](VkCommandBuffer cb) { + vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); + vkCmdPushConstants(cb, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push), &push); + for (int i = 0; i < iterations; ++i) vkCmdDispatch(cb, N / kTile, M / kTile, 1); + }; + + Submit(d, record); // warm-up (compiles or loads the pipeline cache) + const auto start = std::chrono::steady_clock::now(); + Submit(d, record); + const double seconds = + std::chrono::duration(std::chrono::steady_clock::now() - start).count() / static_cast(iterations); + + // The CPU reference is O(M*N*K); only small shapes can afford it. + if (check && static_cast(M) * N * K <= (1ull << 26)) { + std::vector c(static_cast(M) * N); + Download(d, buffer_c, c.data(), c_bytes); + double max_abs = 0.0; + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + double reference = 0.0; + for (uint32_t k = 0; k < K; ++k) { + reference += static_cast(HalfToFloat(a[static_cast(m) * K + k])) * + static_cast(HalfToFloat(b[static_cast(n) * K + k])); + } + max_abs = std::max(max_abs, std::fabs(static_cast(c[static_cast(m) * N + n]) - reference)); + } + } + std::printf("check %ux%ux%u: max|C - CPU| = %.6g\n", M, N, K, max_abs); + } + + const double flops = 2.0 * M * N * K; + std::printf("gemm %ux%ux%u: %8.3f ms %7.1f TFLOPS\n", M, N, K, seconds * 1e3, flops / seconds / 1e12); + + DestroyBuffer(d, buffer_a); + DestroyBuffer(d, buffer_b); + DestroyBuffer(d, buffer_c); +} + +} // namespace + +int main(int argc, char** argv) { + std::string spv = DEFAULT_SPV_PATH; + bool check = true; + struct Shape { + uint32_t M, N, K; + int iterations; + }; + std::vector shapes; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--spv" && i + 1 < argc) { + spv = argv[++i]; + } else if (arg == "--no-check") { + check = 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 [--spv ] [--no-check] [--bench M N K iters]...\n", argv[0]); + return 2; + } + } + if (shapes.empty()) { + shapes = {{64, 64, 64, 1}, {512, 512, 512, 20}, {3072, 3072, 3072, 10}, {3072, 12288, 3072, 5}, + {3072, 3072, 12288, 5}}; + } + for (const Shape& s : shapes) { + if (s.M % kTile != 0 || s.N % kTile != 0 || s.K % 16 != 0) { + std::printf("shapes must be multiples of %u (M, N) and 16 (K): %ux%ux%u\n", kTile, s.M, s.N, s.K); + return 2; + } + } + + Device d = OpenDevice(); + PrintCooperativeMatrixShapes(d); + + const std::vector spirv = ReadFile(spv); + VkShaderModuleCreateInfo module_info{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; + module_info.codeSize = spirv.size() * 4; + module_info.pCode = spirv.data(); + VkShaderModule module = VK_NULL_HANDLE; + Check(vkCreateShaderModule(d.device, &module_info, nullptr, &module), "vkCreateShaderModule"); + + VkPushConstantRange push_range{}; + push_range.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT; + push_range.offset = 0; + push_range.size = 48; + VkPipelineLayoutCreateInfo layout_info{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; + layout_info.pushConstantRangeCount = 1; + layout_info.pPushConstantRanges = &push_range; + VkPipelineLayout layout = VK_NULL_HANDLE; + Check(vkCreatePipelineLayout(d.device, &layout_info, nullptr, &layout), "vkCreatePipelineLayout"); + + VkComputePipelineCreateInfo pipeline_info{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + pipeline_info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; + pipeline_info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT; + pipeline_info.stage.module = module; + pipeline_info.stage.pName = "main"; + pipeline_info.layout = layout; + VkPipeline pipeline = VK_NULL_HANDLE; + Check(vkCreateComputePipelines(d.device, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &pipeline), + "vkCreateComputePipelines"); + + for (const Shape& s : shapes) RunGemm(d, pipeline, layout, s.M, s.N, s.K, s.iterations, check); + + vkDestroyPipeline(d.device, pipeline, nullptr); + vkDestroyPipelineLayout(d.device, layout, nullptr); + vkDestroyShaderModule(d.device, module, nullptr); + vkDestroyCommandPool(d.device, d.pool, nullptr); + vkDestroyDevice(d.device, nullptr); + vkDestroyInstance(d.instance, nullptr); + return 0; +} diff --git a/tools/vulkan_gemm/probe.py b/tools/vulkan_gemm/probe.py new file mode 100644 index 0000000..b7366ab --- /dev/null +++ b/tools/vulkan_gemm/probe.py @@ -0,0 +1,142 @@ +"""Lists a Vulkan device's cooperative-matrix shapes and memory heaps, using only vulkan-1.dll. + +The C++ tool needs Vulkan headers and an import library; this script needs neither (ctypes only), so it works on a +machine with just the GPU driver installed. It is how the capability table in docs/vulkan-backend-research.md was +produced. + + 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 +VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_NV = 1000249001 + +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(): + try: + vk = C.WinDLL("vulkan-1.dll") + except OSError: + vk = 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}:") + # The KHR query is authoritative on current drivers; the NV one is only a fallback for older ones (its enums + # are narrower, so its properties are not decoded here). + khr_shapes = 0 + for name, s_type in (("vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR", + VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR), + ("vkGetPhysicalDeviceCooperativeMatrixPropertiesNV", + VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_NV)): + if name.endswith("NV") and khr_shapes: + continue + query = proc(instance, name, C.c_int32, + [C.c_void_p, C.POINTER(C.c_uint32), C.POINTER(VkCooperativeMatrixProperties)]) + if query is None: + continue + n = C.c_uint32(0) + result = query(device, C.byref(n), None) + if result != VK_SUCCESS or n.value == 0: + continue + shapes = (VkCooperativeMatrixProperties * n.value)() + for shape in shapes: + shape.sType = s_type + if query(device, C.byref(n), shapes) != VK_SUCCESS: + continue + if name.endswith("KHR"): + khr_shapes = n.value + print(f" {name}: {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()) From 6a5272e86f083e1f795053f79e4d917b1addca40 Mon Sep 17 00:00:00 2001 From: xycld Date: Sun, 1 Nov 2026 22:38:50 -0500 Subject: [PATCH 03/41] vulkan: tool fixes from review; notes state the fallback direction 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. --- docs/vulkan-backend-design.md | 54 +++--- docs/vulkan-backend-research.md | 56 +++--- tools/vulkan_gemm/CMakeLists.txt | 13 +- tools/vulkan_gemm/README.md | 80 ++++----- tools/vulkan_gemm/gemm_f16.comp | 4 +- tools/vulkan_gemm/main.cpp | 298 +++++++++++++++++++++---------- tools/vulkan_gemm/probe.py | 50 ++---- 7 files changed, 326 insertions(+), 229 deletions(-) diff --git a/docs/vulkan-backend-design.md b/docs/vulkan-backend-design.md index 95c09b9..b82867d 100644 --- a/docs/vulkan-backend-design.md +++ b/docs/vulkan-backend-design.md @@ -5,14 +5,15 @@ Status: proposal; nothing in this document is implemented yet. ## Scope and goal -Add a **Vulkan compute backend** that can run the whole MossInfer pipeline on a non-CUDA GPU (or on NVIDIA without the -CUDA toolkit), at a usable fraction of the CUDA speed, without regressing the CUDA path. +Add a **Vulkan compute backend** as the fallback for hardware the engine does not otherwise support, able to run the +MossInfer pipeline there without regressing the CUDA path. CUDA stays the path on every GPU it supports. Requirements: - Same model files (official bf16 bundle, ComfyUI INT8/INT4 ConvRot files) and the same public API/CLI/server surface. - CPU side identical: tokenizer, samplers, PNG, planner, reports. -- The backend is selected at run time (or at build time) and the probe reports which one ran. +- The backend for a device follows one fixed rule, never a measurement: CUDA where the engine supports the device, Vulkan + otherwise. A caller can name a backend explicitly, and the probe and the report always say which one ran. - Correctness first: every kernel validated against the CUDA implementation on the same weights; end-to-end images compared statistically (Vulkan numerics differ in accumulation order, like sm_75 does today). @@ -27,7 +28,7 @@ The engine is not CUDA-abstracted today: arenas (`runtime/DeviceArena`). - `Planner`/`pipeline` call `cudaMemGetInfo`-based free-memory helpers, `cudaEvent`s and cuBLASLt handles, and the model code calls kernel functions directly (`kernels::Gemm::linear`, `moss_dit::MossDit`, ...). -- ~120 CUDA translation units implement the kernels (GEMM, attention, norm, RoPE, GDN, VAE, quantized formats). +- About 70 CUDA translation units implement the kernels (GEMM, attention, norm, RoPE, GDN, VAE, quantized formats). A Vulkan backend therefore needs two layers: a **device runtime** (buffers, queues, commands, events, memory budget) and a **kernel set** behind the same call shapes the model code uses. @@ -36,7 +37,7 @@ a **kernel set** behind the same call shapes the model code uses. | Option | What it means | Cost | Risk | |---|---|---|---| -| **A. Backend abstraction inside the engine** | Introduce `mi::runtime::Backend` (device/buffer/stream/event) and `mi::kernels::Backend` (GEMM, attention, norm, ...) interfaces; CUDA and Vulkan implement them; model/pipeline code talks to the interfaces. | High upfront refactor (~120 TUs touched), but one pipeline for both. | Refactor stalls the CUDA path; hard to land incrementally. | +| **A. Backend abstraction inside the engine** | Introduce `mi::runtime::Backend` (device/buffer/stream/event) and `mi::kernels::Backend` (GEMM, attention, norm, ...) interfaces; CUDA and Vulkan implement them; model/pipeline code talks to the interfaces. | High upfront refactor (about 70 kernel files and their callers), but one pipeline for both. | Refactor stalls the CUDA path; hard to land incrementally. | | **B. Separate Vulkan runner first** | A standalone `moss-vulkan` executable (or tool) that loads the same bundle and runs the pipeline with Vulkan kernels; the engine stays untouched. Merge later through A. | Medium; no risk to CUDA; proves kernels and performance with real weights. | Duplicated pipeline logic until merged; second code path to maintain. | | **C. CUDA shim over Vulkan** | Implement the CUDA runtime API subset (`cudaMalloc`, `cudaStream_t`, `cuBLASLt` calls) on top of Vulkan so the engine compiles unchanged. | Tempting but the semantics (streams/events, cuBLASLt descriptors, error model) do not map cleanly; CUTLASS kernels are SASS, not SPIR-V, so "compile the same kernels" is impossible. | Dead end for the kernel layer. | @@ -108,7 +109,7 @@ tools/vulkan_check.cpp bring-up tool (like tools/sm75_check.cu): devi `free_budget_bytes = heap_budget - heap_usage - reserve` to the planner, the same contract `cudaMemGetInfo` provides. - Weight streaming: the engine's DiT block streaming (5.16 GB INT8 at 256², 9.6 GB bf16) becomes: host-side ring buffer → transfer queue → device-local block buffers, with timeline semaphores per block and a per-step upload budget. The - planner's existing `HiddenFraction` model (copy hidden behind compute) applies unchanged. + planner's existing estimate of copies hidden behind compute applies unchanged. ### Dispatch and command recording @@ -123,10 +124,10 @@ tools/vulkan_check.cpp bring-up tool (like tools/sm75_check.cu): devi Legend: CM = cooperative matrix, SIMT = plain compute shader. -| Stage | Operation | Shapes (DiT 3072, FFN 12288, VAE 96–384) | Vulkan kernel | +| Stage | Operation | Shapes (DiT 3072, its FFN 8192; text encoder 4096, its FFN 12288; VAE 96–384) | Vulkan kernel | |---|---|---|---| | DiT block | modulation/norm | [tokens, 3072] | SIMT, vectorized f16 | -| DiT block | QKV / out / FFN up / FFN down / final | 3072×3072, 3072×12288, 12288×3072 | CM f16 16×16×16, f32 acc; epilogues in shader (bias, gated residual) | +| DiT block | QKV / out / FFN up-gate / FFN down / final (N×K, M = tokens) | 9216×3072, 3072×3072, 16384×3072, 3072×8192 | CM f16 16×16×16, f32 acc; epilogues in shader (bias, gated residual) | | DiT block | joint attention | tokens 4096 (1024²), heads 24, head_dim 128, non-causal varlen | tiled FA: CM QK^T, f32 softmax, SIMT PV or CM PV | | DiT block | RoPE, SiTU-GLU, residual | elementwise | SIMT | | Text encoder | RMSNorm, linear, RoPE, causal attention, GDN recurrence | hidden 4096, 32 layers | GEMM CM; attention CM/SIMT; GDN SIMT scan | @@ -134,13 +135,22 @@ Legend: CM = cooperative matrix, SIMT = plain compute shader. | INT8 path | ConvRot GEMM + dequant | same GEMM shapes | CM s8 16×16×32, s32 acc, dequant epilogue | | INT4 path | unpack to INT8 then GEMM | same | SIMT unpack + INT8 CM | +### Devices + +Tile sizes come from the device's limits, not from this document: Vulkan only guarantees 16 KiB of shared memory, and +subgroup sizes differ (the tiles below need 32). A device without a usable cooperative-matrix shape runs a plain compute +GEMM, slower but correct; that path is what makes the backend a fallback rather than a second fast path. + ### GEMM tiling (48 KiB shared memory, f16, f32 accumulate) +M is free: whole tiles always run, A and C are allocated with M rounded up to 64 rows, and row m of C depends only on row m +of A. N is a multiple of 64 in every engine GEMM. + - Workgroup 64×64 output tile, 4 warps (2×2), each warp computes 32×32 via four 16×16×16 cooperative-matrix ops. - K step 32: A tile 64×32 f16 = 4 KiB, B tile 32×64 f16 = 4 KiB; 3-stage pipeline = 24 KiB, leaves room for epilogue staging and the bias/gate vectors. - Registers: 32×32 f32 accumulator per warp = 1024 floats = 32 per thread, plus A/B fragments; well within a 255-register - budget at 256 threads/workgroup. + budget at 128 threads/workgroup. - Split-K for skinny GEMMs (e.g., final layer, modulation) or when M is small (prompt tokens); the engine's CUDA GEMM already chunks rows, so the Vulkan dispatcher can reuse that policy. - Epilogues are fused in the shader: bias add, `ffn_down` residual add, gated residual (multiply by gate), SiTU-GLU @@ -148,13 +158,14 @@ Legend: CM = cooperative matrix, SIMT = plain compute shader. ### Attention kernel -- One workgroup per (block of 64 queries, head, sequence chunk): 4 warps, BM=64, BN=64, head_dim 128 (DiT) or 256 - (Qwen). +- One workgroup per (block of queries, head, sequence chunk). Staged f16 tiles must fit 48 KiB, Q + K + V + P: + head_dim 128 (DiT) with 64 queries and 32 keys is 16 + 8 + 8 + 4 = 36 KiB; head_dim 256 (Qwen) with 32 queries and + 16 keys is 16 + 8 + 8 + 1 = 33 KiB. - QK^T on cooperative matrix (f16 inputs, f32 accumulate), softmax in SIMT f32 with the same online algorithm as the CUDA sm_75 kernel (`src/kernels/sm75/attention_sm75.cu`), PV on cooperative matrix or SIMT. - Varlen/GQA: per-sequence offsets and head mapping in push constants; the DiT joint attention is non-causal over [text + image] tokens, Qwen is causal. -- 4096-token attention at 1024² needs 64×64 tiles × 64 query blocks × 24 heads; the same tiling as the CUDA kernel. +- 4096-token attention at 1024² walks those tiles over 64 query blocks and 24 heads. ### Quantization formats @@ -168,12 +179,12 @@ Legend: CM = cooperative matrix, SIMT = plain compute shader. ## Integration with the engine - **Probe**: `src/backend/probe.*` gains a Vulkan path when built with `MOSSINFER_VULKAN`: enumerate physical devices, - report name/vendor/heap/coopmat shapes, recommend `vulkan` when no CUDA device is usable (or when `--backend vulkan`). + report name/vendor/heap/coopmat shapes and, per device, the backend that the fixed rule above gives it. - **Planner**: the cost model's rates (`bf16_tflops`, PCIe rate) are learned per backend and per device; the tier rules are backend-neutral (reference/high/turbo), except that on Vulkan the 8-bit attention modes may not exist. - **Tier/attention**: Vulkan on Turing-class hardware resolves an unset attention to `exact` (same rule as sm_75); on devices with f16 cooperative matrix a `q8_fp16` port can come later. -- **CLI/server**: `--backend cuda|vulkan|auto`; the report gains `quality.backend = "vulkan"` and the device string; +- **CLI/server**: `--backend cuda|vulkan` overrides the rule; the report gains `quality.backend = "vulkan"` and the device string; `moss-server` passes it through `load`'s new optional `backend` parameter (protocol addition). - **Packaging**: a `-vulkan` package flavour (loader + shaders embedded); the CUDA package stays the default. @@ -194,8 +205,8 @@ Legend: CM = cooperative matrix, SIMT = plain compute shader. | Milestone | Deliverable | Acceptance | Estimate | |---|---|---|---| -| **V0** | Probe tool + microbenchmarks (GEMM, bandwidth, dispatch overhead, attention) | Numbers for all five unknowns in the research doc; go/no-go | 2–3 days | -| **V1** | Vulkan runtime (device, allocator, commands, pipeline cache) + f16 GEMM validated vs CPU/CUDA | GEMM cosine ≥ 0.9999 at engine shapes; ≥ 40 % of cuBLAS | 1–2 weeks | +| **V0** | Probe tool + GEMM microbenchmark (`tools/vulkan_gemm`); done for the GEMM | GEMM correct and measured at the engine's shapes against cuBLASLt | done | +| **V1** | Vulkan runtime (device, allocator, commands, pipeline cache) + f16 GEMM, shared-memory tiles and a plain compute path | GEMM cosine ≥ 0.9999 at engine shapes, on a device without cooperative matrices too | 1–2 weeks | | **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights | Block output matches CUDA block (cosine ≥ 0.999) | 1–2 weeks | | **V3** | Text encoder (GEMM, RMSNorm, RoPE, causal attention, GDN) and VAE decoder | Encode/decode parity; first full Vulkan image | 2–3 weeks | | **V4** | Pipeline integration (runtime abstraction, planner, streaming, CLI/server) | `moss-infer --backend vulkan` and `moss-server` produce images on the 2080 | 2–4 weeks | @@ -209,14 +220,13 @@ the VAE convolutions are the least predictable parts. - **Scope explosion.** B-first keeps the engine shippable; V4 (refactor) is gated on V2/V3 parity. - **Vendor differences.** Cooperative matrix shapes and shared-memory limits vary; the kernel set is selected from a capability table (NVIDIA first, AMD/Intel fallbacks second). -- **Performance shortfall.** If V0 shows < 40 % of cuBLAS on the 2080, the Vulkan backend is positioned as a - portability backend (AMD/Intel), not a replacement on NVIDIA. +- **Performance.** As a fallback it never replaces CUDA, so its speed gate is whether a generation is usable on the + fallback device. That can only be measured on one. - **Maintenance.** Two kernel sets are expensive; the shared parity tests and the runtime abstraction are the mitigation. If the backend cannot be maintained, V1–V3 remain a standalone tool. ## Decisions requested -1. Confirm the target: portability first (AMD/Intel) or NVIDIA parity first? -2. Approve the two-phase plan (standalone runner, then runtime abstraction) or prefer a full abstraction up front? -3. Choose the first device matrix for validation (2080 only, or also an AMD/Intel card). -4. Decide whether the Vulkan path ships as a separate package flavour or inside the default package (loader dependency). +1. Decided: Vulkan is the fallback for hardware the engine does not otherwise support. +2. Choose the first fallback devices for validation. +3. Decide whether the Vulkan path ships as a separate package flavour or inside the default package (loader dependency). diff --git a/docs/vulkan-backend-research.md b/docs/vulkan-backend-research.md index 6742b22..d1de120 100644 --- a/docs/vulkan-backend-research.md +++ b/docs/vulkan-backend-research.md @@ -7,17 +7,15 @@ Vulkan compute backend. ## Why a Vulkan backend MossInfer is CUDA-only: the pipeline calls cuBLASLt, CUTLASS kernels and CUDA kernels directly, and the probe refuses any -GPU whose compute capability is below 8.0 (7.5 with `MOSSINFER_SM75`). That confines the engine to NVIDIA on Windows and -Linux, with the CUDA toolkit's driver and memory-management model. +GPU whose compute capability is below 8.0 (7.5 with `MOSSINFER_SM75`, 7.0 with `MOSSINFER_SM70`). That confines the engine +to NVIDIA on Windows and Linux. The Vulkan backend is the fallback for the hardware that leaves out; CUDA stays the path +on every GPU it supports. A Vulkan compute backend would give: -- **Other GPUs.** AMD (RDNA2/3/4), Intel (Arc/Xe), and NVIDIA on systems where the CUDA toolkit is not desirable. The - engine's model code is portable C++/CUDA; only the kernels and the device runtime are CUDA-specific. -- **One packaging story** for Windows and Linux without the CUDA runtime dependency, and a path to Steam Deck (RADV) - and other low-power devices. -- **An escape hatch from CUDA's memory model.** Vulkan exposes `VK_EXT_memory_budget`, explicit heaps and suballocation, - which fits the engine's planner/streaming design; CUDA's `cudaMemGetInfo` is coarser. +- **Other GPUs.** AMD, Intel and older NVIDIA cards. The engine's model code is portable; only the kernels and the device + runtime are CUDA-specific. +- **No CUDA runtime dependency** in the fallback package. The cost is a second implementation of every tensor-core kernel, plus a device/memory/command abstraction the engine does not have today (see the design document). @@ -35,7 +33,7 @@ and one physical device. | `maxComputeSharedMemorySize` | 48 KiB | 64 KiB is a CUDA opt-in only; Vulkan tiles must fit 48 KiB (Turing's hardware limit is 64 KiB but Vulkan caps at 48 KiB here). | | `maxComputeWorkGroupInvocations` | 1024 | Up to 32 warps per workgroup. | | `subgroupSize` | 32 | Cooperative-matrix `gl_ScopeSubgroup` = one warp. | -| `maxStorageBufferRange` | 2^32-1 bytes | 4 GiB per buffer; weights are 75 MB at most per tensor, so a small number of big buffers is fine. | +| `maxStorageBufferRange` | 2^32-1 bytes | 4 GiB per buffer; the largest DiT weight is 96 MiB as f16 (up/gate), so a small number of big buffers is fine. | | Device-local heap | 7989 MiB | The engine's planner sees the same ~8 GiB. A second device-local heap of 214 MiB exists; host-visible heap is 16 GiB. | ### Extensions a compute backend needs (all present) @@ -132,28 +130,26 @@ attention remains the bottleneck. It is the closest prior art to what MossInfer ## Performance model -Measured CUDA numbers on this machine (RTX 2080 8 GiB, SFT bundle), from `docs/sm75.md`: +Measured on an RTX 4090 Laptop, same GPU and session, with the fragment-from-global-memory kernel in `tools/vulkan_gemm`: -| Tier | 256×256 step | Compute per step (approx.) | -|---|---|---| -| reference (bf16→fp16 GEMMs) | 0.80 s | 42 blocks × ~0.015 s + text encoder ~3 s per prompt | -| high (INT8 DiT) | 0.19 s | 42 blocks × ~0.004 s | -| turbo (INT4 via INT8) | 0.22 s | unpack overhead + INT8 | +| Shape (M x N x K) | Vulkan | cuBLASLt (f16, f32 compute) | Ratio | +|---|---:|---:|---:| +| 8222 x 9216 x 3072 (DiT QKV) | 20.1 TFLOPS | 52.6 | 38 % | +| 8222 x 3072 x 3072 (DiT projection) | 31.1 | 57.4 | 54 % | +| 8222 x 16384 x 3072 (DiT up/gate) | 19.7 | 45.5 | 43 % | +| 8222 x 3072 x 8192 (DiT down) | 13.5 | 70.1 | 19 % | +| 282 x 12288 x 4096 (text encoder) | 25.0 | 59.5 | 42 % | +| 282 x 4096 x 12288 (text encoder) | 12.7 | 52.2 | 24 % | -Turing's fp16 tensor-core peak on the 2080 is roughly 50–60 TFLOPS (NVIDIA's figures for the card are in that range, -depending on fp16 vs fp32 accumulate); cuBLASLt reaches a large fraction of it on the GEMM shapes of this model -(3072×12288×3072 etc.). A Vulkan cooperative-matrix GEMM with 16×16×16 f16×f16→f32 tiles can use -the same instruction; the realistic outcome is **50–80 % of cuBLAS** for large GEMMs, dropping for small ones (the DiT's -per-block GEMMs are 3072-wide, large enough; the modulation/norm kernels are bandwidth-bound and vendor-independent). +The large-K shapes suffer most, which fits the kernel's strided global loads. These are NVIDIA numbers: they show the +kernel works and where it loses time, not how a fallback device will perform. CUDA stays the path wherever it runs. -What that means per step at 256² (INT8 high, 0.19 s): a Vulkan backend at 60 % of the CUDA GEMM throughput would be -~0.30 s/step; the planner's cost model would then choose the Vulkan tier/streaming plan accordingly. The VAE decode is -fp32/conv-heavy and would be the second-largest cost; Vulkan convolution on Turing has no tensor-core help for fp32, so -it stays bandwidth-bound (same as the CUDA fp32 VAE path). +For scale, `docs/sm75.md` measures CUDA on an RTX 2080 Ti at 256x256: 0.175 s per step in `reference`, 0.091 s in `high` +(INT8) and 0.103 s in `turbo` (INT4). The critical unknowns to measure in a prototype (M0): -1. Cooperative-matrix GEMM throughput at the engine's shapes (f16, f32 accumulate, 48 KiB tiles). +1. Cooperative-matrix GEMM throughput at the engine's shapes: measured above, without shared-memory tiles. 2. Dispatch overhead per kernel (engine runs ~20–40 dispatches per block; command-buffer batching must keep the CPU cost per step below the GPU cost). 3. Attention throughput (fp16 QK^T + softmax + PV, 4096-token joint attention at 1024²). @@ -169,7 +165,7 @@ The critical unknowns to measure in a prototype (M0): - **Compute:** f16 inputs with f32 accumulate via cooperative matrix (matching cuBLASLt's `CUBLAS_COMPUTE_32F` use in the engine's fp16 path); INT8 weights use s32 accumulate (exact, same result as CUTLASS on CUDA); the VAE decoder stays fp32 (SIMT), since Vulkan exposes no fp32 tensor core. -- **No FP8 on Turing** (`VK_KHR_shader_fp8` is not exposed), so the engine's `q8_fp8` attention path cannot be ported +- **No FP8 on Turing** (`VK_EXT_shader_float8` is not exposed), so the engine's `q8_fp8` attention path cannot be ported to this GPU; `q8_fp16` would be the 8-bit attention variant on Vulkan, and `exact` remains the default on sm_75-class hardware. - **Denormals:** f16 cooperative matrix on NVIDIA flushes subnormal inputs/outputs? The CUDA sm_75 port relies on fp16 @@ -180,7 +176,7 @@ The critical unknowns to measure in a prototype (M0): | Risk | Impact | Mitigation | |---|---|---| -| Cooperative matrix on AMD/Intel differs (16×16×16 only, no int8) | Kernels must specialize per vendor; Vulkan backend becomes NVIDIA-first | Query properties at startup and select kernels; scalar/DP4A fallbacks | +| Cooperative-matrix shapes, subgroup sizes and shared memory differ per device, or there is no cooperative matrix at all | The fast kernel cannot run everywhere | Query properties at startup and select kernels; a plain compute path is required, not optional | | Shader compilation time and driver bugs | Cold start and hangs | Precompile SPIR-V at build time, ship pipeline cache, validation layers in CI | | 48 KiB shared memory | Smaller tiles than CUDA, lower GEMM efficiency | 3-stage 64×64×32 f16 tiles; split-K for skinny shapes | | Dispatch overhead (CPU-bound at small sizes) | Step time worse than CUDA | Batch dispatches into one command buffer; reuse command buffers; push constants + BDA | @@ -191,9 +187,9 @@ The critical unknowns to measure in a prototype (M0): ## Open questions for the design phase -1. Is the target "NVIDIA parity plus AMD/Intel support", or "works on anything"? That decides how many fallbacks the - kernel layer needs. +1. Decided: Vulkan is the fallback for hardware the engine does not otherwise support. It therefore needs a path for devices + without cooperative matrices, not only the fast one. 2. Should the Vulkan path reuse the engine's existing planner/streaming code through a device abstraction, or start as a separate `moss-vulkan` runner and merge later? (The design document recommends the latter first.) 3. Which attention variant is acceptable on Vulkan: `exact` only (like sm_75), or a `q8_fp16` port? -4. Is a WebGPU/wgpu fallback worth a third backend, or is Vulkan enough for the non-NVIDIA target? +4. Which fallback devices are the first validation targets? No non-NVIDIA device has run any of this. diff --git a/tools/vulkan_gemm/CMakeLists.txt b/tools/vulkan_gemm/CMakeLists.txt index 7c82eb1..4b678b3 100644 --- a/tools/vulkan_gemm/CMakeLists.txt +++ b/tools/vulkan_gemm/CMakeLists.txt @@ -5,6 +5,9 @@ project(mossinfer_vulkan_gemm 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() # Prefer the Vulkan SDK's CMake package; fall back to explicit paths for machines without the SDK (the loader's import # library can be generated from vulkan-1.dll, see README.md). @@ -28,17 +31,17 @@ if(NOT GLSL_COMPILER) message(FATAL_ERROR "glslc not found: it ships with the Vulkan SDK") endif() -set(GERM_SPV ${CMAKE_CURRENT_BINARY_DIR}/gemm_f16.spv) +set(GEMM_SPV ${CMAKE_CURRENT_BINARY_DIR}/gemm_f16.spv) add_custom_command( - OUTPUT ${GERM_SPV} + OUTPUT ${GEMM_SPV} COMMAND ${GLSL_COMPILER} --target-env=vulkan1.3 -fshader-stage=comp -O ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp - -o ${GERM_SPV} + -o ${GEMM_SPV} DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp COMMENT "glslc gemm_f16.comp") -add_custom_target(gemm_spv ALL DEPENDS ${GERM_SPV}) +add_custom_target(gemm_spv ALL DEPENDS ${GEMM_SPV}) add_executable(vulkan_gemm main.cpp) target_include_directories(vulkan_gemm PRIVATE ${VULKAN_INCLUDE}) target_link_libraries(vulkan_gemm PRIVATE ${VULKAN_TARGET}) -target_compile_definitions(vulkan_gemm PRIVATE "DEFAULT_SPV_PATH=\"${GERM_SPV}\"") +target_compile_definitions(vulkan_gemm PRIVATE "DEFAULT_SPV_PATH=\"${GEMM_SPV}\"") add_dependencies(vulkan_gemm gemm_spv) diff --git a/tools/vulkan_gemm/README.md b/tools/vulkan_gemm/README.md index e1a0098..fff750d 100644 --- a/tools/vulkan_gemm/README.md +++ b/tools/vulkan_gemm/README.md @@ -1,14 +1,14 @@ # Vulkan GEMM bring-up -The first artefact of the Vulkan backend work ([research](../docs/vulkan-backend-research.md), -[design](../docs/vulkan-backend-design.md)): a standalone tool that runs one f16 GEMM on a Vulkan device with +The first artefact of the Vulkan backend work ([research](../../docs/vulkan-backend-research.md), +[design](../../docs/vulkan-backend-design.md)): a standalone tool that runs one f16 GEMM on a Vulkan device with cooperative matrices, checks it against a CPU reference and times it at the engine's matrix shapes. It links no engine code and is not part of the engine build. The point is to answer the two questions the design document leaves open before any integration work starts: 1. Does `VK_KHR_cooperative_matrix` give a correct f16 GEMM with f32 accumulation on this device? -2. How fast is it, compared with the cuBLASLt path the engine uses? +2. How fast is it, compared with cuBLASLt on the same GPU? ## What the kernel does @@ -21,7 +21,7 @@ C[M,N] = A[M,K] * B[K,N] with `A` row-major (stride `K`) and `B` column-major (stride `K`), i.e. `B[k,n]` lives at `n*K + k`. That is exactly the engine's weight layout (`[N,K]` row-major), so the GEMM consumes weights as stored, with no transpose. -- One workgroup per 64x64 output tile, 128 threads (four 32-wide subgroups). +- One workgroup per 64x64 output tile, 128 threads: four full subgroups of 32, which the pipeline requires. - Each subgroup owns a 32x32 quarter, computed as four `16x16x16` cooperative-matrix operations (`f16` inputs, `f32` accumulate) per K step, 16 elements at a time. - Output tiles are walked in bands of 8 N-tiles (the same idea as the engine's `BandedSwizzle` for INT8 GEMMs): a band @@ -48,7 +48,7 @@ cmake -S tools/vulkan_gemm -B build-vulkan-gemm \ -DVULKAN_LIBRARY= ``` -Two notes from a machine without administrator rights (where the SDK installer cannot run): +Notes from a machine without administrator rights (where the SDK installer cannot run): - Headers: the [Vulkan-Headers](https://github.com/KhronosGroup/Vulkan-Headers) repository's `include/vulkan` is enough. Its `vulkan_core.h` includes `vk_video/*`, which is only in the SDK; a local copy either adds those headers or @@ -60,56 +60,50 @@ Two notes from a machine without administrator rights (where the SDK installer c ## Run ```sh -build-vulkan-gemm/vulkan_gemm # check 64^3, then the engine shapes -build-vulkan-gemm/vulkan_gemm --bench 3072 3072 3072 10 -build-vulkan-gemm/vulkan_gemm --bench 3072 12288 3072 5 --no-check +build-vulkan-gemm/vulkan_gemm # the checks, then the engine's shapes +build-vulkan-gemm/vulkan_gemm --device 1 # another device; without it, the first one that can run the kernel +build-vulkan-gemm/vulkan_gemm --bench 8222 9216 3072 5 --no-check ``` -`M` and `N` must be multiples of 64 and `K` a multiple of 16 (the kernel has no bounds checks yet). `--spv` selects a -different SPIR-V file. The CPU reference is `O(M*N*K)`, so it only runs for shapes up to 2^26 multiply-adds (64³, not -512³ or the engine shapes); `--no-check` disables it. +`N` must be a multiple of 64 and `K` of 16. `M` is free: the kernel runs whole 64-row tiles, so the tool allocates A and C +with `M` rounded up to 64 rows. The exit code is non-zero when a check fails. The CPU reference is `O(M*N*K)`, so it runs +for shapes up to 2^31 multiply-adds; the defaults check one tile, a partial last band, several bands, and row tails with +NaN in the padding rows. -Output on an RTX 2080 8 GiB, NVIDIA driver 621.01, Vulkan 1.4.357: +The tool refuses a device it cannot run on and says why: no Vulkan 1.3, no `VK_KHR_cooperative_matrix`, no +16x16x16 f16 x f16 -> f32 subgroup shape, or no way to fix the compute subgroup size at 32. -``` -device: NVIDIA GeForce RTX 2080 (Vulkan 1.4.357) -features: shaderFloat16=1 storage16=1 storage8=1 bda=1 coopmat=1 -cooperative matrix shapes (17): - 16x16x16 A=f16 B=f16 C=f16 R=f16 sat=0 - 16x8x16 A=f16 B=f16 C=f16 R=f16 sat=0 - ... - 16x16x16 A=f16 B=f16 C=f32 R=f32 sat=0 <- the shape the kernel uses - ... - 16x16x32 A=s8 B=s8 C=s32 R=s32 sat=0 <- the INT8 path, later - 16x16x16 A=bf16 B=bf16 C=f32 R=f32 sat=0 <- driver-emulated on Turing; not used -check 64x64x64: max|C - CPU| = 1.01342e-05 -gemm 64x64x64: 0.102 ms 0.0 TFLOPS -gemm 512x512x512: 0.038 ms 7.0 TFLOPS -gemm 3072x3072x3072: 6.349 ms 9.1 TFLOPS -gemm 3072x12288x3072: 24.712 ms 9.4 TFLOPS -gemm 3072x3072x12288: 49.082 ms 4.7 TFLOPS -``` +## Results + +RTX 4090 Laptop, NVIDIA driver 615.78, Vulkan 1.4.351, against cuBLASLt f16 inputs with f32 compute on the same GPU in the +same session. The shapes are the DiT's at 1024x1024 with CFG (M = 8222) and the text encoder's: + +| Shape (M x N x K) | Vulkan | cuBLASLt | Vulkan / cuBLASLt | +|---|---:|---:|---:| +| 8222 x 9216 x 3072 (QKV) | 20.1 TFLOPS | 52.6 | 38 % | +| 8222 x 3072 x 3072 (projection) | 31.1 | 57.4 | 54 % | +| 8222 x 16384 x 3072 (up/gate) | 19.7 | 45.5 | 43 % | +| 8222 x 3072 x 8192 (down) | 13.5 | 70.1 | 19 % | +| 282 x 12288 x 4096 | 25.0 | 59.5 | 42 % | +| 282 x 4096 x 12288 | 12.7 | 52.2 | 24 % | -The small shapes mostly measure dispatch overhead (one dispatch is ~0.1 ms here); the large ones are what matters. +RTX 2080 8 GiB, driver 621.01, at the earlier 3072-row shapes: 8-9 TFLOPS. The checks match a double-precision CPU +reference to within fp16 input and f32 accumulator rounding (1e-5 at 64 x 64 x 64, 5e-5 at 512^3). ## Status and limitations -- **Correctness**: the 64³ case matches a double-precision CPU reference to `1e-5` (fp16 rounding of the inputs and the - f32 accumulator), which is the expected error for this data. The same kernel runs the engine's shapes. -- **Performance**: ~8-9 TFLOPS at 3072³, roughly 10-15 % of the card's fp16 tensor peak. The kernel is memory-bound: - the cooperative-matrix fragments load straight from global memory, so every subgroup issues strided per-lane loads and - each A/B element is fetched by several workgroups. Making the tiles larger (64x128) did not help, which confirmed the - diagnosis. -- **Next step**: stage A and B tiles in shared memory with vectorised coalesced copies and double buffering, then load - the fragments from shared memory. The 48 KiB shared-memory limit and the 64x64x32 tile staging are worked out in the - design document; the expected result is 20-30 TFLOPS, which would put a Vulkan DiT step within 2x of the CUDA path. -- No INT8/INT4 path yet, no bias/epilogue fusion, no multi-queue overlap, and one queue only. +- The fragments load straight from global memory, so large K suffers most. Staging A and B tiles in shared memory is the + next step for speed. +- No bias or epilogue fusion, no INT8/INT4 path, no path for devices without cooperative matrices, one queue only. +- N must be a multiple of 64. Every N in the engine's DiT and text encoder GEMMs is. +- Vulkan's synchronization validation does not report on this kernel's accesses: it stays quiet with or without barriers, + because they go through buffer device addresses. Core validation runs clean. ## Files | File | Purpose | |---|---| | `gemm_f16.comp` | the cooperative-matrix GEMM (GLSL 460, compiled by `glslc`) | -| `main.cpp` | device setup, buffers, pipeline, CPU check, timing | -| `probe.py` | optional: lists the device's cooperative-matrix shapes and memory heaps using only `vulkan-1.dll` (no SDK) | +| `main.cpp` | device selection and setup, buffers, pipeline, CPU check, timing | +| `probe.py` | optional: lists the device's cooperative-matrix shapes and memory heaps using only the Vulkan loader (no SDK) | | `CMakeLists.txt` | standalone build | diff --git a/tools/vulkan_gemm/gemm_f16.comp b/tools/vulkan_gemm/gemm_f16.comp index 5385d32..d978719 100644 --- a/tools/vulkan_gemm/gemm_f16.comp +++ b/tools/vulkan_gemm/gemm_f16.comp @@ -1,7 +1,9 @@ // C[M,N] = A[M,K] * B[K,N] with A row-major (stride lda) and B column-major (stride ldb == K), i.e. B[k,n] lives at // n*ldb + k: exactly the engine's [N,K] weight layout, so the GEMM consumes weights as stored (no transpose). // -// One workgroup per 64x64 output tile, 128 threads (4 subgroups of 32). Each subgroup owns a 32x32 quarter, computed as +// One workgroup per 64x64 output tile, 128 threads: exactly 4 full subgroups of 32, which the pipeline must require +// (VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT and requiredSubgroupSize). Whole tiles always run, so A and +// C need M rounded up to 64 rows, and N must be a multiple of 64. Each subgroup owns a 32x32 quarter, computed as // four 16x16x16 cooperative-matrix operations (f16 inputs, f32 accumulate) per K step. Output tiles are walked in bands // of BAND N-tiles (the same idea as the engine's BandedSwizzle for INT8 GEMMs): a band of weight tiles stays in L2 // while every M tile streams past it. diff --git a/tools/vulkan_gemm/main.cpp b/tools/vulkan_gemm/main.cpp index 02b4005..daaf376 100644 --- a/tools/vulkan_gemm/main.cpp +++ b/tools/vulkan_gemm/main.cpp @@ -4,9 +4,11 @@ // // Kernel contract: C[M,N] = A[M,K] * B[K,N] with A row-major (stride K) and B column-major (stride K), i.e. B[k,n] lives // at n*K + k, which is exactly the engine's [N,K] weight layout. That lets the GEMM consume weights as stored, with no -// transposes. M and N must be multiples of 64 (the tool checks). +// transposes. N must be a multiple of 64 and K of 16. M is free: the kernel always runs whole 64-row tiles, so A and C +// must be allocated with M rounded up to 64 rows. Row m of C depends only on row m of A, so the padding rows never reach +// the first M rows (the check fills them with NaN to show it). // -// Usage: vulkan_gemm [--spv ] [--bench M N K iters]... [--no-check] +// Usage: vulkan_gemm [--spv ] [--device N] [--bench M N K iters]... [--no-check] #include #include @@ -29,6 +31,7 @@ namespace { constexpr uint32_t kTile = 64; // BM == BN, one workgroup per tile +constexpr uint32_t kSubgroupSize = 32; // gemm_f16.comp divides its 128 invocations into four subgroups of this size [[noreturn]] void Fail(const char* what, VkResult r) { std::printf("%s failed: VkResult %d\n", what, static_cast(r)); @@ -40,45 +43,37 @@ void Check(VkResult r, const char* what) { } float HalfToFloat(uint16_t h) { - const uint32_t sign = static_cast(h >> 15) << 31; const uint32_t exp = (h >> 10) & 0x1f; const uint32_t mant = h & 0x3ff; - uint32_t bits = 0; + float f = 0.0f; if (exp == 0) { - if (mant == 0) { - bits = sign; - } else { - uint32_t e = 127 - 15 + 1; - uint32_t m = mant; - while ((m & 0x400) == 0) { - m <<= 1; - --e; - } - bits = sign | ((e - 1) << 23) | ((m & 0x3ff) << 13); - } - } else if (exp == 31) { - bits = sign | 0x7f800000u | (mant << 13); + f = std::ldexp(static_cast(mant), -24); // zero or subnormal: mant * 2^-24, exact in fp32 } else { - bits = sign | ((exp + 127 - 15) << 23) | (mant << 13); + const uint32_t bits = exp == 31 ? 0x7f800000u | (mant << 13) : ((exp + 127 - 15) << 23) | (mant << 13); + std::memcpy(&f, &bits, 4); } - float f = 0.0f; - std::memcpy(&f, &bits, 4); - return f; + return (h & 0x8000u) != 0 ? -f : f; } +// Rounds to nearest, ties to even, like the hardware conversion. uint16_t FloatToHalf(float f) { uint32_t x = 0; std::memcpy(&x, &f, 4); - const uint32_t sign = (x >> 16) & 0x8000u; - const int32_t exp = static_cast((x >> 23) & 0xff) - 127 + 15; - uint32_t mant = x & 0x7fffffu; - if (exp <= 0) { - if (exp < -10) return static_cast(sign); - mant = (mant | 0x800000u) >> (1 - exp); - return static_cast(sign | ((mant + 0x1000u) >> 13)); - } - if (exp >= 31) return static_cast(sign | 0x7c00u); - return static_cast(sign | (static_cast(exp) << 10) | ((mant + 0x1000u) >> 13)); + const auto sign = static_cast((x >> 16) & 0x8000u); + const uint32_t magnitude = x & 0x7fffffffu; + if (magnitude > 0x7f800000u) return static_cast(sign | 0x7e00u); // NaN + const 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 uint32_t mant = magnitude & 0x7fffffu; + const uint32_t significand = exp > 0 ? mant : mant | 0x800000u; + const int shift = exp > 0 ? 13 : 14 - exp; + uint32_t rounded = significand >> shift; + const uint32_t rest = significand & ((1u << shift) - 1u); + const 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)); } struct Buffer { @@ -168,6 +163,16 @@ void Submit(Device& d, const std::function& record) { vkFreeCommandBuffers(d.device, d.pool, 1, &command_buffer); } +// A global memory barrier. vkQueueWaitIdle only makes writes available to the host: work in a later submit still needs +// one of these to see them. +void Barrier(VkCommandBuffer cb, VkPipelineStageFlags src_stage, VkAccessFlags src_access, + VkPipelineStageFlags dst_stage, VkAccessFlags dst_access) { + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = src_access; + barrier.dstAccessMask = dst_access; + vkCmdPipelineBarrier(cb, src_stage, dst_stage, 0, 1, &barrier, 0, nullptr, 0, nullptr); +} + void Upload(Device& d, Buffer& dst, const void* data, VkDeviceSize size) { Buffer staging = CreateBuffer(d, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, false); @@ -175,6 +180,8 @@ void Upload(Device& d, Buffer& dst, const void* data, VkDeviceSize size) { Submit(d, [&](VkCommandBuffer cb) { VkBufferCopy copy{0, 0, size}; vkCmdCopyBuffer(cb, staging.buffer, dst.buffer, 1, ©); + Barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_ACCESS_SHADER_READ_BIT); }); DestroyBuffer(d, staging); } @@ -185,6 +192,8 @@ void Download(Device& d, Buffer& src, void* data, VkDeviceSize size) { Submit(d, [&](VkCommandBuffer cb) { VkBufferCopy copy{0, 0, size}; vkCmdCopyBuffer(cb, src.buffer, staging.buffer, 1, ©); + Barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_HOST_BIT, + VK_ACCESS_HOST_READ_BIT); }); std::memcpy(data, staging.mapped, static_cast(size)); DestroyBuffer(d, staging); @@ -224,16 +233,22 @@ const char* ComponentName(int type) { default: return "?"; } } -void PrintCooperativeMatrixShapes(Device& d) { +std::vector CooperativeMatrixShapes(VkInstance instance, VkPhysicalDevice physical) { auto query = reinterpret_cast( - vkGetInstanceProcAddr(d.instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); - if (query == nullptr) return; + vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); + if (query == nullptr) return {}; uint32_t count = 0; - if (query(d.physical, &count, nullptr) != VK_SUCCESS || count == 0) return; + if (query(physical, &count, nullptr) != VK_SUCCESS || count == 0) return {}; std::vector shapes(count); for (auto& s : shapes) s.sType = VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR; - if (query(d.physical, &count, shapes.data()) != VK_SUCCESS) return; - std::printf("cooperative matrix shapes (%u):\n", count); + if (query(physical, &count, shapes.data()) != VK_SUCCESS) return {}; + shapes.resize(count); + return shapes; +} + +void PrintCooperativeMatrixShapes(Device& d) { + const auto shapes = CooperativeMatrixShapes(d.instance, d.physical); + std::printf("cooperative matrix shapes (%zu):\n", shapes.size()); for (const auto& s : shapes) { std::printf(" %ux%ux%u A=%s B=%s C=%s R=%s sat=%u\n", s.MSize, s.NSize, s.KSize, ComponentName(static_cast(s.AType)), ComponentName(static_cast(s.BType)), @@ -242,7 +257,62 @@ void PrintCooperativeMatrixShapes(Device& d) { } } -Device OpenDevice() { +// Why the kernel cannot run on this device; empty when it can. +std::string WhyUnsupported(VkInstance instance, VkPhysicalDevice physical) { + VkPhysicalDeviceVulkan13Properties properties13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_PROPERTIES}; + VkPhysicalDeviceProperties2 properties{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2}; + vkGetPhysicalDeviceProperties(physical, &properties.properties); + if (properties.properties.apiVersion < VK_API_VERSION_1_3) return "needs Vulkan 1.3"; + properties.pNext = &properties13; + vkGetPhysicalDeviceProperties2(physical, &properties); + + uint32_t count = 0; + vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, nullptr); + std::vector extensions(count); + vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, extensions.data()); + const bool has_extension = std::any_of(extensions.begin(), extensions.end(), [](const VkExtensionProperties& e) { + return std::strcmp(e.extensionName, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME) == 0; + }); + if (!has_extension) 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}; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES}; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; + features.pNext = &v11; + v11.pNext = &v12; + v12.pNext = &v13; + v13.pNext = &coop; + vkGetPhysicalDeviceFeatures2(physical, &features); + if (coop.cooperativeMatrix == 0) return "the cooperativeMatrix feature is off"; + if (features.features.shaderInt64 == 0) return "no shaderInt64"; + if (v11.storageBuffer16BitAccess == 0) return "no storageBuffer16BitAccess"; + if (v12.shaderFloat16 == 0) return "no shaderFloat16"; + if (v12.bufferDeviceAddress == 0) return "no bufferDeviceAddress"; + if (v12.vulkanMemoryModel == 0) return "no vulkanMemoryModel"; + if (v13.subgroupSizeControl == 0 || v13.computeFullSubgroups == 0) return "no subgroup size control"; + + if (properties13.minSubgroupSize > kSubgroupSize || properties13.maxSubgroupSize < kSubgroupSize || + (properties13.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT) == 0) { + return "cannot fix the compute subgroup size at " + std::to_string(kSubgroupSize) + " (device range " + + std::to_string(properties13.minSubgroupSize) + "-" + std::to_string(properties13.maxSubgroupSize) + ")"; + } + if (properties13.maxComputeWorkgroupSubgroups < 4) return "needs 4 subgroups per workgroup"; + + const auto shapes = CooperativeMatrixShapes(instance, physical); + const bool has_shape = std::any_of(shapes.begin(), shapes.end(), [](const VkCooperativeMatrixPropertiesKHR& s) { + return s.MSize == 16 && s.NSize == 16 && s.KSize == 16 && s.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && + s.BType == VK_COMPONENT_TYPE_FLOAT16_KHR && s.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + s.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR && s.saturatingAccumulation == VK_FALSE && + s.scope == VK_SCOPE_SUBGROUP_KHR; + }); + if (!has_shape) return "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; + return {}; +} + +// device_index < 0: the first device that can run the kernel. +Device OpenDevice(int device_index) { Device d; VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.pApplicationName = "mossinfer-vulkan-gemm"; @@ -253,13 +323,28 @@ Device OpenDevice() { uint32_t count = 0; Check(vkEnumeratePhysicalDevices(d.instance, &count, nullptr), "vkEnumeratePhysicalDevices"); - if (count == 0) { - std::printf("no Vulkan physical device\n"); - std::exit(1); - } std::vector devices(count); Check(vkEnumeratePhysicalDevices(d.instance, &count, devices.data()), "vkEnumeratePhysicalDevices"); - d.physical = devices[0]; + bool chosen = false; + for (uint32_t i = 0; i < count; ++i) { + VkPhysicalDeviceProperties properties{}; + vkGetPhysicalDeviceProperties(devices[i], &properties); + const std::string why = WhyUnsupported(d.instance, devices[i]); + std::printf("device %u: %s: %s\n", i, properties.deviceName, why.empty() ? "usable" : why.c_str()); + const bool wanted = device_index < 0 ? why.empty() && !chosen : static_cast(device_index) == i; + if (wanted) { + if (!why.empty()) { + std::printf("device %u cannot run this kernel\n", i); + std::exit(1); + } + d.physical = devices[i]; + chosen = true; + } + } + if (!chosen) { + std::printf("no usable Vulkan device\n"); + std::exit(1); + } vkGetPhysicalDeviceProperties(d.physical, &d.properties); vkGetPhysicalDeviceMemoryProperties(d.physical, &d.memory); @@ -280,40 +365,29 @@ Device OpenDevice() { std::exit(1); } - VkPhysicalDeviceShaderFloat16Int8Features float16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES}; - VkPhysicalDevice16BitStorageFeatures storage16{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_16BIT_STORAGE_FEATURES}; - VkPhysicalDevice8BitStorageFeatures storage8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_8BIT_STORAGE_FEATURES}; - VkPhysicalDeviceBufferDeviceAddressFeatures address{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_BUFFER_DEVICE_ADDRESS_FEATURES}; - VkPhysicalDeviceCooperativeMatrixFeaturesKHR coop{ - VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR}; - VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; - features.pNext = &float16; - float16.pNext = &storage16; - storage16.pNext = &storage8; - storage8.pNext = &address; - address.pNext = &coop; - vkGetPhysicalDeviceFeatures2(d.physical, &features); - - std::printf("device: %s (Vulkan %u.%u.%u)\n", d.properties.deviceName, - VK_API_VERSION_MAJOR(d.properties.apiVersion), VK_API_VERSION_MINOR(d.properties.apiVersion), - VK_API_VERSION_PATCH(d.properties.apiVersion)); - std::printf("features: shaderFloat16=%u storage16=%u storage8=%u bda=%u coopmat=%u\n", float16.shaderFloat16, - storage16.storageBuffer16BitAccess, storage8.storageBuffer8BitAccess, address.bufferDeviceAddress, - coop.cooperativeMatrix); - if (float16.shaderFloat16 == 0 || storage16.storageBuffer16BitAccess == 0 || address.bufferDeviceAddress == 0 || - coop.cooperativeMatrix == 0) { - std::printf("the device lacks a feature the kernel needs\n"); - std::exit(1); - } - float16.shaderFloat16 = VK_TRUE; - float16.shaderInt8 = VK_TRUE; - storage16.storageBuffer16BitAccess = VK_TRUE; - storage16.uniformAndStorageBuffer16BitAccess = VK_TRUE; - storage8.storageBuffer8BitAccess = VK_TRUE; - address.bufferDeviceAddress = VK_TRUE; + // Exactly what gemm_f16.spv declares (spirv-dis lists its capabilities), so a missing one fails at vkCreateDevice. + 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}; + v13.pNext = &coop; + v13.subgroupSizeControl = VK_TRUE; + v13.computeFullSubgroups = VK_TRUE; + VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; + v12.pNext = &v13; + v12.shaderFloat16 = VK_TRUE; + v12.bufferDeviceAddress = VK_TRUE; + v12.vulkanMemoryModel = VK_TRUE; + VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES}; + v11.pNext = &v12; + v11.storageBuffer16BitAccess = VK_TRUE; + VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; + features.pNext = &v11; + features.features.shaderInt64 = VK_TRUE; - const char* extensions[] = {"VK_KHR_cooperative_matrix"}; + std::printf("using: %s (Vulkan %u.%u.%u)\n", d.properties.deviceName, VK_API_VERSION_MAJOR(d.properties.apiVersion), + VK_API_VERSION_MINOR(d.properties.apiVersion), VK_API_VERSION_PATCH(d.properties.apiVersion)); + + const char* extensions[] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME}; const float priority = 1.0f; VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; queue_info.queueFamilyIndex = d.queue_family; @@ -334,20 +408,27 @@ Device OpenDevice() { return d; } -// One GEMM run: fills A and B with deterministic data, optionally checks a 64x64 case against a CPU reference, -// then times `iterations` dispatches. -void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M, uint32_t N, uint32_t K, +// The CPU reference is O(M*N*K), so only shapes up to this many multiply-adds are checked. +constexpr uint64_t kMaxCheckedMultiplyAdds = 1ull << 31; + +// One GEMM run: fills A and B with deterministic data, times `iterations` dispatches, then checks shapes small enough +// for the CPU reference. +double RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M, uint32_t N, uint32_t K, int iterations, bool check) { - const VkDeviceSize a_bytes = static_cast(M) * K * 2; + const uint32_t rows = (M + kTile - 1) / kTile * kTile; // the kernel runs whole tiles + const bool checked = check && static_cast(M) * N * K <= kMaxCheckedMultiplyAdds; + const VkDeviceSize a_bytes = static_cast(rows) * K * 2; const VkDeviceSize b_bytes = static_cast(N) * K * 2; - const VkDeviceSize c_bytes = static_cast(M) * N * 4; + const VkDeviceSize c_bytes = static_cast(rows) * N * 4; - std::vector a(static_cast(M) * K); + std::vector a(static_cast(rows) * K); std::vector b(static_cast(N) * K); std::mt19937 rng(1234); std::uniform_real_distribution dist(-1.0f, 1.0f); for (auto& v : a) v = FloatToHalf(dist(rng)); for (auto& v : b) v = FloatToHalf(dist(rng)); + // NaN in the padding rows: any leak into the first M rows of C fails the check. + if (checked) std::fill(a.begin() + static_cast(static_cast(M) * K), a.end(), uint16_t{0x7e00}); Buffer buffer_a = CreateBuffer(d, a_bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); @@ -375,7 +456,12 @@ void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M const auto record = [&](VkCommandBuffer cb) { vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); vkCmdPushConstants(cb, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push), &push); - for (int i = 0; i < iterations; ++i) vkCmdDispatch(cb, N / kTile, M / kTile, 1); + for (int i = 0; i < iterations; ++i) { + vkCmdDispatch(cb, N / kTile, rows / kTile, 1); + Barrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_WRITE_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT); + } }; Submit(d, record); // warm-up (compiles or loads the pipeline cache) @@ -384,9 +470,9 @@ void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M const double seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count() / static_cast(iterations); - // The CPU reference is O(M*N*K); only small shapes can afford it. - if (check && static_cast(M) * N * K <= (1ull << 26)) { - std::vector c(static_cast(M) * N); + double max_error = -1.0; + if (checked) { + std::vector c(static_cast(rows) * N); Download(d, buffer_c, c.data(), c_bytes); double max_abs = 0.0; for (uint32_t m = 0; m < M; ++m) { @@ -396,9 +482,11 @@ void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M reference += static_cast(HalfToFloat(a[static_cast(m) * K + k])) * static_cast(HalfToFloat(b[static_cast(n) * K + k])); } - max_abs = std::max(max_abs, std::fabs(static_cast(c[static_cast(m) * N + n]) - reference)); + const double error = std::fabs(static_cast(c[static_cast(m) * N + n]) - reference); + max_abs = std::isnan(error) ? INFINITY : std::max(max_abs, error); } } + max_error = max_abs; std::printf("check %ux%ux%u: max|C - CPU| = %.6g\n", M, N, K, max_abs); } @@ -408,6 +496,7 @@ void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M DestroyBuffer(d, buffer_a); DestroyBuffer(d, buffer_b); DestroyBuffer(d, buffer_c); + return max_error; } } // namespace @@ -415,6 +504,7 @@ void RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M int main(int argc, char** argv) { std::string spv = DEFAULT_SPV_PATH; bool check = true; + int device_index = -1; struct Shape { uint32_t M, N, K; int iterations; @@ -424,6 +514,8 @@ int main(int argc, char** argv) { const std::string arg = argv[i]; if (arg == "--spv" && i + 1 < argc) { spv = argv[++i]; + } else if (arg == "--device" && i + 1 < argc) { + device_index = std::atoi(argv[++i]); } else if (arg == "--no-check") { check = false; } else if (arg == "--bench" && i + 4 < argc) { @@ -434,22 +526,25 @@ int main(int argc, char** argv) { s.iterations = std::atoi(argv[++i]); shapes.push_back(s); } else { - std::printf("usage: %s [--spv ] [--no-check] [--bench M N K iters]...\n", argv[0]); + std::printf("usage: %s [--spv ] [--device N] [--no-check] [--bench M N K iters]...\n", argv[0]); return 2; } } if (shapes.empty()) { - shapes = {{64, 64, 64, 1}, {512, 512, 512, 20}, {3072, 3072, 3072, 10}, {3072, 12288, 3072, 5}, - {3072, 3072, 12288, 5}}; + // Checked: one tile, a partial last band, several bands, and row tails. Then the engine's real shapes at 1024x1024 + // with CFG: the DiT's QKV, projection, up/gate and down GEMMs, and the text encoder's GatedDeltaNet projections. + shapes = {{64, 64, 64, 1}, {128, 640, 64, 1}, {256, 1024, 64, 1}, {94, 640, 64, 1}, + {1030, 9216, 64, 1}, {512, 512, 512, 20}, {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) { - if (s.M % kTile != 0 || s.N % kTile != 0 || s.K % 16 != 0) { - std::printf("shapes must be multiples of %u (M, N) and 16 (K): %ux%ux%u\n", kTile, s.M, s.N, s.K); + if (s.M == 0 || s.N == 0 || s.K == 0 || s.iterations <= 0 || s.N % kTile != 0 || s.K % 16 != 0) { + std::printf("N must be a multiple of %u and K of 16, and all must be positive: %ux%ux%u\n", kTile, s.M, s.N, s.K); return 2; } } - Device d = OpenDevice(); + Device d = OpenDevice(device_index); PrintCooperativeMatrixShapes(d); const std::vector spirv = ReadFile(spv); @@ -470,7 +565,13 @@ int main(int argc, char** argv) { Check(vkCreatePipelineLayout(d.device, &layout_info, nullptr, &layout), "vkCreatePipelineLayout"); VkComputePipelineCreateInfo pipeline_info{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + // The kernel's tile split assumes exactly four full subgroups of 32. Without this a driver may pick another size. + VkPipelineShaderStageRequiredSubgroupSizeCreateInfo subgroup_size{ + VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; + subgroup_size.requiredSubgroupSize = kSubgroupSize; pipeline_info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; + pipeline_info.stage.pNext = &subgroup_size; + pipeline_info.stage.flags = VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT; pipeline_info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT; pipeline_info.stage.module = module; pipeline_info.stage.pName = "main"; @@ -479,7 +580,16 @@ int main(int argc, char** argv) { Check(vkCreateComputePipelines(d.device, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &pipeline), "vkCreateComputePipelines"); - for (const Shape& s : shapes) RunGemm(d, pipeline, layout, s.M, s.N, s.K, s.iterations, check); + // Inputs are fp16 and the accumulator is fp32 (no fp64 reference is exact), so this bounds rounding, not bits. + constexpr double kMaxError = 1e-2; + bool failed = false; + for (const Shape& s : shapes) { + const double error = RunGemm(d, pipeline, layout, s.M, s.N, s.K, s.iterations, check); + if (error > kMaxError) { + std::printf("FAILED: %ux%ux%u: error %.6g is above %g\n", s.M, s.N, s.K, error, kMaxError); + failed = true; + } + } vkDestroyPipeline(d.device, pipeline, nullptr); vkDestroyPipelineLayout(d.device, layout, nullptr); @@ -487,5 +597,5 @@ int main(int argc, char** argv) { vkDestroyCommandPool(d.device, d.pool, nullptr); vkDestroyDevice(d.device, nullptr); vkDestroyInstance(d.instance, nullptr); - return 0; + return failed ? 1 : 0; } diff --git a/tools/vulkan_gemm/probe.py b/tools/vulkan_gemm/probe.py index b7366ab..1e3e5f9 100644 --- a/tools/vulkan_gemm/probe.py +++ b/tools/vulkan_gemm/probe.py @@ -1,4 +1,4 @@ -"""Lists a Vulkan device's cooperative-matrix shapes and memory heaps, using only vulkan-1.dll. +"""Lists a Vulkan device's cooperative-matrix shapes and memory heaps, using only the Vulkan loader. The C++ tool needs Vulkan headers and an import library; this script needs neither (ctypes only), so it works on a machine with just the GPU driver installed. It is how the capability table in docs/vulkan-backend-research.md was @@ -13,7 +13,6 @@ VK_STRUCTURE_TYPE_APPLICATION_INFO = 0 VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO = 1 VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR = 1000506001 -VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_NV = 1000249001 COMPONENT = { 0: "f16", 1: "f32", 2: "f64", 3: "s8", 4: "s16", 5: "s32", 6: "s64", @@ -64,10 +63,7 @@ class VkPhysicalDeviceMemoryProperties(C.Structure): def main(): - try: - vk = C.WinDLL("vulkan-1.dll") - except OSError: - vk = C.CDLL("libvulkan.so.1") + 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 @@ -98,36 +94,22 @@ def proc(instance, name, restype, argtypes): for index in range(count.value): device = devices[index] print(f"\ndevice {index}:") - # The KHR query is authoritative on current drivers; the NV one is only a fallback for older ones (its enums - # are narrower, so its properties are not decoded here). - khr_shapes = 0 - for name, s_type in (("vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR", - VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_KHR), - ("vkGetPhysicalDeviceCooperativeMatrixPropertiesNV", - VK_STRUCTURE_TYPE_COOPERATIVE_MATRIX_PROPERTIES_NV)): - if name.endswith("NV") and khr_shapes: - continue - query = proc(instance, name, C.c_int32, - [C.c_void_p, C.POINTER(C.c_uint32), C.POINTER(VkCooperativeMatrixProperties)]) - if query is None: - continue - n = C.c_uint32(0) - result = query(device, C.byref(n), None) - if result != VK_SUCCESS or n.value == 0: - continue + 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 = s_type - if query(device, C.byref(n), shapes) != VK_SUCCESS: - continue - if name.endswith("KHR"): - khr_shapes = n.value - print(f" {name}: {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)}") + 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)) From ae542e7f1776a02198f4939c7f339b90dc0674d7 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 00:13:10 -0500 Subject: [PATCH 04/41] vulkan: device runtime, shared-memory coopmat GEMM and a plain compute 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. --- docs/vulkan-backend-design.md | 17 +- docs/vulkan-backend-research.md | 18 +- tools/vulkan_gemm/CMakeLists.txt | 47 -- tools/vulkan_gemm/README.md | 109 ---- tools/vulkan_gemm/gemm_f16.comp | 86 --- tools/vulkan_gemm/main.cpp | 601 ------------------- vulkan/CMakeLists.txt | 35 ++ vulkan/README.md | 67 +++ vulkan/gemm.cpp | 77 +++ vulkan/gemm.h | 45 ++ vulkan/half.h | 59 ++ vulkan/runtime.cpp | 372 ++++++++++++ vulkan/runtime.h | 113 ++++ vulkan/shaders/gemm_coop.comp | 120 ++++ vulkan/shaders/gemm_plain.comp | 47 ++ {tools/vulkan_gemm => vulkan/tools}/probe.py | 0 vulkan/tools/vulkan_check.cpp | 199 ++++++ 17 files changed, 1147 insertions(+), 865 deletions(-) delete mode 100644 tools/vulkan_gemm/CMakeLists.txt delete mode 100644 tools/vulkan_gemm/README.md delete mode 100644 tools/vulkan_gemm/gemm_f16.comp delete mode 100644 tools/vulkan_gemm/main.cpp create mode 100644 vulkan/CMakeLists.txt create mode 100644 vulkan/README.md create mode 100644 vulkan/gemm.cpp create mode 100644 vulkan/gemm.h create mode 100644 vulkan/half.h create mode 100644 vulkan/runtime.cpp create mode 100644 vulkan/runtime.h create mode 100644 vulkan/shaders/gemm_coop.comp create mode 100644 vulkan/shaders/gemm_plain.comp rename {tools/vulkan_gemm => vulkan/tools}/probe.py (100%) create mode 100644 vulkan/tools/vulkan_check.cpp diff --git a/docs/vulkan-backend-design.md b/docs/vulkan-backend-design.md index b82867d..33ea531 100644 --- a/docs/vulkan-backend-design.md +++ b/docs/vulkan-backend-design.md @@ -71,14 +71,15 @@ This keeps `main` shippable at every step and gives measurable checkpoints befor ## Module layout and build ``` -cmake/modules/vulkan.cmake option MOSSINFER_VULKAN, find Vulkan (loader + headers), shader compilation -src/runtime/vulkan/ device, queues, allocator, staging ring, pipeline cache, descriptor/BDA helpers -src/kernels/vulkan/ GEMM, attention, norm, rope, gdn, vae, quantized formats (C++ dispatch) -shaders/vulkan/ *.comp GLSL sources; compiled to SPIR-V at build time by glslc -tests/vulkan/ kernel parity tests, end-to-end image tests -tools/vulkan_check.cpp bring-up tool (like tools/sm75_check.cu): device report, GEMM/attention/VAE checks +vulkan/runtime.* device, capabilities, buffers, pipelines, barriers +vulkan/gemm.* GEMM paths, then the other kernels beside it +vulkan/shaders/ *.comp GLSL sources, compiled to SPIR-V at build time and embedded +vulkan/tools/ vulkan_check (device list, checks, timing), probe.py ``` +It stays outside `src/` while it is standalone, because the engine's build globs that directory. Integration moves it behind +`MOSSINFER_VULKAN`. + - **Shaders**: GLSL 4.60 + `GL_KHR_cooperative_matrix` + `GL_EXT_shader_16bit_storage`/`8bit_storage` + `GL_EXT_shader_explicit_arithmetic_types_*` + `GL_KHR_shader_subgroup_*`, compiled by `glslc` to SPIR-V at build time and embedded as `const uint32_t[]` (a generated header), so the binary has no external files. A pipeline cache file is @@ -205,8 +206,8 @@ of A. N is a multiple of 64 in every engine GEMM. | Milestone | Deliverable | Acceptance | Estimate | |---|---|---|---| -| **V0** | Probe tool + GEMM microbenchmark (`tools/vulkan_gemm`); done for the GEMM | GEMM correct and measured at the engine's shapes against cuBLASLt | done | -| **V1** | Vulkan runtime (device, allocator, commands, pipeline cache) + f16 GEMM, shared-memory tiles and a plain compute path | GEMM cosine ≥ 0.9999 at engine shapes, on a device without cooperative matrices too | 1–2 weeks | +| **V0** | Probe tool + GEMM microbenchmark; done | GEMM correct and measured at the engine's shapes against cuBLASLt | done | +| **V1** | Vulkan runtime + f16 GEMM with shared-memory tiles and a plain compute path; done, except the suballocator and pipeline cache, which wait until the block runner needs them | Checks at engine shapes on a device without the fast path too | done | | **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights | Block output matches CUDA block (cosine ≥ 0.999) | 1–2 weeks | | **V3** | Text encoder (GEMM, RMSNorm, RoPE, causal attention, GDN) and VAE decoder | Encode/decode parity; first full Vulkan image | 2–3 weeks | | **V4** | Pipeline integration (runtime abstraction, planner, streaming, CLI/server) | `moss-infer --backend vulkan` and `moss-server` produce images on the 2080 | 2–4 weeks | diff --git a/docs/vulkan-backend-research.md b/docs/vulkan-backend-research.md index d1de120..5b515ee 100644 --- a/docs/vulkan-backend-research.md +++ b/docs/vulkan-backend-research.md @@ -130,26 +130,16 @@ attention remains the bottleneck. It is the closest prior art to what MossInfer ## Performance model -Measured on an RTX 4090 Laptop, same GPU and session, with the fragment-from-global-memory kernel in `tools/vulkan_gemm`: - -| Shape (M x N x K) | Vulkan | cuBLASLt (f16, f32 compute) | Ratio | -|---|---:|---:|---:| -| 8222 x 9216 x 3072 (DiT QKV) | 20.1 TFLOPS | 52.6 | 38 % | -| 8222 x 3072 x 3072 (DiT projection) | 31.1 | 57.4 | 54 % | -| 8222 x 16384 x 3072 (DiT up/gate) | 19.7 | 45.5 | 43 % | -| 8222 x 3072 x 8192 (DiT down) | 13.5 | 70.1 | 19 % | -| 282 x 12288 x 4096 (text encoder) | 25.0 | 59.5 | 42 % | -| 282 x 4096 x 12288 (text encoder) | 12.7 | 52.2 | 24 % | - -The large-K shapes suffer most, which fits the kernel's strided global loads. These are NVIDIA numbers: they show the -kernel works and where it loses time, not how a fallback device will perform. CUDA stays the path wherever it runs. +Measured on an RTX 4090 Laptop against cuBLASLt on the same GPU. See [`vulkan/README.md`](../vulkan/README.md) for the table: +the cooperative-matrix GEMM reaches 35-58 % of cuBLASLt at the engine's shapes, and the plain compute path 1-1.6 TFLOPS. +These are NVIDIA numbers: they show the kernels work and where they lose time, not how a fallback device will perform. For scale, `docs/sm75.md` measures CUDA on an RTX 2080 Ti at 256x256: 0.175 s per step in `reference`, 0.091 s in `high` (INT8) and 0.103 s in `turbo` (INT4). The critical unknowns to measure in a prototype (M0): -1. Cooperative-matrix GEMM throughput at the engine's shapes: measured above, without shared-memory tiles. +1. Cooperative-matrix GEMM throughput at the engine's shapes: measured above. 2. Dispatch overhead per kernel (engine runs ~20–40 dispatches per block; command-buffer batching must keep the CPU cost per step below the GPU cost). 3. Attention throughput (fp16 QK^T + softmax + PV, 4096-token joint attention at 1024²). diff --git a/tools/vulkan_gemm/CMakeLists.txt b/tools/vulkan_gemm/CMakeLists.txt deleted file mode 100644 index 4b678b3..0000000 --- a/tools/vulkan_gemm/CMakeLists.txt +++ /dev/null @@ -1,47 +0,0 @@ -# Standalone build for the Vulkan cooperative-matrix GEMM bring-up tool. It is deliberately separate from the engine's -# CMake project: no engine code is linked, and nothing here affects a normal build. See README.md. -cmake_minimum_required(VERSION 3.20) -project(mossinfer_vulkan_gemm 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() - -# Prefer the Vulkan SDK's CMake package; fall back to explicit paths for machines without the SDK (the loader's import -# library can be generated from vulkan-1.dll, see README.md). -find_package(Vulkan QUIET) -set(VULKAN_INCLUDE_DIR "" CACHE PATH "Vulkan include directory when no SDK is installed") -set(VULKAN_LIBRARY "" CACHE FILEPATH "Vulkan loader import library when no SDK is installed") - -if(VULKAN_INCLUDE_DIR AND VULKAN_LIBRARY) - set(VULKAN_INCLUDE ${VULKAN_INCLUDE_DIR}) - set(VULKAN_TARGET ${VULKAN_LIBRARY}) -elseif(Vulkan_FOUND) - set(VULKAN_INCLUDE ${Vulkan_INCLUDE_DIRS}) - set(VULKAN_TARGET Vulkan::Vulkan) -else() - message(FATAL_ERROR - "Vulkan not found: install the Vulkan SDK, or pass -DVULKAN_INCLUDE_DIR= -DVULKAN_LIBRARY=") -endif() - -find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN_SDK}/bin") -if(NOT GLSL_COMPILER) - message(FATAL_ERROR "glslc not found: it ships with the Vulkan SDK") -endif() - -set(GEMM_SPV ${CMAKE_CURRENT_BINARY_DIR}/gemm_f16.spv) -add_custom_command( - OUTPUT ${GEMM_SPV} - COMMAND ${GLSL_COMPILER} --target-env=vulkan1.3 -fshader-stage=comp -O ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp - -o ${GEMM_SPV} - DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/gemm_f16.comp - COMMENT "glslc gemm_f16.comp") -add_custom_target(gemm_spv ALL DEPENDS ${GEMM_SPV}) - -add_executable(vulkan_gemm main.cpp) -target_include_directories(vulkan_gemm PRIVATE ${VULKAN_INCLUDE}) -target_link_libraries(vulkan_gemm PRIVATE ${VULKAN_TARGET}) -target_compile_definitions(vulkan_gemm PRIVATE "DEFAULT_SPV_PATH=\"${GEMM_SPV}\"") -add_dependencies(vulkan_gemm gemm_spv) diff --git a/tools/vulkan_gemm/README.md b/tools/vulkan_gemm/README.md deleted file mode 100644 index fff750d..0000000 --- a/tools/vulkan_gemm/README.md +++ /dev/null @@ -1,109 +0,0 @@ -# Vulkan GEMM bring-up - -The first artefact of the Vulkan backend work ([research](../../docs/vulkan-backend-research.md), -[design](../../docs/vulkan-backend-design.md)): a standalone tool that runs one f16 GEMM on a Vulkan device with -cooperative matrices, checks it against a CPU reference and times it at the engine's matrix shapes. - -It links no engine code and is not part of the engine build. The point is to answer the two questions the design -document leaves open before any integration work starts: - -1. Does `VK_KHR_cooperative_matrix` give a correct f16 GEMM with f32 accumulation on this device? -2. How fast is it, compared with cuBLASLt on the same GPU? - -## What the kernel does - -`gemm_f16.comp` computes - -``` -C[M,N] = A[M,K] * B[K,N] -``` - -with `A` row-major (stride `K`) and `B` column-major (stride `K`), i.e. `B[k,n]` lives at `n*K + k`. That is exactly the -engine's weight layout (`[N,K]` row-major), so the GEMM consumes weights as stored, with no transpose. - -- One workgroup per 64x64 output tile, 128 threads: four full subgroups of 32, which the pipeline requires. -- Each subgroup owns a 32x32 quarter, computed as four `16x16x16` cooperative-matrix operations (`f16` inputs, `f32` - accumulate) per K step, 16 elements at a time. -- Output tiles are walked in bands of 8 N-tiles (the same idea as the engine's `BandedSwizzle` for INT8 GEMMs): a band - of weight tiles stays in L2 while every M tile streams past it. -- Arguments (buffer device addresses, sizes, strides) go through push constants; the buffers are addressed with - `GL_EXT_buffer_reference`, so there are no descriptor sets. - -## Build - -With the Vulkan SDK installed: - -```sh -cmake -S tools/vulkan_gemm -B build-vulkan-gemm -DCMAKE_BUILD_TYPE=Release -cmake --build build-vulkan-gemm -``` - -The SDK's `glslc` compiles the shader; the build fails with a clear message when either the SDK or `glslc` is missing. - -Without the SDK (for example a CI machine or a user who only has the driver), pass the two paths explicitly: - -```sh -cmake -S tools/vulkan_gemm -B build-vulkan-gemm \ - -DVULKAN_INCLUDE_DIR=/include \ - -DVULKAN_LIBRARY= -``` - -Notes from a machine without administrator rights (where the SDK installer cannot run): - -- Headers: the [Vulkan-Headers](https://github.com/KhronosGroup/Vulkan-Headers) repository's `include/vulkan` is enough. - Its `vulkan_core.h` includes `vk_video/*`, which is only in the SDK; a local copy either adds those headers or - defines the handful of `StdVideo*` types the core header names (the tool never touches video). -- Import library: `dumpbin /exports C:\Windows\System32\vulkan-1.dll`, keep the `vk*` names in a `.def` file, and - `lib /def:vulkan-1.def /machine:x64 /out:vulkan-1.lib`. -- `glslc` can be extracted from the SDK installer with 7-Zip (`7z x VulkanSDK.exe -o`); it lands in `Bin/`. - -## Run - -```sh -build-vulkan-gemm/vulkan_gemm # the checks, then the engine's shapes -build-vulkan-gemm/vulkan_gemm --device 1 # another device; without it, the first one that can run the kernel -build-vulkan-gemm/vulkan_gemm --bench 8222 9216 3072 5 --no-check -``` - -`N` must be a multiple of 64 and `K` of 16. `M` is free: the kernel runs whole 64-row tiles, so the tool allocates A and C -with `M` rounded up to 64 rows. The exit code is non-zero when a check fails. The CPU reference is `O(M*N*K)`, so it runs -for shapes up to 2^31 multiply-adds; the defaults check one tile, a partial last band, several bands, and row tails with -NaN in the padding rows. - -The tool refuses a device it cannot run on and says why: no Vulkan 1.3, no `VK_KHR_cooperative_matrix`, no -16x16x16 f16 x f16 -> f32 subgroup shape, or no way to fix the compute subgroup size at 32. - -## Results - -RTX 4090 Laptop, NVIDIA driver 615.78, Vulkan 1.4.351, against cuBLASLt f16 inputs with f32 compute on the same GPU in the -same session. The shapes are the DiT's at 1024x1024 with CFG (M = 8222) and the text encoder's: - -| Shape (M x N x K) | Vulkan | cuBLASLt | Vulkan / cuBLASLt | -|---|---:|---:|---:| -| 8222 x 9216 x 3072 (QKV) | 20.1 TFLOPS | 52.6 | 38 % | -| 8222 x 3072 x 3072 (projection) | 31.1 | 57.4 | 54 % | -| 8222 x 16384 x 3072 (up/gate) | 19.7 | 45.5 | 43 % | -| 8222 x 3072 x 8192 (down) | 13.5 | 70.1 | 19 % | -| 282 x 12288 x 4096 | 25.0 | 59.5 | 42 % | -| 282 x 4096 x 12288 | 12.7 | 52.2 | 24 % | - -RTX 2080 8 GiB, driver 621.01, at the earlier 3072-row shapes: 8-9 TFLOPS. The checks match a double-precision CPU -reference to within fp16 input and f32 accumulator rounding (1e-5 at 64 x 64 x 64, 5e-5 at 512^3). - -## Status and limitations - -- The fragments load straight from global memory, so large K suffers most. Staging A and B tiles in shared memory is the - next step for speed. -- No bias or epilogue fusion, no INT8/INT4 path, no path for devices without cooperative matrices, one queue only. -- N must be a multiple of 64. Every N in the engine's DiT and text encoder GEMMs is. -- Vulkan's synchronization validation does not report on this kernel's accesses: it stays quiet with or without barriers, - because they go through buffer device addresses. Core validation runs clean. - -## Files - -| File | Purpose | -|---|---| -| `gemm_f16.comp` | the cooperative-matrix GEMM (GLSL 460, compiled by `glslc`) | -| `main.cpp` | device selection and setup, buffers, pipeline, CPU check, timing | -| `probe.py` | optional: lists the device's cooperative-matrix shapes and memory heaps using only the Vulkan loader (no SDK) | -| `CMakeLists.txt` | standalone build | diff --git a/tools/vulkan_gemm/gemm_f16.comp b/tools/vulkan_gemm/gemm_f16.comp deleted file mode 100644 index d978719..0000000 --- a/tools/vulkan_gemm/gemm_f16.comp +++ /dev/null @@ -1,86 +0,0 @@ -// C[M,N] = A[M,K] * B[K,N] with A row-major (stride lda) and B column-major (stride ldb == K), i.e. B[k,n] lives at -// n*ldb + k: exactly the engine's [N,K] weight layout, so the GEMM consumes weights as stored (no transpose). -// -// One workgroup per 64x64 output tile, 128 threads: exactly 4 full subgroups of 32, which the pipeline must require -// (VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT and requiredSubgroupSize). Whole tiles always run, so A and -// C need M rounded up to 64 rows, and N must be a multiple of 64. Each subgroup owns a 32x32 quarter, computed as -// four 16x16x16 cooperative-matrix operations (f16 inputs, f32 accumulate) per K step. Output tiles are walked in bands -// of BAND N-tiles (the same idea as the engine's BandedSwizzle for INT8 GEMMs): a band of weight tiles stays in L2 -// while every M tile streams past it. -// -// The fragments load straight from global memory in this version. That is the correctness baseline; the shared-memory -// staged version (vectorised coalesced copies, double buffering) is the next step and is expected to be 3-4x faster. -#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 = 2) buffer F16Buf { float16_t d[]; }; -layout(buffer_reference, std430, buffer_reference_align = 4) buffer F32Buf { 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 BAND = 8u; - -void main() { - F16Buf A = F16Buf(pc.a); - F16Buf B = F16Buf(pc.b); - F32Buf C = F32Buf(pc.c); - - 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_tile = pos / width; - const uint n_tile = band * BAND + (pos - m_tile * width); - - const uint warp = gl_SubgroupID; - const uint m0 = m_tile * BM + (warp >> 1) * 32u; - const uint n0 = n_tile * BN + (warp & 1u) * 32u; - - coopmat fa0, fa1; - coopmat fb0, fb1; - coopmat acc00, acc01, acc10, acc11; - const float zero = 0.0; - acc00 = coopmat(zero); - acc01 = acc00; - acc10 = acc00; - acc11 = acc00; - - for (uint k = 0u; k < pc.K; k += 16u) { - coopMatLoad(fa0, A.d, m0 * pc.lda + k, pc.lda, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(fa1, A.d, (m0 + 16u) * pc.lda + k, pc.lda, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(fb0, B.d, n0 * pc.ldb + k, pc.ldb, gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad(fb1, B.d, (n0 + 16u) * pc.ldb + k, pc.ldb, gl_CooperativeMatrixLayoutColumnMajor); - acc00 = coopMatMulAdd(fa0, fb0, acc00); - acc01 = coopMatMulAdd(fa0, fb1, acc01); - acc10 = coopMatMulAdd(fa1, fb0, acc10); - acc11 = coopMatMulAdd(fa1, fb1, acc11); - } - - coopMatStore(acc00, C.d, m0 * pc.ldc + n0, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(acc01, C.d, m0 * pc.ldc + n0 + 16u, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(acc10, C.d, (m0 + 16u) * pc.ldc + n0, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); - coopMatStore(acc11, C.d, (m0 + 16u) * pc.ldc + n0 + 16u, pc.ldc, gl_CooperativeMatrixLayoutRowMajor); -} diff --git a/tools/vulkan_gemm/main.cpp b/tools/vulkan_gemm/main.cpp deleted file mode 100644 index daaf376..0000000 --- a/tools/vulkan_gemm/main.cpp +++ /dev/null @@ -1,601 +0,0 @@ -// Vulkan cooperative-matrix GEMM bring-up: one f16 GEMM kernel, checked against a CPU reference and timed at the -// engine's matrix shapes. This is the first artefact of the Vulkan backend work in docs/vulkan-backend-*.md; it is a -// standalone tool (no engine code is linked), so it can run before any backend integration exists. -// -// Kernel contract: C[M,N] = A[M,K] * B[K,N] with A row-major (stride K) and B column-major (stride K), i.e. B[k,n] lives -// at n*K + k, which is exactly the engine's [N,K] weight layout. That lets the GEMM consume weights as stored, with no -// transposes. N must be a multiple of 64 and K of 16. M is free: the kernel always runs whole 64-row tiles, so A and C -// must be allocated with M rounded up to 64 rows. Row m of C depends only on row m of A, so the padding rows never reach -// the first M rows (the check fills them with NaN to show it). -// -// Usage: vulkan_gemm [--spv ] [--device N] [--bench M N K iters]... [--no-check] -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifndef DEFAULT_SPV_PATH -#define DEFAULT_SPV_PATH "gemm_f16.spv" -#endif - -namespace { - -constexpr uint32_t kTile = 64; // BM == BN, one workgroup per tile -constexpr uint32_t kSubgroupSize = 32; // gemm_f16.comp divides its 128 invocations into four subgroups of this size - -[[noreturn]] void Fail(const char* what, VkResult r) { - std::printf("%s failed: VkResult %d\n", what, static_cast(r)); - std::exit(1); -} - -void Check(VkResult r, const char* what) { - if (r != VK_SUCCESS) Fail(what, r); -} - -float HalfToFloat(uint16_t h) { - const uint32_t exp = (h >> 10) & 0x1f; - const 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 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. -uint16_t FloatToHalf(float f) { - uint32_t x = 0; - std::memcpy(&x, &f, 4); - const auto sign = static_cast((x >> 16) & 0x8000u); - const uint32_t magnitude = x & 0x7fffffffu; - if (magnitude > 0x7f800000u) return static_cast(sign | 0x7e00u); // NaN - const 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 uint32_t mant = magnitude & 0x7fffffu; - const uint32_t significand = exp > 0 ? mant : mant | 0x800000u; - const int shift = exp > 0 ? 13 : 14 - exp; - uint32_t rounded = significand >> shift; - const uint32_t rest = significand & ((1u << shift) - 1u); - const 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)); -} - -struct Buffer { - VkBuffer buffer = VK_NULL_HANDLE; - VkDeviceMemory memory = VK_NULL_HANDLE; - VkDeviceSize size = 0; - VkDeviceAddress address = 0; - void* mapped = nullptr; -}; - -// Everything the tool needs from one device. Not a general runtime: the backend's runtime (buffers, arenas, queues) -// replaces this when the engine integration starts (docs/vulkan-backend-design.md). -struct Device { - VkInstance instance = VK_NULL_HANDLE; - VkPhysicalDevice physical = VK_NULL_HANDLE; - VkDevice device = VK_NULL_HANDLE; - VkQueue queue = VK_NULL_HANDLE; - uint32_t queue_family = 0; - VkCommandPool pool = VK_NULL_HANDLE; - VkPhysicalDeviceMemoryProperties memory{}; - VkPhysicalDeviceProperties properties{}; -}; - -uint32_t FindMemoryType(const Device& d, uint32_t type_bits, VkMemoryPropertyFlags flags) { - for (uint32_t i = 0; i < d.memory.memoryTypeCount; ++i) { - if ((type_bits & (1u << i)) != 0 && (d.memory.memoryTypes[i].propertyFlags & flags) == flags) return i; - } - std::printf("no memory type with flags 0x%x\n", static_cast(flags)); - std::exit(1); -} - -Buffer CreateBuffer(Device& d, VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags flags, - bool device_address) { - Buffer b; - b.size = size; - VkBufferCreateInfo info{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; - info.size = size; - info.usage = usage; - if (device_address) info.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT; - info.sharingMode = VK_SHARING_MODE_EXCLUSIVE; - Check(vkCreateBuffer(d.device, &info, nullptr, &b.buffer), "vkCreateBuffer"); - - VkMemoryRequirements requirements{}; - vkGetBufferMemoryRequirements(d.device, b.buffer, &requirements); - VkMemoryAllocateFlagsInfo allocate_flags{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO}; - allocate_flags.flags = device_address ? VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT : 0; - VkMemoryAllocateInfo allocate{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; - allocate.pNext = &allocate_flags; - allocate.allocationSize = requirements.size; - allocate.memoryTypeIndex = FindMemoryType(d, requirements.memoryTypeBits, flags); - Check(vkAllocateMemory(d.device, &allocate, nullptr, &b.memory), "vkAllocateMemory"); - Check(vkBindBufferMemory(d.device, b.buffer, b.memory, 0), "vkBindBufferMemory"); - if ((flags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) != 0) { - Check(vkMapMemory(d.device, b.memory, 0, VK_WHOLE_SIZE, 0, &b.mapped), "vkMapMemory"); - } - if (device_address) { - VkBufferDeviceAddressInfo info{VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO}; - info.buffer = b.buffer; - b.address = vkGetBufferDeviceAddress(d.device, &info); - } - return b; -} - -void DestroyBuffer(Device& d, Buffer& b) { - if (b.buffer != VK_NULL_HANDLE) vkDestroyBuffer(d.device, b.buffer, nullptr); - if (b.memory != VK_NULL_HANDLE) vkFreeMemory(d.device, b.memory, nullptr); - b = Buffer{}; -} - -void Submit(Device& d, const std::function& record) { - VkCommandBufferAllocateInfo allocate{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; - allocate.commandPool = d.pool; - allocate.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - allocate.commandBufferCount = 1; - VkCommandBuffer command_buffer = VK_NULL_HANDLE; - Check(vkAllocateCommandBuffers(d.device, &allocate, &command_buffer), "vkAllocateCommandBuffers"); - VkCommandBufferBeginInfo begin{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; - begin.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT; - Check(vkBeginCommandBuffer(command_buffer, &begin), "vkBeginCommandBuffer"); - record(command_buffer); - Check(vkEndCommandBuffer(command_buffer), "vkEndCommandBuffer"); - VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO}; - submit.commandBufferCount = 1; - submit.pCommandBuffers = &command_buffer; - Check(vkQueueSubmit(d.queue, 1, &submit, VK_NULL_HANDLE), "vkQueueSubmit"); - Check(vkQueueWaitIdle(d.queue), "vkQueueWaitIdle"); - vkFreeCommandBuffers(d.device, d.pool, 1, &command_buffer); -} - -// A global memory barrier. vkQueueWaitIdle only makes writes available to the host: work in a later submit still needs -// one of these to see them. -void Barrier(VkCommandBuffer cb, VkPipelineStageFlags src_stage, VkAccessFlags src_access, - VkPipelineStageFlags dst_stage, VkAccessFlags dst_access) { - VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; - barrier.srcAccessMask = src_access; - barrier.dstAccessMask = dst_access; - vkCmdPipelineBarrier(cb, src_stage, dst_stage, 0, 1, &barrier, 0, nullptr, 0, nullptr); -} - -void Upload(Device& d, Buffer& dst, const void* data, VkDeviceSize size) { - Buffer staging = CreateBuffer(d, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, false); - std::memcpy(staging.mapped, data, static_cast(size)); - Submit(d, [&](VkCommandBuffer cb) { - VkBufferCopy copy{0, 0, size}; - vkCmdCopyBuffer(cb, staging.buffer, dst.buffer, 1, ©); - Barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - VK_ACCESS_SHADER_READ_BIT); - }); - DestroyBuffer(d, staging); -} - -void Download(Device& d, Buffer& src, void* data, VkDeviceSize size) { - Buffer staging = CreateBuffer(d, size, VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, false); - Submit(d, [&](VkCommandBuffer cb) { - VkBufferCopy copy{0, 0, size}; - vkCmdCopyBuffer(cb, src.buffer, staging.buffer, 1, ©); - Barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_HOST_BIT, - VK_ACCESS_HOST_READ_BIT); - }); - std::memcpy(data, staging.mapped, static_cast(size)); - DestroyBuffer(d, staging); -} - -std::vector ReadFile(const std::string& path) { - std::ifstream file(path, std::ios::binary | std::ios::ate); - if (!file) { - std::printf("cannot open %s\n", path.c_str()); - std::exit(1); - } - const auto size = static_cast(file.tellg()); - if (size == 0 || size % 4 != 0) { - std::printf("%s is not a SPIR-V blob\n", path.c_str()); - std::exit(1); - } - std::vector words(size / 4); - file.seekg(0); - file.read(reinterpret_cast(words.data()), static_cast(size)); - return words; -} - -const char* ComponentName(int type) { - switch (type) { - case 0: return "f16"; - case 1: return "f32"; - case 2: return "f64"; - case 3: return "s8"; - case 4: return "s16"; - case 5: return "s32"; - case 6: return "s64"; - case 7: return "u8"; - case 8: return "u16"; - case 9: return "u32"; - case 10: return "u64"; - case 1000141000: return "bf16"; - default: return "?"; - } -} -std::vector CooperativeMatrixShapes(VkInstance instance, VkPhysicalDevice physical) { - auto query = reinterpret_cast( - vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); - if (query == nullptr) return {}; - uint32_t count = 0; - if (query(physical, &count, nullptr) != VK_SUCCESS || count == 0) 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 {}; - shapes.resize(count); - return shapes; -} - -void PrintCooperativeMatrixShapes(Device& d) { - const auto shapes = CooperativeMatrixShapes(d.instance, d.physical); - std::printf("cooperative matrix shapes (%zu):\n", shapes.size()); - for (const auto& s : shapes) { - std::printf(" %ux%ux%u A=%s B=%s C=%s R=%s sat=%u\n", s.MSize, s.NSize, s.KSize, - ComponentName(static_cast(s.AType)), ComponentName(static_cast(s.BType)), - ComponentName(static_cast(s.CType)), ComponentName(static_cast(s.ResultType)), - s.saturatingAccumulation); - } -} - -// Why the kernel cannot run on this device; empty when it can. -std::string WhyUnsupported(VkInstance instance, VkPhysicalDevice physical) { - VkPhysicalDeviceVulkan13Properties properties13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_PROPERTIES}; - VkPhysicalDeviceProperties2 properties{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2}; - vkGetPhysicalDeviceProperties(physical, &properties.properties); - if (properties.properties.apiVersion < VK_API_VERSION_1_3) return "needs Vulkan 1.3"; - properties.pNext = &properties13; - vkGetPhysicalDeviceProperties2(physical, &properties); - - uint32_t count = 0; - vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, nullptr); - std::vector extensions(count); - vkEnumerateDeviceExtensionProperties(physical, nullptr, &count, extensions.data()); - const bool has_extension = std::any_of(extensions.begin(), extensions.end(), [](const VkExtensionProperties& e) { - return std::strcmp(e.extensionName, VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME) == 0; - }); - if (!has_extension) 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}; - VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; - VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES}; - VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; - features.pNext = &v11; - v11.pNext = &v12; - v12.pNext = &v13; - v13.pNext = &coop; - vkGetPhysicalDeviceFeatures2(physical, &features); - if (coop.cooperativeMatrix == 0) return "the cooperativeMatrix feature is off"; - if (features.features.shaderInt64 == 0) return "no shaderInt64"; - if (v11.storageBuffer16BitAccess == 0) return "no storageBuffer16BitAccess"; - if (v12.shaderFloat16 == 0) return "no shaderFloat16"; - if (v12.bufferDeviceAddress == 0) return "no bufferDeviceAddress"; - if (v12.vulkanMemoryModel == 0) return "no vulkanMemoryModel"; - if (v13.subgroupSizeControl == 0 || v13.computeFullSubgroups == 0) return "no subgroup size control"; - - if (properties13.minSubgroupSize > kSubgroupSize || properties13.maxSubgroupSize < kSubgroupSize || - (properties13.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT) == 0) { - return "cannot fix the compute subgroup size at " + std::to_string(kSubgroupSize) + " (device range " + - std::to_string(properties13.minSubgroupSize) + "-" + std::to_string(properties13.maxSubgroupSize) + ")"; - } - if (properties13.maxComputeWorkgroupSubgroups < 4) return "needs 4 subgroups per workgroup"; - - const auto shapes = CooperativeMatrixShapes(instance, physical); - const bool has_shape = std::any_of(shapes.begin(), shapes.end(), [](const VkCooperativeMatrixPropertiesKHR& s) { - return s.MSize == 16 && s.NSize == 16 && s.KSize == 16 && s.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && - s.BType == VK_COMPONENT_TYPE_FLOAT16_KHR && s.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && - s.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR && s.saturatingAccumulation == VK_FALSE && - s.scope == VK_SCOPE_SUBGROUP_KHR; - }); - if (!has_shape) return "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; - return {}; -} - -// device_index < 0: the first device that can run the kernel. -Device OpenDevice(int device_index) { - Device d; - VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; - app.pApplicationName = "mossinfer-vulkan-gemm"; - app.apiVersion = VK_API_VERSION_1_3; - VkInstanceCreateInfo instance_info{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; - instance_info.pApplicationInfo = &app; - Check(vkCreateInstance(&instance_info, nullptr, &d.instance), "vkCreateInstance"); - - uint32_t count = 0; - Check(vkEnumeratePhysicalDevices(d.instance, &count, nullptr), "vkEnumeratePhysicalDevices"); - std::vector devices(count); - Check(vkEnumeratePhysicalDevices(d.instance, &count, devices.data()), "vkEnumeratePhysicalDevices"); - bool chosen = false; - for (uint32_t i = 0; i < count; ++i) { - VkPhysicalDeviceProperties properties{}; - vkGetPhysicalDeviceProperties(devices[i], &properties); - const std::string why = WhyUnsupported(d.instance, devices[i]); - std::printf("device %u: %s: %s\n", i, properties.deviceName, why.empty() ? "usable" : why.c_str()); - const bool wanted = device_index < 0 ? why.empty() && !chosen : static_cast(device_index) == i; - if (wanted) { - if (!why.empty()) { - std::printf("device %u cannot run this kernel\n", i); - std::exit(1); - } - d.physical = devices[i]; - chosen = true; - } - } - if (!chosen) { - std::printf("no usable Vulkan device\n"); - std::exit(1); - } - vkGetPhysicalDeviceProperties(d.physical, &d.properties); - vkGetPhysicalDeviceMemoryProperties(d.physical, &d.memory); - - uint32_t families = 0; - vkGetPhysicalDeviceQueueFamilyProperties(d.physical, &families, nullptr); - std::vector queue_families(families); - vkGetPhysicalDeviceQueueFamilyProperties(d.physical, &families, queue_families.data()); - bool found = false; - for (uint32_t i = 0; i < families; ++i) { - if ((queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) != 0) { - d.queue_family = i; - found = true; - break; - } - } - if (!found) { - std::printf("no compute queue\n"); - std::exit(1); - } - - // Exactly what gemm_f16.spv declares (spirv-dis lists its capabilities), so a missing one fails at vkCreateDevice. - 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}; - v13.pNext = &coop; - v13.subgroupSizeControl = VK_TRUE; - v13.computeFullSubgroups = VK_TRUE; - VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES}; - v12.pNext = &v13; - v12.shaderFloat16 = VK_TRUE; - v12.bufferDeviceAddress = VK_TRUE; - v12.vulkanMemoryModel = VK_TRUE; - VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES}; - v11.pNext = &v12; - v11.storageBuffer16BitAccess = VK_TRUE; - VkPhysicalDeviceFeatures2 features{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2}; - features.pNext = &v11; - features.features.shaderInt64 = VK_TRUE; - - std::printf("using: %s (Vulkan %u.%u.%u)\n", d.properties.deviceName, VK_API_VERSION_MAJOR(d.properties.apiVersion), - VK_API_VERSION_MINOR(d.properties.apiVersion), VK_API_VERSION_PATCH(d.properties.apiVersion)); - - const char* extensions[] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME}; - const float priority = 1.0f; - VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; - queue_info.queueFamilyIndex = d.queue_family; - queue_info.queueCount = 1; - queue_info.pQueuePriorities = &priority; - VkDeviceCreateInfo device_info{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO}; - device_info.pNext = &features; - device_info.queueCreateInfoCount = 1; - device_info.pQueueCreateInfos = &queue_info; - device_info.enabledExtensionCount = 1; - device_info.ppEnabledExtensionNames = extensions; - Check(vkCreateDevice(d.physical, &device_info, nullptr, &d.device), "vkCreateDevice"); - vkGetDeviceQueue(d.device, d.queue_family, 0, &d.queue); - - VkCommandPoolCreateInfo pool_info{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; - pool_info.queueFamilyIndex = d.queue_family; - Check(vkCreateCommandPool(d.device, &pool_info, nullptr, &d.pool), "vkCreateCommandPool"); - return d; -} - -// The CPU reference is O(M*N*K), so only shapes up to this many multiply-adds are checked. -constexpr uint64_t kMaxCheckedMultiplyAdds = 1ull << 31; - -// One GEMM run: fills A and B with deterministic data, times `iterations` dispatches, then checks shapes small enough -// for the CPU reference. -double RunGemm(Device& d, VkPipeline pipeline, VkPipelineLayout layout, uint32_t M, uint32_t N, uint32_t K, - int iterations, bool check) { - const uint32_t rows = (M + kTile - 1) / kTile * kTile; // the kernel runs whole tiles - const bool checked = check && static_cast(M) * N * K <= kMaxCheckedMultiplyAdds; - const VkDeviceSize a_bytes = static_cast(rows) * K * 2; - const VkDeviceSize b_bytes = static_cast(N) * K * 2; - const VkDeviceSize c_bytes = static_cast(rows) * N * 4; - - std::vector a(static_cast(rows) * K); - std::vector b(static_cast(N) * K); - std::mt19937 rng(1234); - std::uniform_real_distribution dist(-1.0f, 1.0f); - for (auto& v : a) v = FloatToHalf(dist(rng)); - for (auto& v : b) v = FloatToHalf(dist(rng)); - // NaN in the padding rows: any leak into the first M rows of C fails the check. - if (checked) std::fill(a.begin() + static_cast(static_cast(M) * K), a.end(), uint16_t{0x7e00}); - - Buffer buffer_a = CreateBuffer(d, a_bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); - Buffer buffer_b = CreateBuffer(d, b_bytes, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); - Buffer buffer_c = CreateBuffer(d, c_bytes, - VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | - VK_BUFFER_USAGE_TRANSFER_SRC_BIT, - VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, true); - Upload(d, buffer_a, a.data(), a_bytes); - Upload(d, buffer_b, b.data(), b_bytes); - - struct Push { - uint64_t a; - uint64_t b; - uint64_t c; - uint32_t M; - uint32_t N; - uint32_t K; - uint32_t lda; - uint32_t ldb; - uint32_t ldc; - } push{buffer_a.address, buffer_b.address, buffer_c.address, M, N, K, K, K, N}; - - const auto record = [&](VkCommandBuffer cb) { - vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline); - vkCmdPushConstants(cb, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(push), &push); - for (int i = 0; i < iterations; ++i) { - vkCmdDispatch(cb, N / kTile, rows / kTile, 1); - Barrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_WRITE_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT); - } - }; - - Submit(d, record); // warm-up (compiles or loads the pipeline cache) - const auto start = std::chrono::steady_clock::now(); - Submit(d, record); - const double seconds = - std::chrono::duration(std::chrono::steady_clock::now() - start).count() / static_cast(iterations); - - double max_error = -1.0; - if (checked) { - std::vector c(static_cast(rows) * N); - Download(d, buffer_c, c.data(), c_bytes); - double max_abs = 0.0; - for (uint32_t m = 0; m < M; ++m) { - for (uint32_t n = 0; n < N; ++n) { - double reference = 0.0; - for (uint32_t k = 0; k < K; ++k) { - reference += static_cast(HalfToFloat(a[static_cast(m) * K + k])) * - static_cast(HalfToFloat(b[static_cast(n) * K + k])); - } - const double error = std::fabs(static_cast(c[static_cast(m) * N + n]) - reference); - max_abs = std::isnan(error) ? INFINITY : std::max(max_abs, error); - } - } - max_error = max_abs; - std::printf("check %ux%ux%u: max|C - CPU| = %.6g\n", M, N, K, max_abs); - } - - const double flops = 2.0 * M * N * K; - std::printf("gemm %ux%ux%u: %8.3f ms %7.1f TFLOPS\n", M, N, K, seconds * 1e3, flops / seconds / 1e12); - - DestroyBuffer(d, buffer_a); - DestroyBuffer(d, buffer_b); - DestroyBuffer(d, buffer_c); - return max_error; -} - -} // namespace - -int main(int argc, char** argv) { - std::string spv = DEFAULT_SPV_PATH; - bool check = true; - int device_index = -1; - struct Shape { - uint32_t M, N, K; - int iterations; - }; - std::vector shapes; - for (int i = 1; i < argc; ++i) { - const std::string arg = argv[i]; - if (arg == "--spv" && i + 1 < argc) { - spv = argv[++i]; - } else if (arg == "--device" && i + 1 < argc) { - device_index = std::atoi(argv[++i]); - } else if (arg == "--no-check") { - check = 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 [--spv ] [--device N] [--no-check] [--bench M N K iters]...\n", argv[0]); - return 2; - } - } - if (shapes.empty()) { - // Checked: one tile, a partial last band, several bands, and row tails. Then the engine's real shapes at 1024x1024 - // with CFG: the DiT's QKV, projection, up/gate and down GEMMs, and the text encoder's GatedDeltaNet projections. - shapes = {{64, 64, 64, 1}, {128, 640, 64, 1}, {256, 1024, 64, 1}, {94, 640, 64, 1}, - {1030, 9216, 64, 1}, {512, 512, 512, 20}, {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) { - if (s.M == 0 || s.N == 0 || s.K == 0 || s.iterations <= 0 || s.N % kTile != 0 || s.K % 16 != 0) { - std::printf("N must be a multiple of %u and K of 16, and all must be positive: %ux%ux%u\n", kTile, s.M, s.N, s.K); - return 2; - } - } - - Device d = OpenDevice(device_index); - PrintCooperativeMatrixShapes(d); - - const std::vector spirv = ReadFile(spv); - VkShaderModuleCreateInfo module_info{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; - module_info.codeSize = spirv.size() * 4; - module_info.pCode = spirv.data(); - VkShaderModule module = VK_NULL_HANDLE; - Check(vkCreateShaderModule(d.device, &module_info, nullptr, &module), "vkCreateShaderModule"); - - VkPushConstantRange push_range{}; - push_range.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT; - push_range.offset = 0; - push_range.size = 48; - VkPipelineLayoutCreateInfo layout_info{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; - layout_info.pushConstantRangeCount = 1; - layout_info.pPushConstantRanges = &push_range; - VkPipelineLayout layout = VK_NULL_HANDLE; - Check(vkCreatePipelineLayout(d.device, &layout_info, nullptr, &layout), "vkCreatePipelineLayout"); - - VkComputePipelineCreateInfo pipeline_info{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; - // The kernel's tile split assumes exactly four full subgroups of 32. Without this a driver may pick another size. - VkPipelineShaderStageRequiredSubgroupSizeCreateInfo subgroup_size{ - VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; - subgroup_size.requiredSubgroupSize = kSubgroupSize; - pipeline_info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; - pipeline_info.stage.pNext = &subgroup_size; - pipeline_info.stage.flags = VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT; - pipeline_info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT; - pipeline_info.stage.module = module; - pipeline_info.stage.pName = "main"; - pipeline_info.layout = layout; - VkPipeline pipeline = VK_NULL_HANDLE; - Check(vkCreateComputePipelines(d.device, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &pipeline), - "vkCreateComputePipelines"); - - // Inputs are fp16 and the accumulator is fp32 (no fp64 reference is exact), so this bounds rounding, not bits. - constexpr double kMaxError = 1e-2; - bool failed = false; - for (const Shape& s : shapes) { - const double error = RunGemm(d, pipeline, layout, s.M, s.N, s.K, s.iterations, check); - if (error > kMaxError) { - std::printf("FAILED: %ux%ux%u: error %.6g is above %g\n", s.M, s.N, s.K, error, kMaxError); - failed = true; - } - } - - vkDestroyPipeline(d.device, pipeline, nullptr); - vkDestroyPipelineLayout(d.device, layout, nullptr); - vkDestroyShaderModule(d.device, module, nullptr); - vkDestroyCommandPool(d.device, d.pool, nullptr); - vkDestroyDevice(d.device, nullptr); - vkDestroyInstance(d.instance, nullptr); - return failed ? 1 : 0; -} diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt new file mode 100644 index 0000000..01ea2ac --- /dev/null +++ b/vulkan/CMakeLists.txt @@ -0,0 +1,35 @@ +# The Vulkan backend, the fallback for hardware the engine does not otherwise support. A standalone project for now: it +# links no engine code, and nothing here affects the engine's build. See README.md. +cmake_minimum_required(VERSION 3.20) +project(mossinfer_vulkan 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() + +find_package(Vulkan REQUIRED) +find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN_SDK}/bin" REQUIRED) + +# 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 "") +foreach(shader IN ITEMS gemm_coop gemm_plain) + 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.3 -fshader-stage=comp -O -mfmt=num ${_source} -o ${_output} + DEPENDS ${_source} + COMMENT "glslc ${shader}.comp") + list(APPEND _spirv_files ${_output}) +endforeach() + +add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp ${_spirv_files}) +target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) +target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) + +add_executable(vulkan_check tools/vulkan_check.cpp) +target_link_libraries(vulkan_check PRIVATE mossinfer_vulkan) diff --git a/vulkan/README.md b/vulkan/README.md new file mode 100644 index 0000000..a75452e --- /dev/null +++ b/vulkan/README.md @@ -0,0 +1,67 @@ +# Vulkan backend + +The fallback for hardware the engine does not otherwise support. CUDA stays the path on every GPU it supports. This is a +standalone project for now: it links no engine code and nothing here affects the engine's build. Design and research: +[design](../docs/vulkan-backend-design.md), [research](../docs/vulkan-backend-research.md). + +What exists is the device runtime and the GEMM. It cannot generate images yet. + +## Build and check + +```sh +cmake -S vulkan -B build-vulkan +cmake --build build-vulkan +build-vulkan/vulkan_check --list # each device, and which GEMM it gets +build-vulkan/vulkan_check # checks against an fp64 CPU reference, then the engine's shapes +build-vulkan/vulkan_check --path plain --no-bench +``` + +It needs the Vulkan headers and loader, and `glslc`. Shaders compile into the binary. The exit code is non-zero when a check +fails. A check poisons C and uses strides larger than the rows, so a write past a row or past M fails it. + +## GEMM + +`C[M,N] = A[M,K] * B[N,K]^T`, with B in the engine's weight layout, f32 accumulation and f32 output. M and N are free. + +| Path | Used when | Inputs | K | +|---|---|---|---| +| Cooperative matrix (`gemm_coop.comp`) | the device has 16x16x16 f16 x f16 -> f32 in subgroup scope, a compute subgroup size fixable at 32, and 24 KiB of shared memory | f16 | multiple of 16 | +| Plain compute (`gemm_plain.comp`) | every other device with Vulkan 1.2, buffer device addresses and 64-bit integers | f16 or bf16 | multiple of 4 | + +`--list` says why a device gets the plain path. The plain shader declares only `Shader`, `Int64` and +`PhysicalStorageBufferAddresses`: it reads 32-bit words, so it needs no 16-bit support and does not depend on the subgroup +size. + +## Results + +TFLOPS on an RTX 4090 Laptop, driver 615.78, with cuBLASLt (f16 inputs, f32 compute) on the same GPU. M = 8222 is the DiT at +1024x1024 with CFG. + +| Shape (M x N x K) | Cooperative matrix | Plain | cuBLASLt | +|---|---:|---:|---:| +| 8222 x 9216 x 3072 (QKV) | 27.0 | 1.39 | 52.6 | +| 8222 x 3072 x 3072 (projection) | 33.1 | 1.62 | 57.4 | +| 8222 x 16384 x 3072 (up/gate) | 23.5 | 1.14 | 45.5 | +| 8222 x 3072 x 8192 (down) | 24.3 | 1.52 | 70.1 | +| 282 x 12288 x 4096 (text encoder) | 27.3 | 1.02 | 59.5 | +| 282 x 4096 x 12288 (text encoder) | 18.1 | | 52.2 | + +Checks pass on that GPU on both paths, and through the plain path on llvmpipe (Mesa's CPU driver, 8-wide subgroups, 32 KiB of +shared memory), with core validation clean. + +## Numerics + +The cooperative-matrix path rounds like the tensor cores: its error matches cuBLASLt on the same data up to K = 1040 and is +biased toward zero. At K = 4096 it reaches 8.1e-4 against cuBLASLt's 1.3e-4, on values near 80. The plain path has no such +bias. + +## Files + +| File | Purpose | +|---|---| +| `runtime.h`, `runtime.cpp` | device selection, capabilities, buffers, pipelines, barriers | +| `gemm.h`, `gemm.cpp` | path selection, argument checks, dispatch | +| `half.h` | host fp16 and bf16 conversions | +| `shaders/` | compute shaders | +| `tools/vulkan_check.cpp` | device list, checks, timing | +| `tools/probe.py` | a device's cooperative-matrix shapes and heaps, using only the Vulkan loader (no headers) | diff --git a/vulkan/gemm.cpp b/vulkan/gemm.cpp new file mode 100644 index 0000000..38bbdcd --- /dev/null +++ b/vulkan/gemm.cpp @@ -0,0 +1,77 @@ +#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 kPlainTile = 8; +constexpr std::uint32_t kCoopSubgroup = 32; + +struct Push { + std::uint64_t a, b, c; + std::uint32_t m, n, k, lda, ldb, ldc; +}; + +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(Push), 0, + std::vector{0}); + plain_bf16_ = std::make_unique(device, kPlainSpirv, sizeof(kPlainSpirv), sizeof(Push), 0, + std::vector{1}); + } +} + +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"); + + const Push push{args.a->address(), args.b->address(), args.c->address(), args.m, args.n, args.k, args.lda, args.ldb, + args.ldc}; + if (path_ == GemmPath::kCoop) { + coop_->dispatch(cb, &push, tiles(args.n, kCoopTile), tiles(args.m, kCoopTile)); + } else { + const Pipeline& pipeline = input == InputType::kBF16 ? *plain_bf16_ : *plain_f16_; + pipeline.dispatch(cb, &push, tiles(args.n, kPlainTile), tiles(args.m, kPlainTile)); + } + 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/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/runtime.cpp b/vulkan/runtime.cpp new file mode 100644 index 0000000..ac27c08 --- /dev/null +++ b/vulkan/runtime.cpp @@ -0,0 +1,372 @@ +#include "runtime.h" + +#include +#include +#include + +namespace mi::vk { +namespace { + +// 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 = 4; + +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() { + 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; }); +} + +std::string why_no_coop(VkInstance instance, 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 4 subgroups per workgroup"; + if (properties.limits.maxComputeSharedMemorySize < kCoopSharedBytes) { + return "needs " + std::to_string(kCoopSharedBytes) + " bytes of shared memory"; + } + + auto query = reinterpret_cast( + vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); + if (query == nullptr) return "no cooperative matrix query"; + std::uint32_t count = 0; + if (query(physical, &count, nullptr) != VK_SUCCESS) return "cooperative matrix query failed"; + 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 "cooperative matrix query failed"; + const bool has_shape = std::any_of(shapes.begin(), shapes.end(), [](const VkCooperativeMatrixPropertiesKHR& s) { + return s.MSize == 16 && s.NSize == 16 && s.KSize == 16 && s.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && + s.BType == VK_COMPONENT_TYPE_FLOAT16_KHR && s.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && + s.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR && s.saturatingAccumulation == VK_FALSE && + s.scope == VK_SCOPE_SUBGROUP_KHR; + }); + if (!has_shape) return "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; + return {}; +} + +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; + 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"; + info.caps.coop_reason = why_no_coop(instance, physical, properties, info.caps); + info.caps.coop = info.caps.coop_reason.empty(); + 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) { + instance_ = create_instance(); + const std::vector devices = physical_devices(instance_); + 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); + + // 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; + std::vector extensions; + if (info_.caps.coop) { + 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 + } + + const float priority = 1.0f; + VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; + queue_info.queueFamilyIndex = family; + queue_info.queueCount = 1; + queue_info.pQueuePriorities = &priority; + VkDeviceCreateInfo device_info{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO}; + device_info.pNext = &features; + device_info.queueCreateInfoCount = 1; + device_info.pQueueCreateInfos = &queue_info; + 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"); +} + +Device::~Device() { + 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::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; + check(vkBeginCommandBuffer(cb, &begin), "vkBeginCommandBuffer"); + record(cb); + check(vkEndCommandBuffer(cb), "vkEndCommandBuffer"); + VkSubmitInfo submit{VK_STRUCTURE_TYPE_SUBMIT_INFO}; + submit.commandBufferCount = 1; + submit.pCommandBuffers = &cb; + const VkResult submitted = vkQueueSubmit(queue_, 1, &submit, VK_NULL_HANDLE); + const VkResult waited = submitted == VK_SUCCESS ? vkQueueWaitIdle(queue_) : submitted; + vkFreeCommandBuffers(device_, pool_, 1, &cb); + check(submitted, "vkQueueSubmit"); + check(waited, "vkQueueWaitIdle"); +} + +Buffer::Buffer(Device& device, VkDeviceSize bytes, Kind kind) : device_(device), bytes_(bytes) { + 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; + 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) { + Buffer staging(device_, bytes, Kind::kStaging); + std::memcpy(staging.mapped(), data, static_cast(bytes)); + device_.run([&](VkCommandBuffer cb) { + const VkBufferCopy copy{0, 0, 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) { + Buffer staging(device_, bytes, Kind::kStaging); + device_.run([&](VkCommandBuffer cb) { + const VkBufferCopy copy{0, 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_; + check(vkCreateComputePipelines(device_.handle(), VK_NULL_HANDLE, 1, &info, nullptr, &pipeline_), + "vkCreateComputePipelines"); +} + +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); +} + +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); +} + +} // namespace mi::vk diff --git a/vulkan/runtime.h b/vulkan/runtime.h new file mode 100644 index 0000000..263cc0b --- /dev/null +++ b/vulkan/runtime.h @@ -0,0 +1,113 @@ +// 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 + +#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 4 per workgroup, and its shared memory. Otherwise `coop_reason` says why not. + bool coop = false; + std::string coop_reason; + std::uint32_t shared_memory_bytes = 0; + std::uint32_t min_subgroup = 0, max_subgroup = 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_; } + + // Records into one command buffer, submits it and waits. + void run(const std::function& record); + + VkDevice handle() const { return device_; } + 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_; +}; + +// A device-local storage buffer with a device address, or a host-visible staging buffer. +class Buffer { + public: + enum class Kind { kDevice, kStaging }; + Buffer(Device& device, VkDeviceSize bytes, Kind kind = Kind::kDevice); + ~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_; } + + void upload(const void* data, VkDeviceSize bytes); + void download(void* data, VkDeviceSize bytes); + + 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: + 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); + +} // namespace mi::vk 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..3698abc --- /dev/null +++ b/vulkan/shaders/gemm_plain.comp @@ -0,0 +1,47 @@ +// C[M,N] = A[M,K] * B[N,K]^T without cooperative matrices, for any device with buffer device addresses and 64-bit +// integers: 16-bit inputs (f16, or bf16 when INPUT_BF16 is set), 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 4. +// +// Inputs are read as 32-bit words, so the shader needs neither 16-bit storage nor 16-bit arithmetic, and it does not +// depend on the subgroup size. Each invocation computes one output element of an 8x8 tile. +#version 460 +#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require +#extension GL_EXT_buffer_reference : require + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 0) const bool INPUT_BF16 = false; + +layout(buffer_reference, std430, buffer_reference_align = 4) readonly buffer Words { uvec2 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 TILE = 8u; + +vec4 decode(uvec2 words) { + if (INPUT_BF16) { + return uintBitsToFloat(uvec4(words.x << 16, words.x & 0xffff0000u, words.y << 16, words.y & 0xffff0000u)); + } + return vec4(unpackHalf2x16(words.x), unpackHalf2x16(words.y)); +} + +void main() { + const uint m = gl_WorkGroupID.y * TILE + gl_LocalInvocationIndex / TILE; + const uint n = gl_WorkGroupID.x * TILE + gl_LocalInvocationIndex % TILE; + if (m >= pc.M || n >= pc.N) return; + Words A = Words(pc.a + (uint64_t(m) * pc.lda) * 2ul); + Words B = Words(pc.b + (uint64_t(n) * pc.ldb) * 2ul); + float sum = 0.0; + for (uint i = 0u; i < pc.K / 4u; ++i) sum += dot(decode(A.d[i]), decode(B.d[i])); + F32(pc.c + (uint64_t(m) * pc.ldc + n) * 4ul).d[0] = sum; +} diff --git a/tools/vulkan_gemm/probe.py b/vulkan/tools/probe.py similarity index 100% rename from tools/vulkan_gemm/probe.py rename to vulkan/tools/probe.py diff --git a/vulkan/tools/vulkan_check.cpp b/vulkan/tools/vulkan_check.cpp new file mode 100644 index 0000000..6164446 --- /dev/null +++ b/vulkan/tools/vulkan_check.cpp @@ -0,0 +1,199 @@ +// 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 "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()); + } +} + +} // 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, 1, 16, 1}, {3, 8, 16, 1}, {64, 64, 64, 1}, {65, 65, 48, 1}, + {67, 136, 1040, 1}, {129, 256, 4096, 1}, {128, 640, 64, 1}, {94, 643, 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; + } +} From bb04f52fadfc58dea48a0536b902c7fbe3009773 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 00:19:39 -0500 Subject: [PATCH 05/41] vulkan: the DiT block's plain compute kernels, checked against host replays 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/CMakeLists.txt | 21 +- vulkan/block.cpp | 166 +++++++++++ vulkan/block.h | 77 ++++++ vulkan/gemm.cpp | 23 +- vulkan/shaders/attention_plain.comp | 135 +++++++++ vulkan/shaders/common.glsl | 26 ++ vulkan/shaders/gemm_plain.comp | 109 ++++++-- vulkan/shaders/norm_modulate.comp | 63 +++++ vulkan/shaders/qk_norm_rope.comp | 68 +++++ vulkan/tools/vulkan_block_check.cpp | 413 ++++++++++++++++++++++++++++ vulkan/tools/vulkan_check.cpp | 4 +- 11 files changed, 1064 insertions(+), 41 deletions(-) create mode 100644 vulkan/block.cpp create mode 100644 vulkan/block.h create mode 100644 vulkan/shaders/attention_plain.comp create mode 100644 vulkan/shaders/common.glsl create mode 100644 vulkan/shaders/norm_modulate.comp create mode 100644 vulkan/shaders/qk_norm_rope.comp create mode 100644 vulkan/tools/vulkan_block_check.cpp diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 01ea2ac..8c6f164 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -16,20 +16,31 @@ find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN set(_spirv_dir ${CMAKE_CURRENT_BINARY_DIR}/spirv) file(MAKE_DIRECTORY ${_spirv_dir}) set(_spirv_files "") -foreach(shader IN ITEMS gemm_coop gemm_plain) +# The cooperative-matrix shader needs Vulkan 1.3; the rest must also load on Vulkan 1.2 devices, which reject newer SPIR-V. +function(_mi_add_shader shader target_env) 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.3 -fshader-stage=comp -O -mfmt=num ${_source} -o ${_output} - DEPENDS ${_source} + COMMAND ${GLSL_COMPILER} --target-env=${target_env} -fshader-stage=comp -O -mfmt=num + -I${CMAKE_CURRENT_SOURCE_DIR}/shaders ${_source} -o ${_output} + DEPENDS ${_source} ${CMAKE_CURRENT_SOURCE_DIR}/shaders/common.glsl COMMENT "glslc ${shader}.comp") - list(APPEND _spirv_files ${_output}) + set(_spirv_files ${_spirv_files} ${_output} PARENT_SCOPE) +endfunction() +_mi_add_shader(gemm_coop vulkan1.3) +foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain) + _mi_add_shader(${shader} vulkan1.2) endforeach() -add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp ${_spirv_files}) +add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp ${_spirv_files}) target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) 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 -ffp-contract=off) +target_link_libraries(vulkan_block_check PRIVATE mossinfer_vulkan) diff --git a/vulkan/block.cpp b/vulkan/block.cpp new file mode 100644 index 0000000..c436f67 --- /dev/null +++ b/vulkan/block.cpp @@ -0,0 +1,166 @@ +#include "block.h" + +#include +#include +#include + +namespace mi::vk { +namespace { + +const std::uint32_t kGemmSpirv[] = { +#include "gemm_plain.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 kAttentionSpirv[] = { +#include "attention_plain.spv.inc" +}; + +struct GemmPush { + std::uint64_t a, b, out, gate; + std::uint32_t m, n, k, lda, ldb, ldo; +}; +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 { + std::uint64_t q, k, v, o; + std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_row0, k_row0, seqlen_k; + float log2_scale; +}; + +enum GemmMode : std::uint32_t { kF32 = 0, kBF16 = 1, kGated = 2, kSitu = 3 }; + +template +std::unique_ptr gemm_pipeline(Device& device, const Spirv& spirv, GemmMode mode) { + return std::make_unique(device, spirv, sizeof(spirv), sizeof(GemmPush), 0, + std::vector{1u, static_cast(mode)}); +} + +void require(bool condition, const std::string& message) { + if (!condition) throw std::invalid_argument("vulkan block: " + message); +} + +constexpr std::uint32_t kGemmTileRows = 8, kGemmTileCols = 32; + +void dispatch_gemm(const Pipeline& pipeline, VkCommandBuffer cb, const GemmPush& push) { + require(push.m > 0 && push.n % 4 == 0 && push.n > 0 && push.k % 4 == 0 && push.k > 0, "bad GEMM shape"); + require(push.lda % 4 == 0 && push.ldb % 4 == 0 && push.ldo % 4 == 0, "GEMM strides must be multiples of 4"); + pipeline.dispatch(cb, &push, (push.n + kGemmTileCols - 1) / kGemmTileCols, (push.m + kGemmTileRows - 1) / kGemmTileRows); + compute_barrier(cb); +} + +} // namespace + +BlockKernels::BlockKernels(Device& device) { + gemm_bf16_ = gemm_pipeline(device, kGemmSpirv, kBF16); + gemm_gated_ = gemm_pipeline(device, kGemmSpirv, kGated); + gemm_situ_ = gemm_pipeline(device, kGemmSpirv, kSitu); + norm_modulate_ = std::make_unique(device, kNormSpirv, sizeof(kNormSpirv), sizeof(NormPush)); + qk_norm_rope_ = std::make_unique(device, kQkSpirv, sizeof(kQkSpirv), sizeof(QkPush)); + 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_, cb, GemmPush{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_, cb, GemmPush{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_, cb, GemmPush{x, w_up_gate, out, 0, m, f, k, ldx, ldw, ldo}); +} + +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); + 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; + 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_k), log2_scale}; + attention_->dispatch(cb, &push, a.heads, static_cast(seqlen_q)); + } + 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..9ad082b --- /dev/null +++ b/vulkan/block.h @@ -0,0 +1,77 @@ +// 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) +// It runs on any device the runtime accepts, needs no cooperative matrices, and does not depend on the subgroup size. +// Operands are device addresses with row strides that are multiples of 4 elements. Each record_* adds its own barrier. +#pragma once + +#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 +}; + +class BlockKernels { + public: + explicit BlockKernels(Device& device); + ~BlockKernels(); + + 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; + 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: + std::unique_ptr gemm_bf16_, gemm_gated_, gemm_situ_, norm_modulate_, qk_norm_rope_, attention_; +}; + +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/gemm.cpp b/vulkan/gemm.cpp index 38bbdcd..93ad70e 100644 --- a/vulkan/gemm.cpp +++ b/vulkan/gemm.cpp @@ -14,13 +14,17 @@ const std::uint32_t kPlainSpirv[] = { }; constexpr std::uint32_t kCoopTile = 64; -constexpr std::uint32_t kPlainTile = 8; +constexpr std::uint32_t kPlainTileRows = 8, kPlainTileCols = 32; 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); @@ -42,10 +46,10 @@ Gemm::Gemm(Device& device, GemmPath path) : path_(path) { } coop_ = std::make_unique(device, kCoopSpirv, sizeof(kCoopSpirv), sizeof(Push), kCoopSubgroup); } else { - plain_f16_ = std::make_unique(device, kPlainSpirv, sizeof(kPlainSpirv), sizeof(Push), 0, - std::vector{0}); - plain_bf16_ = std::make_unique(device, kPlainSpirv, sizeof(kPlainSpirv), sizeof(Push), 0, - std::vector{1}); + 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}); } } @@ -62,14 +66,17 @@ void Gemm::record(VkCommandBuffer cb, const GemmArgs& args, InputType input) con 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"); - const Push push{args.a->address(), args.b->address(), args.c->address(), args.m, args.n, args.k, args.lda, args.ldb, - args.ldc}; 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, kPlainTile), tiles(args.m, kPlainTile)); + pipeline.dispatch(cb, &push, tiles(args.n, kPlainTileCols), tiles(args.m, kPlainTileRows)); } compute_barrier(cb); } diff --git a/vulkan/shaders/attention_plain.comp b/vulkan/shaders/attention_plain.comp new file mode 100644 index 0000000..aa9ca54 --- /dev/null +++ b/vulkan/shaders/attention_plain.comp @@ -0,0 +1,135 @@ +// Scaled-dot-product attention over one sequence, head_dim 128, non-causal, bf16 in and out, fp32 online softmax. One query +// row of one head per workgroup, 32 lanes. Key tokens walk from the last step of 32 down to step 0, as FlashAttention's +// reverse iteration: +// S = q * K_step^T +// row_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 +// acc_o += P * V_step^T +// and at the end out = bf16(acc_o / row_sum). Lane c owns key c of each step for S and P, and dims 4c..4c+3 for acc_o. +#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 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_k; + float log2_scale; +} pc; + +const uint STEP = 32u; + +shared vec4 q_shared[32]; +shared float s_shared[STEP]; +shared float w_shared[STEP]; +shared float p_shared[STEP]; +shared float max_scaled_shared; +shared float rescale_shared; +shared float inv_sum_shared; + +// 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 q_row = uint64_t(pc.q_row0) + gl_WorkGroupID.y; + const uint64_t k_row0 = pc.k_row0; + + q_shared[lane] = f4(Words(pc.q + (q_row * pc.ldq + uint64_t(head) * pc.hsq) * 2ul).d[lane]); + sync_shared(); + + const uint64_t k_head = pc.k + (k_row0 * pc.ldk + uint64_t(head) * pc.hsk) * 2ul; + const uint64_t v_head = pc.v + (k_row0 * pc.ldv + uint64_t(head) * pc.hsv) * 2ul; + const uint64_t ldk_bytes = uint64_t(pc.ldk) * 2ul, ldv_bytes = uint64_t(pc.ldv) * 2ul; + + precise vec4 acc = vec4(0.0); + float row_max = 0.0, row_sum = 0.0; // lane 0's + const uint steps = (pc.seqlen_k + STEP - 1u) / STEP; + 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 (lane < cols) { + Words key = Words(k_head + uint64_t(n0 + lane) * ldk_bytes); + precise float s = 0.0; + for (uint g = 0u; g < 32u; ++g) { + const vec4 q4 = q_shared[g]; + const vec4 k4 = f4(key.d[g]); + s = s + q4.x * k4.x; + s = s + q4.y * k4.y; + s = s + q4.z * k4.z; + s = s + q4.w * k4.w; + } + s_shared[lane] = s; + } + sync_shared(); + + if (lane == 0u) { + const float prev = row_max; + float cur = first ? s_shared[0] : max(prev, s_shared[0]); + for (uint c = 1u; c < cols; ++c) cur = max(cur, s_shared[c]); + row_max = cur; + if (first) { + rescale_shared = 1.0; + row_sum = 0.0; + } else { + precise const float diff = (prev - cur) * pc.log2_scale; + const float rescale = exp2(diff); + rescale_shared = rescale; + precise const float scaled = row_sum * rescale; + row_sum = scaled; + } + precise const float max_scaled = isinf(cur) && cur < 0.0 ? 0.0 : cur * pc.log2_scale; + max_scaled_shared = max_scaled; + } + sync_shared(); + + float w = 0.0; + if (lane < cols) { + // FlashAttention's one fused step; a device without a fused fma rounds it twice, within the bf16 weight's error. + w = exp2(fma(s_shared[lane], pc.log2_scale, -max_scaled_shared)); + } + w_shared[lane] = w; + p_shared[lane] = uintBitsToFloat(bf(w) << 16); + if (!first) { + precise const vec4 rescaled = acc * rescale_shared; + acc = rescaled; + } + sync_shared(); + + if (lane == 0u) { + precise float sum = row_sum; + for (uint c = 0u; c < cols; ++c) sum = sum + w_shared[c]; + row_sum = sum; + } + + for (uint c = 0u; c < cols; ++c) { + const vec4 v4 = f4(Words(v_head + uint64_t(n0 + c) * ldv_bytes).d[lane]); + acc = acc + p_shared[c] * v4; + } + sync_shared(); + } + + if (lane == 0u) inv_sum_shared = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + sync_shared(); + precise const vec4 result = acc * inv_sum_shared; + Words(pc.o + (q_row * pc.ldo + uint64_t(head) * pc.hso) * 2ul).d[lane] = bf4(result); +} diff --git a/vulkan/shaders/common.glsl b/vulkan/shaders/common.glsl new file mode 100644 index 0000000..322f9a2 --- /dev/null +++ b/vulkan/shaders/common.glsl @@ -0,0 +1,26 @@ +// 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)); } + +#endif diff --git a/vulkan/shaders/gemm_plain.comp b/vulkan/shaders/gemm_plain.comp index 3698abc..2f90d5a 100644 --- a/vulkan/shaders/gemm_plain.comp +++ b/vulkan/shaders/gemm_plain.comp @@ -1,47 +1,104 @@ -// C[M,N] = A[M,K] * B[N,K]^T without cooperative matrices, for any device with buffer device addresses and 64-bit -// integers: 16-bit inputs (f16, or bf16 when INPUT_BF16 is set), 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 4. +// 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 accumulators 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. // -// Inputs are read as 32-bit words, so the shader needs neither 16-bit storage nor 16-bit arithmetic, and it does not -// depend on the subgroup size. Each invocation computes one output element of an 8x8 tile. +// Each invocation computes 4 adjacent output columns of one row, in tiles of 8 rows by 32 columns. 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 #version 460 -#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require -#extension GL_EXT_buffer_reference : require +#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 bool INPUT_BF16 = false; - -layout(buffer_reference, std430, buffer_reference_align = 4) readonly buffer Words { uvec2 d[]; }; -layout(buffer_reference, std430, buffer_reference_align = 4) writeonly buffer F32 { float d[]; }; +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 c; + uint64_t out_; + uint64_t gate; uint M; uint N; uint K; uint lda; uint ldb; - uint ldc; + uint ldo; } pc; -const uint TILE = 8u; +vec4 decode(uvec2 w) { return INPUT_BF16 ? f4(w) : vec4(unpackHalf2x16(w.x), unpackHalf2x16(w.y)); } -vec4 decode(uvec2 words) { - if (INPUT_BF16) { - return uintBitsToFloat(uvec4(words.x << 16, words.x & 0xffff0000u, words.y << 16, words.y & 0xffff0000u)); - } - return vec4(unpackHalf2x16(words.x), unpackHalf2x16(words.y)); +// exp(-2 |x|) stays in range where GLSL's tanh may divide infinities; beyond 12 the result rounds to 1 in fp32 anyway. +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); + return sign(x) * t; +} + +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; } void main() { - const uint m = gl_WorkGroupID.y * TILE + gl_LocalInvocationIndex / TILE; - const uint n = gl_WorkGroupID.x * TILE + gl_LocalInvocationIndex % TILE; + const uint inv = gl_LocalInvocationIndex; + const uint m = gl_WorkGroupID.y * 8u + inv / 8u; + const uint n = (gl_WorkGroupID.x * 8u + inv % 8u) * 4u; if (m >= pc.M || n >= pc.N) return; - Words A = Words(pc.a + (uint64_t(m) * pc.lda) * 2ul); - Words B = Words(pc.b + (uint64_t(n) * pc.ldb) * 2ul); - float sum = 0.0; - for (uint i = 0u; i < pc.K / 4u; ++i) sum += dot(decode(A.d[i]), decode(B.d[i])); - F32(pc.c + (uint64_t(m) * pc.ldc + n) * 4ul).d[0] = sum; + + Words A = Words(pc.a + uint64_t(m) * pc.lda * 2ul); + const uint64_t ldb_bytes = uint64_t(pc.ldb) * 2ul; + const uint64_t b0 = pc.b + uint64_t(n) * ldb_bytes; + Words B0 = Words(b0), B1 = Words(b0 + ldb_bytes), B2 = Words(b0 + 2ul * ldb_bytes), B3 = Words(b0 + 3ul * ldb_bytes); + const uint64_t g0 = pc.b + uint64_t(pc.N + n) * ldb_bytes; + Words G0 = Words(g0), G1 = Words(g0 + ldb_bytes), G2 = Words(g0 + 2ul * ldb_bytes), G3 = Words(g0 + 3ul * ldb_bytes); + + precise vec4 acc = vec4(0.0), acc_gate = vec4(0.0); + const uint words = pc.K / 4u; + for (uint i = 0u; i < words; ++i) { + const vec4 x = decode(A.d[i]); + const vec4 c0 = decode(B0.d[i]), c1 = decode(B1.d[i]), c2 = decode(B2.d[i]), c3 = decode(B3.d[i]); + acc = acc + x.x * vec4(c0.x, c1.x, c2.x, c3.x); + acc = acc + x.y * vec4(c0.y, c1.y, c2.y, c3.y); + acc = acc + x.z * vec4(c0.z, c1.z, c2.z, c3.z); + acc = acc + x.w * vec4(c0.w, c1.w, c2.w, c3.w); + if (MODE == 3) { + const vec4 d0 = decode(G0.d[i]), d1 = decode(G1.d[i]), d2 = decode(G2.d[i]), d3 = decode(G3.d[i]); + acc_gate = acc_gate + x.x * vec4(d0.x, d1.x, d2.x, d3.x); + acc_gate = acc_gate + x.y * vec4(d0.y, d1.y, d2.y, d3.y); + acc_gate = acc_gate + x.z * vec4(d0.z, d1.z, d2.z, d3.z); + acc_gate = acc_gate + x.w * vec4(d0.w, d1.w, d2.w, d3.w); + } + } + + if (MODE == 0) { + Floats out_f = Floats(pc.out_ + (uint64_t(m) * pc.ldo + n) * 4ul); + out_f.d[0] = acc.x; + out_f.d[1] = acc.y; + out_f.d[2] = acc.z; + out_f.d[3] = acc.w; + return; + } + Words out_w = Words(pc.out_ + (uint64_t(m) * pc.ldo + n) * 2ul); + if (MODE == 1) { + out_w.d[0] = bf4(acc); + } else if (MODE == 2) { + const vec4 gate = f4(Words(pc.gate + uint64_t(n) * 2ul).d[0]); + precise const vec4 gated = rbf4(acc) * gate; + precise const vec4 sum = f4(out_w.d[0]) + rbf4(gated); + out_w.d[0] = bf4(sum); + } else { + out_w.d[0] = bf4(situ(rbf4(acc), rbf4(acc_gate))); + } } diff --git a/vulkan/shaders/norm_modulate.comp b/vulkan/shaders/norm_modulate.comp new file mode 100644 index 0000000..08cc9df --- /dev/null +++ b/vulkan/shaders/norm_modulate.comp @@ -0,0 +1,63 @@ +// out[rows,d] = bf16(bf16(bf16(bf16(rms_norm(x)) * w) * bf16(1 + scale)) + shift); w, scale, shift: [d]. 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]); + 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/tools/vulkan_block_check.cpp b/vulkan/tools/vulkan_block_check.cpp new file mode 100644 index 0000000..f0109cf --- /dev/null +++ b/vulkan/tools/vulkan_block_check.cpp @@ -0,0 +1,413 @@ +// 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; +}; + +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)); + } + 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}); + + 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; +} + +bool check_gemms(Device& device, const BlockKernels& kernels) { + constexpr int m = 37, n = 104, k = 40, f = 24, ld = 3 * k; + const auto x_h = synth_bf16(21, static_cast(m) * ld, 0), 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); + 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); + } + } + 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}); + 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}) && 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}) && 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::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 == "--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] [--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()); + const BlockKernels kernels(device); + bool ok = check_norms(device, kernels); + ok = check_gemms(device, kernels) && 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_check.cpp b/vulkan/tools/vulkan_check.cpp index 6164446..937ea33 100644 --- a/vulkan/tools/vulkan_check.cpp +++ b/vulkan/tools/vulkan_check.cpp @@ -174,8 +174,8 @@ int main(int argc, char** argv) { : 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, 1, 16, 1}, {3, 8, 16, 1}, {64, 64, 64, 1}, {65, 65, 48, 1}, - {67, 136, 1040, 1}, {129, 256, 4096, 1}, {128, 640, 64, 1}, {94, 643, 64, 1}, + 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; From 68a985262ae700c4a46d1e4c7abe715db8d31ebb Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 00:22:05 -0500 Subject: [PATCH 06/41] vulkan: check real DiT blocks against the engine's output 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. --- docs/vulkan-backend-design.md | 2 +- vulkan/CMakeLists.txt | 4 + vulkan/README.md | 26 +++++- vulkan/tools/vulkan_block_real.cpp | 145 +++++++++++++++++++++++++++++ 4 files changed, 174 insertions(+), 3 deletions(-) create mode 100644 vulkan/tools/vulkan_block_real.cpp diff --git a/docs/vulkan-backend-design.md b/docs/vulkan-backend-design.md index 33ea531..3c0d50d 100644 --- a/docs/vulkan-backend-design.md +++ b/docs/vulkan-backend-design.md @@ -208,7 +208,7 @@ of A. N is a multiple of 64 in every engine GEMM. |---|---|---|---| | **V0** | Probe tool + GEMM microbenchmark; done | GEMM correct and measured at the engine's shapes against cuBLASLt | done | | **V1** | Vulkan runtime + f16 GEMM with shared-memory tiles and a plain compute path; done, except the suballocator and pipeline cache, which wait until the block runner needs them | Checks at engine shapes on a device without the fast path too | done | -| **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights | Block output matches CUDA block (cosine ≥ 0.999) | 1–2 weeks | +| **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights; done on the plain compute path | Block output matches the engine's block (cosine ≥ 0.999) | done | | **V3** | Text encoder (GEMM, RMSNorm, RoPE, causal attention, GDN) and VAE decoder | Encode/decode parity; first full Vulkan image | 2–3 weeks | | **V4** | Pipeline integration (runtime abstraction, planner, streaming, CLI/server) | `moss-infer --backend vulkan` and `moss-server` produce images on the 2080 | 2–4 weeks | | **V5** | Quantized paths (INT8/INT4), packaging, CI, docs | `high`/`turbo` tiers on Vulkan; `-vulkan` package; tests in CI | 1–2 weeks | diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 8c6f164..6003f30 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -44,3 +44,7 @@ target_link_libraries(vulkan_check PRIVATE mossinfer_vulkan) add_executable(vulkan_block_check tools/vulkan_block_check.cpp) target_compile_options(vulkan_block_check PRIVATE -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) diff --git a/vulkan/README.md b/vulkan/README.md index a75452e..2ad6f07 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -4,7 +4,7 @@ The fallback for hardware the engine does not otherwise support. CUDA stays the standalone project for now: it links no engine code and nothing here affects the engine's build. Design and research: [design](../docs/vulkan-backend-design.md), [research](../docs/vulkan-backend-research.md). -What exists is the device runtime and the GEMM. It cannot generate images yet. +What exists is the device runtime, the GEMM and one DiT block's kernels. It cannot generate images yet: the text encoder, the VAE, the DiT's outer layers and the pipeline are still to come. ## Build and check @@ -49,6 +49,25 @@ TFLOPS on an RTX 4090 Laptop, driver 615.78, with cuBLASLt (f16 inputs, f32 comp Checks pass on that GPU on both paths, and through the plain path on llvmpipe (Mesa's CPU driver, 8-wide subgroups, 32 KiB of shared memory), with core validation clean. +## DiT block + +`block.h` is the plain compute path of one bf16 DiT block: RMSNorm with modulation, the GEMMs with their gated residual and SiTU-GLU epilogues, QK norm with RoPE, and varlen attention with an online softmax. It has the engine's rounding points and needs no cooperative matrices, 16-bit support or particular subgroup size. + +```sh +build-vulkan/vulkan_block_check # each kernel against a host replay of its op order, then one block +build-vulkan/vulkan_block_real DIR # real blocks against the engine's output, from files captured from its forward +``` + +On real weights and activations at the 1024x1024 golden case, against the engine: + +| Block | Same bits | p99 scaled ULPs | Max scaled ULPs | Cosine | +|---:|---:|---:|---:|---:| +| 1 | 83.98 % | 0.06 | 2.00 | 0.999999977 | +| 22 | 72.36 % | 0.25 | 4.00 | 0.999999867 | +| 41 | 69.23 % | 1.00 | 96.00 | 0.999997331 | + +Bits differ because fp32 scalar and tensor-core accumulation are not bit-compatible. Known-answer checks pass on an RTX 4090 Laptop and on llvmpipe. One block of 8484 tokens takes about 2 s on that GPU through this path, against 31 ms for the engine, so about 86 s per CFG step at 1024x1024. + ## Numerics The cooperative-matrix path rounds like the tensor cores: its error matches cuBLASLt on the same data up to K = 1040 and is @@ -61,7 +80,10 @@ bias. |---|---| | `runtime.h`, `runtime.cpp` | device selection, capabilities, buffers, pipelines, barriers | | `gemm.h`, `gemm.cpp` | path selection, argument checks, dispatch | +| `block.h`, `block.cpp` | one DiT block's plain compute kernels | | `half.h` | host fp16 and bf16 conversions | | `shaders/` | compute shaders | -| `tools/vulkan_check.cpp` | device list, checks, timing | +| `tools/vulkan_check.cpp` | device list, GEMM checks, timing | +| `tools/vulkan_block_check.cpp` | the block's kernels against host replays, one block's timing | +| `tools/vulkan_block_real.cpp` | real blocks against the engine's captured output | | `tools/probe.py` | a device's cooperative-matrix shapes and heaps, using only the Vulkan loader (no headers) | diff --git a/vulkan/tools/vulkan_block_real.cpp b/vulkan/tools/vulkan_block_real.cpp new file mode 100644 index 0000000..65643e5 --- /dev/null +++ b/vulkan/tools/vulkan_block_real.cpp @@ -0,0 +1,145 @@ +// Runs one DiT block's plain compute 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]\n", argv[0]); + return 2; + } + const fs::path dir = argv[1]; + const int device_index = argc >= 4 && std::string(argv[2]) == "--device" ? std::atoi(argv[3]) : -1; + try { + Device device(device_index); + std::printf("device %u: %s\n", device.info().index, device.info().name.c_str()); + const BlockKernels kernels(device); + + 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; + } +} From 36d68e9fc27afeff681def0972a7b35dfbcec61e Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:04:46 -0500 Subject: [PATCH 07/41] vulkan: the DiT block on bf16 cooperative matrices, GEMMs with epilogues 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. --- vulkan/CMakeLists.txt | 2 + vulkan/block.cpp | 89 ++++++++-- vulkan/block.h | 13 +- vulkan/runtime.cpp | 87 +++++++--- vulkan/runtime.h | 6 +- vulkan/shaders/attention_coop.comp | 252 ++++++++++++++++++++++++++++ vulkan/shaders/common.glsl | 21 +++ vulkan/shaders/gemm_bf16_coop.comp | 178 ++++++++++++++++++++ vulkan/shaders/gemm_plain.comp | 21 --- vulkan/tools/vulkan_block_check.cpp | 23 ++- 10 files changed, 627 insertions(+), 65 deletions(-) create mode 100644 vulkan/shaders/attention_coop.comp create mode 100644 vulkan/shaders/gemm_bf16_coop.comp diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 6003f30..9166f7e 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -29,6 +29,8 @@ function(_mi_add_shader shader target_env) set(_spirv_files ${_spirv_files} ${_output} 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) foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain) _mi_add_shader(${shader} vulkan1.2) endforeach() diff --git a/vulkan/block.cpp b/vulkan/block.cpp index c436f67..e6c4f62 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -10,6 +10,9 @@ 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" }; @@ -19,6 +22,9 @@ const std::uint32_t kQkSpirv[] = { const std::uint32_t kAttentionSpirv[] = { #include "attention_plain.spv.inc" }; +const std::uint32_t kAttentionCoopSpirv[] = { +#include "attention_coop.spv.inc" +}; struct GemmPush { std::uint64_t a, b, out, gate; @@ -40,12 +46,26 @@ struct AttentionPush { float log2_scale; }; +struct CoopAttentionPush { + 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; + enum GemmMode : std::uint32_t { kF32 = 0, kBF16 = 1, kGated = 2, kSitu = 3 }; -template -std::unique_ptr gemm_pipeline(Device& device, const Spirv& spirv, GemmMode mode) { - return std::make_unique(device, spirv, sizeof(spirv), sizeof(GemmPush), 0, - std::vector{1u, static_cast(mode)}); +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}); + } + return std::make_unique(device, kGemmSpirv, sizeof(kGemmSpirv), sizeof(GemmPush), 0, + std::vector{1u, m}); } void require(bool condition, const std::string& message) { @@ -54,22 +74,45 @@ void require(bool condition, const std::string& message) { constexpr std::uint32_t kGemmTileRows = 8, kGemmTileCols = 32; -void dispatch_gemm(const Pipeline& pipeline, VkCommandBuffer cb, const GemmPush& push) { - require(push.m > 0 && push.n % 4 == 0 && push.n > 0 && push.k % 4 == 0 && push.k > 0, "bad GEMM shape"); - require(push.lda % 4 == 0 && push.ldb % 4 == 0 && push.ldo % 4 == 0, "GEMM strides must be multiples of 4"); - pipeline.dispatch(cb, &push, (push.n + kGemmTileCols - 1) / kGemmTileCols, (push.m + kGemmTileRows - 1) / kGemmTileRows); +void dispatch_gemm(const Pipeline& pipeline, BlockPath path, GemmMode mode, VkCommandBuffer cb, const GemmPush& push) { + const std::uint32_t align = path == BlockPath::kCoop ? 8 : 4; + require(push.m > 0 && push.n > 0 && push.k > 0, "bad GEMM shape"); + require(push.n % align == 0 && push.k % align == 0 && push.lda % align == 0 && push.ldb % align == 0 && + push.ldo % align == 0, + "GEMM N, K and strides must be multiples of " + std::to_string(align)); + if (path == BlockPath::kCoop) { + // 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, (push.n + cols - 1) / cols, (push.m + kCoopTileRows - 1) / kCoopTileRows); + } else { + pipeline.dispatch(cb, &push, (push.n + kGemmTileCols - 1) / kGemmTileCols, (push.m + kGemmTileRows - 1) / kGemmTileRows); + } compute_barrier(cb); } } // namespace -BlockKernels::BlockKernels(Device& device) { - gemm_bf16_ = gemm_pipeline(device, kGemmSpirv, kBF16); - gemm_gated_ = gemm_pipeline(device, kGemmSpirv, kGated); - gemm_situ_ = gemm_pipeline(device, kGemmSpirv, kSitu); +const char* block_path_name(BlockPath path) { return path == BlockPath::kCoop ? "coop" : "plain"; } + +BlockKernels::BlockKernels(Device& device) + : BlockKernels(device, device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) {} + +BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { + if (path == BlockPath::kCoop && !device.info().caps.coop_bf16) { + throw std::runtime_error(device.info().name + " cannot run the cooperative-matrix DiT block: " + + device.info().caps.coop_bf16_reason); + } + gemm_bf16_ = gemm_pipeline(device, path, kBF16); + gemm_gated_ = gemm_pipeline(device, path, kGated); + gemm_situ_ = gemm_pipeline(device, path, kSitu); norm_modulate_ = std::make_unique(device, kNormSpirv, sizeof(kNormSpirv), sizeof(NormPush)); qk_norm_rope_ = std::make_unique(device, kQkSpirv, sizeof(kQkSpirv), sizeof(QkPush)); - attention_ = std::make_unique(device, kAttentionSpirv, sizeof(kAttentionSpirv), sizeof(AttentionPush)); + if (path == BlockPath::kCoop) { + attention_ = std::make_unique(device, kAttentionCoopSpirv, sizeof(kAttentionCoopSpirv), + sizeof(CoopAttentionPush), kCoopSubgroup); + } else { + attention_ = std::make_unique(device, kAttentionSpirv, sizeof(kAttentionSpirv), sizeof(AttentionPush)); + } } BlockKernels::~BlockKernels() = default; @@ -89,20 +132,20 @@ void BlockKernels::record_norm_modulate(VkCommandBuffer cb, VkDeviceAddress x, s 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_, cb, GemmPush{x, w, out, 0, m, n, k, ldx, ldw, ldo}); + dispatch_gemm(*gemm_bf16_, path_, kBF16, cb, GemmPush{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_, cb, GemmPush{x, w, residual, gate, m, n, k, ldx, ldw, ldr}); + dispatch_gemm(*gemm_gated_, path_, kGated, cb, GemmPush{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_, cb, GemmPush{x, w_up_gate, out, 0, m, f, k, ldx, ldw, ldo}); + dispatch_gemm(*gemm_situ_, path_, kSitu, cb, GemmPush{x, w_up_gate, out, 0, m, f, k, ldx, ldw, ldo}); } void BlockKernels::record_qk_norm_rope(VkCommandBuffer cb, VkDeviceAddress qkv, std::uint32_t ld, std::uint32_t tokens, @@ -122,9 +165,23 @@ void BlockKernels::record_attention(VkCommandBuffer cb, const AttentionArgs& a) 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 (path_ == BlockPath::kCoop) { + 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"); + if (path_ == BlockPath::kCoop) { + const CoopAttentionPush 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}; + attention_->dispatch(cb, &push, a.heads, + (static_cast(seqlen_q) + kCoopAttentionRows - 1) / kCoopAttentionRows); + continue; + } 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]), diff --git a/vulkan/block.h b/vulkan/block.h index 9ad082b..899703f 100644 --- a/vulkan/block.h +++ b/vulkan/block.h @@ -7,8 +7,9 @@ // 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) -// It runs on any device the runtime accepts, needs no cooperative matrices, and does not depend on the subgroup size. -// Operands are device addresses with row strides that are multiples of 4 elements. Each record_* adds its own barrier. +// Two 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 on bf16 cooperative +// matrices (Capabilities::coop_bf16) and needs strides, N and K that are multiples of 8. Each record_* adds its own barrier. #pragma once #include @@ -33,11 +34,18 @@ struct AttentionArgs { float scale; // SDPA's 1 / sqrt(head_dim) for the DiT }; +enum class BlockPath { kCoop, kPlain }; +const char* block_path_name(BlockPath path); + 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; @@ -55,6 +63,7 @@ class BlockKernels { void record_attention(VkCommandBuffer cb, const AttentionArgs& args) const; private: + BlockPath path_; std::unique_ptr gemm_bf16_, gemm_gated_, gemm_situ_, norm_modulate_, qk_norm_rope_, attention_; }; diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index ac27c08..7e99a62 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -10,7 +10,9 @@ namespace { // 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 = 4; +constexpr std::uint32_t kCoopSubgroupsPerWorkgroup = 8; +// The bf16 DiT kernels' largest shared memory: 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)); @@ -44,8 +46,9 @@ bool has_extension(VkPhysicalDevice physical, const char* name) { [&](const VkExtensionProperties& e) { return std::strcmp(e.extensionName, name) == 0; }); } -std::string why_no_coop(VkInstance instance, VkPhysicalDevice physical, const VkPhysicalDeviceProperties& properties, - Capabilities& caps) { +// 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}; @@ -70,27 +73,65 @@ std::string why_no_coop(VkInstance instance, VkPhysicalDevice physical, const Vk 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 4 subgroups per workgroup"; - if (properties.limits.maxComputeSharedMemorySize < kCoopSharedBytes) { - return "needs " + std::to_string(kCoopSharedBytes) + " bytes of shared memory"; - } + if (p13.maxComputeWorkgroupSubgroups < kCoopSubgroupsPerWorkgroup) return "needs 8 subgroups per workgroup"; + return {}; +} +std::vector coop_shapes(VkInstance instance, VkPhysicalDevice physical) { auto query = reinterpret_cast( vkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR")); - if (query == nullptr) return "no cooperative matrix query"; + if (query == nullptr) return {}; std::uint32_t count = 0; - if (query(physical, &count, nullptr) != VK_SUCCESS) return "cooperative matrix query failed"; + 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 "cooperative matrix query failed"; - const bool has_shape = std::any_of(shapes.begin(), shapes.end(), [](const VkCooperativeMatrixPropertiesKHR& s) { - return s.MSize == 16 && s.NSize == 16 && s.KSize == 16 && s.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && - s.BType == VK_COMPONENT_TYPE_FLOAT16_KHR && s.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && - s.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR && s.saturatingAccumulation == VK_FALSE && - s.scope == VK_SCOPE_SUBGROUP_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; }); - if (!has_shape) return "no 16x16x16 f16 x f16 -> f32 cooperative matrix in subgroup scope"; - return {}; +} + +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(); } DeviceInfo describe(VkInstance instance, VkPhysicalDevice physical, std::uint32_t index) { @@ -112,8 +153,7 @@ DeviceInfo describe(VkInstance instance, VkPhysicalDevice physical, std::uint32_ vkGetPhysicalDeviceFeatures2(physical, &features); if (v12.bufferDeviceAddress == 0) info.unusable_reason = "no bufferDeviceAddress"; if (features.features.shaderInt64 == 0) info.unusable_reason = "no shaderInt64"; - info.caps.coop_reason = why_no_coop(instance, physical, properties, info.caps); - info.caps.coop = info.caps.coop_reason.empty(); + describe_coop(instance, physical, properties, info.caps); return info; } @@ -171,8 +211,15 @@ Device::Device(int index) { 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) { + 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) { v12.pNext = &v13; v12.shaderFloat16 = VK_TRUE; v12.vulkanMemoryModel = VK_TRUE; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 263cc0b..83f9bea 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -15,9 +15,13 @@ 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 4 per workgroup, and its shared memory. Otherwise `coop_reason` says why not. + // 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; std::uint32_t shared_memory_bytes = 0; std::uint32_t min_subgroup = 0, max_subgroup = 0; }; diff --git a/vulkan/shaders/attention_coop.comp b/vulkan/shaders/attention_coop.comp new file mode 100644 index 0000000..faa4f49 --- /dev/null +++ b/vulkan/shaders/attention_coop.comp @@ -0,0 +1,252 @@ +// 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. +#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_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] = rk[j]; + smem[SV + at] = rv[j]; + } +} + +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] = v; + } + 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] = uvec4(p_bits[0], p_bits[1], p_bits[2], p_bits[3]); + smem[sp + frow * P_ROW + half_ * 2u + 1u] = uvec4(p_bits[4], p_bits[5], p_bits[6], p_bits[7]); + + // 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. + const float inv_sum = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + 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/common.glsl b/vulkan/shaders/common.glsl index 322f9a2..ccac00b 100644 --- a/vulkan/shaders/common.glsl +++ b/vulkan/shaders/common.glsl @@ -23,4 +23,25 @@ uint bf(float v) { 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)); } +// exp(-2 |x|) stays in range where GLSL's tanh may divide infinities; beyond 12 the result rounds to 1 in fp32 anyway. +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); + return sign(x) * t; +} + +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/gemm_bf16_coop.comp b/vulkan/shaders/gemm_bf16_coop.comp new file mode 100644 index 0000000..3e174c0 --- /dev/null +++ b/vulkan/shaders/gemm_bf16_coop.comp @@ -0,0 +1,178 @@ +// 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 +// +// 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. +#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 int MODE = 1; + +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 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 = 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; } + +// 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]; + } +} + +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); + } + + 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; + [[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 vec4 lo = staged(base0, frow, fcol / 4u), hi = staged(base0, frow, fcol / 4u + 1u); + 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 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 { + const vec4 g_lo = staged(base1, frow, fcol / 4u), g_hi = staged(base1, frow, fcol / 4u + 1u); + Vec4s(pc.out_ + at * 2ul).d[0] = + uvec4(bf4(situ(rbf4(lo), rbf4(g_lo))), bf4(situ(rbf4(hi), rbf4(g_hi)))); + } + } + sync_subgroup(); + } + } +} diff --git a/vulkan/shaders/gemm_plain.comp b/vulkan/shaders/gemm_plain.comp index 2f90d5a..4802620 100644 --- a/vulkan/shaders/gemm_plain.comp +++ b/vulkan/shaders/gemm_plain.comp @@ -30,27 +30,6 @@ layout(push_constant) uniform PC { vec4 decode(uvec2 w) { return INPUT_BF16 ? f4(w) : vec4(unpackHalf2x16(w.x), unpackHalf2x16(w.y)); } -// exp(-2 |x|) stays in range where GLSL's tanh may divide infinities; beyond 12 the result rounds to 1 in fp32 anyway. -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); - return sign(x) * t; -} - -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; -} - void main() { const uint inv = gl_LocalInvocationIndex; const uint m = gl_WorkGroupID.y * 8u + inv / 8u; diff --git a/vulkan/tools/vulkan_block_check.cpp b/vulkan/tools/vulkan_block_check.cpp index f0109cf..4887d01 100644 --- a/vulkan/tools/vulkan_block_check.cpp +++ b/vulkan/tools/vulkan_block_check.cpp @@ -262,8 +262,9 @@ bool check_norms(Device& device, const BlockKernels& kernels) { return report("qk_norm_rope", read(*qkv, want.size()), rope_ref, {3.0, 1.0}) && ok; } -bool check_gemms(Device& device, const BlockKernels& kernels) { - constexpr int m = 37, n = 104, k = 40, f = 24, ld = 3 * k; +// m, n and k: tails on every axis of both paths' tiles; the larger shape also crosses tiles and bands of tiles. +bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int k, int f) { + const int ld = 3 * k; const auto x_h = synth_bf16(21, static_cast(m) * ld, 0), 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); auto x = device_buffer(device, x_h), w = device_buffer(device, w_h), gate = device_buffer(device, gate_h), @@ -282,6 +283,7 @@ bool check_gemms(Device& device, const BlockKernels& kernels) { situ[static_cast(r) * f + c] = host_situ(h, g); } } + std::printf(" GEMMs at M %d, N %d, K %d, SiTU F %d\n", m, n, k, f); 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}); out->upload(res_h.data(), res_h.size() * 2); @@ -374,6 +376,7 @@ bool run_block(Device& device, const BlockKernels& kernels, const std::vector rows = {4378, 4106}; for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; @@ -381,6 +384,8 @@ int main(int argc, char** argv) { 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) { @@ -389,16 +394,24 @@ int main(int argc, char** argv) { std::string item; while (std::getline(ss, item, ',')) rows.push_back(std::atoi(item.c_str())); } else { - std::printf("usage: %s [--device N] [--rows L1,L2,...] [--iters N] [--no-block]\n", argv[0]); + std::printf("usage: %s [--device N] [--path coop|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()); - const BlockKernels kernels(device); + if (!path_name.empty() && path_name != "coop" && path_name != "plain") { + std::printf("--path takes coop or plain\n"); + return 2; + } + const BlockKernels kernels = path_name.empty() ? BlockKernels(device) + : BlockKernels(device, path_name == "coop" ? BlockPath::kCoop : BlockPath::kPlain); + std::printf("path: %s%s%s\n", block_path_name(kernels.path()), device.info().caps.coop_bf16 ? "" : " (no coop: ", + device.info().caps.coop_bf16 ? "" : (device.info().caps.coop_bf16_reason + ")").c_str()); bool ok = check_norms(device, kernels); - ok = check_gemms(device, kernels) && ok; + ok = check_gemms(device, kernels, 37, 104, 40, 24) && ok; + ok = check_gemms(device, kernels, 300, 1048, 264, 520) && ok; ok = check_attention(device, kernels) && ok; if (!ok) { std::printf("known-answer check FAILED: not timing.\n"); From d94623fd3b9847d2da8f18c4a219444ce1a6ec22 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:09:59 -0500 Subject: [PATCH 08/41] vulkan: safetensors and npy readers, buffer copies at an offset --- vulkan/CMakeLists.txt | 2 +- vulkan/npy.h | 97 +++++++++++++++ vulkan/runtime.cpp | 10 +- vulkan/runtime.h | 5 +- vulkan/safetensors.cpp | 262 +++++++++++++++++++++++++++++++++++++++++ vulkan/safetensors.h | 45 +++++++ 6 files changed, 414 insertions(+), 7 deletions(-) create mode 100644 vulkan/npy.h create mode 100644 vulkan/safetensors.cpp create mode 100644 vulkan/safetensors.h diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 9166f7e..458a221 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -35,7 +35,7 @@ foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain) _mi_add_shader(${shader} vulkan1.2) endforeach() -add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp ${_spirv_files}) +add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp safetensors.cpp ${_spirv_files}) target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) diff --git a/vulkan/npy.h b/vulkan/npy.h new file mode 100644 index 0000000..988c24d --- /dev/null +++ b/vulkan/npy.h @@ -0,0 +1,97 @@ +// Minimal .npy reader and writer (format 1-3 headers, little-endian, C order) for the Vulkan backend's tools. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mi::vk::npy { + +struct Array { + std::string descr; // e.g. " shape; + std::vector data; + + std::int64_t numel() const { + std::int64_t n = 1; + for (const std::int64_t d : shape) n *= d; + return n; + } + template + std::vector as(const char* want) const { + if (descr != want) throw std::runtime_error("npy: dtype " + descr + ", expected " + want); + std::vector out(data.size() / sizeof(T)); + if (!out.empty()) std::memcpy(out.data(), data.data(), out.size() * sizeof(T)); + return out; + } +}; + +inline Array load(const std::filesystem::path& path) { + std::ifstream in(path, std::ios::binary); + if (!in) throw std::runtime_error("npy: cannot open " + path.string()); + char magic[8]; + if (!in.read(magic, 8) || std::memcmp(magic, "\x93NUMPY", 6) != 0) throw std::runtime_error("npy: bad magic in " + path.string()); + std::uint32_t header_len = 0; + if (magic[6] == 1) { + unsigned char b[2]; + in.read(reinterpret_cast(b), 2); + header_len = b[0] | (b[1] << 8); + } else { + unsigned char b[4]; + in.read(reinterpret_cast(b), 4); + header_len = b[0] | (b[1] << 8) | (b[2] << 16) | (static_cast(b[3]) << 24); + } + std::string header(header_len, '\0'); + in.read(header.data(), header_len); + Array a; + const auto field = [&](const char* key) { + const std::size_t at = header.find(key); + if (at == std::string::npos) throw std::runtime_error(std::string("npy: no ") + key + " in " + path.string()); + return at + std::strlen(key); + }; + std::size_t at = header.find('\'', field("'descr':")) + 1; + a.descr = header.substr(at, header.find('\'', at) - at); + if (header.find("True", field("'fortran_order':")) == field("'fortran_order':") + 1) { + throw std::runtime_error("npy: Fortran order in " + path.string()); + } + at = header.find('(', field("'shape':")) + 1; + const std::size_t end = header.find(')', at); + std::string dims = header.substr(at, end - at); + std::size_t pos = 0; + while (pos < dims.size()) { + const std::size_t comma = dims.find(',', pos); + const std::string item = dims.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + if (item.find_first_of("0123456789") != std::string::npos) a.shape.push_back(std::stoll(item)); + if (comma == std::string::npos) break; + pos = comma + 1; + } + const std::size_t item_bytes = static_cast(std::stoi(a.descr.substr(2))); + a.data.resize(static_cast(a.numel()) * item_bytes); + if (!in.read(reinterpret_cast(a.data.data()), static_cast(a.data.size()))) { + throw std::runtime_error("npy: short data in " + path.string()); + } + return a; +} + +inline void save(const std::filesystem::path& path, const char* descr, const std::vector& 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 index 7e99a62..eb6ad28 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -323,11 +323,12 @@ Buffer::~Buffer() { vkFreeMemory(device_.handle(), memory_, nullptr); } -void Buffer::upload(const void* data, VkDeviceSize bytes) { +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, 0, bytes}; + 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; @@ -337,10 +338,11 @@ void Buffer::upload(const void* data, VkDeviceSize bytes) { }); } -void Buffer::download(void* data, VkDeviceSize bytes) { +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{0, 0, bytes}; + 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; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 83f9bea..772e61b 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -80,8 +80,9 @@ class Buffer { VkDeviceSize bytes() const { return bytes_; } void* mapped() const { return mapped_; } - void upload(const void* data, VkDeviceSize bytes); - void download(void* data, VkDeviceSize bytes); + // 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_; diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp new file mode 100644 index 0000000..644061b --- /dev/null +++ b/vulkan/safetensors.cpp @@ -0,0 +1,262 @@ +#include "safetensors.h" + +#include +#include +#include +#include + +namespace mi::vk::st { +namespace { + +[[noreturn]] void fail(const std::string& what) { throw std::runtime_error("safetensors: " + what); } + +// Just enough JSON for safetensors headers and index files: objects, arrays, strings without \u escapes beyond ASCII +// passthrough, integers, true/false/null. +struct Json { + enum class Kind { kNull, kBool, kNumber, kString, kArray, kObject } kind = Kind::kNull; + std::string text; // string value, or the number's spelling + 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() { + Json j = value(); + space(); + if (i_ != s_.size()) fail("trailing characters in JSON"); + return j; + } + + private: + void space() { + while (i_ < s_.size() && (s_[i_] == ' ' || s_[i_] == '\n' || s_[i_] == '\r' || s_[i_] == '\t')) ++i_; + } + char peek() { + space(); + 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_; + } + std::string string() { + expect('"'); + std::string out; + while (true) { + if (i_ >= s_.size()) fail("unterminated JSON string"); + const char c = s_[i_++]; + if (c == '"') return out; + if (c != '\\') { + out += c; + continue; + } + if (i_ >= s_.size()) fail("unterminated JSON escape"); + const char e = s_[i_++]; + switch (e) { + case 'n': out += '\n'; break; + case 't': out += '\t'; break; + case 'r': out += '\r'; break; + case 'b': out += '\b'; break; + case 'f': out += '\f'; break; + case 'u': { + if (i_ + 4 > s_.size()) fail("short \\u escape"); + const unsigned code = static_cast(std::stoul(s_.substr(i_, 4), nullptr, 16)); + i_ += 4; + // Names and dtypes are ASCII; anything else only has to survive as some bytes. + if (code < 0x80) out += static_cast(code); + else out += '?'; + break; + } + default: out += e; + } + } + } + Json value() { + Json j; + const char c = peek(); + if (c == '{') { + ++i_; + j.kind = Json::Kind::kObject; + if (peek() == '}') { + ++i_; + return j; + } + while (true) { + std::string key = string(); + expect(':'); + j.members.emplace_back(std::move(key), value()); + if (peek() == ',') { + ++i_; + continue; + } + expect('}'); + return j; + } + } + if (c == '[') { + ++i_; + j.kind = Json::Kind::kArray; + if (peek() == ']') { + ++i_; + return j; + } + while (true) { + j.items.push_back(value()); + if (peek() == ',') { + ++i_; + continue; + } + expect(']'); + return j; + } + } + if (c == '"') { + j.kind = Json::Kind::kString; + j.text = string(); + return j; + } + if (s_.compare(i_, 4, "true") == 0 || s_.compare(i_, 5, "false") == 0) { + j.kind = Json::Kind::kBool; + j.text = s_[i_] == 't' ? "true" : "false"; + i_ += j.text.size(); + return j; + } + if (s_.compare(i_, 4, "null") == 0) { + i_ += 4; + return j; + } + const std::size_t start = i_; + while (i_ < s_.size() && (std::strchr("+-0123456789.eE", s_[i_]) != nullptr)) ++i_; + if (start == i_) fail("bad JSON value"); + j.kind = Json::Kind::kNumber; + j.text = s_.substr(start, i_ - start); + return j; + } + + const std::string& s_; + std::size_t i_ = 0; +}; + +std::string read_text(const std::filesystem::path& path) { + std::ifstream in(path, std::ios::binary); + if (!in) fail("cannot open " + path.string()); + return {std::istreambuf_iterator(in), std::istreambuf_iterator()}; +} + +std::uint64_t to_u64(const Json& j) { + if (j.kind != Json::Kind::kNumber) fail("expected a number"); + return std::stoull(j.text); +} + +std::size_t dtype_bytes(const std::string& dtype) { + if (dtype == "F64" || dtype == "I64" || dtype == "U64") return 8; + if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; + if (dtype == "F16" || dtype == "BF16" || dtype == "I16" || dtype == "U16") return 2; + if (dtype == "I8" || dtype == "U8" || dtype == "BOOL" || dtype == "F8_E4M3" || dtype == "F8_E5M2") return 1; + fail("unknown dtype " + dtype); +} + +} // namespace + +Index::Index(const std::filesystem::path& path) { + if (!std::filesystem::is_directory(path)) { + add_file(path); + return; + } + std::vector indexes, files; + for (const auto& entry : std::filesystem::directory_iterator(path)) { + const std::string name = entry.path().filename().string(); + if (name.size() > 22 && 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) fail(path.string() + " holds more than one index file"); + if (indexes.empty()) { + if (files.size() != 1) fail(path.string() + " holds no index file and not exactly one .safetensors file"); + add_file(files.front()); + return; + } + const Json index = Parser(read_text(indexes.front())).parse(); + const Json* map = index.find("weight_map"); + if (map == nullptr || map->kind != Json::Kind::kObject) fail(indexes.front().string() + " has no weight_map"); + std::set shards; + for (const auto& [name, file] : map->members) shards.insert(file.text); + for (const std::string& shard : shards) add_file(path / shard); + for (const auto& [name, file] : map->members) { + if (!has(name)) fail(name + " is in the index but in no shard"); + } +} + +void Index::add_file(const std::filesystem::path& file) { + std::ifstream in(file, std::ios::binary); + if (!in) fail("cannot open " + file.string()); + std::uint64_t header_bytes = 0; + unsigned char size[8]; + if (!in.read(reinterpret_cast(size), 8)) fail(file.string() + " is too short"); + for (int i = 7; i >= 0; --i) header_bytes = (header_bytes << 8) | size[i]; + if (header_bytes > (std::uint64_t{100} << 20)) fail(file.string() + " has an implausible header size"); + std::string header(header_bytes, '\0'); + if (!in.read(header.data(), static_cast(header_bytes))) fail(file.string() + " has a short header"); + const std::uint64_t data_start = 8 + header_bytes; + const std::uint64_t file_bytes = std::filesystem::file_size(file); + + const Json root = Parser(header).parse(); + if (root.kind != Json::Kind::kObject) fail(file.string() + ": the header is not an object"); + for (const auto& [name, t] : root.members) { + if (name == "__metadata__") continue; + const Json* dtype = t.find("dtype"); + const Json* shape = t.find("shape"); + const Json* offsets = t.find("data_offsets"); + if (dtype == nullptr || shape == nullptr || offsets == nullptr || offsets->items.size() != 2) { + fail(file.string() + ": malformed entry " + name); + } + TensorInfo info; + info.dtype = dtype->text; + std::uint64_t count = 1; + for (const Json& d : shape->items) { + info.shape.push_back(static_cast(to_u64(d))); + count *= to_u64(d); + } + const std::uint64_t begin = to_u64(offsets->items[0]), end = to_u64(offsets->items[1]); + if (end < begin || end - begin != count * dtype_bytes(info.dtype) || data_start + end > file_bytes) { + fail(file.string() + ": bad extent for " + name); + } + info.offset = data_start + begin; + info.bytes = end - begin; + info.file = file; + if (!tensors_.emplace(name, std::move(info)).second) fail(name + " appears twice"); + } +} + +const TensorInfo& Index::at(const std::string& name) const { + const auto it = tensors_.find(name); + if (it == tensors_.end()) fail("no tensor " + name); + return it->second; +} + +std::vector Index::read(const std::string& name) const { + std::vector out(at(name).bytes); + read_into(name, out.data(), out.size()); + return out; +} + +void Index::read_into(const std::string& name, void* out, std::uint64_t bytes) const { + const TensorInfo& info = at(name); + if (bytes != info.bytes) fail(name + ": read of " + std::to_string(bytes) + " bytes, tensor has " + std::to_string(info.bytes)); + std::ifstream in(info.file, std::ios::binary); + if (!in) fail("cannot open " + info.file.string()); + in.seekg(static_cast(info.offset)); + if (!in.read(static_cast(out), static_cast(bytes))) fail("short read of " + name); +} + +} // namespace mi::vk::st diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h new file mode 100644 index 0000000..df6a3d1 --- /dev/null +++ b/vulkan/safetensors.h @@ -0,0 +1,45 @@ +// Minimal safetensors reader for the Vulkan backend: a file's header, and one tensor's bytes on demand. A sharded model is +// a set of files named by its model.safetensors.index.json (or diffusion_pytorch_model.safetensors.index.json). +// Reads with std::ifstream, so it is the same on every platform. Failures throw std::runtime_error. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace mi::vk::st { + +struct TensorInfo { + std::string dtype; // as the header spells it: "BF16", "F32", "F16", "I64", ... + std::vector shape; + std::uint64_t offset = 0; // from the start of the file + std::uint64_t bytes = 0; + std::filesystem::path file; +}; + +// The tensors of one file, or of every shard of a model directory. +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); + + const std::map& tensors() const { return tensors_; } + bool has(const std::string& name) const { return tensors_.count(name) != 0; } + // Throws if the tensor is missing. + const TensorInfo& at(const std::string& name) const; + + // Reads a tensor's bytes, or throws. + std::vector read(const std::string& name) const; + // Reads into caller memory of exactly info.bytes. + void read_into(const std::string& name, void* out, std::uint64_t bytes) const; + + private: + void add_file(const std::filesystem::path& file); + std::map tensors_; +}; + +} // namespace mi::vk::st From f143f0dedb7e2875f19dc7f4d6cfab16a502e23a Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:16:32 -0500 Subject: [PATCH 09/41] vulkan: the whole DiT forward and the text refiner, and a sampling loop 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. --- vulkan/CMakeLists.txt | 11 +- vulkan/block.cpp | 52 ++++- vulkan/block.h | 17 +- vulkan/dit.cpp | 297 ++++++++++++++++++++++++++++ vulkan/dit.h | 87 ++++++++ vulkan/runtime.cpp | 20 ++ vulkan/runtime.h | 5 + vulkan/shaders/elementwise.comp | 29 +++ vulkan/shaders/gemm_bf16_coop.comp | 15 +- vulkan/shaders/gemm_plain.comp | 12 +- vulkan/shaders/norm_modulate.comp | 7 +- vulkan/tools/vulkan_block_check.cpp | 44 +++++ vulkan/tools/vulkan_dit_run.cpp | 163 +++++++++++++++ 13 files changed, 751 insertions(+), 8 deletions(-) create mode 100644 vulkan/dit.cpp create mode 100644 vulkan/dit.h create mode 100644 vulkan/shaders/elementwise.comp create mode 100644 vulkan/tools/vulkan_dit_run.cpp diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 458a221..afa71a9 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -31,11 +31,11 @@ endfunction() _mi_add_shader(gemm_coop vulkan1.3) _mi_add_shader(gemm_bf16_coop vulkan1.3) _mi_add_shader(attention_coop vulkan1.3) -foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain) +foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain elementwise) _mi_add_shader(${shader} vulkan1.2) endforeach() -add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp safetensors.cpp ${_spirv_files}) +add_library(mossinfer_vulkan STATIC 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}) target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) @@ -44,9 +44,14 @@ 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 -ffp-contract=off) +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) diff --git a/vulkan/block.cpp b/vulkan/block.cpp index e6c4f62..f6ed424 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -19,6 +19,9 @@ const std::uint32_t kNormSpirv[] = { 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" }; @@ -53,7 +56,12 @@ struct CoopAttentionPush { }; constexpr std::uint32_t kCoopAttentionRows = 64; -enum GemmMode : std::uint32_t { kF32 = 0, kBF16 = 1, kGated = 2, kSitu = 3 }; +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; @@ -105,6 +113,12 @@ BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { 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) { @@ -148,6 +162,42 @@ void BlockKernels::record_gemm_situ_glu(VkCommandBuffer cb, VkDeviceAddress out, dispatch_gemm(*gemm_situ_, path_, kSitu, cb, GemmPush{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_, path_, kBias, cb, GemmPush{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_, path_, kBiasResidual, cb, GemmPush{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; + p.dispatch(cb, &push, (push.n / 4 + 63) / 64, 1); + 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 { diff --git a/vulkan/block.h b/vulkan/block.h index 899703f..0369892 100644 --- a/vulkan/block.h +++ b/vulkan/block.h @@ -57,6 +57,20 @@ class BlockKernels { 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; @@ -64,7 +78,8 @@ class BlockKernels { private: BlockPath path_; - std::unique_ptr gemm_bf16_, gemm_gated_, gemm_situ_, norm_modulate_, qk_norm_rope_, attention_; + std::unique_ptr gemm_bf16_, gemm_gated_, gemm_situ_, gemm_bias_, gemm_bias_residual_, norm_modulate_, + qk_norm_rope_, attention_, silu_, situ_; }; struct BlockWeights { diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp new file mode 100644 index 0000000..f042c13 --- /dev/null +++ b/vulkan/dit.cpp @@ -0,0 +1,297 @@ +#include "dit.h" + +#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) + : device_(device), kernels_(kernels) { + const st::Index index(transformer_dir); + std::vector host; + // Uploads tensors into one new buffer, each at a 256-byte aligned offset; returns their addresses. + const auto load = [&](const std::vector& names, const std::vector>& shapes) { + 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; + } + owned_.push_back(std::make_unique(device_, total)); + Buffer& buf = *owned_.back(); + weight_bytes_ += total; + std::vector addresses; + for (std::size_t i = 0; i < names.size(); ++i) { + host.resize(index.at(names[i]).bytes); + index.read_into(names[i], host.data(), host.size()); + buf.upload(host.data(), host.size(), offsets[i]); + addresses.push_back(buf.address() + offsets[i]); + } + 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]}); + } + + for (int l = 0; l < kDiTLayers; ++l) { + const std::string p = "layers." + std::to_string(l) + "."; + // q, k and v load back to back, 256-aligned sizes of 18 MiB each, so they form the [9216, 3072] QKV weight. + auto t = load({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"}, + {{d}, {d}, {kHead}, {kHead}, {d, d}, {d, d}, {d, d}, {d, d}, {2 * f, d}, {d, f}}); + if (t[5] != t[4] + kDim * kDim * kBf16 || t[6] != t[5] + kDim * kDim * kBf16) fail("q, k and v are not contiguous"); + blocks_.push_back({t[0], t[1], t[2], t[3], t[4], t[7], t[8], t[9]}); + } +} + +Dit::~Dit() = default; + +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 float scale = static_cast(1.0 / std::sqrt(static_cast(kHead))); + const std::vector cu = {0, static_cast(length)}; + + device_.run([&](VkCommandBuffer cb) { + kernels_.record_gemm_bias(cb, x, kDim, e.address(), kDiTTextDim, g_.text_w, kDiTTextDim, g_.text_b, length, kDim, + kDiTTextDim); + }); + for (const TextWeights& b : text_blocks_) { + device_.run([&](VkCommandBuffer cb) { + kernels_.record_norm(cb, x, kDim, b.norm1, y.address(), kDim, length, kDim); + kernels_.record_gemm_bias(cb, qkv.address(), kQkv, y.address(), kDim, b.qkv_w, kDim, b.qkv_b, length, kQkv, kDim); + kernels_.record_qk_norm_rope(cb, qkv.address(), kQkv, length, kDiTHeads, b.q_norm, b.k_norm, rope.address(), kHead); + }); + device_.run([&](VkCommandBuffer cb) { + const VkDeviceAddress k = at(qkv, kDim * kBf16), v = at(qkv, 2 * kDim * kBf16); + kernels_.record_attention(cb, AttentionArgs{qkv.address(), k, v, attn.address(), kQkv, kQkv, kQkv, kDim, kHead, kHead, + kHead, kHead, cu, cu, kDiTHeads, scale}); + kernels_.record_gemm_bias_residual(cb, x, kDim, attn.address(), kDim, b.proj_w, kDim, b.proj_b, length, kDim, kDim); + kernels_.record_norm(cb, x, kDim, b.norm2, y.address(), kDim, length, kDim); + }); + device_.run([&](VkCommandBuffer cb) { + kernels_.record_gemm_bias(cb, up.address(), kFfn, y.address(), kDim, b.up_w, kDim, b.up_b, length, kFfn, kDim); + kernels_.record_gemm_bias(cb, gate.address(), kFfn, y.address(), kDim, b.gate_w, kDim, b.gate_b, length, kFfn, kDim); + kernels_.record_situ(cb, up.address(), gate.address(), ffn.address(), length * kFfn); + kernels_.record_gemm_bias_residual(cb, x, kDim, ffn.address(), kFfn, b.down_w, kFfn, b.down_b, length, 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(); + + 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) { + // 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(); + device_.run([&](VkCommandBuffer cb) { record_block_op(kernels_, cb, op, blocks_[static_cast(l)], b); }); + op_ms_[static_cast(op)] += now_ms() - t0; + } + r0 = r1; + } + } + + // 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; + return out; +} + +} // namespace mi::vk diff --git a/vulkan/dit.h b/vulkan/dit.h new file mode 100644 index 0000000..a3cc35a --- /dev/null +++ b/vulkan/dit.h @@ -0,0 +1,87 @@ +// 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, so no single submit runs long on a slow device. +#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; + +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); + ~Dit(); + Dit(const Dit&) = delete; + Dit& operator=(const Dit&) = delete; + + std::uint64_t weight_bytes() const { return weight_bytes_; } + + // 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. + const std::vector& last_op_ms() const { return op_ms_; } + + private: + struct Weight; + struct TextBlock; + struct Scratch; + + Buffer& scratch(std::unique_ptr& slot, VkDeviceSize bytes); + void time_modulation(float timestep); + + Device& device_; + const BlockKernels& kernels_; + std::uint64_t weight_bytes_ = 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/runtime.cpp b/vulkan/runtime.cpp index eb6ad28..6549727 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -418,4 +418,24 @@ void compute_barrier(VkCommandBuffer cb) { 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 index 772e61b..c66f0a3 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -115,4 +115,9 @@ class Pipeline { // 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/shaders/elementwise.comp b/vulkan/shaders/elementwise.comp new file mode 100644 index 0000000..9ac56e7 --- /dev/null +++ b/vulkan/shaders/elementwise.comp @@ -0,0 +1,29 @@ +// 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() { + const uint i = 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 index 3e174c0..8b8e95b 100644 --- a/vulkan/shaders/gemm_bf16_coop.comp +++ b/vulkan/shaders/gemm_bf16_coop.comp @@ -5,6 +5,8 @@ // 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 @@ -166,10 +168,21 @@ void main() { 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 { + } else if (MODE == 3) { const vec4 g_lo = staged(base1, frow, fcol / 4u), g_hi = staged(base1, frow, fcol / 4u + 1u); 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_plain.comp b/vulkan/shaders/gemm_plain.comp index 4802620..29053d8 100644 --- a/vulkan/shaders/gemm_plain.comp +++ b/vulkan/shaders/gemm_plain.comp @@ -7,6 +7,8 @@ // 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 #include "common.glsl" @@ -77,7 +79,15 @@ void main() { precise const vec4 gated = rbf4(acc) * gate; precise const vec4 sum = f4(out_w.d[0]) + rbf4(gated); out_w.d[0] = bf4(sum); - } else { + } else if (MODE == 3) { out_w.d[0] = bf4(situ(rbf4(acc), rbf4(acc_gate))); + } else { + precise const vec4 biased = acc + 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 index 08cc9df..32dd1d3 100644 --- a/vulkan/shaders/norm_modulate.comp +++ b/vulkan/shaders/norm_modulate.comp @@ -1,4 +1,5 @@ -// out[rows,d] = bf16(bf16(bf16(bf16(rms_norm(x)) * w) * bf16(1 + scale)) + shift); w, scale, shift: [d]. x and out are +// 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 @@ -55,6 +56,10 @@ void main() { 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]); diff --git a/vulkan/tools/vulkan_block_check.cpp b/vulkan/tools/vulkan_block_check.cpp index 4887d01..6c50fdf 100644 --- a/vulkan/tools/vulkan_block_check.cpp +++ b/vulkan/tools/vulkan_block_check.cpp @@ -240,6 +240,31 @@ bool check_norms(Device& device, const BlockKernels& kernels) { 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), @@ -291,6 +316,25 @@ bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int 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}) && 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}) && 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}) && 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}) && ok; diff --git a/vulkan/tools/vulkan_dit_run.cpp b/vulkan/tools/vulkan_dit_run.cpp new file mode 100644 index 0000000..b773c58 --- /dev/null +++ b/vulkan/tools/vulkan_dit_run.cpp @@ -0,0 +1,163 @@ +// 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|plain] [--device N] +// 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).as("& 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|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; + 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 { + 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.as(" 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, path_name == "plain" ? BlockPath::kPlain : BlockPath::kCoop); + 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(); + Dit dit(device, kernels, transformer); + std::printf("weights: %.2f GiB in %.1f s\n", static_cast(dit.weight_bytes()) / (1u << 30), 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.as(" 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)\n", ms[kBlockOps]); + 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, " Date: Mon, 2 Nov 2026 06:26:13 -0500 Subject: [PATCH 10/41] vulkan: the plain GEMM in 128x128 register-blocked tiles, and pipeline 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. --- vulkan/block.cpp | 5 +- vulkan/gemm.cpp | 2 +- vulkan/runtime.cpp | 53 ++++++++++ vulkan/runtime.h | 3 + vulkan/shaders/gemm_plain.comp | 171 ++++++++++++++++++++++----------- 5 files changed, 177 insertions(+), 57 deletions(-) diff --git a/vulkan/block.cpp b/vulkan/block.cpp index f6ed424..10d251c 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -80,7 +80,7 @@ void require(bool condition, const std::string& message) { if (!condition) throw std::invalid_argument("vulkan block: " + message); } -constexpr std::uint32_t kGemmTileRows = 8, kGemmTileCols = 32; +constexpr std::uint32_t kGemmTileRows = 128, kGemmTileCols = 128; void dispatch_gemm(const Pipeline& pipeline, BlockPath path, GemmMode mode, VkCommandBuffer cb, const GemmPush& push) { const std::uint32_t align = path == BlockPath::kCoop ? 8 : 4; @@ -93,7 +93,8 @@ void dispatch_gemm(const Pipeline& pipeline, BlockPath path, GemmMode mode, VkCo const std::uint32_t cols = mode == kSitu ? kCoopTileCols / 2 : kCoopTileCols; pipeline.dispatch(cb, &push, (push.n + cols - 1) / cols, (push.m + kCoopTileRows - 1) / kCoopTileRows); } else { - pipeline.dispatch(cb, &push, (push.n + kGemmTileCols - 1) / kGemmTileCols, (push.m + kGemmTileRows - 1) / kGemmTileRows); + const std::uint32_t cols = mode == kSitu ? kGemmTileCols / 2 : kGemmTileCols; + pipeline.dispatch(cb, &push, (push.n + cols - 1) / cols, (push.m + kGemmTileRows - 1) / kGemmTileRows); } compute_barrier(cb); } diff --git a/vulkan/gemm.cpp b/vulkan/gemm.cpp index 93ad70e..2b69096 100644 --- a/vulkan/gemm.cpp +++ b/vulkan/gemm.cpp @@ -14,7 +14,7 @@ const std::uint32_t kPlainSpirv[] = { }; constexpr std::uint32_t kCoopTile = 64; -constexpr std::uint32_t kPlainTileRows = 8, kPlainTileCols = 32; +constexpr std::uint32_t kPlainTileRows = 128, kPlainTileCols = 128; constexpr std::uint32_t kCoopSubgroup = 32; struct Push { diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index 6549727..49e90ca 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -1,12 +1,21 @@ #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; @@ -230,6 +239,16 @@ Device::Device(int index) { 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_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; queue_info.queueFamilyIndex = family; @@ -393,8 +412,42 @@ Pipeline::Pipeline(Device& device, const std::uint32_t* spirv, std::size_t spirv 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() { diff --git a/vulkan/runtime.h b/vulkan/runtime.h index c66f0a3..269a57e 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -54,6 +54,7 @@ class Device { void run(const std::function& record); 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: @@ -64,6 +65,7 @@ class Device { VkCommandPool pool_ = VK_NULL_HANDLE; VkPhysicalDeviceMemoryProperties memory_{}; DeviceInfo info_; + bool pipeline_stats_ = false; }; // A device-local storage buffer with a device address, or a host-visible staging buffer. @@ -105,6 +107,7 @@ class Pipeline { 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; diff --git a/vulkan/shaders/gemm_plain.comp b/vulkan/shaders/gemm_plain.comp index 29053d8..f344fff 100644 --- a/vulkan/shaders/gemm_plain.comp +++ b/vulkan/shaders/gemm_plain.comp @@ -1,8 +1,12 @@ // 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 accumulators 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. +// 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. // -// Each invocation computes 4 adjacent output columns of one row, in tiles of 8 rows by 32 columns. MODE: +// 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 @@ -11,9 +15,10 @@ // 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 = 64, local_size_y = 1, local_size_z = 1) in; +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; @@ -30,64 +35,122 @@ layout(push_constant) uniform PC { 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 inv = gl_LocalInvocationIndex; - const uint m = gl_WorkGroupID.y * 8u + inv / 8u; - const uint n = (gl_WorkGroupID.x * 8u + inv % 8u) * 4u; - if (m >= pc.M || n >= pc.N) return; + 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; - Words A = Words(pc.a + uint64_t(m) * pc.lda * 2ul); - const uint64_t ldb_bytes = uint64_t(pc.ldb) * 2ul; - const uint64_t b0 = pc.b + uint64_t(n) * ldb_bytes; - Words B0 = Words(b0), B1 = Words(b0 + ldb_bytes), B2 = Words(b0 + 2ul * ldb_bytes), B3 = Words(b0 + 3ul * ldb_bytes); - const uint64_t g0 = pc.b + uint64_t(pc.N + n) * ldb_bytes; - Words G0 = Words(g0), G1 = Words(g0 + ldb_bytes), G2 = Words(g0 + 2ul * ldb_bytes), G3 = Words(g0 + 3ul * ldb_bytes); + 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); + } - precise vec4 acc = vec4(0.0), acc_gate = vec4(0.0); - const uint words = pc.K / 4u; - for (uint i = 0u; i < words; ++i) { - const vec4 x = decode(A.d[i]); - const vec4 c0 = decode(B0.d[i]), c1 = decode(B1.d[i]), c2 = decode(B2.d[i]), c3 = decode(B3.d[i]); - acc = acc + x.x * vec4(c0.x, c1.x, c2.x, c3.x); - acc = acc + x.y * vec4(c0.y, c1.y, c2.y, c3.y); - acc = acc + x.z * vec4(c0.z, c1.z, c2.z, c3.z); - acc = acc + x.w * vec4(c0.w, c1.w, c2.w, c3.w); - if (MODE == 3) { - const vec4 d0 = decode(G0.d[i]), d1 = decode(G1.d[i]), d2 = decode(G2.d[i]), d3 = decode(G3.d[i]); - acc_gate = acc_gate + x.x * vec4(d0.x, d1.x, d2.x, d3.x); - acc_gate = acc_gate + x.y * vec4(d0.y, d1.y, d2.y, d3.y); - acc_gate = acc_gate + x.z * vec4(d0.z, d1.z, d2.z, d3.z); - acc_gate = acc_gate + x.w * vec4(d0.w, d1.w, d2.w, d3.w); + 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(); } - if (MODE == 0) { - Floats out_f = Floats(pc.out_ + (uint64_t(m) * pc.ldo + n) * 4ul); - out_f.d[0] = acc.x; - out_f.d[1] = acc.y; - out_f.d[2] = acc.z; - out_f.d[3] = acc.w; - return; - } - Words out_w = Words(pc.out_ + (uint64_t(m) * pc.ldo + n) * 2ul); - if (MODE == 1) { - out_w.d[0] = bf4(acc); - } else if (MODE == 2) { - const vec4 gate = f4(Words(pc.gate + uint64_t(n) * 2ul).d[0]); - precise const vec4 gated = rbf4(acc) * 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(acc), rbf4(acc_gate))); - } else { - precise const vec4 biased = acc + 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); + [[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); + } + } } } } From 91a03d8d450daa36a96903c10f619556723f15c6 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:28:19 -0500 Subject: [PATCH 11/41] vulkan: plain attention over 32-row tiles in shared memory 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). --- vulkan/block.cpp | 29 +--- vulkan/shaders/attention_plain.comp | 228 ++++++++++++++++++---------- 2 files changed, 159 insertions(+), 98 deletions(-) diff --git a/vulkan/block.cpp b/vulkan/block.cpp index 10d251c..cc7914d 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -43,18 +43,13 @@ struct QkPush { std::uint32_t ld, table_ld, heads; float factor; }; -struct AttentionPush { - std::uint64_t q, k, v, o; - std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_row0, k_row0, seqlen_k; - float log2_scale; -}; -struct CoopAttentionPush { +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; +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 }; @@ -124,7 +119,7 @@ BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { 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(CoopAttentionPush), kCoopSubgroup); + sizeof(AttentionPush), kCoopSubgroup); } else { attention_ = std::make_unique(device, kAttentionSpirv, sizeof(kAttentionSpirv), sizeof(AttentionPush)); } @@ -225,19 +220,11 @@ void BlockKernels::record_attention(VkCommandBuffer cb, const AttentionArgs& a) 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"); - if (path_ == BlockPath::kCoop) { - const CoopAttentionPush 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}; - attention_->dispatch(cb, &push, a.heads, - (static_cast(seqlen_q) + kCoopAttentionRows - 1) / kCoopAttentionRows); - continue; - } - 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_k), log2_scale}; - attention_->dispatch(cb, &push, a.heads, static_cast(seqlen_q)); + 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 = path_ == BlockPath::kCoop ? kCoopAttentionRows : kPlainAttentionRows; + attention_->dispatch(cb, &push, a.heads, (static_cast(seqlen_q) + rows - 1) / rows); } compute_barrier(cb); } diff --git a/vulkan/shaders/attention_plain.comp b/vulkan/shaders/attention_plain.comp index aa9ca54..973c026 100644 --- a/vulkan/shaders/attention_plain.comp +++ b/vulkan/shaders/attention_plain.comp @@ -1,16 +1,23 @@ -// Scaled-dot-product attention over one sequence, head_dim 128, non-causal, bf16 in and out, fp32 online softmax. One query -// row of one head per workgroup, 32 lanes. Key tokens walk from the last step of 32 down to step 0, as FlashAttention's -// reverse iteration: +// 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 -// row_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 -// acc_o += P * V_step^T -// and at the end out = bf16(acc_o / row_sum). Lane c owns key c of each step for S and P, and dims 4c..4c+3 for acc_o. +// 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 = 32, local_size_y = 1, local_size_z = 1) in; +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; @@ -27,109 +34,176 @@ layout(push_constant) uniform PC { 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 vec4 q_shared[32]; -shared float s_shared[STEP]; -shared float w_shared[STEP]; -shared float p_shared[STEP]; -shared float max_scaled_shared; -shared float rescale_shared; -shared float inv_sum_shared; +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]; -// 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(); } +// 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 lane = gl_LocalInvocationIndex; + const uint t = gl_LocalInvocationIndex; const uint head = gl_WorkGroupID.x; - const uint64_t q_row = uint64_t(pc.q_row0) + gl_WorkGroupID.y; - const uint64_t k_row0 = pc.k_row0; - - q_shared[lane] = f4(Words(pc.q + (q_row * pc.ldq + uint64_t(head) * pc.hsq) * 2ul).d[lane]); - sync_shared(); + 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 - const uint64_t k_head = pc.k + (k_row0 * pc.ldk + uint64_t(head) * pc.hsk) * 2ul; - const uint64_t v_head = pc.v + (k_row0 * pc.ldv + uint64_t(head) * pc.hsv) * 2ul; - const uint64_t ldk_bytes = uint64_t(pc.ldk) * 2ul, ldv_bytes = uint64_t(pc.ldv) * 2ul; + // 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; + } - precise vec4 acc = vec4(0.0); - float row_max = 0.0, row_sum = 0.0; // lane 0's + 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); - if (lane < cols) { - Words key = Words(k_head + uint64_t(n0 + lane) * ldk_bytes); - precise float s = 0.0; - for (uint g = 0u; g < 32u; ++g) { - const vec4 q4 = q_shared[g]; - const vec4 k4 = f4(key.d[g]); - s = s + q4.x * k4.x; - s = s + q4.y * k4.y; - s = s + q4.z * k4.z; - s = s + q4.w * k4.w; - } - s_shared[lane] = s; + // 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; } - sync_shared(); - - if (lane == 0u) { - const float prev = row_max; - float cur = first ? s_shared[0] : max(prev, s_shared[0]); - for (uint c = 1u; c < cols; ++c) cur = max(cur, s_shared[c]); - row_max = cur; - if (first) { - rescale_shared = 1.0; - row_sum = 0.0; - } else { - precise const float diff = (prev - cur) * pc.log2_scale; - const float rescale = exp2(diff); - rescale_shared = rescale; - precise const float scaled = row_sum * rescale; - row_sum = scaled; + 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])); } - precise const float max_scaled = isinf(cur) && cur < 0.0 ? 0.0 : cur * pc.log2_scale; - max_scaled_shared = max_scaled; + } + [[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(); - float w = 0.0; - if (lane < cols) { - // FlashAttention's one fused step; a device without a fused fma rounds it twice, within the bf16 weight's error. - w = exp2(fma(s_shared[lane], pc.log2_scale, -max_scaled_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]); } - w_shared[lane] = w; - p_shared[lane] = uintBitsToFloat(bf(w) << 16); + 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 vec4 rescaled = acc * rescale_shared; - acc = rescaled; + 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; - if (lane == 0u) { - precise float sum = row_sum; - for (uint c = 0u; c < cols; ++c) sum = sum + w_shared[c]; - row_sum = 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 c = 0u; c < cols; ++c) { - const vec4 v4 = f4(Words(v_head + uint64_t(n0 + c) * ldv_bytes).d[lane]); - acc = acc + p_shared[c] * v4; + 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 (lane == 0u) inv_sum_shared = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + if (spart == 0u) row_inv[srow] = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; sync_shared(); - precise const vec4 result = acc * inv_sum_shared; - Words(pc.o + (q_row * pc.ldo + uint64_t(head) * pc.hso) * 2ul).d[lane] = bf4(result); + [[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); + } + } } From cca0a3ebfc315eeadf88d6a7ebf499d64d57566f Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:31:15 -0500 Subject: [PATCH 12/41] vulkan: the FLUX.2 VAE decoder, TF32-rounded convolutions on bf16 cooperative 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. --- vulkan/CMakeLists.txt | 19 + vulkan/safetensors.cpp | 287 ++++++++++ vulkan/safetensors.h | 39 ++ vulkan/shaders/vae_conv_coop.comp | 265 +++++++++ vulkan/shaders/vae_conv_plain.comp | 152 +++++ vulkan/shaders/vae_gn_apply.comp | 45 ++ vulkan/shaders/vae_gn_finalize.comp | 50 ++ vulkan/shaders/vae_gn_stats.comp | 85 +++ vulkan/shaders/vae_latent_unpack.comp | 30 + vulkan/shaders/vae_rgb8.comp | 38 ++ vulkan/shaders/vae_softmax.comp | 58 ++ vulkan/shaders/vae_transpose.comp | 32 ++ vulkan/tools/vulkan_vae_check.cpp | 774 ++++++++++++++++++++++++++ vulkan/vae.cpp | 727 ++++++++++++++++++++++++ vulkan/vae.h | 143 +++++ 15 files changed, 2744 insertions(+) create mode 100644 vulkan/safetensors.cpp create mode 100644 vulkan/safetensors.h create mode 100644 vulkan/shaders/vae_conv_coop.comp create mode 100644 vulkan/shaders/vae_conv_plain.comp create mode 100644 vulkan/shaders/vae_gn_apply.comp create mode 100644 vulkan/shaders/vae_gn_finalize.comp create mode 100644 vulkan/shaders/vae_gn_stats.comp create mode 100644 vulkan/shaders/vae_latent_unpack.comp create mode 100644 vulkan/shaders/vae_rgb8.comp create mode 100644 vulkan/shaders/vae_softmax.comp create mode 100644 vulkan/shaders/vae_transpose.comp create mode 100644 vulkan/tools/vulkan_vae_check.cpp create mode 100644 vulkan/vae.cpp create mode 100644 vulkan/vae.h diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 9166f7e..b736d37 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -50,3 +50,22 @@ 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 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_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}) +add_library(mossinfer_vulkan_vae STATIC vae.cpp safetensors.cpp ${_spirv_vae}) +target_include_directories(mossinfer_vulkan_vae PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) +target_link_libraries(mossinfer_vulkan_vae PUBLIC 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 -ffp-contract=off) +target_link_libraries(vulkan_vae_check PRIVATE mossinfer_vulkan_vae ZLIB::ZLIB) diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp new file mode 100644 index 0000000..a72b48a --- /dev/null +++ b/vulkan/safetensors.cpp @@ -0,0 +1,287 @@ +#include "safetensors.h" + +#include +#include +#include +#include + +namespace mi::vk::st { +namespace { + +// The header is one JSON object: {"name": {"dtype": "F32", "shape": [..], "data_offsets": [begin, end]}, ...}, with an +// optional "__metadata__" object of strings. The parser accepts general JSON and keeps only what the entries need. +class Parser { + public: + Parser(const std::string& text, const std::string& path) : s_(text), path_(path) {} + + std::vector parse_header() { + std::vector out; + expect('{'); + if (peek() == '}') { + ++i_; + return out; + } + for (;;) { + const std::string name = string(); + expect(':'); + if (name == "__metadata__") { + skip_value(); + } else { + out.push_back(entry(name)); + } + const char c = next(); + if (c == '}') break; + if (c != ',') fail("expected , or } after an entry"); + } + ws(); + if (i_ != s_.size()) fail("trailing bytes after the header object"); + return out; + } + + private: + [[noreturn]] void fail(const std::string& what) const { + throw std::runtime_error("safetensors " + path_ + ": header: " + what + " at byte " + std::to_string(i_)); + } + void ws() { + while (i_ < s_.size() && (s_[i_] == ' ' || s_[i_] == '\n' || s_[i_] == '\r' || s_[i_] == '\t')) ++i_; + } + char peek() { + ws(); + if (i_ >= s_.size()) fail("unexpected end"); + return s_[i_]; + } + char next() { + const char c = peek(); + ++i_; + return c; + } + void expect(char c) { + if (next() != c) fail(std::string("expected '") + c + "'"); + } + + std::string string() { + expect('"'); + std::string out; + for (;;) { + if (i_ >= s_.size()) fail("unterminated string"); + const char c = s_[i_++]; + if (c == '"') return out; + if (c != '\\') { + out.push_back(c); + continue; + } + if (i_ >= s_.size()) fail("unterminated escape"); + const char e = s_[i_++]; + switch (e) { + case '"': case '\\': case '/': out.push_back(e); break; + case 'b': out.push_back('\b'); break; + case 'f': out.push_back('\f'); break; + case 'n': out.push_back('\n'); break; + case 'r': out.push_back('\r'); break; + case 't': out.push_back('\t'); break; + case 'u': { + // Tensor names are ASCII in practice; other code points are kept as UTF-8 of the BMP value. + if (i_ + 4 > s_.size()) fail("short \\u escape"); + const unsigned v = static_cast(std::stoul(s_.substr(i_, 4), nullptr, 16)); + i_ += 4; + if (v < 0x80) { + out.push_back(static_cast(v)); + } else if (v < 0x800) { + out.push_back(static_cast(0xc0 | (v >> 6))); + out.push_back(static_cast(0x80 | (v & 0x3f))); + } else { + out.push_back(static_cast(0xe0 | (v >> 12))); + out.push_back(static_cast(0x80 | ((v >> 6) & 0x3f))); + out.push_back(static_cast(0x80 | (v & 0x3f))); + } + break; + } + default: fail("bad escape"); + } + } + } + + std::uint64_t integer() { + ws(); + if (i_ >= s_.size() || s_[i_] < '0' || s_[i_] > '9') fail("expected a non-negative integer"); + std::uint64_t v = 0; + while (i_ < s_.size() && s_[i_] >= '0' && s_[i_] <= '9') { + const auto d = static_cast(s_[i_++] - '0'); + if (v > (std::numeric_limits::max() - d) / 10) fail("integer overflow"); + v = v * 10 + d; + } + return v; + } + + std::vector integer_array() { + std::vector out; + expect('['); + if (peek() == ']') { + ++i_; + return out; + } + for (;;) { + out.push_back(integer()); + const char c = next(); + if (c == ']') return out; + if (c != ',') fail("expected , or ] in an array"); + } + } + + void skip_value() { + const char c = peek(); + if (c == '"') { + string(); + } else if (c == '{' || c == '[') { + const char close = c == '{' ? '}' : ']'; + ++i_; + if (peek() == close) { + ++i_; + return; + } + for (;;) { + if (c == '{') { + string(); + expect(':'); + } + skip_value(); + const char d = next(); + if (d == close) return; + if (d != ',') fail("expected , in a container"); + } + } else { + // number, true, false, null + ws(); + const std::size_t start = i_; + while (i_ < s_.size() && std::strchr(",}] \n\r\t", s_[i_]) == nullptr) ++i_; + if (i_ == start) fail("expected a value"); + } + } + + TensorInfo entry(const std::string& name) { + TensorInfo t; + t.name = name; + bool has_dtype = false, has_shape = false, has_offsets = false; + std::vector offsets; + expect('{'); + for (;;) { + const std::string key = string(); + expect(':'); + if (key == "dtype") { + t.dtype = string(); + has_dtype = true; + } else if (key == "shape") { + for (const std::uint64_t d : integer_array()) { + if (d > static_cast(std::numeric_limits::max())) fail("dimension too large"); + t.shape.push_back(static_cast(d)); + } + has_shape = true; + } else if (key == "data_offsets") { + offsets = integer_array(); + has_offsets = true; + } else { + skip_value(); + } + const char c = next(); + if (c == '}') break; + if (c != ',') fail("expected , or } in a tensor entry"); + } + if (!has_dtype || !has_shape || !has_offsets || offsets.size() != 2 || offsets[1] < offsets[0]) { + fail("tensor '" + name + "' lacks dtype, shape or a valid data_offsets pair"); + } + t.offset = offsets[0]; + t.bytes = offsets[1] - offsets[0]; + return t; + } + + const std::string& s_; + const std::string& path_; + std::size_t i_ = 0; +}; + +std::uint64_t dtype_bytes(const std::string& dtype) { + if (dtype == "F64" || dtype == "I64" || dtype == "U64") return 8; + if (dtype == "F32" || dtype == "I32" || dtype == "U32") return 4; + if (dtype == "F16" || dtype == "BF16" || dtype == "I16" || dtype == "U16") return 2; + if (dtype == "I8" || dtype == "U8" || dtype == "BOOL" || dtype == "F8_E4M3" || dtype == "F8_E5M2") return 1; + return 0; // unknown: the size check is skipped +} + +} // namespace + +File::File(const std::string& path) : path_(path), in_(path, std::ios::binary) { + if (!in_) throw std::runtime_error("safetensors: cannot open " + path); + in_.seekg(0, std::ios::end); + file_bytes_ = static_cast(in_.tellg()); + in_.seekg(0); + unsigned char len_bytes[8] = {}; + if (file_bytes_ < 8 || !in_.read(reinterpret_cast(len_bytes), 8)) { + throw std::runtime_error("safetensors " + path + ": shorter than its 8-byte header length"); + } + std::uint64_t header_bytes = 0; + for (int i = 7; i >= 0; --i) header_bytes = header_bytes << 8 | len_bytes[i]; + // The format caps the header at 100 MB. + if (header_bytes > 100'000'000 || header_bytes > file_bytes_ - 8) { + throw std::runtime_error("safetensors " + path + ": header length " + std::to_string(header_bytes) + " is out of range"); + } + std::string header(static_cast(header_bytes), '\0'); + if (!in_.read(header.data(), static_cast(header_bytes))) { + throw std::runtime_error("safetensors " + path + ": cannot read the header"); + } + tensors_ = Parser(header, path_).parse_header(); + const std::uint64_t data_start = 8 + header_bytes, data_bytes = file_bytes_ - data_start; + for (TensorInfo& t : tensors_) { + if (t.offset > data_bytes || t.bytes > data_bytes - t.offset) { + throw std::runtime_error("safetensors " + path + ": tensor '" + t.name + "' lies outside the file"); + } + std::uint64_t count = 1; + for (const std::int64_t d : t.shape) { + const auto u = static_cast(d); + if (u != 0 && count > std::numeric_limits::max() / u) { + throw std::runtime_error("safetensors " + path + ": tensor '" + t.name + "' has too many elements"); + } + count *= u; + } + const std::uint64_t element = dtype_bytes(t.dtype); + if (element != 0 && (count > std::numeric_limits::max() / element || count * element != t.bytes)) { + throw std::runtime_error("safetensors " + path + ": tensor '" + t.name + "' byte range does not match its shape"); + } + t.offset += data_start; + } +} + +const TensorInfo* File::find(const std::string& name) const { + const auto it = std::find_if(tensors_.begin(), tensors_.end(), [&](const TensorInfo& t) { return t.name == name; }); + return it == tensors_.end() ? nullptr : &*it; +} + +void File::read(const TensorInfo& info, void* dst) { + in_.clear(); + in_.seekg(static_cast(info.offset)); + if (!in_.read(static_cast(dst), static_cast(info.bytes))) { + throw std::runtime_error("safetensors " + path_ + ": cannot read tensor '" + info.name + "'"); + } +} + +std::vector File::read_f32(const std::string& name) { + const TensorInfo* t = find(name); + if (t == nullptr) throw std::runtime_error("safetensors " + path_ + ": no tensor '" + 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"); +} + +} // namespace mi::vk::st diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h new file mode 100644 index 0000000..0614081 --- /dev/null +++ b/vulkan/safetensors.h @@ -0,0 +1,39 @@ +// A minimal safetensors reader: the JSON header's tensor entries (dtype, shape, byte range) and raw reads of one +// tensor's bytes. Links nothing beyond the standard library. Failures throw std::runtime_error. +#pragma once + +#include +#include +#include +#include + +namespace mi::vk::st { + +struct TensorInfo { + std::string name; + std::string dtype; // as stored: "F32", "BF16", "F16", "I64", ... + std::vector shape; + std::uint64_t offset = 0; // from the start of the file + std::uint64_t bytes = 0; +}; + +class File { + public: + explicit File(const std::string& path); + + const std::vector& tensors() const { return tensors_; } + // nullptr when the file has no tensor of that name. + const TensorInfo* find(const std::string& name) const; + // Reads the tensor's bytes into dst, which must hold info.bytes. + void read(const TensorInfo& info, void* dst); + // A tensor of dtype F32 or BF16 as fp32 values; throws for other dtypes. + std::vector read_f32(const std::string& name); + + private: + std::string path_; + std::ifstream in_; + std::uint64_t file_bytes_ = 0; + std::vector tensors_; +}; + +} // namespace mi::vk::st 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_plain.comp b/vulkan/shaders/vae_conv_plain.comp new file mode 100644 index 0000000..61c1017 --- /dev/null +++ b/vulkan/shaders/vae_conv_plain.comp @@ -0,0 +1,152 @@ +// 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 16 TM rows x 16 TN columns, invocation (x, y) the rows +// y + 16 i and columns x + 16 j. K walks in steps of 16 through shared memory, the next step loading into registers. +#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 TM = 4; // 1..8 +layout(constant_id = 4) const uint TN = 4; // 1..4 + +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 BM = 16u * TM; +const uint BN = 16u * TN; +const uint BK = 16u; +// [k][row], rows padded by 4 so the transposed stores of 4 k values spread over banks. +shared float As[BK * (BM + 4u)]; +shared float Bs[BK * (BN + 4u)]; + +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: BM * 4 of A (invocation i takes chunks i and i + 256) and BN * 4 <= 256 of B per step. +const uint A_CHUNKS = (BM * 4u + 255u) / 256u; + +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[8][4]; + [[unroll]] for (uint i = 0u; i < 8u; ++i) { + [[unroll]] for (uint j = 0u; j < 4u; ++j) acc[i][j] = 0.0; + } + + vec4 ra[2], rb[1]; + const uint steps = (pc.K + BK - 1u) / BK; + for (uint step = 0u; step < steps; ++step) { + const uint k0 = step * BK; + [[unroll]] for (uint q = 0u; q < 2u; ++q) { + const uint chunk = t + q * 256u; + if (q < A_CHUNKS && chunk < BM * 4u) ra[q] = load_a(m_base + chunk / 4u, k0 + (chunk % 4u) * 4u); + } + if (t < BN * 4u) rb[0] = load_b(n_base + t / 4u, k0 + (t % 4u) * 4u); + barrier(); // the previous step's reads of As and Bs are done + [[unroll]] for (uint q = 0u; q < 2u; ++q) { + const uint chunk = t + q * 256u; + if (q < A_CHUNKS && chunk < BM * 4u) { + const uint row = chunk / 4u, kk = (chunk % 4u) * 4u; + As[(kk + 0u) * (BM + 4u) + row] = ra[q].x; + As[(kk + 1u) * (BM + 4u) + row] = ra[q].y; + As[(kk + 2u) * (BM + 4u) + row] = ra[q].z; + As[(kk + 3u) * (BM + 4u) + row] = ra[q].w; + } + } + if (t < BN * 4u) { + const uint row = t / 4u, kk = (t % 4u) * 4u; + Bs[(kk + 0u) * (BN + 4u) + row] = rb[0].x; + Bs[(kk + 1u) * (BN + 4u) + row] = rb[0].y; + Bs[(kk + 2u) * (BN + 4u) + row] = rb[0].z; + Bs[(kk + 3u) * (BN + 4u) + row] = rb[0].w; + } + barrier(); + for (uint k = 0u; k < BK; ++k) { + float av[8], bv[4]; + [[unroll]] for (uint i = 0u; i < 8u; ++i) av[i] = i < TM ? As[k * (BM + 4u) + ty + 16u * i] : 0.0; + [[unroll]] for (uint j = 0u; j < 4u; ++j) bv[j] = j < TN ? Bs[k * (BN + 4u) + tx + 16u * j] : 0.0; + [[unroll]] for (uint i = 0u; i < 8u; ++i) { + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + if (i < TM && j < TN) acc[i][j] += av[i] * bv[j]; + } + } + } + } + + [[unroll]] for (uint i = 0u; i < 8u; ++i) { + const uint m = m_base + ty + 16u * i; + if (i >= TM || m >= pc.M) continue; + [[unroll]] for (uint j = 0u; j < 4u; ++j) { + const uint n = n_base + tx + 16u * 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/tools/vulkan_vae_check.cpp b/vulkan/tools/vulkan_vae_check.cpp new file mode 100644 index 0000000..8d94c3b --- /dev/null +++ b/vulkan/tools/vulkan_vae_check.cpp @@ -0,0 +1,774 @@ +// 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] +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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(" %-34s %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; +} + +double rms(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 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(" %-34s 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}, + {"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(" %-34s %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(" %-34s %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(" %-34s %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); + } +} + +// ---- Real decode ------------------------------------------------------------------------------------------------------ + +struct Options { + int device = -1; + std::string path; + bool kat = true, decode = true, bench = 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); + 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)); + 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" : ""); + } + 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 == "--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]\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..a16dbd9 --- /dev/null +++ b/vulkan/vae.cpp @@ -0,0 +1,727 @@ +#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 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 +// kernel, whose operands then get the same rounding. +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; + const std::uint32_t tm = a.n <= 16 ? 8 : 4, tn = a.n <= 16 ? 1 : 4; + tile_rows = 16 * tm; + tile_cols = 16 * tn; + const std::vector key = {1, taps, up, tf32, tm, tn}; + auto& p = conv_[key]; + if (!p) { + p = std::make_unique(device_, kConvPlainSpirv, sizeof(kConvPlainSpirv), sizeof(ConvPush), 0, + std::vector{taps, up, tf32, tm, 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::string& path) + : VaeDecoder(device, path, device.info().caps.coop_bf16 ? VaePath::kCoop : VaePath::kPlain) {} + +VaeDecoder::VaeDecoder(Device& device, const std::string& 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: " + 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). +class Submitter { + public: + static constexpr std::uint64_t kElementCost = 32; + Submitter(Device& device, std::uint64_t budget) : device_(device), budget_(budget) {} + void add(std::uint64_t cost, std::function op) { + 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> ops_; + std::uint64_t cost_ = 0; + int submits_ = 0; +}; + +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(ConvArgs a, VaePrecision precision) { + const std::uint64_t per_row = static_cast(a.n) * a.k; + const std::uint64_t rows = 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((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; + gemm(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; + s_.add(cost, [this, x, pixels, c = n.c](VkCommandBuffer cb) { k_.record_gn_stats(cb, x, pixels, c, part_, stats_, kResnetEps); }); + s_.add(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(p, VaePrecision::kFp32); + s_.add(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(s, VaePrecision::kSplit16); + s_.add(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(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(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(); +} + +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_; + + Submitter sub(device_, max_macs_); + 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(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(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); + } + VaeDecodeStats st; + 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..e6ab571 --- /dev/null +++ b/vulkan/vae.h @@ -0,0 +1,143 @@ +// 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, chosen as BlockKernels does: 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 "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; +}; + +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::string& safetensors_path); + VaeDecoder(Device& device, const std::string& 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; + void set_max_macs_per_submit(std::uint64_t macs) { max_macs_ = macs; } + + // 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; +}; + +} // namespace mi::vk From f5bb7c8edca6ff73b93555aacba92716d40dd946 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:33:17 -0500 Subject: [PATCH 13/41] moss-infer generate --backend vulkan: the engine side 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. --- cmake/modules/vulkan.cmake | 11 +++ src/tools/cli/cmd_generate.cpp | 48 ++++++++++ src/tools/cli/vulkan_backend.cpp | 150 +++++++++++++++++++++++++++++++ src/tools/cli/vulkan_backend.h | 29 ++++++ vulkan/CMakeLists.txt | 8 ++ vulkan/generate.cpp | 103 +++++++++++++++++++++ vulkan/generate.h | 44 +++++++++ 7 files changed, 393 insertions(+) create mode 100644 cmake/modules/vulkan.cmake create mode 100644 src/tools/cli/vulkan_backend.cpp create mode 100644 src/tools/cli/vulkan_backend.h create mode 100644 vulkan/generate.cpp create mode 100644 vulkan/generate.h diff --git a/cmake/modules/vulkan.cmake b/cmake/modules/vulkan.cmake new file mode 100644 index 0000000..300436f --- /dev/null +++ b/cmake/modules/vulkan.cmake @@ -0,0 +1,11 @@ +# The Vulkan backend (vulkan/), the fallback for GPUs the CUDA engine does not support: `moss-infer generate --backend +# vulkan`. Off by default; it needs the Vulkan headers (1.4.313 or newer, for VK_KHR_shader_bfloat16), the loader and +# glslc. Without it the CLI refuses --backend vulkan with a message saying so. +option(MOSSINFER_VULKAN "Build the Vulkan backend into the engine (moss-infer generate --backend vulkan)" OFF) +if(MOSSINFER_VULKAN) + set(MOSSINFER_VULKAN_TOOLS OFF CACHE BOOL "Build the Vulkan backend's check and benchmark tools") + add_subdirectory(${CMAKE_SOURCE_DIR}/vulkan ${CMAKE_BINARY_DIR}/vulkan) + set_target_properties(mossinfer_vulkan PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_link_libraries(mossinfer_core PUBLIC mossinfer_vulkan) + target_compile_definitions(mossinfer_core PUBLIC MOSSINFER_VULKAN=1) +endif() diff --git a/src/tools/cli/cmd_generate.cpp b/src/tools/cli/cmd_generate.cpp index 4737abf..28f44ec 100644 --- a/src/tools/cli/cmd_generate.cpp +++ b/src/tools/cli/cmd_generate.cpp @@ -27,6 +27,7 @@ #include "pipeline/report_json.h" #include "planner/planner.h" #include "tools/cli/commands.h" +#include "tools/cli/vulkan_backend.h" #include "util/json.h" #include "util/path.h" @@ -56,10 +57,15 @@ constexpr const char* kUsage = " [--cfg-truncation X] [--step-cache off|fbcache] [--step-cache-threshold X]\n" " [--control [--control-scale S] [--control-image =\n" " [--control-strength X] [--control-start X] [--control-end X]]...]\n" + " [--backend cuda|vulkan [--vulkan-device N] [--vulkan-kernels coop|plain]]\n" " [--reference ]... [--reference-max-edge N]\n" " [--lora-placement default|merge|side_branch]\n" "`moss-infer check` takes the same options (-o optional): it checks them and plans the request as generate would,\n" "without loading weights or writing anything, and prints the verdict as JSON (exit 0 accepted, 1 refused).\n" + "--backend vulkan (builds with -DMOSSINFER_VULKAN=ON) runs on Vulkan instead of CUDA, for GPUs the CUDA engine does\n" + "not support: txt2img with the bf16 model files and the official sampler and schedule only, much slower than CUDA.\n" + "--vulkan-device picks a device (default: the first usable one, discrete first); --vulkan-kernels picks the bf16\n" + "cooperative-matrix kernels or the plain ones (default: coop where the device has them). Both are reported.\n" "Flags override values read from --params. At least one of -o / --latent-out is required;\n" "without -o the VAE decode is skipped. --vram-budget caps the process's GPU memory (CUDA context included).\n" "Each --then-prompt generates one more image in the same engine session (same settings, warm models and\n" @@ -343,6 +349,8 @@ int generate_or_check(const std::vector& args, bool check) { std::optional strength; std::string hires; bool distill = false, distill_dataspace = false; + std::string backend = "cuda"; + VulkanRun vulkan; json::Value pdoc; const json::Value* pobj = nullptr; // the parameters object of --params int distill_steps = 0; @@ -489,6 +497,12 @@ int generate_or_check(const std::vector& args, bool check) { else if (f == "--reference") reference_images.push_back(path_arg(v, f)); else if (f == "--reference-max-edge") req.references.max_edge = to_int32(f, v); else if (f == "--lora-placement") req.lora_placement = lora::parse_placement(v); + else if (f == "--backend") { + if (v != "cuda" && v != "vulkan") throw std::invalid_argument("--backend takes cuda or vulkan"); + backend = v; + } + else if (f == "--vulkan-device") vulkan.device = to_int32(f, v); + else if (f == "--vulkan-kernels") vulkan.kernels = v; else throw std::invalid_argument("unknown flag " + f); } } catch (const std::invalid_argument& e) { @@ -580,6 +594,10 @@ int generate_or_check(const std::vector& args, bool check) { return 1; } req.decode = !out_png.empty() || (check && out_latent.empty()); // check: a PNG unless only --latent-out is given + if (check && backend != "cuda") { + std::cerr << "moss-infer check: --backend vulkan is not supported by check\n"; + return 2; + } if (check) return run_check(options, req, then_prompts, !preview_png.empty(), quiet); pipeline::validate(req); for (const std::string& p : then_prompts) { @@ -588,6 +606,36 @@ int generate_or_check(const std::vector& args, bool check) { pipeline::validate(f); } + if ((given("--vulkan-device") || given("--vulkan-kernels")) && backend != "vulkan") { + std::cerr << "moss-infer generate: --vulkan-device and --vulkan-kernels need --backend vulkan\n" << kUsage; + return 2; + } + if (backend == "vulkan") { + // CUDA-only settings are refused rather than ignored. + for (const char* f : {"--mode", "--vram-budget", "--vram-reserve", "--host-ram-budget", "--cache-dir", "--then-prompt", + "--direct-io", "--buffered-io", "--te-host-pin", "--cert-dir", "--dit-blocks", "--dit", "--te", + "--vae", "--preview", "--tier", "--attention", "--json-report", "--time-embedding", "--distill", + "--lora-placement"}) { + if (given(f)) { + std::cerr << "moss-infer generate: " << f << " is not supported with --backend vulkan\n"; + return 2; + } + } + try { + vulkan.out_png = out_png; + vulkan.quiet = quiet; + const std::vector latents = run_vulkan_generate(req, vulkan); + if (!out_latent.empty()) write_npy_f32(out_latent, latents, {moss_dit::kLatentChannels, req.height / 16, req.width / 16}); + } catch (const std::invalid_argument& e) { + std::cerr << "moss-infer generate: " << e.what() << "\n"; + return 2; + } catch (const std::exception& e) { + std::cerr << "moss-infer generate: " << e.what() << "\n"; + return 1; + } + return 0; + } + std::signal(SIGINT, on_sigint); pipeline::Pipeline pipe(options); pipeline::GenerateControl control; diff --git a/src/tools/cli/vulkan_backend.cpp b/src/tools/cli/vulkan_backend.cpp new file mode 100644 index 0000000..e0bbf92 --- /dev/null +++ b/src/tools/cli/vulkan_backend.cpp @@ -0,0 +1,150 @@ +#include "tools/cli/vulkan_backend.h" + +#include +#include +#include +#include +#include +#include + +#include "mossinfer/mossinfer.h" +#include "pipeline/png_writer.h" +#include "pipeline/samplers/initial_noise.h" +#include "pipeline/samplers/schedules.h" +#include "tokenizer/prompt_tokens.h" +#include "tokenizer/qwen_bpe.h" + +#if defined(MOSSINFER_VULKAN) +#include "generate.h" +#include "runtime.h" +#endif + +namespace mi::cli { +namespace { + +// What the Vulkan backend runs: txt2img with the official sampler and schedule. Anything else is refused by name. +void check_supported(const pipeline::Request& req) { + const auto refuse = [](const char* what) { + throw std::invalid_argument(std::string("the Vulkan backend does not support ") + what); + }; + if (req.sampler != pipeline::SamplerKind::kEuler) refuse("--sampler (it runs the official euler sampler)"); + if (req.scheduler != pipeline::SchedulerKind::kMoss) refuse("--scheduler (it runs the official moss schedule)"); + if (req.flow_multiplier != 1.0) refuse("--flow-multiplier"); + if (req.noise_layout != pipeline::NoiseLayout::kOfficial) refuse("--noise-layout"); + if (!req.slg_layers.empty() && req.slg_scale != 0.0) refuse("--slg-layers"); + if (!req.loras.empty()) refuse("LoRAs"); + if (req.init_image || req.mask || req.denoise_strength != 1.0) refuse("img2img (--init-image, --mask, --strength)"); + if (req.hires.enabled()) refuse("--hires"); + if (req.control.enabled()) refuse("--control"); + if (req.references.enabled()) refuse("reference images"); + if (req.accel.active(req.cfg != 1.0)) refuse("accelerators (--cfg-truncation, --step-cache)"); + if (req.width % 16 != 0 || req.height % 16 != 0) refuse("sizes that are not multiples of 16"); +} + +} // namespace + +#if defined(MOSSINFER_VULKAN) + +bool vulkan_backend_built() { return true; } + +std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run) { + check_supported(req); + if (!run.kernels.empty() && run.kernels != "coop" && run.kernels != "plain") { + throw std::invalid_argument("--vulkan-kernels takes coop or plain"); + } + vk::Device device(run.device); + const vk::Capabilities& caps = device.info().caps; + const vk::BlockPath path = run.kernels.empty() ? (caps.coop_bf16 ? vk::BlockPath::kCoop : vk::BlockPath::kPlain) + : run.kernels == "coop" ? vk::BlockPath::kCoop + : vk::BlockPath::kPlain; + // The choice is always reported: the plain path is much slower, and the user should know which one ran. + std::fprintf(stderr, "vulkan: device %u, %s; kernels: %s%s\n", device.info().index, device.info().name.c_str(), + vk::block_path_name(path), + path == vk::BlockPath::kPlain && !caps.coop_bf16 ? (" (no bf16 cooperative matrices: " + caps.coop_bf16_reason + ")").c_str() + : ""); + + const tok::Tokenizer tokenizer = tok::Tokenizer::load(req.model / "tokenizer" / "tokenizer.json"); + vk::GenerateRequest g; + g.model_dir = req.model; + tok::PromptTokenReport cond_report, uncond_report; + g.cond_ids = tok::encode_moss_prompt_report(tokenizer, req.prompt, static_cast(req.max_sequence_length), + cond_report); + const bool cfg = req.cfg != 1.0; + if (cfg) { + g.uncond_ids = tok::encode_moss_prompt_report(tokenizer, req.negative_prompt, + static_cast(req.max_sequence_length), 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 (req.seed >= 0) { + seed = static_cast(req.seed); + } else { + std::random_device rd; + seed = std::uniform_int_distribution(0, (std::uint64_t{1} << 32) - 1)(rd); + } + g.width = static_cast(req.width); + g.height = static_cast(req.height); + g.noise = pipeline::initial_noise(pipeline::NoiseLayout::kOfficial, seed, vk::kDiTLatentChannels, req.height / 16, + req.width / 16); + const pipeline::Schedule schedule = pipeline::moss_schedule(req.steps, req.shift); + g.sigmas = schedule.sigmas; + g.timesteps = schedule.timesteps; + g.cfg = static_cast(req.cfg); + for (const int l : req.uncond_skip_layers) g.uncond_skip.set(static_cast(l)); + if (!run.quiet) { + g.progress = [](const char* stage, int done, int total) { + std::fprintf(stderr, "\r%s %d/%d ", stage, done, total); + if (done == total) std::fprintf(stderr, "\n"); + }; + } + + const vk::GenerateResult r = vk::generate(device, path, g); + if (!run.out_png.empty()) { + std::vector text; + if (req.save_metadata) { + pipeline::GenerationMetadata m; + m.engine_version = mi_version(); + m.tier = "reference"; + m.models.push_back({"bundle", req.model.filename().string(), ""}); + m.prompt = req.prompt; + m.negative_prompt = req.negative_prompt; + m.sampler = "euler"; + m.scheduler = pipeline::scheduler_name(req.scheduler); + m.steps = req.steps; + m.cfg = req.cfg; + m.shift = req.shift; + m.seed = seed; + m.width = req.width; + m.height = req.height; + m.extra.emplace_back("backend", std::string("vulkan ") + vk::block_path_name(path)); + text = pipeline::generation_text_chunks(m); + } + pipeline::write_png_rgb8(run.out_png, r.rgb, static_cast(r.width), static_cast(r.height), text); + } + if (!run.quiet) { + std::fprintf(stderr, "seed %llu; text encoder %.1f s, DiT load %.1f s, denoise %.1f s (%.2f s/step), VAE %.1f s\n", + static_cast(seed), r.text_encoder_s, r.dit_load_s, r.denoise_s, + r.denoise_s / static_cast(g.timesteps.size()), r.vae_s); + } + return r.latents; +} + +#else + +bool vulkan_backend_built() { return false; } + +std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun&) { + check_supported(req); + throw std::runtime_error("this build has no Vulkan backend (configure with -DMOSSINFER_VULKAN=ON)"); +} + +#endif + +} // namespace mi::cli diff --git a/src/tools/cli/vulkan_backend.h b/src/tools/cli/vulkan_backend.h new file mode 100644 index 0000000..9fcaaef --- /dev/null +++ b/src/tools/cli/vulkan_backend.h @@ -0,0 +1,29 @@ +// `moss-infer generate --backend vulkan`: text to image on the Vulkan backend (vulkan/), the fallback for GPUs the CUDA +// engine does not support. The engine's host code supplies what does not depend on the backend: the tokenizer and chat +// template, the seed's initial noise, the schedule and the PNG writer, so a seed gives the same inputs as on CUDA. +#pragma once + +#include +#include +#include + +#include "pipeline/pipeline.h" + +namespace mi::cli { + +struct VulkanRun { + int device = -1; // -1: the first usable device, a discrete one first + std::string kernels; // "coop", "plain", or empty: coop where the device has bf16 cooperative matrices + std::filesystem::path out_png; // empty: no image written (the decode still runs) + bool quiet = false; +}; + +// Whether this build has the Vulkan backend (MOSSINFER_VULKAN). +bool vulkan_backend_built(); + +// Runs txt2img for `req` (already validated) on Vulkan and writes the PNG. Throws std::invalid_argument for a request the +// backend does not support (it names the option), std::runtime_error for a device or model failure. Returns the final +// latents, f32 [128, height / 16, width / 16]. +std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run); + +} // namespace mi::cli diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index afa71a9..629d095 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -39,6 +39,14 @@ add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp dit.cpp safet target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} PRIVATE ${_spirv_dir}) target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) +# Built into the engine (cmake/modules/vulkan.cmake), the tools are off unless asked for. +if(NOT DEFINED MOSSINFER_VULKAN_TOOLS) + set(MOSSINFER_VULKAN_TOOLS ON) +endif() +if(NOT MOSSINFER_VULKAN_TOOLS) + return() +endif() + add_executable(vulkan_check tools/vulkan_check.cpp) target_link_libraries(vulkan_check PRIVATE mossinfer_vulkan) diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp new file mode 100644 index 0000000..c330918 --- /dev/null +++ b/vulkan/generate.cpp @@ -0,0 +1,103 @@ +#include "generate.h" + +#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; + + // 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").string(), 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); + Dit dit(device, kernels, r.model_dir / "transformer"); + 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. + { + const double t0 = now_s(); + report("vae", 0, 1); + VaeDecoder vae(device, (r.model_dir / "vae" / "diffusion_pytorch_model.safetensors").string(), + path == BlockPath::kCoop ? VaePath::kCoop : VaePath::kPlain); + 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..08c5a4f --- /dev/null +++ b/vulkan/generate.h @@ -0,0 +1,44 @@ +// 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; + // 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] + 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; +}; + +// path: the DiT's and text encoder's kernels (kCoop needs Capabilities::coop_bf16); the VAE follows it. +GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& request); + +} // namespace mi::vk From f2f2ab7ecc4adfd4349a78c9312adcd559b2ce32 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:39:32 -0500 Subject: [PATCH 14/41] vulkan: DiT blocks past the device's memory stream from host memory 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. --- vulkan/dit.cpp | 94 ++++++++++++++++++++++++++------- vulkan/dit.h | 31 ++++++++--- vulkan/generate.cpp | 10 +++- vulkan/generate.h | 1 + vulkan/runtime.cpp | 17 ++++++ vulkan/runtime.h | 4 ++ vulkan/tools/vulkan_dit_run.cpp | 13 +++-- 7 files changed, 140 insertions(+), 30 deletions(-) diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index f042c13..48ac07d 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -62,12 +62,14 @@ std::vector rope_table(std::uint32_t height, std::uint32_t width, return out; } -Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::path& transformer_dir) +Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::path& transformer_dir, const DitOptions& options) : device_(device), kernels_(kernels) { const st::Index index(transformer_dir); std::vector host; - // Uploads tensors into one new buffer, each at a 256-byte aligned offset; returns their addresses. - const auto load = [&](const std::vector& names, const std::vector>& shapes) { + // 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) { @@ -77,16 +79,26 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat offsets.push_back(total); total += (t.bytes + 255) / 256 * 256; } - owned_.push_back(std::make_unique(device_, total)); + owned_.push_back(std::make_unique(device_, total, kind)); Buffer& buf = *owned_.back(); - weight_bytes_ += total; - std::vector addresses; for (std::size_t i = 0; i < names.size(); ++i) { - host.resize(index.at(names[i]).bytes); - index.read_into(names[i], host.data(), host.size()); - buf.upload(host.data(), host.size(), offsets[i]); - addresses.push_back(buf.address() + offsets[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; @@ -109,20 +121,58 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat 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]}); } - for (int l = 0; l < kDiTLayers; ++l) { + // 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) + "."; - // q, k and v load back to back, 256-aligned sizes of 18 MiB each, so they form the [9216, 3072] QKV weight. - auto t = load({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"}, - {{d}, {d}, {kHead}, {kHead}, {d, d}, {d, d}, {d, d}, {d, d}, {2 * f, d}, {d, f}}); - if (t[5] != t[4] + kDim * kDim * kBf16 || t[6] != t[5] + kDim * kDim * kBf16) fail("q, k and v are not contiguous"); - blocks_.push_back({t[0], t[1], t[2], t[3], t[4], t[7], t[8], t[9]}); + 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; + // A streamed block needs the slot too. + resident = free >= kDiTLayers * block_bytes ? kDiTLayers + : static_cast(free > block_bytes ? free / block_bytes - 1 : 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); + } else { + if (!slot_) { + slot_ = std::make_unique(device_, block_bytes); + weight_bytes_ += block_bytes; + } + blocks_.push_back(block_weights(slot_->address(), o)); + host_blocks_.push_back(owned_.back().get()); + } + } + resident_blocks_ = resident; } Dit::~Dit() = default; +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(); @@ -218,6 +268,7 @@ std::vector> Dit::forward(const std::vector& latents, 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); @@ -264,6 +315,12 @@ std::vector> Dit::forward(const std::vector& latents, 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])}; + if (const Buffer* host = host_blocks_[static_cast(l)]; host != nullptr && streamed_layer_ != l) { + const double t0 = now_ms(); + device_.run([&](VkCommandBuffer cb) { record_copy(cb, *host, 0, *slot_, 0, host->bytes()); }); + streamed_layer_ = l; + copy_ms += now_ms() - t0; + } for (int op = 0; op < kBlockOps; ++op) { const double t0 = now_ms(); device_.run([&](VkCommandBuffer cb) { record_block_op(kernels_, cb, op, blocks_[static_cast(l)], b); }); @@ -291,6 +348,7 @@ std::vector> Dit::forward(const std::vector& latents, 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; } diff --git a/vulkan/dit.h b/vulkan/dit.h index a3cc35a..fb2b284 100644 --- a/vulkan/dit.h +++ b/vulkan/dit.h @@ -26,15 +26,27 @@ 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; +}; + 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); + 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_; } + 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); @@ -49,20 +61,23 @@ class Dit { 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. + // 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 streamed blocks' copies). const std::vector& last_op_ms() const { return op_ms_; } + double last_copy_ms() const { return last_copy_ms_; } private: - struct Weight; - struct TextBlock; - struct Scratch; - Buffer& scratch(std::unique_ptr& slot, VkDeviceSize bytes); void time_modulation(float timestep); Device& device_; const BlockKernels& kernels_; - std::uint64_t weight_bytes_ = 0; + std::uint64_t weight_bytes_ = 0, host_weight_bytes_ = 0; + int resident_blocks_ = 0; + std::vector host_blocks_; // per block: its host copy when streamed, else null + std::unique_ptr slot_; // the device slot streamed blocks run from + int streamed_layer_ = -1; // the block the slot holds + double last_copy_ms_ = 0.0; std::vector> owned_; std::vector blocks_; struct Globals { diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp index c330918..c1d461d 100644 --- a/vulkan/generate.cpp +++ b/vulkan/generate.cpp @@ -51,7 +51,15 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r { double t0 = now_s(); const BlockKernels kernels(device, path); - Dit dit(device, kernels, r.model_dir / "transformer"); + // 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) { diff --git a/vulkan/generate.h b/vulkan/generate.h index 08c5a4f..e9be3b7 100644 --- a/vulkan/generate.h +++ b/vulkan/generate.h @@ -36,6 +36,7 @@ struct GenerateResult { std::vector rgb; // [height, width, 3] 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); the VAE follows it. diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index 49e90ca..68291bc 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -274,6 +274,23 @@ Device::~Device() { 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; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 269a57e..15ff87f 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -50,6 +50,10 @@ class Device { 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); diff --git a/vulkan/tools/vulkan_dit_run.cpp b/vulkan/tools/vulkan_dit_run.cpp index b773c58..0cfc56a 100644 --- a/vulkan/tools/vulkan_dit_run.cpp +++ b/vulkan/tools/vulkan_dit_run.cpp @@ -2,6 +2,7 @@ // 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|plain] [--device N] +// [--dit-budget GiB] (device memory for the weights; the blocks past it stream) // 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 @@ -73,6 +74,7 @@ int main(int argc, char** argv) { 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; fs::path out_path; std::string path_name; for (int i = 4; i < argc; ++i) { @@ -81,6 +83,7 @@ int main(int argc, char** argv) { 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 { std::printf("unknown argument %s\n", arg.c_str()); return 2; @@ -104,8 +107,12 @@ int main(int argc, char** argv) { : BlockKernels(device, path_name == "plain" ? BlockPath::kPlain : BlockPath::kCoop); 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(); - Dit dit(device, kernels, transformer); - std::printf("weights: %.2f GiB in %.1f s\n", static_cast(dit.weight_bytes()) / (1u << 30), now_s() - t0); + DitOptions options; + options.device_budget = static_cast(budget_gib * (1u << 30)); + 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; @@ -144,7 +151,7 @@ int main(int argc, char** argv) { 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)\n", ms[kBlockOps]); + 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) { From f098f76ff483409b0ec69640ad4b49026af9099b Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:40:23 -0500 Subject: [PATCH 15/41] vulkan: the Qwen3.5 text encoder, streamed layer by layer 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. --- vulkan/CMakeLists.txt | 26 + vulkan/npy.h | 158 +++++ vulkan/safetensors.cpp | 317 +++++++++++ vulkan/safetensors.h | 46 ++ vulkan/shaders/te_attention.comp | 134 +++++ vulkan/shaders/te_common.glsl | 57 ++ vulkan/shaders/te_conv_silu.comp | 38 ++ vulkan/shaders/te_gate_mul.comp | 44 ++ vulkan/shaders/te_gates.comp | 55 ++ vulkan/shaders/te_gdn_chunk.comp | 116 ++++ vulkan/shaders/te_gdn_qk.comp | 61 ++ vulkan/shaders/te_gdn_solve.comp | 62 ++ vulkan/shaders/te_gdn_state.comp | 110 ++++ vulkan/shaders/te_rms_norm.comp | 112 ++++ vulkan/te_kernels.cpp | 243 ++++++++ vulkan/te_kernels.h | 91 +++ vulkan/text_encoder.cpp | 570 ++++++++++++++++++ vulkan/text_encoder.h | 84 +++ vulkan/tools/vulkan_te_check.cpp | 165 ++++++ vulkan/tools/vulkan_te_kernel_check.cpp | 729 ++++++++++++++++++++++++ 20 files changed, 3218 insertions(+) create mode 100644 vulkan/npy.h create mode 100644 vulkan/safetensors.cpp create mode 100644 vulkan/safetensors.h create mode 100644 vulkan/shaders/te_attention.comp create mode 100644 vulkan/shaders/te_common.glsl create mode 100644 vulkan/shaders/te_conv_silu.comp create mode 100644 vulkan/shaders/te_gate_mul.comp create mode 100644 vulkan/shaders/te_gates.comp create mode 100644 vulkan/shaders/te_gdn_chunk.comp create mode 100644 vulkan/shaders/te_gdn_qk.comp create mode 100644 vulkan/shaders/te_gdn_solve.comp create mode 100644 vulkan/shaders/te_gdn_state.comp create mode 100644 vulkan/shaders/te_rms_norm.comp create mode 100644 vulkan/te_kernels.cpp create mode 100644 vulkan/te_kernels.h create mode 100644 vulkan/text_encoder.cpp create mode 100644 vulkan/text_encoder.h create mode 100644 vulkan/tools/vulkan_te_check.cpp create mode 100644 vulkan/tools/vulkan_te_kernel_check.cpp diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 9166f7e..ee24ad1 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -50,3 +50,29 @@ 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 Qwen3.5 text encoder (text_encoder.h): its own kernels on top of the DiT block's GEMMs, and a safetensors reader for +# streaming its weights. +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 safetensors.cpp ${_te_spirv_files}) + +# 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 -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) diff --git a/vulkan/npy.h b/vulkan/npy.h new file mode 100644 index 0000000..145377c --- /dev/null +++ b/vulkan/npy.h @@ -0,0 +1,158 @@ +// 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; +} + +} // namespace mi::vk::npy diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp new file mode 100644 index 0000000..c1b2ac6 --- /dev/null +++ b/vulkan/safetensors.cpp @@ -0,0 +1,317 @@ +#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 + +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::string& path) : path_(path) { +#ifdef _WIN32 + fd_ = _open(path.c_str(), _O_RDONLY | _O_BINARY); +#else + fd_ = ::open(path.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::string& path) { + std::ifstream in(path, 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; +} + +} // namespace mi::vk::st diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h new file mode 100644 index 0000000..0821058 --- /dev/null +++ b/vulkan/safetensors.h @@ -0,0 +1,46 @@ +// 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 + +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::string& path); + ~File(); + File(const File&) = delete; + File& operator=(const File&) = delete; + + const std::string& path() const { return path_; } + 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); } + + 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::string& path); + +} // namespace mi::vk::st diff --git a/vulkan/shaders/te_attention.comp b/vulkan/shaders/te_attention.comp new file mode 100644 index 0000000..76e457d --- /dev/null +++ b/vulkan/shaders/te_attention.comp @@ -0,0 +1,134 @@ +// 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. One query row of one head per +// workgroup of 32 lanes; row i attends to keys 0 .. i, in steps of 32 walked from the last down to the first, as +// FlashAttention's reverse iteration: +// S = q * K_step^T (lane c: key c of the step, dims ascending, each product and sum rounded) +// 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 +// acc_o += P * V_step (lane l: dims 4l .. 4l + 3 and 128 + 4l .. 128 + 4l + 3) +// and out = bf16(acc_o / row_sum). Rows q_first .. q_first + gl_NumWorkGroups.y - 1 of the sequence. +#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 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; // within the sequence; q, k, v and o point at its first row + uint group; // query heads per key/value head + float log2_scale; +} pc; + +const uint STEP = 32u; + +shared vec4 q_shared[64]; +shared float s_shared[STEP]; +shared float w_shared[STEP]; +shared float p_shared[STEP]; +shared float max_scaled_shared; +shared float rescale_shared; +shared float inv_sum_shared; + +void main() { + const uint lane = gl_LocalInvocationIndex; + const uint head = gl_WorkGroupID.x, kv_head = head / pc.group; + const uint row = pc.q_first + gl_WorkGroupID.y; + const uint seqlen_k = row + 1u; + + Words qrow = Words(pc.q + (uint64_t(row) * pc.ldq + uint64_t(head) * pc.hsq) * 2ul); + q_shared[lane] = f4(qrow.d[lane]); + q_shared[lane + 32u] = f4(qrow.d[lane + 32u]); + barrier(); + + 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 uint64_t ldk_bytes = uint64_t(pc.ldk) * 2ul, ldv_bytes = uint64_t(pc.ldv) * 2ul; + + precise vec4 acc_lo = vec4(0.0), acc_hi = vec4(0.0); + float row_max = 0.0, row_sum = 0.0; // lane 0's + const uint steps = (seqlen_k + STEP - 1u) / STEP; + 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, seqlen_k - n0); + + if (lane < cols) { + Words key = Words(k_head + uint64_t(n0 + lane) * ldk_bytes); + precise float s = 0.0; + for (uint g = 0u; g < 64u; ++g) { + const vec4 q4 = q_shared[g]; + const vec4 k4 = f4(key.d[g]); + s = s + q4.x * k4.x; + s = s + q4.y * k4.y; + s = s + q4.z * k4.z; + s = s + q4.w * k4.w; + } + s_shared[lane] = s; + } + barrier(); + + if (lane == 0u) { + const float prev = row_max; + float cur = first ? s_shared[0] : max(prev, s_shared[0]); + for (uint c = 1u; c < cols; ++c) cur = max(cur, s_shared[c]); + row_max = cur; + if (first) { + rescale_shared = 1.0; + row_sum = 0.0; + } else { + precise const float diff = (prev - cur) * pc.log2_scale; + const float rescale = exp2(diff); + rescale_shared = rescale; + precise const float scaled = row_sum * rescale; + row_sum = scaled; + } + precise const float max_scaled = cur * pc.log2_scale; + max_scaled_shared = max_scaled; + } + barrier(); + + float w = 0.0; + if (lane < cols) w = exp2(fma(s_shared[lane], pc.log2_scale, -max_scaled_shared)); + w_shared[lane] = w; + p_shared[lane] = uintBitsToFloat(bf(w) << 16); + if (!first) { + precise const vec4 lo = acc_lo * rescale_shared, hi = acc_hi * rescale_shared; + acc_lo = lo; + acc_hi = hi; + } + barrier(); + + if (lane == 0u) { + precise float sum = row_sum; + for (uint c = 0u; c < cols; ++c) sum = sum + w_shared[c]; + row_sum = sum; + } + for (uint c = 0u; c < cols; ++c) { + Words val = Words(v_head + uint64_t(n0 + c) * ldv_bytes); + const float p = p_shared[c]; + acc_lo = acc_lo + p * f4(val.d[lane]); + acc_hi = acc_hi + p * f4(val.d[lane + 32u]); + } + barrier(); + } + + if (lane == 0u) inv_sum_shared = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + barrier(); + precise const vec4 lo = acc_lo * inv_sum_shared, hi = acc_hi * inv_sum_shared; + Words out_row = Words(pc.o + (uint64_t(row) * pc.ldo + uint64_t(head) * pc.hso) * 2ul); + out_row.d[lane] = bf4(lo); + out_row.d[lane + 32u] = bf4(hi); +} 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/te_kernels.cpp b/vulkan/te_kernels.cpp new file mode 100644 index 0000000..97b6e18 --- /dev/null +++ b/vulkan/te_kernels.cpp @@ -0,0 +1,243 @@ +#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, 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.rows <= kMaxGroups, "attention: bad head grouping 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.group, log2_scale}; + attention_->dispatch(cb, &push, a.heads, a.rows); + 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..4d853ef --- /dev/null +++ b/vulkan/te_kernels.h @@ -0,0 +1,91 @@ +// 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. +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..0ef176a --- /dev/null +++ b/vulkan/text_encoder.cpp @@ -0,0 +1,570 @@ +#include "text_encoder.h" + +#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 computes it on the CPU (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, then 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::string& 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 in one group are placed back to back; their sizes are multiples of 256 bytes, so the 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::kCoop ? 200.0 : 20.0); + return gflop * 1e9; + } + + std::vector> encode(const std::vector>& seqs, const TeLayerHook& hook); +}; + +TextEncoder::TextEncoder(Device& device, const std::string& dir) + : TextEncoder(device, dir, device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) {} + +TextEncoder::TextEncoder(Device& device, const std::string& 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; + + Buffer weights(device, layer_bytes); + 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 = weights.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), + std::make_unique(device, layer_bytes, Buffer::Kind::kStaging)}; + 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(); + auto* dst = static_cast(staging[s]->mapped()); + for (const Piece& piece : layers[static_cast(k)].pieces) piece.file->read(*piece.tensor, dst + piece.offset); + 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; + 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(), wa = weights.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); }); + }); + }; + + for (int k = 0; k < kTeLayers; ++k) { + const LayerPlan& l = layers[static_cast(k)]; + const int s = k % 2; + { + const auto t0 = Clock::now(); + std::unique_lock lock(handoff.mu); + handoff.cv.wait(lock, [&] { return handoff.error != nullptr || handoff.filled[s] == k; }); + if (handoff.error) std::rethrow_exception(handoff.error); + stats.read_wait_s += seconds_since(t0); + } + { + const auto t0 = Clock::now(); + device.run([&](VkCommandBuffer cb) { + const VkBufferCopy region{0, 0, l.bytes}; + vkCmdCopyBuffer(cb, staging[s]->handle(), weights.handle(), 1, ®ion); + VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; + barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; + vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, + nullptr, 0, nullptr); + }); + stats.upload_s += seconds_since(t0); + ++stats.submits; + std::lock_guard lock(handoff.mu); + handoff.filled[s] = -1; + handoff.cv.notify_all(); + } + 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 = static_cast(std::max(1.0, std::floor(sub.budget() / per_row))); + 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; + // Per token and value head: the chunk products (2 x 64 x 128), the solve (256 x 64) and the scan (about 6 x 128 x 128). + 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 (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(); + 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..78c8213 --- /dev/null +++ b/vulkan/text_encoder.h @@ -0,0 +1,84 @@ +// 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 + 1 into one of two host staging buffers while layer k runs, and each layer is copied into one device +// buffer before it runs. 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 "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; + // 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; // staging -> device copies + 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::string& dir); + TextEncoder(Device& device, const std::string& 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/vulkan_te_check.cpp b/vulkan/tools/vulkan_te_check.cpp new file mode 100644 index 0000000..d5ef436 --- /dev/null +++ b/vulkan/tools/vulkan_te_check.cpp @@ -0,0 +1,165 @@ +// Runs the Vulkan text encoder on the golden token ids of the official PyTorch pipeline (tests/golden) and compares the +// raw outputs of decoder layers 0, 3 and 31 (te_layerNN, before the final norm) and last_hidden_state (te_last_hidden) +// with them: cosine over the sequence, worst row cosine, row relative L2, share of identical bf16 bits, worst bf16 ULPs +// scaled to max(|golden|, rms(golden) / 64), NaN count. The cond and uncond prompts of a case are encoded in one call; +// --iters repeats the call for timing; --profile submits each op alone and prints seconds per kind of op. Exit status: 0 +// ran without NaN, 1 otherwise. +// Usage: vulkan_te_check [--device N] [--path coop|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] [--profile] +#include +#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|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 != "plain") { + std::printf("--path takes coop or plain\n"); + return 2; + } + try { + Device device(device_index); + const BlockPath path = path_name.empty() ? (device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) + : (path_name == "coop" ? BlockPath::kCoop : BlockPath::kPlain); + 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..368b6be --- /dev/null +++ b/vulkan/tools/vulkan_te_kernel_check.cpp @@ -0,0 +1,729 @@ +// 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; + ok = report_f32("gates g", read(g, g_want.size()), g_want, 8.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 = 32; + 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; + 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; + } +} From e5358f8dbc7f05018679b3b94158754c2d2a4714 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 06:56:48 -0500 Subject: [PATCH 16/41] vulkan: bounded submits, and the Vulkan loader opened at run time 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. --- vulkan/CMakeLists.txt | 7 ++- vulkan/dit.cpp | 102 +++++++++++++++++++++++++------- vulkan/dit.h | 8 ++- vulkan/runtime.cpp | 2 + vulkan/runtime.h | 2 +- vulkan/tools/vulkan_dit_run.cpp | 5 +- vulkan/vk_api.cpp | 65 ++++++++++++++++++++ vulkan/vk_api.h | 69 +++++++++++++++++++++ 8 files changed, 234 insertions(+), 26 deletions(-) create mode 100644 vulkan/vk_api.cpp create mode 100644 vulkan/vk_api.h diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 629d095..7ab540b 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -9,7 +9,7 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) set(CMAKE_BUILD_TYPE Release) # the CPU reference is unusably slow unoptimized endif() -find_package(Vulkan REQUIRED) +find_package(Vulkan REQUIRED COMPONENTS glslc) find_program(GLSL_COMPILER NAMES glslc HINTS "$ENV{VULKAN_SDK}/Bin" "$ENV{VULKAN_SDK}/bin" REQUIRED) # Shaders become SPIR-V word lists compiled into the binary, so nothing beyond it ships. @@ -35,9 +35,10 @@ foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain el _mi_add_shader(${shader} vulkan1.2) endforeach() -add_library(mossinfer_vulkan STATIC runtime.cpp gemm.cpp block.cpp dit.cpp safetensors.cpp ${_spirv_files}) +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}) -target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Vulkan) +# Headers only: vk_api.cpp opens the loader at run time, so nothing links against it. +target_link_libraries(mossinfer_vulkan PUBLIC Vulkan::Headers ${CMAKE_DL_LIBS}) # Built into the engine (cmake/modules/vulkan.cmake), the tools are off unless asked for. if(NOT DEFINED MOSSINFER_VULKAN_TOOLS) diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index 48ac07d..7ef4d27 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -64,6 +65,7 @@ std::vector rope_table(std::uint32_t height, std::uint32_t width, 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::kCoop ? 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 @@ -164,6 +166,50 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat Dit::~Dit() = default; +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 @@ -197,31 +243,47 @@ std::unique_ptr Dit::encode_text(const std::vector& embed rope.upload(table.data(), table.size() * kBf16); auto out = std::make_unique(device_, s * kDim * kBf16); const VkDeviceAddress x = out->address(); - const float scale = static_cast(1.0 / std::sqrt(static_cast(kHead))); - const std::vector cu = {0, static_cast(length)}; + 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; }; - device_.run([&](VkCommandBuffer cb) { - kernels_.record_gemm_bias(cb, x, kDim, e.address(), kDiTTextDim, g_.text_w, kDiTTextDim, g_.text_b, length, kDim, - kDiTTextDim); + 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, y.address(), kDim, length, kDim); - kernels_.record_gemm_bias(cb, qkv.address(), kQkv, y.address(), kDim, b.qkv_w, kDim, b.qkv_b, length, kQkv, kDim); - kernels_.record_qk_norm_rope(cb, qkv.address(), kQkv, length, kDiTHeads, b.q_norm, b.k_norm, rope.address(), kHead); + 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) { - const VkDeviceAddress k = at(qkv, kDim * kBf16), v = at(qkv, 2 * kDim * kBf16); - kernels_.record_attention(cb, AttentionArgs{qkv.address(), k, v, attn.address(), kQkv, kQkv, kQkv, kDim, kHead, kHead, - kHead, kHead, cu, cu, kDiTHeads, scale}); - kernels_.record_gemm_bias_residual(cb, x, kDim, attn.address(), kDim, b.proj_w, kDim, b.proj_b, length, kDim, kDim); - kernels_.record_norm(cb, x, kDim, b.norm2, y.address(), kDim, length, kDim); + kernels_.record_qk_norm_rope(cb, qa, kQkv, length, kDiTHeads, b.q_norm, b.k_norm, rope.address(), kHead); }); - device_.run([&](VkCommandBuffer cb) { - kernels_.record_gemm_bias(cb, up.address(), kFfn, y.address(), kDim, b.up_w, kDim, b.up_b, length, kFfn, kDim); - kernels_.record_gemm_bias(cb, gate.address(), kFfn, y.address(), kDim, b.gate_w, kDim, b.gate_b, length, kFfn, kDim); - kernels_.record_situ(cb, up.address(), gate.address(), ffn.address(), length * kFfn); - kernels_.record_gemm_bias_residual(cb, x, kDim, ffn.address(), kFfn, b.down_w, kFfn, b.down_b, length, kDim, kFfn); + 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; @@ -323,7 +385,7 @@ std::vector> Dit::forward(const std::vector& latents, } for (int op = 0; op < kBlockOps; ++op) { const double t0 = now_ms(); - device_.run([&](VkCommandBuffer cb) { record_block_op(kernels_, cb, op, blocks_[static_cast(l)], b); }); + run_block_op(op, blocks_[static_cast(l)], b); op_ms_[static_cast(op)] += now_ms() - t0; } r0 = r1; diff --git a/vulkan/dit.h b/vulkan/dit.h index fb2b284..ba8ab3c 100644 --- a/vulkan/dit.h +++ b/vulkan/dit.h @@ -5,7 +5,7 @@ // 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, so no single submit runs long on a slow device. +// Work is submitted op by op, and an op over many tokens in parts (DitOptions::submit_gflop). #pragma once #include @@ -31,6 +31,10 @@ struct DitOptions { // 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 { @@ -68,12 +72,14 @@ class Dit { 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 std::unique_ptr slot_; // the device slot streamed blocks run from int streamed_layer_ = -1; // the block the slot holds diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index 68291bc..fec3a64 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -28,6 +28,8 @@ void check(VkResult result, const char* what) { } 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; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 15ff87f..ad9278d 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -2,7 +2,7 @@ // embedded in the binary. Failures throw std::runtime_error. #pragma once -#include +#include "vk_api.h" #include #include diff --git a/vulkan/tools/vulkan_dit_run.cpp b/vulkan/tools/vulkan_dit_run.cpp index 0cfc56a..1b492bf 100644 --- a/vulkan/tools/vulkan_dit_run.cpp +++ b/vulkan/tools/vulkan_dit_run.cpp @@ -3,6 +3,7 @@ // 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|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 @@ -74,7 +75,7 @@ int main(int argc, char** argv) { 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; + double budget_gib = 0, submit_gflop = 0; fs::path out_path; std::string path_name; for (int i = 4; i < argc; ++i) { @@ -84,6 +85,7 @@ int main(int argc, char** argv) { 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; @@ -109,6 +111,7 @@ int main(int argc, char** argv) { 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), diff --git a/vulkan/vk_api.cpp b/vulkan/vk_api.cpp new file mode 100644 index 0000000..323fa3d --- /dev/null +++ b/vulkan/vk_api.cpp @@ -0,0 +1,65 @@ +#include "vk_api.h" + +#include + +#if defined(_WIN32) +#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..00fd274 --- /dev/null +++ b/vulkan/vk_api.h @@ -0,0 +1,69 @@ +// 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) + +#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 From d21c401917e588b9c96e1180fc7b1d7a090753a5 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:04:58 -0500 Subject: [PATCH 17/41] vulkan: VAE plain convolutions in 128x96 tiles, a narrow kernel for conv_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. --- vulkan/CMakeLists.txt | 2 +- vulkan/shaders/vae_conv_narrow.comp | 113 ++++++++++++++++++++++++++++ vulkan/shaders/vae_conv_plain.comp | 109 +++++++++++++++------------ vulkan/tools/vulkan_vae_check.cpp | 27 +++++-- vulkan/vae.cpp | 86 +++++++++++++++------ vulkan/vae.h | 6 ++ 6 files changed, 261 insertions(+), 82 deletions(-) create mode 100644 vulkan/shaders/vae_conv_narrow.comp diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index b736d37..b37de0b 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -54,7 +54,7 @@ target_link_libraries(vulkan_block_real PRIVATE mossinfer_vulkan) # 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_gn_stats vae_gn_finalize vae_gn_apply vae_softmax vae_transpose +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() 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 index 61c1017..ec1919e 100644 --- a/vulkan/shaders/vae_conv_plain.comp +++ b/vulkan/shaders/vae_conv_plain.comp @@ -6,8 +6,9 @@ // 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 16 TM rows x 16 TN columns, invocation (x, y) the rows -// y + 16 i and columns x + 16 j. K walks in steps of 16 through shared memory, the next step loading into registers. +// 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 @@ -17,8 +18,7 @@ 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 TM = 4; // 1..8 -layout(constant_id = 4) const uint TN = 4; // 1..4 +layout(constant_id = 3) const uint TN = 6; // 4 or 6 layout(buffer_reference, std430, buffer_reference_align = 16) buffer Vec4s { vec4 d[]; }; @@ -40,12 +40,14 @@ layout(push_constant) uniform PC { uint m0; } pc; -const uint BM = 16u * TM; +const uint TM = 8u; +const uint BM = 128u; const uint BN = 16u * TN; const uint BK = 16u; -// [k][row], rows padded by 4 so the transposed stores of 4 k values spread over banks. -shared float As[BK * (BM + 4u)]; -shared float Bs[BK * (BN + 4u)]; +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); @@ -82,8 +84,20 @@ vec4 load_b(uint n, uint k) { return round4(Vec4s(pc.b + (uint64_t(n) * pc.ldb + k) * 4ul).d[0]); } -// Chunks of 4 k values: BM * 4 of A (invocation i takes chunks i and i + 256) and BN * 4 <= 256 of B per step. -const uint A_CHUNKS = (BM * 4u + 255u) / 256u; +// 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; @@ -91,56 +105,53 @@ void main() { const uint m_base = pc.m0 + gl_WorkGroupID.y * BM; const uint n_base = gl_WorkGroupID.x * BN; - float acc[8][4]; - [[unroll]] for (uint i = 0u; i < 8u; ++i) { - [[unroll]] for (uint j = 0u; j < 4u; ++j) acc[i][j] = 0.0; + 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 ra[2], rb[1]; + 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) { - const uint k0 = step * BK; - [[unroll]] for (uint q = 0u; q < 2u; ++q) { - const uint chunk = t + q * 256u; - if (q < A_CHUNKS && chunk < BM * 4u) ra[q] = load_a(m_base + chunk / 4u, k0 + (chunk % 4u) * 4u); - } - if (t < BN * 4u) rb[0] = load_b(n_base + t / 4u, k0 + (t % 4u) * 4u); barrier(); // the previous step's reads of As and Bs are done - [[unroll]] for (uint q = 0u; q < 2u; ++q) { - const uint chunk = t + q * 256u; - if (q < A_CHUNKS && chunk < BM * 4u) { - const uint row = chunk / 4u, kk = (chunk % 4u) * 4u; - As[(kk + 0u) * (BM + 4u) + row] = ra[q].x; - As[(kk + 1u) * (BM + 4u) + row] = ra[q].y; - As[(kk + 2u) * (BM + 4u) + row] = ra[q].z; - As[(kk + 3u) * (BM + 4u) + row] = ra[q].w; - } - } - if (t < BN * 4u) { - const uint row = t / 4u, kk = (t % 4u) * 4u; - Bs[(kk + 0u) * (BN + 4u) + row] = rb[0].x; - Bs[(kk + 1u) * (BN + 4u) + row] = rb[0].y; - Bs[(kk + 2u) * (BN + 4u) + row] = rb[0].z; - Bs[(kk + 3u) * (BN + 4u) + row] = rb[0].w; - } + 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(); - for (uint k = 0u; k < BK; ++k) { - float av[8], bv[4]; - [[unroll]] for (uint i = 0u; i < 8u; ++i) av[i] = i < TM ? As[k * (BM + 4u) + ty + 16u * i] : 0.0; - [[unroll]] for (uint j = 0u; j < 4u; ++j) bv[j] = j < TN ? Bs[k * (BN + 4u) + tx + 16u * j] : 0.0; - [[unroll]] for (uint i = 0u; i < 8u; ++i) { - [[unroll]] for (uint j = 0u; j < 4u; ++j) { - if (i < TM && j < TN) acc[i][j] += av[i] * bv[j]; + 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 < 8u; ++i) { - const uint m = m_base + ty + 16u * i; - if (i >= TM || m >= pc.M) continue; - [[unroll]] for (uint j = 0u; j < 4u; ++j) { - const uint n = n_base + tx + 16u * 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]; diff --git a/vulkan/tools/vulkan_vae_check.cpp b/vulkan/tools/vulkan_vae_check.cpp index 8d94c3b..f7fc259 100644 --- a/vulkan/tools/vulkan_vae_check.cpp +++ b/vulkan/tools/vulkan_vae_check.cpp @@ -4,7 +4,7 @@ // 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] +// [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile] #include #include @@ -238,7 +238,7 @@ bool report(const char* name, const std::vector& got, const std::vector(got[i]) - ref[i]) / (std::abs(ref[i]) + floor[i])); } const bool ok = nans == 0 && worst <= limit.max_err; - std::printf(" %-34s %s %zu values, %.2f %% same as the fp32-rounded replay, worst scaled error %.3g (limit %.3g: %s), %zu NaN\n", + 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; @@ -365,7 +365,7 @@ bool check_conv(Device& device, const VaeKernels& k, const ConvCase& cc) { : Limit{4e-5, "16-bit operand pieces + fp32 accumulation"}; bool ok = report(cc.name, got, want, fl, limit); if (clobbered != 0) { - std::printf(" %-34s FAILED %zu values outside the output columns were written\n", cc.name, clobbered); + std::printf(" %-40s FAILED %zu values outside the output columns were written\n", cc.name, clobbered); ok = false; } return ok; @@ -378,6 +378,8 @@ bool check_convs(Device& device, const VaeKernels& k) { {"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}, @@ -489,7 +491,7 @@ bool check_exact(Device& device, const VaeKernels& k) { bad += std::bit_cast(got[static_cast(i) * ld_out + j]) != std::bit_cast(want); } } - std::printf(" %-34s %s %zu of %zu values differ\n", "transpose", bad == 0 ? "ok " : "FAILED", bad, got.size()); + std::printf(" %-40s %s %zu of %zu values differ\n", "transpose", bad == 0 ? "ok " : "FAILED", bad, got.size()); ok = ok && bad == 0; } { @@ -514,7 +516,7 @@ bool check_exact(Device& device, const VaeKernels& k) { } } } - std::printf(" %-34s %s %zu of %zu values differ\n", "latent unpack", bad == 0 ? "ok " : "FAILED", bad, got.size()); + std::printf(" %-40s %s %zu of %zu values differ\n", "latent unpack", bad == 0 ? "ok " : "FAILED", bad, got.size()); ok = ok && bad == 0; } { @@ -538,7 +540,7 @@ bool check_exact(Device& device, const VaeKernels& k) { } bad += got[i] != want; } - std::printf(" %-34s %s %zu of %zu bytes differ\n", "rgb8", bad == 0 ? "ok " : "FAILED", bad, got.size()); + std::printf(" %-40s %s %zu of %zu bytes differ\n", "rgb8", bad == 0 ? "ok " : "FAILED", bad, got.size()); ok = ok && bad == 0; } return ok; @@ -618,7 +620,7 @@ void run_bench(Device& device, const VaeKernels& k) { struct Options { int device = -1; std::string path; - bool kat = true, decode = true, bench = false; + bool kat = true, decode = true, bench = false, profile = 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"; @@ -651,6 +653,7 @@ bool run_decode(Device& device, const Options& o, VaePath path) { } 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)); @@ -663,6 +666,12 @@ bool run_decode(Device& device, const Options& o, VaePath path) { 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.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()); @@ -716,6 +725,8 @@ int main(int argc, char** argv) { } else if (arg == "--bench") { o.bench = true; o.decode = false; + } else if (arg == "--profile") { + o.profile = true; } else if (arg == "--no-kat") { o.kat = false; } else if (arg == "--vae") { @@ -740,7 +751,7 @@ int main(int argc, char** argv) { } } 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]\n", + " [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile]\n", e.what(), argv[0]); return 2; } diff --git a/vulkan/vae.cpp b/vulkan/vae.cpp index a16dbd9..ff3c96a 100644 --- a/vulkan/vae.cpp +++ b/vulkan/vae.cpp @@ -20,6 +20,9 @@ const std::uint32_t kConvCoopSpirv[] = { 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" }; @@ -122,7 +125,7 @@ VaeKernels::VaeKernels(Device& device, VaePath path) : device_(device), path_(pa 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 -// kernel, whose operands then get the same rounding. +// 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; @@ -145,14 +148,26 @@ const Pipeline& VaeKernels::conv_pipeline(const ConvArgs& a, VaePrecision precis return *p; } const std::uint32_t tf32 = precision == VaePrecision::kTf32 ? 1 : 0; - const std::uint32_t tm = a.n <= 16 ? 8 : 4, tn = a.n <= 16 ? 1 : 4; - tile_rows = 16 * tm; + 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, tm, 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, tm, tn}); + std::vector{taps, up, tf32, tn}); } return *p; } @@ -478,11 +493,26 @@ Layout layout(std::uint32_t c, std::uint32_t h, std::uint32_t w, std::uint64_t m } // 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) : device_(device), budget_(budget) {} - void add(std::uint64_t cost, std::function op) { + 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; @@ -501,26 +531,31 @@ class Submitter { 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(ConvArgs a, VaePrecision precision) { + void gemm(const std::string& name, ConvArgs a, VaePrecision precision) { const std::uint64_t per_row = static_cast(a.n) * a.k; - const std::uint64_t rows = std::max(VaeKernels::kRowAlign, max_macs_ / per_row / VaeKernels::kRowAlign * - VaeKernels::kRowAlign); + // 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((a.m - a.m0) * per_row, [this, a, precision](VkCommandBuffer cb) { k_.record_conv(cb, a, precision); }); + 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, @@ -542,12 +577,15 @@ class Run { a.w = w; a.c = cv.c; a.b_bf16 = cv.bf16; - gemm(a, VaePrecision::kTf32); + 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; - s_.add(cost, [this, x, pixels, c = n.c](VkCommandBuffer cb) { k_.record_gn_stats(cb, x, pixels, c, part_, stats_, kResnetEps); }); - s_.add(cost, [this, n, x, y, pixels, silu](VkCommandBuffer cb) { + 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); }); } @@ -588,8 +626,8 @@ class Run { p.ldb = mid; p.ldo = 3 * mid; p.b_bf16 = w.qkv.bf16; - gemm(p, VaePrecision::kFp32); - s_.add(static_cast(l_tokens) * mid * Submitter::kElementCost, [this, qkv, vt, l_tokens, mid](VkCommandBuffer cb) { + 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. @@ -607,8 +645,8 @@ class Run { s.lda = 3 * mid; s.ldb = 3 * mid; s.ldo = l_tokens; - gemm(s, VaePrecision::kSplit16); - s_.add(static_cast(rows) * l_tokens * Submitter::kElementCost, + 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; @@ -620,7 +658,7 @@ class Run { pv.lda = l_tokens; pv.ldb = l_tokens; pv.ldo = mid; - gemm(pv, VaePrecision::kSplit16); + gemm("attention P V", pv, VaePrecision::kSplit16); } // Reference order: round(to_out(o) + bias), then + residual. ConvArgs out; @@ -636,7 +674,7 @@ class Run { out.ldb = mid; out.ldo = mid; out.b_bf16 = w.attn_out.bf16; - gemm(out, VaePrecision::kFp32); + gemm("attention out", out, VaePrecision::kFp32); } private: @@ -666,13 +704,14 @@ VaeDecodeStats VaeDecoder::decode(const float* latents, std::uint32_t h, std::ui const std::uint64_t half = l.region / 2; const Weights& wt = *weights_; - Submitter sub(device_, max_macs_); + 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(128ull * h * w * Submitter::kElementCost, [&, unpacked](VkCommandBuffer cb) { + 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); @@ -702,7 +741,7 @@ VaeDecodeStats VaeDecoder::decode(const float* latents, std::uint32_t h, std::ui 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(static_cast(n) * Submitter::kElementCost, + sub.add("rgb8", static_cast(n) * Submitter::kElementCost, [&, q, n](VkCommandBuffer cb) { kernels_->record_rgb8(cb, q, out.address(), n); }); sub.flush(); @@ -717,7 +756,6 @@ VaeDecodeStats VaeDecoder::decode(const float* latents, std::uint32_t h, std::ui image->resize(n); (q == r0.address() ? r0 : r1).download(image->data(), static_cast(n) * 4); } - VaeDecodeStats st; st.seconds = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); st.device_bytes = l.total(); st.submits = sub.submits(); diff --git a/vulkan/vae.h b/vulkan/vae.h index e6ab571..a18a4bb 100644 --- a/vulkan/vae.h +++ b/vulkan/vae.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include "runtime.h" @@ -105,6 +106,8 @@ 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 { @@ -124,6 +127,8 @@ class VaeDecoder { // 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; 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. @@ -138,6 +143,7 @@ class VaeDecoder { std::unique_ptr weights_; int base_ = 0; std::uint64_t max_macs_ = 8ull << 30; + bool profile_ = false; }; } // namespace mi::vk From 82ef6f3f17e554686d2865d10be69411b26fb8a8 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:10:41 -0500 Subject: [PATCH 18/41] vulkan: text encoder weights on the transfer queue, read by four threads 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. --- vulkan/runtime.cpp | 123 +++++++++++++++++++++++++++++++++++++--- vulkan/runtime.h | 31 +++++++++- vulkan/text_encoder.cpp | 119 ++++++++++++++++++++++++++++---------- vulkan/text_encoder.h | 8 ++- 4 files changed, 239 insertions(+), 42 deletions(-) diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index 7e99a62..7dae95b 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -199,6 +199,14 @@ Device::Device(int index) { 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}; @@ -231,14 +239,18 @@ Device::Device(int index) { } const float priority = 1.0f; - VkDeviceQueueCreateInfo queue_info{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; - queue_info.queueFamilyIndex = family; - queue_info.queueCount = 1; - queue_info.pQueuePriorities = &priority; + 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 = 1; - device_info.pQueueCreateInfos = &queue_info; + 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"); @@ -247,9 +259,21 @@ Device::Device(int index) { 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); @@ -277,20 +301,39 @@ void Device::run(const std::function& record) { 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) : device_(device), bytes_(bytes) { +Buffer::Buffer(Device& device, VkDeviceSize bytes, Kind kind, bool transfer) : device_(device), bytes_(bytes) { 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{}; @@ -407,6 +450,72 @@ void Pipeline::dispatch(VkCommandBuffer cb, const void* push, std::uint32_t x, s 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; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 83f9bea..775dcfb 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -53,6 +54,19 @@ class Device { // 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_; } std::uint32_t find_memory_type(std::uint32_t type_bits, VkMemoryPropertyFlags flags) const; @@ -64,13 +78,28 @@ class Device { VkCommandPool pool_ = VK_NULL_HANDLE; VkPhysicalDeviceMemoryProperties memory_{}; DeviceInfo info_; + 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 }; - Buffer(Device& device, VkDeviceSize bytes, Kind kind = Kind::kDevice); + // `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; diff --git a/vulkan/text_encoder.cpp b/vulkan/text_encoder.cpp index 0ef176a..4b5c8ea 100644 --- a/vulkan/text_encoder.cpp +++ b/vulkan/text_encoder.cpp @@ -1,6 +1,7 @@ #include "text_encoder.h" #include +#include #include #include #include @@ -248,6 +249,41 @@ struct TextEncoder::Impl { } 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::string& dir) @@ -291,12 +327,14 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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; - Buffer weights(device, layer_bytes); + // 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 = weights.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), - std::make_unique(device, layer_bytes, Buffer::Kind::kStaging)}; + 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 @@ -323,8 +361,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec handoff.busy[s] = true; } const auto t0 = Clock::now(); - auto* dst = static_cast(staging[s]->mapped()); - for (const Piece& piece : layers[static_cast(k)].pieces) piece.file->read(*piece.tensor, dst + piece.offset); + 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; @@ -348,6 +385,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec }; 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. { @@ -364,7 +402,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec } Submitter sub(device, budget(), options.profile, stats); - const VkDeviceAddress xa = x.address(), ha = h.address(), sa = scratch.address(), wa = weights.address(); + 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. @@ -412,33 +450,45 @@ std::vector> TextEncoder::Impl::encode(const std::vec }); }; + 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)]; - const int s = k % 2; - { - const auto t0 = Clock::now(); - std::unique_lock lock(handoff.mu); - handoff.cv.wait(lock, [&] { return handoff.error != nullptr || handoff.filled[s] == k; }); - if (handoff.error) std::rethrow_exception(handoff.error); - stats.read_wait_s += seconds_since(t0); - } - { - const auto t0 = Clock::now(); - device.run([&](VkCommandBuffer cb) { - const VkBufferCopy region{0, 0, l.bytes}; - vkCmdCopyBuffer(cb, staging[s]->handle(), weights.handle(), 1, ®ion); - VkMemoryBarrier barrier{VK_STRUCTURE_TYPE_MEMORY_BARRIER}; - barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; - vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &barrier, 0, - nullptr, 0, nullptr); - }); - stats.upload_s += seconds_since(t0); - ++stats.submits; - std::lock_guard lock(handoff.mu); - handoff.filled[s] = -1; - handoff.cv.notify_all(); + 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); @@ -547,6 +597,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec }); 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(); @@ -560,6 +611,12 @@ std::vector> TextEncoder::Impl::encode(const std::vec } } 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; diff --git a/vulkan/text_encoder.h b/vulkan/text_encoder.h index 78c8213..344a778 100644 --- a/vulkan/text_encoder.h +++ b/vulkan/text_encoder.h @@ -6,8 +6,8 @@ // 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 + 1 into one of two host staging buffers while layer k runs, and each layer is copied into one device -// buffer before it runs. Work goes to the device in submits of a bounded size (TextEncoderOptions::submit_gflop), so no +// 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 @@ -35,6 +35,8 @@ 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; }; @@ -46,7 +48,7 @@ struct TextEncoderStats { 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; // staging -> device copies + 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) From 222a8a0e2cd1a95c4b183ca234f50b368320ca42 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:11:51 -0500 Subject: [PATCH 19/41] vulkan: format the text encoder's sources at 120 columns --- vulkan/npy.h | 15 ++- vulkan/safetensors.cpp | 24 +++-- vulkan/safetensors.h | 2 +- vulkan/te_kernels.cpp | 37 ++++++-- vulkan/te_kernels.h | 19 ++-- vulkan/text_encoder.cpp | 63 +++++++----- vulkan/text_encoder.h | 35 +++---- vulkan/tools/vulkan_te_check.cpp | 28 +++--- vulkan/tools/vulkan_te_kernel_check.cpp | 121 ++++++++++++++---------- 9 files changed, 212 insertions(+), 132 deletions(-) diff --git a/vulkan/npy.h b/vulkan/npy.h index 145377c..f51be9c 100644 --- a/vulkan/npy.h +++ b/vulkan/npy.h @@ -1,5 +1,5 @@ -// 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. +// 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 @@ -33,7 +33,10 @@ struct Array { 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'; + 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))); } @@ -130,7 +133,8 @@ inline Array load(const std::filesystem::path& path) { 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::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(' '); @@ -151,7 +155,8 @@ inline Array load(const std::filesystem::path& path) { } 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"); + if (!in.read(reinterpret_cast(arr.data.data()), static_cast(bytes))) + fail(path, "truncated data"); return arr; } diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp index c1b2ac6..7442fb8 100644 --- a/vulkan/safetensors.cpp +++ b/vulkan/safetensors.cpp @@ -135,11 +135,21 @@ class Parser { 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 '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"); @@ -157,7 +167,9 @@ class Parser { } continue; } - default: c = e; break; + default: + c = e; + break; } } out += c; diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h index 0821058..c4f4a63 100644 --- a/vulkan/safetensors.h +++ b/vulkan/safetensors.h @@ -27,7 +27,7 @@ class File { const std::string& path() const { return path_; } 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 + 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; diff --git a/vulkan/te_kernels.cpp b/vulkan/te_kernels.cpp index 97b6e18..5c4af60 100644 --- a/vulkan/te_kernels.cpp +++ b/vulkan/te_kernels.cpp @@ -119,7 +119,8 @@ 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); + chunk = align(kGdnValueHeads * t * kGdnChunk * 4), + value = align(kGdnValueHeads * t * kGdnHeadDim * 4); GdnWorkspace ws{}; ws.qn = base; ws.kn = ws.qn + key; @@ -158,9 +159,22 @@ void TeKernels::record_rms_norm(VkCommandBuffer cb, const RmsNormArgs& a) const 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}; + 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); } @@ -175,7 +189,8 @@ void TeKernels::record_conv_silu(VkCommandBuffer cb, VkDeviceAddress x, std::uin } 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 { + 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}; @@ -199,7 +214,8 @@ void TeKernels::record_gated_delta_rule(VkCommandBuffer cb, const GdnArgs& a) co } }; 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}; + 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); @@ -209,7 +225,8 @@ void TeKernels::record_gated_delta_rule(VkCommandBuffer cb, const GdnArgs& a) co }); 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}; + 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); @@ -219,11 +236,13 @@ void TeKernels::record_attention(VkCommandBuffer cb, const TeAttentionArgs& a) c 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.rows <= kMaxGroups, "attention: bad head grouping or too many rows"); + require(a.group >= 1 && a.heads % a.group == 0 && a.rows <= kMaxGroups, + "attention: bad head grouping 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.group, log2_scale}; + 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.group, log2_scale}; attention_->dispatch(cb, &push, a.heads, a.rows); compute_barrier(cb); } diff --git a/vulkan/te_kernels.h b/vulkan/te_kernels.h index 4d853ef..34fd07e 100644 --- a/vulkan/te_kernels.h +++ b/vulkan/te_kernels.h @@ -1,7 +1,7 @@ -// 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. +// 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 @@ -18,11 +18,12 @@ 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. +// 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. +// 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; @@ -80,8 +81,8 @@ class TeKernels { // (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; + 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_, diff --git a/vulkan/text_encoder.cpp b/vulkan/text_encoder.cpp index 4b5c8ea..dfb9926 100644 --- a/vulkan/text_encoder.cpp +++ b/vulkan/text_encoder.cpp @@ -26,25 +26,26 @@ double seconds_since(Clock::time_point t0) { return std::chrono::duration rope_table(std::uint32_t len) { std::vector out(static_cast(len) * kRopeDim); for (std::uint32_t p = 0; p < len; ++p) { @@ -198,7 +199,8 @@ struct TextEncoder::Impl { 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"); + if (it->second.second->dtype != "BF16" || it->second.second->shape != shape) + fail("tensor " + name + " has an unexpected dtype or shape"); return it->second; } @@ -206,7 +208,7 @@ struct TextEncoder::Impl { LayerPlan l; l.full_attention = i % 4 == 3; const std::string p = "language_model.layers." + std::to_string(i) + "."; - // Tensors in one group are placed back to back; their sizes are multiples of 256 bytes, so the alignment adds no gap. + // 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"); @@ -248,10 +250,11 @@ struct TextEncoder::Impl { return gflop * 1e9; } - std::vector> encode(const std::vector>& seqs, const TeLayerHook& hook); + 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. + // 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 { @@ -260,7 +263,8 @@ struct TextEncoder::Impl { }; 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)}); + 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))); @@ -311,9 +315,11 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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 (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"); + 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())); } @@ -412,17 +418,19 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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 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); + 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 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); @@ -464,7 +472,8 @@ std::vector> TextEncoder::Impl::encode(const std::vec }; // 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)); + 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. @@ -546,8 +555,8 @@ std::vector> TextEncoder::Impl::encode(const std::vec }); 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; + 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) { @@ -564,7 +573,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec ga.out = core; ga.ws = GdnWorkspace::at(sa + gs.ws, T); ga.cu = cu; - // Per token and value head: the chunk products (2 x 64 x 128), the solve (256 x 64) and the scan (about 6 x 128 x 128). + // 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); @@ -583,7 +592,8 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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); }); + 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); } @@ -607,7 +617,8 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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); + result.emplace_back(out.begin() + std::ptrdiff_t{cu[i]} * kTeHidden, + out.begin() + std::ptrdiff_t{cu[i + 1]} * kTeHidden); } } catch (...) { stop_reader(); diff --git a/vulkan/text_encoder.h b/vulkan/text_encoder.h index 344a778..38fb2c5 100644 --- a/vulkan/text_encoder.h +++ b/vulkan/text_encoder.h @@ -1,13 +1,14 @@ // 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. +// .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 +// 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 @@ -32,8 +33,8 @@ 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. + // 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; @@ -45,14 +46,14 @@ struct TextEncoderOptions { 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 + 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::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 }; diff --git a/vulkan/tools/vulkan_te_check.cpp b/vulkan/tools/vulkan_te_check.cpp index d5ef436..71bd561 100644 --- a/vulkan/tools/vulkan_te_check.cpp +++ b/vulkan/tools/vulkan_te_check.cpp @@ -2,9 +2,9 @@ // raw outputs of decoder layers 0, 3 and 31 (te_layerNN, before the final norm) and last_hidden_state (te_last_hidden) // with them: cosine over the sequence, worst row cosine, row relative L2, share of identical bf16 bits, worst bf16 ULPs // scaled to max(|golden|, rms(golden) / 64), NaN count. The cond and uncond prompts of a case are encoded in one call; -// --iters repeats the call for timing; --profile submits each op alone and prints seconds per kind of op. Exit status: 0 -// ran without NaN, 1 otherwise. -// Usage: vulkan_te_check [--device N] [--path coop|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] [--profile] +// --iters repeats the call for timing; --profile submits each op alone and prints seconds per kind of op. Exit status: +// 0 ran without NaN, 1 otherwise. Usage: vulkan_te_check [--device N] [--path coop|plain] [--weights DIR] [--golden +// DIR] [--case ID]... [--iters N] [--profile] #include #include #include @@ -64,8 +64,9 @@ Metrics compare(const std::uint16_t* got, const std::vector& ref) { } 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::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) { @@ -100,7 +101,10 @@ int main(int argc, char** argv) { } else if (arg == "--iters" && i + 1 < argc) { iters = std::max(1, std::atoi(argv[++i])); } else { - std::printf("usage: %s [--device N] [--path coop|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] [--profile]\n", argv[0]); + std::printf( + "usage: %s [--device N] [--path coop|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] " + "[--profile]\n", + argv[0]); return 2; } } @@ -114,7 +118,8 @@ int main(int argc, char** argv) { const BlockPath path = path_name.empty() ? (device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) : (path_name == "coop" ? BlockPath::kCoop : BlockPath::kPlain); 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())); + 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"}; @@ -134,10 +139,11 @@ int main(int argc, char** argv) { 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)); + 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; diff --git a/vulkan/tools/vulkan_te_kernel_check.cpp b/vulkan/tools/vulkan_te_kernel_check.cpp index 368b6be..9cbc6fd 100644 --- a/vulkan/tools/vulkan_te_kernel_check.cpp +++ b/vulkan/tools/vulkan_te_kernel_check.cpp @@ -1,8 +1,8 @@ -// 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] +// 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 @@ -106,8 +106,8 @@ bool report_f32(const char* name, const std::vector& got, const std::vect } 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); + 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; } @@ -146,9 +146,9 @@ float host_inv_rms(const bits* row, int d, int bw, float factor, float 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) { +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); @@ -201,8 +201,9 @@ bool check_norms(Device& device, const TeKernels& k) { 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 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; @@ -232,7 +233,8 @@ bool check_norms(Device& device, const TeKernels& k) { } } 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)); + 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; @@ -255,8 +257,9 @@ bool check_norms(Device& device, const TeKernels& k) { 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); + 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); @@ -267,9 +270,11 @@ bool check_norms(Device& device, const TeKernels& k) { 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]; + 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); + 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; } @@ -277,8 +282,8 @@ bool check_norms(Device& device, const TeKernels& k) { // 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); + 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(); @@ -309,7 +314,8 @@ bool check_conv_gates(Device& device, const TeKernels& k) { 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); + 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); @@ -317,8 +323,8 @@ bool check_conv_gates(Device& device, const TeKernels& k) { 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)); + 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) { @@ -326,7 +332,8 @@ bool check_conv_gates(Device& device, const TeKernels& k) { 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); + 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)))); } @@ -342,12 +349,15 @@ bool check_conv_gates(Device& device, const TeKernels& k) { 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); }); + 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]); + 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)); @@ -362,32 +372,36 @@ bool check_conv_gates(Device& device, const TeKernels& k) { { 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)); + 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); + 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]); + 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); + 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); + 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))); + 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; @@ -400,8 +414,8 @@ bool check_conv_gates(Device& device, const TeKernels& k) { 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) { +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()); @@ -420,7 +434,8 @@ std::vector replay_gdn(const std::vector& qkv, int ld, const std:: } 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; + for (int d = 0; d < D; ++d) + dst[d] = which == 0 ? (from_bf16(x[d]) * inv) * 0.08838834764831845f : from_bf16(x[d]) * inv; } } } @@ -443,7 +458,8 @@ std::vector replay_gdn(const std::vector& qkv, int ld, const std:: } 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)]); + 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) { @@ -488,7 +504,8 @@ std::vector replay_gdn(const std::vector& qkv, int ld, const std:: 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); + 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; } } @@ -496,7 +513,9 @@ std::vector replay_gdn(const std::vector& qkv, int ld, const std:: 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); + 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); } } @@ -507,8 +526,8 @@ std::vector replay_gdn(const std::vector& qkv, int ld, const std:: } // 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) { +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) { @@ -567,7 +586,8 @@ bool check_gdn(Device& device, const TeKernels& k) { 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))); + 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; @@ -596,14 +616,16 @@ bool check_gdn(Device& device, const TeKernels& k) { 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); + 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) { +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 = 32; const float log2_scale = static_cast(static_cast(scale) * 1.4426950408889634074); std::vector acc(D, 0.0f); @@ -659,7 +681,8 @@ bool check_attention(Device& device, const TeKernels& k) { 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)); + 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]; @@ -691,7 +714,8 @@ bool check_attention(Device& device, const TeKernels& k) { 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); + 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)); } } @@ -714,7 +738,8 @@ int main(int argc, char** argv) { } 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()); + 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; From d1579a166ccf725c9579d0de48c43a12f3e38bb4 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:30:30 -0500 Subject: [PATCH 20/41] vulkan: streamed DiT blocks copy on the transfer queue while the previous 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. --- vulkan/dit.cpp | 65 ++++++++++++++++++++++++++++++++++++++++++-------- vulkan/dit.h | 16 ++++++++++--- 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index 7ef4d27..1a2f655 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -81,7 +81,8 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat offsets.push_back(total); total += (t.bytes + 255) / 256 * 256; } - owned_.push_back(std::make_unique(device_, total, kind)); + // 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; @@ -142,9 +143,9 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat int resident = kDiTLayers; if (options.device_budget > 0) { const std::uint64_t free = options.device_budget > weight_bytes_ ? options.device_budget - weight_bytes_ : 0; - // A streamed block needs the slot too. + // Streamed blocks need two slots: the next one copies while the current one runs. resident = free >= kDiTLayers * block_bytes ? kDiTLayers - : static_cast(free > block_bytes ? free / block_bytes - 1 : 0); + : static_cast(free > 2 * block_bytes ? free / block_bytes - 2 : 0); } for (int l = 0; l < kDiTLayers; ++l) { const bool on_device = l < resident; @@ -152,19 +153,50 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat if (on_device) { blocks_.push_back(block_weights(owned_.back()->address(), o)); host_blocks_.push_back(nullptr); + slot_of_.push_back(-1); } else { - if (!slot_) { - slot_ = std::make_unique(device_, block_bytes); + 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(slot_->address(), o)); + 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() = default; +Dit::~Dit() { + // Copies still in flight read host buffers this destructor frees. + for (Slot& s : slot_state_) { + if (s.ticket != 0) device_.host_wait(s.ticket); + } +} + +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. @@ -377,10 +409,19 @@ std::vector> Dit::forward(const std::vector& latents, 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])}; - if (const Buffer* host = host_blocks_[static_cast(l)]; host != nullptr && streamed_layer_ != l) { + Slot* slot = nullptr; + if (host_blocks_[static_cast(l)] != nullptr) { + // This block's copy (issued while the previous streamed block ran, or now), then the next one's into the other + // slot, which runs on the transfer queue while this block computes. const double t0 = now_ms(); - device_.run([&](VkCommandBuffer cb) { record_copy(cb, *host, 0, *slot_, 0, host->bytes()); }); - streamed_layer_ = l; + issue_copy(l); + slot = &slot_state_[slot_of_[static_cast(l)]]; + if (!slot->gpu_waited) { + device_.gpu_wait(slot->ticket); + slot->gpu_waited = true; + } + const int next = next_streamed(l); + if (next != l) issue_copy(next); copy_ms += now_ms() - t0; } for (int op = 0; op < kBlockOps; ++op) { @@ -388,6 +429,10 @@ std::vector> Dit::forward(const std::vector& latents, run_block_op(op, blocks_[static_cast(l)], b); op_ms_[static_cast(op)] += now_ms() - t0; } + if (slot != nullptr && slot->ticket != 0) { + device_.host_wait(slot->ticket); // completed: the block's first op waited for it on the device + slot->ticket = 0; + } r0 = r1; } } diff --git a/vulkan/dit.h b/vulkan/dit.h index ba8ab3c..46c64b8 100644 --- a/vulkan/dit.h +++ b/vulkan/dit.h @@ -66,7 +66,7 @@ class Dit { 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 streamed blocks' copies). + // 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_; } @@ -81,8 +81,18 @@ class Dit { int resident_blocks_ = 0; double max_gflop_ = 0; std::vector host_blocks_; // per block: its host copy when streamed, else null - std::unique_ptr slot_; // the device slot streamed blocks run from - int streamed_layer_ = -1; // the block the slot holds + // 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_; From a2b2f47186ff17a0c48bf3228e2c20d64d2620ea Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:43:19 -0500 Subject: [PATCH 21/41] moss-infer generate --backend vulkan: --tier reference, the one tier 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). --- src/tools/cli/cmd_generate.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/tools/cli/cmd_generate.cpp b/src/tools/cli/cmd_generate.cpp index 28f44ec..5d71a03 100644 --- a/src/tools/cli/cmd_generate.cpp +++ b/src/tools/cli/cmd_generate.cpp @@ -63,7 +63,8 @@ constexpr const char* kUsage = "`moss-infer check` takes the same options (-o optional): it checks them and plans the request as generate would,\n" "without loading weights or writing anything, and prints the verdict as JSON (exit 0 accepted, 1 refused).\n" "--backend vulkan (builds with -DMOSSINFER_VULKAN=ON) runs on Vulkan instead of CUDA, for GPUs the CUDA engine does\n" - "not support: txt2img with the bf16 model files and the official sampler and schedule only, much slower than CUDA.\n" + "not support: txt2img at --tier reference with the bf16 model files and the official sampler and schedule only, much\n" + "slower than CUDA.\n" "--vulkan-device picks a device (default: the first usable one, discrete first); --vulkan-kernels picks the bf16\n" "cooperative-matrix kernels or the plain ones (default: coop where the device has them). Both are reported.\n" "Flags override values read from --params. At least one of -o / --latent-out is required;\n" @@ -614,13 +615,18 @@ int generate_or_check(const std::vector& args, bool check) { // CUDA-only settings are refused rather than ignored. for (const char* f : {"--mode", "--vram-budget", "--vram-reserve", "--host-ram-budget", "--cache-dir", "--then-prompt", "--direct-io", "--buffered-io", "--te-host-pin", "--cert-dir", "--dit-blocks", "--dit", "--te", - "--vae", "--preview", "--tier", "--attention", "--json-report", "--time-embedding", "--distill", + "--vae", "--preview", "--attention", "--json-report", "--time-embedding", "--distill", "--lora-placement"}) { if (given(f)) { std::cerr << "moss-infer generate: " << f << " is not supported with --backend vulkan\n"; return 2; } } + // The backend runs the bf16 model files exactly as the reference tier does; the lossy tiers need the CUDA kernels. + if (options.tier != planner::Tier::kReference) { + std::cerr << "moss-infer generate: --backend vulkan runs --tier reference only\n"; + return 2; + } try { vulkan.out_png = out_png; vulkan.quiet = quiet; From 173d707f05a180ece590d5e3a1f6d6919e0935fa Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:48:35 -0500 Subject: [PATCH 22/41] vulkan: vulkan_vae_check --calibrate times the block's GEMM next to a decode; bands within one dispatch's rows --- vulkan/tools/vulkan_vae_check.cpp | 33 ++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/vulkan/tools/vulkan_vae_check.cpp b/vulkan/tools/vulkan_vae_check.cpp index f7fc259..a688bb7 100644 --- a/vulkan/tools/vulkan_vae_check.cpp +++ b/vulkan/tools/vulkan_vae_check.cpp @@ -4,7 +4,7 @@ // 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] +// [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile] [--calibrate] #include #include @@ -24,6 +24,7 @@ #include #include +#include "block.h" #include "runtime.h" #include "vae.h" @@ -615,12 +616,31 @@ void run_bench(Device& device, const VaeKernels& k) { } } +// 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; + 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"; @@ -657,6 +677,10 @@ bool run_decode(Device& device, const Options& o, VaePath path) { 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; @@ -666,6 +690,7 @@ bool run_decode(Device& device, const Options& o, VaePath path) { 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; @@ -725,6 +750,8 @@ int main(int argc, char** argv) { } 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") { @@ -751,7 +778,7 @@ int main(int argc, char** argv) { } } 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]\n", + " [--golden PNG] [--out PNG] [--crop H,W] [--iters N] [--max-macs N] [--profile] [--calibrate]\n", e.what(), argv[0]); return 2; } From 1d986472a85d82ce5e2b551d3b761fd81625fcfc Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:53:00 -0500 Subject: [PATCH 23/41] vulkan: README section on the text encoder and its results against the goldens --- vulkan/README.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/vulkan/README.md b/vulkan/README.md index 2ad6f07..163f47d 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -87,3 +87,37 @@ bias. | `tools/vulkan_block_check.cpp` | the block's kernels against host replays, one block's timing | | `tools/vulkan_block_real.cpp` | real blocks against the engine's captured output | | `tools/probe.py` | a device's cooperative-matrix shapes and heaps, using only the Vulkan loader (no headers) | + +## 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 both paths share 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.9998972 | 0.9998746 | 0.9998949 | 1.29e-2 | 1.32e-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.9998692 | 0.9998715 | 0.9999026 | 1.47e-2 | 1.50e-2 | 1.33e-2 | +| g_layout cond | 172 | 0.9998756 | 0.9998911 | 0.9998857 | 1.45e-2 | 1.38e-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.9985798 | | 0.9992133 | 2.19e-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). From 21fb05691b53a5a3fd75a29890baba3619772338 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 07:59:42 -0500 Subject: [PATCH 24/41] vulkan: vulkan_vae_check without an unused helper --- vulkan/tools/vulkan_vae_check.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vulkan/tools/vulkan_vae_check.cpp b/vulkan/tools/vulkan_vae_check.cpp index a688bb7..984d3a9 100644 --- a/vulkan/tools/vulkan_vae_check.cpp +++ b/vulkan/tools/vulkan_vae_check.cpp @@ -245,12 +245,6 @@ bool report(const char* name, const std::vector& got, 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 ConvCase { const char* name; std::uint32_t taps; // 0 GEMM, 1, 9 From d02dd05ef117f5ace337567581b6306e590b9cf9 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 08:45:52 -0500 Subject: [PATCH 25/41] vulkan: text encoder attention in tiles of 16 query rows and 16 keys 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. --- vulkan/README.md | 19 ++- vulkan/shaders/te_attention.comp | 177 +++++++++++++----------- vulkan/te_kernels.cpp | 12 +- vulkan/te_kernels.h | 2 +- vulkan/text_encoder.cpp | 2 +- vulkan/tools/vulkan_te_kernel_check.cpp | 4 +- 6 files changed, 122 insertions(+), 94 deletions(-) diff --git a/vulkan/README.md b/vulkan/README.md index 163f47d..29462c7 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -112,12 +112,23 @@ deviation from the same goldens when one of its kernels changes (the worse of SD | Prompt | Tokens | Cosine, coop | Cosine, plain | Cosine, torch | Mean row rel. L2, coop | plain | torch | |---|---:|---:|---:|---:|---:|---:|---:| -| g_fridge cond | 62 | 0.9998972 | 0.9998746 | 0.9998949 | 1.29e-2 | 1.32e-2 | 1.29e-2 | +| 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.9998692 | 0.9998715 | 0.9999026 | 1.47e-2 | 1.50e-2 | 1.33e-2 | -| g_layout cond | 172 | 0.9998756 | 0.9998911 | 0.9998857 | 1.45e-2 | 1.38e-2 | 1.44e-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.9985798 | | 0.9992133 | 2.19e-2 | | 1.89e-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`). diff --git a/vulkan/shaders/te_attention.comp b/vulkan/shaders/te_attention.comp index 76e457d..058ebe2 100644 --- a/vulkan/shaders/te_attention.comp +++ b/vulkan/shaders/te_attention.comp @@ -1,17 +1,20 @@ // 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. One query row of one head per -// workgroup of 32 lanes; row i attends to keys 0 .. i, in steps of 32 walked from the last down to the first, as -// FlashAttention's reverse iteration: -// S = q * K_step^T (lane c: key c of the step, dims ascending, each product and sum rounded) +// 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 -// acc_o += P * V_step (lane l: dims 4l .. 4l + 3 and 128 + 4l .. 128 + 4l + 3) -// and out = bf16(acc_o / row_sum). Rows q_first .. q_first + gl_NumWorkGroups.y - 1 of the sequence. +// 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 = 32, local_size_y = 1, local_size_z = 1) in; +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; layout(push_constant) uniform PC { uint64_t q; @@ -26,109 +29,123 @@ layout(push_constant) uniform PC { uint hsk; uint hsv; uint hso; - uint q_first; // within the sequence; q, k, v and o point at its first row - uint group; // query heads per key/value head + uint q_first; + uint row_end; + uint group; // query heads per key/value head float log2_scale; } pc; -const uint STEP = 32u; +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 vec4 q_shared[64]; -shared float s_shared[STEP]; -shared float w_shared[STEP]; -shared float p_shared[STEP]; -shared float max_scaled_shared; -shared float rescale_shared; -shared float inv_sum_shared; +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]; + +const float NEG_INF = uintBitsToFloat(0xff800000u); void main() { - const uint lane = gl_LocalInvocationIndex; + const uint t = gl_LocalInvocationIndex; const uint head = gl_WorkGroupID.x, kv_head = head / pc.group; - const uint row = pc.q_first + gl_WorkGroupID.y; - const uint seqlen_k = row + 1u; + 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 - Words qrow = Words(pc.q + (uint64_t(row) * pc.ldq + uint64_t(head) * pc.hsq) * 2ul); - q_shared[lane] = f4(qrow.d[lane]); - q_shared[lane + 32u] = f4(qrow.d[lane + 32u]); - barrier(); + 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 uint64_t ldk_bytes = uint64_t(pc.ldk) * 2ul, ldv_bytes = uint64_t(pc.ldv) * 2ul; + const uint my_row = t / 4u, my_cols = t % 4u; - precise vec4 acc_lo = vec4(0.0), acc_hi = vec4(0.0); - float row_max = 0.0, row_sum = 0.0; // lane 0's - const uint steps = (seqlen_k + STEP - 1u) / STEP; - 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, seqlen_k - n0); + 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 - if (lane < cols) { - Words key = Words(k_head + uint64_t(n0 + lane) * ldk_bytes); + 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 g = 0u; g < 64u; ++g) { - const vec4 q4 = q_shared[g]; - const vec4 k4 = f4(key.d[g]); + 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; } - s_shared[lane] = s; + const bool masked = j0 + c > i0 + my_row || j0 + c >= n_keys; + s_tile[my_row * SP + c] = masked ? NEG_INF : s; } barrier(); - if (lane == 0u) { - const float prev = row_max; - float cur = first ? s_shared[0] : max(prev, s_shared[0]); - for (uint c = 1u; c < cols; ++c) cur = max(cur, s_shared[c]); - row_max = cur; - if (first) { - rescale_shared = 1.0; - row_sum = 0.0; - } else { - precise const float diff = (prev - cur) * pc.log2_scale; - const float rescale = exp2(diff); - rescale_shared = rescale; + 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; - row_sum = scaled; + 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); } - precise const float max_scaled = cur * pc.log2_scale; - max_scaled_shared = max_scaled; + row_sum = sum; + rescale_shared[t] = rescale; } barrier(); - float w = 0.0; - if (lane < cols) w = exp2(fma(s_shared[lane], pc.log2_scale, -max_scaled_shared)); - w_shared[lane] = w; - p_shared[lane] = uintBitsToFloat(bf(w) << 16); if (!first) { - precise const vec4 lo = acc_lo * rescale_shared, hi = acc_hi * rescale_shared; - acc_lo = lo; - acc_hi = hi; - } - barrier(); - - if (lane == 0u) { - precise float sum = row_sum; - for (uint c = 0u; c < cols; ++c) sum = sum + w_shared[c]; - row_sum = sum; + [[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 < cols; ++c) { - Words val = Words(v_head + uint64_t(n0 + c) * ldv_bytes); - const float p = p_shared[c]; - acc_lo = acc_lo + p * f4(val.d[lane]); - acc_hi = acc_hi + p * f4(val.d[lane + 32u]); + 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; } - barrier(); } - if (lane == 0u) inv_sum_shared = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + if (t < BR) inv_sum_shared[t] = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; barrier(); - precise const vec4 lo = acc_lo * inv_sum_shared, hi = acc_hi * inv_sum_shared; - Words out_row = Words(pc.o + (uint64_t(row) * pc.ldo + uint64_t(head) * pc.hso) * 2ul); - out_row.d[lane] = bf4(lo); - out_row.d[lane + 32u] = bf4(hi); + [[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/te_kernels.cpp b/vulkan/te_kernels.cpp index 5c4af60..b8c7ec2 100644 --- a/vulkan/te_kernels.cpp +++ b/vulkan/te_kernels.cpp @@ -67,7 +67,7 @@ struct GdnStatePush { }; struct AttentionPush { std::uint64_t q, k, v, o; - std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_first, group; + std::uint32_t ldq, ldk, ldv, ldo, hsq, hsk, hsv, hso, q_first, row_end, group; float log2_scale; }; struct GateMulPush { @@ -236,14 +236,14 @@ void TeKernels::record_attention(VkCommandBuffer cb, const TeAttentionArgs& a) c 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.rows <= kMaxGroups, - "attention: bad head grouping or too many rows"); + 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.group, log2_scale}; - attention_->dispatch(cb, &push, a.heads, a.rows); + 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); } diff --git a/vulkan/te_kernels.h b/vulkan/te_kernels.h index 34fd07e..7e10916 100644 --- a/vulkan/te_kernels.h +++ b/vulkan/te_kernels.h @@ -53,7 +53,7 @@ struct GdnArgs { }; // 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 + 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; diff --git a/vulkan/text_encoder.cpp b/vulkan/text_encoder.cpp index dfb9926..cd024de 100644 --- a/vulkan/text_encoder.cpp +++ b/vulkan/text_encoder.cpp @@ -542,7 +542,7 @@ std::vector> TextEncoder::Impl::encode(const std::vec 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 = static_cast(std::max(1.0, std::floor(sub.budget() / per_row))); + 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); diff --git a/vulkan/tools/vulkan_te_kernel_check.cpp b/vulkan/tools/vulkan_te_kernel_check.cpp index 9cbc6fd..d443b61 100644 --- a/vulkan/tools/vulkan_te_kernel_check.cpp +++ b/vulkan/tools/vulkan_te_kernel_check.cpp @@ -626,7 +626,7 @@ bool check_gdn(Device& device, const TeKernels& k) { 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 = 32; + 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{}; @@ -698,7 +698,7 @@ bool check_attention(Device& device, const TeKernels& k) { 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; + const std::uint32_t split = static_cast(len) / 2 / 16 * 16; device.run([&](VkCommandBuffer cb) { if (split > 0) { a.q_first = 0; From d0db2b316011cdec018fcf91a76efaf9cd83a0a4 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 08:56:53 -0500 Subject: [PATCH 26/41] vulkan: fixes from review - 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. --- src/tools/cli/cmd_generate.cpp | 8 +++-- src/tools/cli/vulkan_backend.cpp | 14 +++++--- src/tools/cli/vulkan_backend.h | 10 +++--- vulkan/CMakeLists.txt | 3 ++ vulkan/block.cpp | 4 ++- vulkan/dit.cpp | 54 ++++++++++++++++++------------ vulkan/generate.cpp | 36 ++++++++++++++++---- vulkan/generate.h | 5 ++- vulkan/runtime.cpp | 42 +++++++++++++++++++++-- vulkan/runtime.h | 2 ++ vulkan/safetensors.cpp | 29 ++++++++++------ vulkan/safetensors.h | 6 ++-- vulkan/shaders/attention_coop.comp | 2 ++ vulkan/shaders/elementwise.comp | 3 +- vulkan/text_encoder.cpp | 10 +++--- vulkan/text_encoder.h | 5 +-- vulkan/vae.cpp | 10 ++++-- vulkan/vae.h | 7 ++-- 18 files changed, 181 insertions(+), 69 deletions(-) diff --git a/src/tools/cli/cmd_generate.cpp b/src/tools/cli/cmd_generate.cpp index 5d71a03..ad9c9da 100644 --- a/src/tools/cli/cmd_generate.cpp +++ b/src/tools/cli/cmd_generate.cpp @@ -630,8 +630,12 @@ int generate_or_check(const std::vector& args, bool check) { try { vulkan.out_png = out_png; vulkan.quiet = quiet; - const std::vector latents = run_vulkan_generate(req, vulkan); - if (!out_latent.empty()) write_npy_f32(out_latent, latents, {moss_dit::kLatentChannels, req.height / 16, req.width / 16}); + if (!out_latent.empty()) { + vulkan.on_latents = [&](const std::vector& latents) { + write_npy_f32(out_latent, latents, {moss_dit::kLatentChannels, req.height / 16, req.width / 16}); + }; + } + run_vulkan_generate(req, vulkan); } catch (const std::invalid_argument& e) { std::cerr << "moss-infer generate: " << e.what() << "\n"; return 2; diff --git a/src/tools/cli/vulkan_backend.cpp b/src/tools/cli/vulkan_backend.cpp index e0bbf92..75ed4f2 100644 --- a/src/tools/cli/vulkan_backend.cpp +++ b/src/tools/cli/vulkan_backend.cpp @@ -47,7 +47,7 @@ void check_supported(const pipeline::Request& req) { bool vulkan_backend_built() { return true; } -std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run) { +void run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run) { check_supported(req); if (!run.kernels.empty() && run.kernels != "coop" && run.kernels != "plain") { throw std::invalid_argument("--vulkan-kernels takes coop or plain"); @@ -105,7 +105,14 @@ std::vector run_vulkan_generate(const pipeline::Request& req, const Vulka }; } - const vk::GenerateResult r = vk::generate(device, path, g); + g.decode = !run.out_png.empty(); + g.on_latents = run.on_latents; + 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 (!run.out_png.empty()) { std::vector text; if (req.save_metadata) { @@ -133,14 +140,13 @@ std::vector run_vulkan_generate(const pipeline::Request& req, const Vulka static_cast(seed), r.text_encoder_s, r.dit_load_s, r.denoise_s, r.denoise_s / static_cast(g.timesteps.size()), r.vae_s); } - return r.latents; } #else bool vulkan_backend_built() { return false; } -std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun&) { +void run_vulkan_generate(const pipeline::Request& req, const VulkanRun&) { check_supported(req); throw std::runtime_error("this build has no Vulkan backend (configure with -DMOSSINFER_VULKAN=ON)"); } diff --git a/src/tools/cli/vulkan_backend.h b/src/tools/cli/vulkan_backend.h index 9fcaaef..16e98ef 100644 --- a/src/tools/cli/vulkan_backend.h +++ b/src/tools/cli/vulkan_backend.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include @@ -14,16 +15,17 @@ namespace mi::cli { struct VulkanRun { int device = -1; // -1: the first usable device, a discrete one first std::string kernels; // "coop", "plain", or empty: coop where the device has bf16 cooperative matrices - std::filesystem::path out_png; // empty: no image written (the decode still runs) + std::filesystem::path out_png; // empty: no decode bool quiet = false; + // Called with the final latents, f32 [128, height / 16, width / 16], before the decode. + std::function&)> on_latents; }; // Whether this build has the Vulkan backend (MOSSINFER_VULKAN). bool vulkan_backend_built(); // Runs txt2img for `req` (already validated) on Vulkan and writes the PNG. Throws std::invalid_argument for a request the -// backend does not support (it names the option), std::runtime_error for a device or model failure. Returns the final -// latents, f32 [128, height / 16, width / 16]. -std::vector run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run); +// backend does not support (it names the option), std::runtime_error for anything else. +void run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run); } // namespace mi::cli diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 3a3adc3..3b60f4c 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -37,6 +37,9 @@ 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 Vulkan::Headers Threads::Threads ${CMAKE_DL_LIBS}) diff --git a/vulkan/block.cpp b/vulkan/block.cpp index cc7914d..1bac856 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -1,5 +1,6 @@ #include "block.h" +#include #include #include #include @@ -180,7 +181,8 @@ 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; - p.dispatch(cb, &push, (push.n / 4 + 63) / 64, 1); + 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 diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index 1a2f655..7b0cb93 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -169,9 +169,14 @@ Dit::Dit(Device& device, const BlockKernels& kernels, const std::filesystem::pat } Dit::~Dit() { - // Copies still in flight read host buffers this destructor frees. + // 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) device_.host_wait(s.ticket); + if (s.ticket == 0) continue; + try { + device_.host_wait(s.ticket); + } catch (const std::exception&) { + } } } @@ -394,6 +399,27 @@ std::vector> Dit::forward(const std::vector& latents, }); 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))) { @@ -409,32 +435,18 @@ std::vector> Dit::forward(const std::vector& latents, 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])}; - Slot* slot = nullptr; - if (host_blocks_[static_cast(l)] != nullptr) { - // This block's copy (issued while the previous streamed block ran, or now), then the next one's into the other - // slot, which runs on the transfer queue while this block computes. - 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; - } - const int next = next_streamed(l); - if (next != l) issue_copy(next); - copy_ms += now_ms() - t0; - } 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; } - if (slot != nullptr && slot->ticket != 0) { - device_.host_wait(slot->ticket); // completed: the block's first op waited for it on the device - slot->ticket = 0; - } 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. diff --git a/vulkan/generate.cpp b/vulkan/generate.cpp index c1d461d..6ef986c 100644 --- a/vulkan/generate.cpp +++ b/vulkan/generate.cpp @@ -1,7 +1,9 @@ #include "generate.h" +#include #include #include +#include #include #include "text_encoder.h" @@ -34,12 +36,35 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r 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").string(), path); + 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); @@ -95,13 +120,12 @@ GenerateResult generate(Device& device, BlockPath path, const GenerateRequest& r out.latents = std::move(latents); } - // VAE decode. - { + // 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); - VaeDecoder vae(device, (r.model_dir / "vae" / "diffusion_pytorch_model.safetensors").string(), - path == BlockPath::kCoop ? VaePath::kCoop : VaePath::kPlain); - vae.decode(out.latents.data(), lh, lw, out.rgb); + vae->decode(out.latents.data(), lh, lw, out.rgb); out.vae_s = now_s() - t0; report("vae", 1, 1); } diff --git a/vulkan/generate.h b/vulkan/generate.h index e9be3b7..c34750e 100644 --- a/vulkan/generate.h +++ b/vulkan/generate.h @@ -27,13 +27,16 @@ struct GenerateRequest { 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] + 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 diff --git a/vulkan/runtime.cpp b/vulkan/runtime.cpp index 58539ff..5b09453 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -85,6 +85,14 @@ std::string why_no_coop_base(VkPhysicalDevice physical, const VkPhysicalDevicePr 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 {}; } @@ -155,6 +163,16 @@ DeviceInfo describe(VkInstance instance, VkPhysicalDevice physical, std::uint32_ 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; @@ -180,8 +198,14 @@ std::vector list_devices() { } 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); @@ -333,9 +357,14 @@ void Device::run(const std::function& record) { 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; - check(vkBeginCommandBuffer(cb, &begin), "vkBeginCommandBuffer"); - record(cb); - check(vkEndCommandBuffer(cb), "vkEndCommandBuffer"); + 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; @@ -360,6 +389,13 @@ void Device::run(const std::function& record) { } 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; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 5d9bf5a..39f407b 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -25,6 +25,8 @@ struct Capabilities { std::string coop_bf16_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 { diff --git a/vulkan/safetensors.cpp b/vulkan/safetensors.cpp index a6d0445..56c524c 100644 --- a/vulkan/safetensors.cpp +++ b/vulkan/safetensors.cpp @@ -211,6 +211,11 @@ std::uint64_t dtype_bytes(const std::string& dtype) { 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) { @@ -240,11 +245,12 @@ void close_fd(int fd) { } // namespace -File::File(const std::string& path) : path_(path) { +File::File(const std::filesystem::path& file) : path_(utf8(file)) { + const std::string& path = path_; #ifdef _WIN32 - fd_ = _open(path.c_str(), _O_RDONLY | _O_BINARY); + fd_ = _wopen(file.c_str(), _O_RDONLY | _O_BINARY); #else - fd_ = ::open(path.c_str(), O_RDONLY | O_CLOEXEC); + fd_ = ::open(file.c_str(), O_RDONLY | O_CLOEXEC); #endif if (fd_ < 0) fail("cannot open " + path); try { @@ -309,8 +315,9 @@ void File::read(const TensorInfo& t, void* dst, std::uint64_t rel, std::uint64_t read_at(fd_, dst, t.offset + rel, bytes, path_); } -std::map read_index(const std::string& path) { - std::ifstream in(path, std::ios::binary); +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(); @@ -348,11 +355,11 @@ std::vector File::read_f32(const std::string& name) const { Index::Index(const std::filesystem::path& path) { const auto add = [&](const std::filesystem::path& file) { - files_.push_back(std::make_unique(file.string())); + 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 " + path.string()); + throw std::runtime_error("safetensors: " + t.name + " appears in more than one file under " + utf8(path)); } } }; @@ -362,20 +369,20 @@ Index::Index(const std::filesystem::path& path) { } std::vector indexes, files; for (const auto& entry : std::filesystem::directory_iterator(path)) { - const std::string name = entry.path().filename().string(); + 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: " + path.string() + " holds more than one index file"); + 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: " + path.string() + " holds no index file and not exactly one .safetensors file"); + 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().string()); + 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) { diff --git a/vulkan/safetensors.h b/vulkan/safetensors.h index abc7394..3752d89 100644 --- a/vulkan/safetensors.h +++ b/vulkan/safetensors.h @@ -22,12 +22,12 @@ struct TensorInfo { class File { public: - explicit File(const std::string& path); + 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_; } + 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 @@ -46,7 +46,7 @@ class File { }; // model.safetensors.index.json's weight_map: tensor name -> shard file name (relative to the index's directory). -std::map read_index(const std::string& path); +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 { diff --git a/vulkan/shaders/attention_coop.comp b/vulkan/shaders/attention_coop.comp index faa4f49..2be652b 100644 --- a/vulkan/shaders/attention_coop.comp +++ b/vulkan/shaders/attention_coop.comp @@ -188,6 +188,8 @@ void main() { smem[sp + frow * P_ROW + half_ * 2u] = uvec4(p_bits[0], p_bits[1], p_bits[2], p_bits[3]); smem[sp + frow * P_ROW + half_ * 2u + 1u] = uvec4(p_bits[4], p_bits[5], p_bits[6], p_bits[7]); + // 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) { diff --git a/vulkan/shaders/elementwise.comp b/vulkan/shaders/elementwise.comp index 9ac56e7..e35d985 100644 --- a/vulkan/shaders/elementwise.comp +++ b/vulkan/shaders/elementwise.comp @@ -16,7 +16,8 @@ layout(push_constant) uniform PC { } pc; void main() { - const uint i = gl_GlobalInvocationID.x; + // 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) { diff --git a/vulkan/text_encoder.cpp b/vulkan/text_encoder.cpp index dfb9926..1ab47b8 100644 --- a/vulkan/text_encoder.cpp +++ b/vulkan/text_encoder.cpp @@ -178,13 +178,13 @@ struct TextEncoder::Impl { std::uint64_t layer_bytes = 0; TextEncoderStats stats; - Impl(Device& d, const std::string& dir, BlockPath p, TextEncoderOptions o) + 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"); + 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); + 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}); @@ -290,10 +290,10 @@ struct TextEncoder::Impl { } }; -TextEncoder::TextEncoder(Device& device, const std::string& dir) +TextEncoder::TextEncoder(Device& device, const std::filesystem::path& dir) : TextEncoder(device, dir, device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) {} -TextEncoder::TextEncoder(Device& device, const std::string& dir, BlockPath path, TextEncoderOptions options) +TextEncoder::TextEncoder(Device& device, const std::filesystem::path& dir, BlockPath path, TextEncoderOptions options) : impl_(std::make_unique(device, dir, path, options)) {} TextEncoder::~TextEncoder() = default; diff --git a/vulkan/text_encoder.h b/vulkan/text_encoder.h index 38fb2c5..e70a53f 100644 --- a/vulkan/text_encoder.h +++ b/vulkan/text_encoder.h @@ -12,6 +12,7 @@ // submit runs long enough to trip a display driver's watchdog. #pragma once +#include #include #include #include @@ -65,8 +66,8 @@ 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::string& dir); - TextEncoder(Device& device, const std::string& dir, BlockPath path, TextEncoderOptions options = {}); + 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; diff --git a/vulkan/vae.cpp b/vulkan/vae.cpp index ff3c96a..08cdd3d 100644 --- a/vulkan/vae.cpp +++ b/vulkan/vae.cpp @@ -385,15 +385,15 @@ void relocate(Resnet& r, VkDeviceAddress base) { } // namespace -VaeDecoder::VaeDecoder(Device& device, const std::string& path) +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::string& path, VaePath vae_path) +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: " + path + " has no FLUX.2 decoder.conv_in.weight [4c, 32, 3, 3]"); + 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"); @@ -690,6 +690,10 @@ std::uint64_t VaeDecoder::workspace_bytes(std::uint32_t h, std::uint32_t w) cons 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(); diff --git a/vulkan/vae.h b/vulkan/vae.h index a18a4bb..6f3ee38 100644 --- a/vulkan/vae.h +++ b/vulkan/vae.h @@ -21,6 +21,7 @@ // 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 @@ -115,8 +116,8 @@ class VaeDecoder { // 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::string& safetensors_path); - VaeDecoder(Device& device, const std::string& safetensors_path, VaePath path); + 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; @@ -126,6 +127,8 @@ class VaeDecoder { 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; } From 1d35e2f440a00bd37cf7e8fe4c3e72a47b750121 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 08:57:17 -0500 Subject: [PATCH 27/41] docs: the Vulkan backend (docs/vulkan.md) --- cmake/modules/vulkan.cmake | 4 +-- docs/vulkan.md | 65 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) create mode 100644 docs/vulkan.md diff --git a/cmake/modules/vulkan.cmake b/cmake/modules/vulkan.cmake index 300436f..1ffc0af 100644 --- a/cmake/modules/vulkan.cmake +++ b/cmake/modules/vulkan.cmake @@ -1,6 +1,6 @@ # The Vulkan backend (vulkan/), the fallback for GPUs the CUDA engine does not support: `moss-infer generate --backend -# vulkan`. Off by default; it needs the Vulkan headers (1.4.313 or newer, for VK_KHR_shader_bfloat16), the loader and -# glslc. Without it the CLI refuses --backend vulkan with a message saying so. +# vulkan`. Off by default; it needs Vulkan headers that define VK_KHR_shader_bfloat16 and a +# glslc that compiles GL_EXT_bfloat16 (a current Vulkan SDK). The loader is opened at run time. Without it the CLI refuses --backend vulkan with a message saying so. option(MOSSINFER_VULKAN "Build the Vulkan backend into the engine (moss-infer generate --backend vulkan)" OFF) if(MOSSINFER_VULKAN) set(MOSSINFER_VULKAN_TOOLS OFF CACHE BOOL "Build the Vulkan backend's check and benchmark tools") diff --git a/docs/vulkan.md b/docs/vulkan.md new file mode 100644 index 0000000..55a8db0 --- /dev/null +++ b/docs/vulkan.md @@ -0,0 +1,65 @@ +# Vulkan backend + +`moss-infer generate --backend vulkan` runs MOSS-Image on Vulkan instead of CUDA. It is the fallback for GPUs the CUDA +engine does not support (AMD, Intel, older NVIDIA cards) and is much slower than CUDA. On a supported NVIDIA GPU, use +the default CUDA backend. + +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 generate --backend vulkan --tier reference --model MOSS-Image/base \ + --prompt "a red apple on a wooden table" --size 1024x1024 -o apple.png +``` + +`--vulkan-device N` picks a device (default: the first usable one, a discrete GPU first). `--vulkan-kernels coop|plain` +picks the kernels (default: `coop` where the device has them). The device and kernels are printed at the start of +every run. + +## 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 | +| `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 | + +## 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 | + +A 512x512, 20-step image with seed 7 is 45.97 dB PSNR from the CUDA backend's image of the same command. + +## Building + +The backend is off by default. It needs Vulkan headers that define `VK_KHR_shader_bfloat16` and a `glslc` that compiles +`GL_EXT_bfloat16`, as in a current Vulkan SDK. The binaries load the Vulkan loader at run time, so a build with the +backend still starts on a machine without Vulkan and only `--backend vulkan` reports it missing. + +```sh +cmake --preset linux-cuda13 -DMOSSINFER_VULKAN=ON +cmake --build --preset linux-cuda13 +``` + +The backend's own checks and benchmarks build from `vulkan/` alone, without CUDA; see [vulkan/README.md](../vulkan/README.md). From 8a79f2dba25e6b52cd21952a2834b195aecbe9d3 Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 08:58:04 -0500 Subject: [PATCH 28/41] vulkan/README: the backend as it stands, from bring-up to text to image --- vulkan/README.md | 135 ++++++++++++++++++++++------------------------- 1 file changed, 64 insertions(+), 71 deletions(-) diff --git a/vulkan/README.md b/vulkan/README.md index 29462c7..02b3118 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -1,92 +1,69 @@ # Vulkan backend -The fallback for hardware the engine does not otherwise support. CUDA stays the path on every GPU it supports. This is a -standalone project for now: it links no engine code and nothing here affects the engine's build. Design and research: -[design](../docs/vulkan-backend-design.md), [research](../docs/vulkan-backend-research.md). +The fallback for GPUs the CUDA engine does not support: `moss-infer generate --backend vulkan` (user guide: +[docs/vulkan.md](../docs/vulkan.md)). CUDA stays the path on every GPU it supports. This directory builds on its own, +without CUDA, for the backend's checks and benchmarks; the engine builds it in with `-DMOSSINFER_VULKAN=ON` +(`cmake/modules/vulkan.cmake`, tools off). -What exists is the device runtime, the GEMM and one DiT block's kernels. It cannot generate images yet: the text encoder, the VAE, the DiT's outer layers and the pipeline are still to come. +Every model runs on one of two kernel paths, chosen per device 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` | +| `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 +cmake -S vulkan -B build-vulkan -G Ninja cmake --build build-vulkan -build-vulkan/vulkan_check --list # each device, and which GEMM it gets -build-vulkan/vulkan_check # checks against an fp64 CPU reference, then the engine's shapes -build-vulkan/vulkan_check --path plain --no-bench +build-vulkan/vulkan_check --list # each device, and why a device gets the plain path +build-vulkan/vulkan_block_check [--path plain] # 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 and loader, and `glslc`. Shaders compile into the binary. The exit code is non-zero when a check -fails. A check poisons C and uses strides larger than the rows, so a write past a row or past M fails it. - -## GEMM - -`C[M,N] = A[M,K] * B[N,K]^T`, with B in the engine's weight layout, f32 accumulation and f32 output. M and N are free. +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. -| Path | Used when | Inputs | K | -|---|---|---|---| -| Cooperative matrix (`gemm_coop.comp`) | the device has 16x16x16 f16 x f16 -> f32 in subgroup scope, a compute subgroup size fixable at 32, and 24 KiB of shared memory | f16 | multiple of 16 | -| Plain compute (`gemm_plain.comp`) | every other device with Vulkan 1.2, buffer device addresses and 64-bit integers | f16 or bf16 | multiple of 4 | +## DiT -`--list` says why a device gets the plain path. The plain shader declares only `Shader`, `Int64` and -`PhysicalStorageBufferAddresses`: it reads 32-bit words, so it needs no 16-bit support and does not depend on the subgroup -size. +`block.h` holds one DiT block's kernels on both paths, 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. -## Results - -TFLOPS on an RTX 4090 Laptop, driver 615.78, with cuBLASLt (f16 inputs, f32 compute) on the same GPU. M = 8222 is the DiT at -1024x1024 with CFG. - -| Shape (M x N x K) | Cooperative matrix | Plain | cuBLASLt | -|---|---:|---:|---:| -| 8222 x 9216 x 3072 (QKV) | 27.0 | 1.39 | 52.6 | -| 8222 x 3072 x 3072 (projection) | 33.1 | 1.62 | 57.4 | -| 8222 x 16384 x 3072 (up/gate) | 23.5 | 1.14 | 45.5 | -| 8222 x 3072 x 8192 (down) | 24.3 | 1.52 | 70.1 | -| 282 x 12288 x 4096 (text encoder) | 27.3 | 1.02 | 59.5 | -| 282 x 4096 x 12288 (text encoder) | 18.1 | | 52.2 | - -Checks pass on that GPU on both paths, and through the plain path on llvmpipe (Mesa's CPU driver, 8-wide subgroups, 32 KiB of -shared memory), with core validation clean. - -## DiT block - -`block.h` is the plain compute path of one bf16 DiT block: RMSNorm with modulation, the GEMMs with their gated residual and SiTU-GLU epilogues, QK norm with RoPE, and varlen attention with an online softmax. It has the engine's rounding points and needs no cooperative matrices, 16-bit support or particular subgroup size. - -```sh -build-vulkan/vulkan_block_check # each kernel against a host replay of its op order, then one block -build-vulkan/vulkan_block_real DIR # real blocks against the engine's output, from files captured from its forward -``` +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. -On real weights and activations at the 1024x1024 golden case, against the engine: +Against the engine's own output on real blocks of the 1024x1024 golden case (captured from its forward): -| Block | Same bits | p99 scaled ULPs | Max scaled ULPs | Cosine | -|---:|---:|---:|---:|---:| -| 1 | 83.98 % | 0.06 | 2.00 | 0.999999977 | -| 22 | 72.36 % | 0.25 | 4.00 | 0.999999867 | -| 41 | 69.23 % | 1.00 | 96.00 | 0.999997331 | +| 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 | -Bits differ because fp32 scalar and tensor-core accumulation are not bit-compatible. Known-answer checks pass on an RTX 4090 Laptop and on llvmpipe. One block of 8484 tokens takes about 2 s on that GPU through this path, against 31 ms for the engine, so about 86 s per CFG step at 1024x1024. +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. -## Numerics +## VAE -The cooperative-matrix path rounds like the tensor cores: its error matches cuBLASLt on the same data up to K = 1040 and is -biased toward zero. At K = 4096 it reaches 8.1e-4 against cuBLASLt's 1.3e-4, on values near 80. The plain path has no such -bias. - -## Files - -| File | Purpose | -|---|---| -| `runtime.h`, `runtime.cpp` | device selection, capabilities, buffers, pipelines, barriers | -| `gemm.h`, `gemm.cpp` | path selection, argument checks, dispatch | -| `block.h`, `block.cpp` | one DiT block's plain compute kernels | -| `half.h` | host fp16 and bf16 conversions | -| `shaders/` | compute shaders | -| `tools/vulkan_check.cpp` | device list, GEMM checks, timing | -| `tools/vulkan_block_check.cpp` | the block's kernels against host replays, one block's timing | -| `tools/vulkan_block_real.cpp` | real blocks against the engine's captured output | -| `tools/probe.py` | a device's cooperative-matrix shapes and heaps, using only the Vulkan loader (no headers) | +`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 @@ -132,3 +109,19 @@ Times on that GPU with nothing else running, the weights read from a drive that 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 both paths | +| `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/` | checks and benchmarks | From 3684915ccba8bc5033074780594971ced734655d Mon Sep 17 00:00:00 2001 From: xycld Date: Mon, 2 Nov 2026 08:58:15 -0500 Subject: [PATCH 29/41] vulkan_block_real: --path --- vulkan/tools/vulkan_block_real.cpp | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/vulkan/tools/vulkan_block_real.cpp b/vulkan/tools/vulkan_block_real.cpp index 65643e5..17ebce7 100644 --- a/vulkan/tools/vulkan_block_real.cpp +++ b/vulkan/tools/vulkan_block_real.cpp @@ -1,4 +1,4 @@ -// Runs one DiT block's plain compute kernels on real weights and activations, and compares each output with the engine's +// 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 @@ -78,15 +78,22 @@ bool compare(const std::vector& got, const std::vector& engine) { int main(int argc, char** argv) { if (argc < 2) { - std::printf("usage: %s DIR [--device N]\n", argv[0]); + std::printf("usage: %s DIR [--device N] [--path coop|plain]\n", argv[0]); return 2; } const fs::path dir = argv[1]; - const int device_index = argc >= 4 && std::string(argv[2]) == "--device" ? std::atoi(argv[3]) : -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); - std::printf("device %u: %s\n", device.info().index, device.info().name.c_str()); - const BlockKernels kernels(device); + const BlockKernels kernels = path_name.empty() ? BlockKernels(device) + : BlockKernels(device, path_name == "plain" ? BlockPath::kPlain : BlockPath::kCoop); + 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 { From 0bb12aefe3e0e32b9fd6c4063d1b202a683082d4 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:35:35 -0500 Subject: [PATCH 30/41] docs/vulkan.md: whole-image accuracy at 1024x1024 and what bounds it --- docs/vulkan.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/vulkan.md b/docs/vulkan.md index 55a8db0..780ba84 100644 --- a/docs/vulkan.md +++ b/docs/vulkan.md @@ -49,7 +49,12 @@ Against the official PyTorch pipeline's golden data at 1024x1024 (`tests/golden/ | 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 | -A 512x512, 20-step image with seed 7 is 45.97 dB PSNR from the CUDA backend's image of the same command. +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 run 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 45.97 dB PSNR from the CUDA backend's image of the same command. ## Building From fba035a57b87bb50b465c5b967dc3f27bd4bdfd3 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:47:51 -0500 Subject: [PATCH 31/41] moss-infer-vulkan: the Vulkan backend as a standalone program; the engine'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. --- cmake/modules/vulkan.cmake | 11 - src/tools/cli/cmd_generate.cpp | 58 ----- src/tools/cli/vulkan_backend.cpp | 156 -------------- src/tools/cli/vulkan_backend.h | 31 --- vulkan/CMakeLists.txt | 58 ++++- vulkan/tools/moss_infer_vulkan.cpp | 336 +++++++++++++++++++++++++++++ 6 files changed, 388 insertions(+), 262 deletions(-) delete mode 100644 cmake/modules/vulkan.cmake delete mode 100644 src/tools/cli/vulkan_backend.cpp delete mode 100644 src/tools/cli/vulkan_backend.h create mode 100644 vulkan/tools/moss_infer_vulkan.cpp diff --git a/cmake/modules/vulkan.cmake b/cmake/modules/vulkan.cmake deleted file mode 100644 index 1ffc0af..0000000 --- a/cmake/modules/vulkan.cmake +++ /dev/null @@ -1,11 +0,0 @@ -# The Vulkan backend (vulkan/), the fallback for GPUs the CUDA engine does not support: `moss-infer generate --backend -# vulkan`. Off by default; it needs Vulkan headers that define VK_KHR_shader_bfloat16 and a -# glslc that compiles GL_EXT_bfloat16 (a current Vulkan SDK). The loader is opened at run time. Without it the CLI refuses --backend vulkan with a message saying so. -option(MOSSINFER_VULKAN "Build the Vulkan backend into the engine (moss-infer generate --backend vulkan)" OFF) -if(MOSSINFER_VULKAN) - set(MOSSINFER_VULKAN_TOOLS OFF CACHE BOOL "Build the Vulkan backend's check and benchmark tools") - add_subdirectory(${CMAKE_SOURCE_DIR}/vulkan ${CMAKE_BINARY_DIR}/vulkan) - set_target_properties(mossinfer_vulkan PROPERTIES POSITION_INDEPENDENT_CODE ON) - target_link_libraries(mossinfer_core PUBLIC mossinfer_vulkan) - target_compile_definitions(mossinfer_core PUBLIC MOSSINFER_VULKAN=1) -endif() diff --git a/src/tools/cli/cmd_generate.cpp b/src/tools/cli/cmd_generate.cpp index ad9c9da..4737abf 100644 --- a/src/tools/cli/cmd_generate.cpp +++ b/src/tools/cli/cmd_generate.cpp @@ -27,7 +27,6 @@ #include "pipeline/report_json.h" #include "planner/planner.h" #include "tools/cli/commands.h" -#include "tools/cli/vulkan_backend.h" #include "util/json.h" #include "util/path.h" @@ -57,16 +56,10 @@ constexpr const char* kUsage = " [--cfg-truncation X] [--step-cache off|fbcache] [--step-cache-threshold X]\n" " [--control [--control-scale S] [--control-image =\n" " [--control-strength X] [--control-start X] [--control-end X]]...]\n" - " [--backend cuda|vulkan [--vulkan-device N] [--vulkan-kernels coop|plain]]\n" " [--reference ]... [--reference-max-edge N]\n" " [--lora-placement default|merge|side_branch]\n" "`moss-infer check` takes the same options (-o optional): it checks them and plans the request as generate would,\n" "without loading weights or writing anything, and prints the verdict as JSON (exit 0 accepted, 1 refused).\n" - "--backend vulkan (builds with -DMOSSINFER_VULKAN=ON) runs on Vulkan instead of CUDA, for GPUs the CUDA engine does\n" - "not support: txt2img at --tier reference with the bf16 model files and the official sampler and schedule only, much\n" - "slower than CUDA.\n" - "--vulkan-device picks a device (default: the first usable one, discrete first); --vulkan-kernels picks the bf16\n" - "cooperative-matrix kernels or the plain ones (default: coop where the device has them). Both are reported.\n" "Flags override values read from --params. At least one of -o / --latent-out is required;\n" "without -o the VAE decode is skipped. --vram-budget caps the process's GPU memory (CUDA context included).\n" "Each --then-prompt generates one more image in the same engine session (same settings, warm models and\n" @@ -350,8 +343,6 @@ int generate_or_check(const std::vector& args, bool check) { std::optional strength; std::string hires; bool distill = false, distill_dataspace = false; - std::string backend = "cuda"; - VulkanRun vulkan; json::Value pdoc; const json::Value* pobj = nullptr; // the parameters object of --params int distill_steps = 0; @@ -498,12 +489,6 @@ int generate_or_check(const std::vector& args, bool check) { else if (f == "--reference") reference_images.push_back(path_arg(v, f)); else if (f == "--reference-max-edge") req.references.max_edge = to_int32(f, v); else if (f == "--lora-placement") req.lora_placement = lora::parse_placement(v); - else if (f == "--backend") { - if (v != "cuda" && v != "vulkan") throw std::invalid_argument("--backend takes cuda or vulkan"); - backend = v; - } - else if (f == "--vulkan-device") vulkan.device = to_int32(f, v); - else if (f == "--vulkan-kernels") vulkan.kernels = v; else throw std::invalid_argument("unknown flag " + f); } } catch (const std::invalid_argument& e) { @@ -595,10 +580,6 @@ int generate_or_check(const std::vector& args, bool check) { return 1; } req.decode = !out_png.empty() || (check && out_latent.empty()); // check: a PNG unless only --latent-out is given - if (check && backend != "cuda") { - std::cerr << "moss-infer check: --backend vulkan is not supported by check\n"; - return 2; - } if (check) return run_check(options, req, then_prompts, !preview_png.empty(), quiet); pipeline::validate(req); for (const std::string& p : then_prompts) { @@ -607,45 +588,6 @@ int generate_or_check(const std::vector& args, bool check) { pipeline::validate(f); } - if ((given("--vulkan-device") || given("--vulkan-kernels")) && backend != "vulkan") { - std::cerr << "moss-infer generate: --vulkan-device and --vulkan-kernels need --backend vulkan\n" << kUsage; - return 2; - } - if (backend == "vulkan") { - // CUDA-only settings are refused rather than ignored. - for (const char* f : {"--mode", "--vram-budget", "--vram-reserve", "--host-ram-budget", "--cache-dir", "--then-prompt", - "--direct-io", "--buffered-io", "--te-host-pin", "--cert-dir", "--dit-blocks", "--dit", "--te", - "--vae", "--preview", "--attention", "--json-report", "--time-embedding", "--distill", - "--lora-placement"}) { - if (given(f)) { - std::cerr << "moss-infer generate: " << f << " is not supported with --backend vulkan\n"; - return 2; - } - } - // The backend runs the bf16 model files exactly as the reference tier does; the lossy tiers need the CUDA kernels. - if (options.tier != planner::Tier::kReference) { - std::cerr << "moss-infer generate: --backend vulkan runs --tier reference only\n"; - return 2; - } - try { - vulkan.out_png = out_png; - vulkan.quiet = quiet; - if (!out_latent.empty()) { - vulkan.on_latents = [&](const std::vector& latents) { - write_npy_f32(out_latent, latents, {moss_dit::kLatentChannels, req.height / 16, req.width / 16}); - }; - } - run_vulkan_generate(req, vulkan); - } catch (const std::invalid_argument& e) { - std::cerr << "moss-infer generate: " << e.what() << "\n"; - return 2; - } catch (const std::exception& e) { - std::cerr << "moss-infer generate: " << e.what() << "\n"; - return 1; - } - return 0; - } - std::signal(SIGINT, on_sigint); pipeline::Pipeline pipe(options); pipeline::GenerateControl control; diff --git a/src/tools/cli/vulkan_backend.cpp b/src/tools/cli/vulkan_backend.cpp deleted file mode 100644 index 75ed4f2..0000000 --- a/src/tools/cli/vulkan_backend.cpp +++ /dev/null @@ -1,156 +0,0 @@ -#include "tools/cli/vulkan_backend.h" - -#include -#include -#include -#include -#include -#include - -#include "mossinfer/mossinfer.h" -#include "pipeline/png_writer.h" -#include "pipeline/samplers/initial_noise.h" -#include "pipeline/samplers/schedules.h" -#include "tokenizer/prompt_tokens.h" -#include "tokenizer/qwen_bpe.h" - -#if defined(MOSSINFER_VULKAN) -#include "generate.h" -#include "runtime.h" -#endif - -namespace mi::cli { -namespace { - -// What the Vulkan backend runs: txt2img with the official sampler and schedule. Anything else is refused by name. -void check_supported(const pipeline::Request& req) { - const auto refuse = [](const char* what) { - throw std::invalid_argument(std::string("the Vulkan backend does not support ") + what); - }; - if (req.sampler != pipeline::SamplerKind::kEuler) refuse("--sampler (it runs the official euler sampler)"); - if (req.scheduler != pipeline::SchedulerKind::kMoss) refuse("--scheduler (it runs the official moss schedule)"); - if (req.flow_multiplier != 1.0) refuse("--flow-multiplier"); - if (req.noise_layout != pipeline::NoiseLayout::kOfficial) refuse("--noise-layout"); - if (!req.slg_layers.empty() && req.slg_scale != 0.0) refuse("--slg-layers"); - if (!req.loras.empty()) refuse("LoRAs"); - if (req.init_image || req.mask || req.denoise_strength != 1.0) refuse("img2img (--init-image, --mask, --strength)"); - if (req.hires.enabled()) refuse("--hires"); - if (req.control.enabled()) refuse("--control"); - if (req.references.enabled()) refuse("reference images"); - if (req.accel.active(req.cfg != 1.0)) refuse("accelerators (--cfg-truncation, --step-cache)"); - if (req.width % 16 != 0 || req.height % 16 != 0) refuse("sizes that are not multiples of 16"); -} - -} // namespace - -#if defined(MOSSINFER_VULKAN) - -bool vulkan_backend_built() { return true; } - -void run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run) { - check_supported(req); - if (!run.kernels.empty() && run.kernels != "coop" && run.kernels != "plain") { - throw std::invalid_argument("--vulkan-kernels takes coop or plain"); - } - vk::Device device(run.device); - const vk::Capabilities& caps = device.info().caps; - const vk::BlockPath path = run.kernels.empty() ? (caps.coop_bf16 ? vk::BlockPath::kCoop : vk::BlockPath::kPlain) - : run.kernels == "coop" ? vk::BlockPath::kCoop - : vk::BlockPath::kPlain; - // The choice is always reported: the plain path is much slower, and the user should know which one ran. - std::fprintf(stderr, "vulkan: device %u, %s; kernels: %s%s\n", device.info().index, device.info().name.c_str(), - vk::block_path_name(path), - path == vk::BlockPath::kPlain && !caps.coop_bf16 ? (" (no bf16 cooperative matrices: " + caps.coop_bf16_reason + ")").c_str() - : ""); - - const tok::Tokenizer tokenizer = tok::Tokenizer::load(req.model / "tokenizer" / "tokenizer.json"); - vk::GenerateRequest g; - g.model_dir = req.model; - tok::PromptTokenReport cond_report, uncond_report; - g.cond_ids = tok::encode_moss_prompt_report(tokenizer, req.prompt, static_cast(req.max_sequence_length), - cond_report); - const bool cfg = req.cfg != 1.0; - if (cfg) { - g.uncond_ids = tok::encode_moss_prompt_report(tokenizer, req.negative_prompt, - static_cast(req.max_sequence_length), 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 (req.seed >= 0) { - seed = static_cast(req.seed); - } else { - std::random_device rd; - seed = std::uniform_int_distribution(0, (std::uint64_t{1} << 32) - 1)(rd); - } - g.width = static_cast(req.width); - g.height = static_cast(req.height); - g.noise = pipeline::initial_noise(pipeline::NoiseLayout::kOfficial, seed, vk::kDiTLatentChannels, req.height / 16, - req.width / 16); - const pipeline::Schedule schedule = pipeline::moss_schedule(req.steps, req.shift); - g.sigmas = schedule.sigmas; - g.timesteps = schedule.timesteps; - g.cfg = static_cast(req.cfg); - for (const int l : req.uncond_skip_layers) g.uncond_skip.set(static_cast(l)); - if (!run.quiet) { - g.progress = [](const char* stage, int done, int total) { - std::fprintf(stderr, "\r%s %d/%d ", stage, done, total); - if (done == total) std::fprintf(stderr, "\n"); - }; - } - - g.decode = !run.out_png.empty(); - g.on_latents = run.on_latents; - 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 (!run.out_png.empty()) { - std::vector text; - if (req.save_metadata) { - pipeline::GenerationMetadata m; - m.engine_version = mi_version(); - m.tier = "reference"; - m.models.push_back({"bundle", req.model.filename().string(), ""}); - m.prompt = req.prompt; - m.negative_prompt = req.negative_prompt; - m.sampler = "euler"; - m.scheduler = pipeline::scheduler_name(req.scheduler); - m.steps = req.steps; - m.cfg = req.cfg; - m.shift = req.shift; - m.seed = seed; - m.width = req.width; - m.height = req.height; - m.extra.emplace_back("backend", std::string("vulkan ") + vk::block_path_name(path)); - text = pipeline::generation_text_chunks(m); - } - pipeline::write_png_rgb8(run.out_png, r.rgb, static_cast(r.width), static_cast(r.height), text); - } - if (!run.quiet) { - std::fprintf(stderr, "seed %llu; text encoder %.1f s, DiT load %.1f s, denoise %.1f s (%.2f s/step), VAE %.1f s\n", - static_cast(seed), r.text_encoder_s, r.dit_load_s, r.denoise_s, - r.denoise_s / static_cast(g.timesteps.size()), r.vae_s); - } -} - -#else - -bool vulkan_backend_built() { return false; } - -void run_vulkan_generate(const pipeline::Request& req, const VulkanRun&) { - check_supported(req); - throw std::runtime_error("this build has no Vulkan backend (configure with -DMOSSINFER_VULKAN=ON)"); -} - -#endif - -} // namespace mi::cli diff --git a/src/tools/cli/vulkan_backend.h b/src/tools/cli/vulkan_backend.h deleted file mode 100644 index 16e98ef..0000000 --- a/src/tools/cli/vulkan_backend.h +++ /dev/null @@ -1,31 +0,0 @@ -// `moss-infer generate --backend vulkan`: text to image on the Vulkan backend (vulkan/), the fallback for GPUs the CUDA -// engine does not support. The engine's host code supplies what does not depend on the backend: the tokenizer and chat -// template, the seed's initial noise, the schedule and the PNG writer, so a seed gives the same inputs as on CUDA. -#pragma once - -#include -#include -#include -#include - -#include "pipeline/pipeline.h" - -namespace mi::cli { - -struct VulkanRun { - int device = -1; // -1: the first usable device, a discrete one first - std::string kernels; // "coop", "plain", or empty: coop where the device has bf16 cooperative matrices - std::filesystem::path out_png; // empty: no decode - bool quiet = false; - // Called with the final latents, f32 [128, height / 16, width / 16], before the decode. - std::function&)> on_latents; -}; - -// Whether this build has the Vulkan backend (MOSSINFER_VULKAN). -bool vulkan_backend_built(); - -// Runs txt2img for `req` (already validated) on Vulkan and writes the PNG. Throws std::invalid_argument for a request the -// backend does not support (it names the option), std::runtime_error for anything else. -void run_vulkan_generate(const pipeline::Request& req, const VulkanRun& run); - -} // namespace mi::cli diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 3b60f4c..53cfab3 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -1,7 +1,8 @@ -# The Vulkan backend, the fallback for hardware the engine does not otherwise support. A standalone project for now: it -# links no engine code, and nothing here affects the engine's build. See README.md. +# 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) -project(mossinfer_vulkan LANGUAGES CXX) +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) @@ -73,10 +74,55 @@ 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>") -# Built into the engine (cmake/modules/vulkan.cmake), the tools are off unless asked for. -if(NOT DEFINED MOSSINFER_VULKAN_TOOLS) - set(MOSSINFER_VULKAN_TOOLS ON) +# 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() diff --git a/vulkan/tools/moss_infer_vulkan.cpp b/vulkan/tools/moss_infer_vulkan.cpp new file mode 100644 index 0000000..1c8bdaf --- /dev/null +++ b/vulkan/tools/moss_infer_vulkan.cpp @@ -0,0 +1,336 @@ +// 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|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 bf16 cooperative-matrix kernels or the\n" + "plain ones (default: coop where the device has them). The device and kernels are printed at the start. Progress goes\n" + "to stderr one line per event: `text encoder`, `denoise / `, `vae`, then 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 != "plain") usage("--kernels takes coop 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 if (d.caps.coop_bf16) { + std::printf(" kernels: coop\n"); + } else { + std::printf(" kernels: plain (no coop: %s)\n", d.caps.coop_bf16_reason.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() ? (caps.coop_bf16 ? vk::BlockPath::kCoop : vk::BlockPath::kPlain) + : a.kernels == "coop" ? vk::BlockPath::kCoop + : vk::BlockPath::kPlain; + // Always reported: the plain kernels are much slower, and which ones ran should be visible. + std::fprintf(stderr, "vulkan: device %u, %s; kernels: %s%s\n", device.info().index, device.info().name.c_str(), + vk::block_path_name(path), + path == vk::BlockPath::kPlain && !caps.coop_bf16 ? (" (no coop: " + caps.coop_bf16_reason + ")").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 From 9cff95c4e502c0a174920b00886ca997d2280d1a Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:49:19 -0500 Subject: [PATCH 32/41] ci: build-vulkan.yml, moss-infer-vulkan's reusable build --- .github/workflows/build-vulkan.yml | 117 +++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 .github/workflows/build-vulkan.yml diff --git a/.github/workflows/build-vulkan.yml b/.github/workflows/build-vulkan.yml new file mode 100644 index 0000000..5178289 --- /dev/null +++ b/.github/workflows/build-vulkan.yml @@ -0,0 +1,117 @@ +# 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" + + - name: Vulkan SDK (Linux) + if: matrix.platform == 'linux' + run: | + f="vulkansdk-linux-x86_64-$VULKAN_SDK_VERSION.tar.xz" + curl -fsSL --retry 3 -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" + sdk="$RUNNER_TEMP/vulkan-sdk/$VULKAN_SDK_VERSION/x86_64" + echo "VULKAN_SDK=$sdk" >> "$GITHUB_ENV" + echo "$sdk/bin" >> "$GITHUB_PATH" + "$sdk/bin/glslc" --version + + - name: Vulkan SDK (Windows) + if: matrix.platform == 'windows' + shell: pwsh + run: | + $f = "vulkansdk-windows-X64-$env:VULKAN_SDK_VERSION.exe" + $exe = Join-Path $env:RUNNER_TEMP $f + Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/$env:VULKAN_SDK_VERSION/windows/$f" -OutFile $exe + $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 "VulkanSDK" + Start-Process -FilePath $exe -Wait -NoNewWindow -ArgumentList @("--root", $root, "--accept-licenses", "--default-answer", "--confirm-command", "install") + "VULKAN_SDK=$root" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "$root\Bin" | Out-File -FilePath $env:GITHUB_PATH -Append -Encoding utf8 + & "$root\Bin\glslc.exe" --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 + build-vulkan/vulkan_block_check --no-block + build-vulkan/vulkan_te_kernel_check + build-vulkan/vulkan_vae_check --kat-only + + - 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 From 4a98fed77a6759466e59b753f0160b0b4659d8f4 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:49:51 -0500 Subject: [PATCH 33/41] ci, release: moss-infer-vulkan for Linux and Windows 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. --- .github/workflows/ci.yml | 4 ++++ .github/workflows/release.yml | 18 ++++++++++++++---- 2 files changed, 18 insertions(+), 4 deletions(-) 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 From ecedde80fb7692ecd9bcc0c207226e4929c9c90f Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:50:07 -0500 Subject: [PATCH 34/41] docs: moss-infer-vulkan, the standalone Vulkan program (docs/vulkan.md, README, vulkan/README.md) --- README.md | 3 ++- docs/vulkan.md | 40 ++++++++++++++++++++++++---------------- vulkan/README.md | 9 +++++---- 3 files changed, 31 insertions(+), 21 deletions(-) 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 index 780ba84..c3a121b 100644 --- a/docs/vulkan.md +++ b/docs/vulkan.md @@ -1,8 +1,9 @@ # Vulkan backend -`moss-infer generate --backend vulkan` runs MOSS-Image on Vulkan instead of CUDA. It is the fallback for GPUs the CUDA -engine does not support (AMD, Intel, older NVIDIA cards) and is much slower than CUDA. On a supported NVIDIA GPU, use -the default CUDA 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 @@ -11,13 +12,17 @@ the lossy accelerators, quantized or GGUF files, and the CUDA memory options. Th and the schedule are the engine's own, so a seed gives the same inputs as on CUDA. ```sh -moss-infer generate --backend vulkan --tier reference --model MOSS-Image/base \ - --prompt "a red apple on a wooden table" --size 1024x1024 -o apple.png +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 ``` -`--vulkan-device N` picks a device (default: the first usable one, a discrete GPU first). `--vulkan-kernels coop|plain` -picks the kernels (default: `coop` where the device has them). The device and kernels are printed at the start of -every run. +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|plain` the kernels (default: `coop` where the device has them); 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 @@ -52,19 +57,22 @@ Against the official PyTorch pipeline's golden data at 1024x1024 (`tests/golden/ 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 run on torch's own text embeddings instead of the Vulkan +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 45.97 dB PSNR from the CUDA backend's image of the same command. +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 -The backend is off by default. It needs Vulkan headers that define `VK_KHR_shader_bfloat16` and a `glslc` that compiles -`GL_EXT_bfloat16`, as in a current Vulkan SDK. The binaries load the Vulkan loader at run time, so a build with the -backend still starts on a machine without Vulkan and only `--backend vulkan` reports it missing. +`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 --preset linux-cuda13 -DMOSSINFER_VULKAN=ON -cmake --build --preset linux-cuda13 +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 ``` -The backend's own checks and benchmarks build from `vulkan/` alone, without CUDA; see [vulkan/README.md](../vulkan/README.md). +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/README.md b/vulkan/README.md index 02b3118..0dab00c 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -1,9 +1,9 @@ # Vulkan backend -The fallback for GPUs the CUDA engine does not support: `moss-infer generate --backend vulkan` (user guide: -[docs/vulkan.md](../docs/vulkan.md)). CUDA stays the path on every GPU it supports. This directory builds on its own, -without CUDA, for the backend's checks and benchmarks; the engine builds it in with `-DMOSSINFER_VULKAN=ON` -(`cmake/modules/vulkan.cmake`, tools off). +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 two kernel paths, chosen per device and reported: @@ -124,4 +124,5 @@ At 4112 tokens on the cooperative-matrix path the GEMMs take 5.6 s of the 7.2 s, | `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 | From 536b01fdcef292ea21c5bd5224b4e4276d81e262 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 01:50:44 -0500 Subject: [PATCH 35/41] vk_api.cpp: NOMINMAX before windows.h --- vulkan/vk_api.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vulkan/vk_api.cpp b/vulkan/vk_api.cpp index 323fa3d..5ba76b4 100644 --- a/vulkan/vk_api.cpp +++ b/vulkan/vk_api.cpp @@ -3,6 +3,9 @@ #include #if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif #ifndef WIN32_LEAN_AND_MEAN #define WIN32_LEAN_AND_MEAN #endif From 7b3b5fa5e4687d603c48447a001514b75f9d46a6 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 04:25:18 -0500 Subject: [PATCH 36/41] docs: drop the Vulkan research and design notes docs/vulkan.md and vulkan/README.md describe what exists; the notes stay readable in PR #8's history. --- docs/vulkan-backend-design.md | 233 -------------------------------- docs/vulkan-backend-research.md | 185 ------------------------- 2 files changed, 418 deletions(-) delete mode 100644 docs/vulkan-backend-design.md delete mode 100644 docs/vulkan-backend-research.md diff --git a/docs/vulkan-backend-design.md b/docs/vulkan-backend-design.md deleted file mode 100644 index 3c0d50d..0000000 --- a/docs/vulkan-backend-design.md +++ /dev/null @@ -1,233 +0,0 @@ -# Vulkan backend: design proposal - -Companion to [`vulkan-backend-research.md`](vulkan-backend-research.md) (device audit, prior art, performance model). -Status: proposal; nothing in this document is implemented yet. - -## Scope and goal - -Add a **Vulkan compute backend** as the fallback for hardware the engine does not otherwise support, able to run the -MossInfer pipeline there without regressing the CUDA path. CUDA stays the path on every GPU it supports. - -Requirements: - -- Same model files (official bf16 bundle, ComfyUI INT8/INT4 ConvRot files) and the same public API/CLI/server surface. -- CPU side identical: tokenizer, samplers, PNG, planner, reports. -- The backend for a device follows one fixed rule, never a measurement: CUDA where the engine supports the device, Vulkan - otherwise. A caller can name a backend explicitly, and the probe and the report always say which one ran. -- Correctness first: every kernel validated against the CUDA implementation on the same weights; end-to-end images - compared statistically (Vulkan numerics differ in accumulation order, like sm_75 does today). - -Non-goals for the first iteration: macOS/MoltenVK, mobile, multi-GPU, Vulkan ray tracing/video, FP8 (not on Turing), -NVFP4 (Blackwell only), and training. - -## Where the CUDA coupling is - -The engine is not CUDA-abstracted today: - -- `Tensor` carries `Device::kCUDA` and a raw `void*`; every kernel takes `cudaStream_t` and uses `cudaMalloc`-style - arenas (`runtime/DeviceArena`). -- `Planner`/`pipeline` call `cudaMemGetInfo`-based free-memory helpers, `cudaEvent`s and cuBLASLt handles, and the - model code calls kernel functions directly (`kernels::Gemm::linear`, `moss_dit::MossDit`, ...). -- About 70 CUDA translation units implement the kernels (GEMM, attention, norm, RoPE, GDN, VAE, quantized formats). - -A Vulkan backend therefore needs two layers: a **device runtime** (buffers, queues, commands, events, memory budget) and -a **kernel set** behind the same call shapes the model code uses. - -## Options - -| Option | What it means | Cost | Risk | -|---|---|---|---| -| **A. Backend abstraction inside the engine** | Introduce `mi::runtime::Backend` (device/buffer/stream/event) and `mi::kernels::Backend` (GEMM, attention, norm, ...) interfaces; CUDA and Vulkan implement them; model/pipeline code talks to the interfaces. | High upfront refactor (about 70 kernel files and their callers), but one pipeline for both. | Refactor stalls the CUDA path; hard to land incrementally. | -| **B. Separate Vulkan runner first** | A standalone `moss-vulkan` executable (or tool) that loads the same bundle and runs the pipeline with Vulkan kernels; the engine stays untouched. Merge later through A. | Medium; no risk to CUDA; proves kernels and performance with real weights. | Duplicated pipeline logic until merged; second code path to maintain. | -| **C. CUDA shim over Vulkan** | Implement the CUDA runtime API subset (`cudaMalloc`, `cudaStream_t`, `cuBLASLt` calls) on top of Vulkan so the engine compiles unchanged. | Tempting but the semantics (streams/events, cuBLASLt descriptors, error model) do not map cleanly; CUTLASS kernels are SASS, not SPIR-V, so "compile the same kernels" is impossible. | Dead end for the kernel layer. | - -**Recommendation: B first, then A.** Build a standalone Vulkan runner that exercises the real weights and produces -images; once the kernels and the memory model are proven, introduce the runtime/kernel interfaces (A) and move the -pipeline behind them, with CUDA as the default and the Vulkan runner's kernels reused as the backend implementation. -This keeps `main` shippable at every step and gives measurable checkpoints before the refactor. - -## Architecture (target state) - -``` - pipeline / planner / models (unchanged, backend-agnostic) - | - mi::runtime (Device, Buffer, Stream, Event, Budget) - / \ - runtime/cuda runtime/vulkan - (DeviceArena, StreamSet, ...) (VkDevice, queues, allocator, ...) - | | - kernels (CUDA) kernels (Vulkan) - gemm / attention / norm / rope / gdn / vae same interfaces, SPIR-V shaders -``` - -- `mi::runtime::Device`: allocation (`Buffer` with size/usage/alignment), upload/download, `Stream` (queue + command - buffer), `Event` (timeline semaphore value), `free_budget_bytes()`, `properties()` (name, vendor, subgroup, coopmat - shapes). -- `mi::kernels::*`: free functions in the existing headers, with the CUDA implementation in `src/kernels/cuda/...` and the - Vulkan one in `src/kernels/vulkan/...`; both built into their own targets and linked according to the selected backend. -- The planner's memory model is backend-neutral: it consumes `free_budget_bytes()` and the transfer-rate profile it - already measures. - -## Module layout and build - -``` -vulkan/runtime.* device, capabilities, buffers, pipelines, barriers -vulkan/gemm.* GEMM paths, then the other kernels beside it -vulkan/shaders/ *.comp GLSL sources, compiled to SPIR-V at build time and embedded -vulkan/tools/ vulkan_check (device list, checks, timing), probe.py -``` - -It stays outside `src/` while it is standalone, because the engine's build globs that directory. Integration moves it behind -`MOSSINFER_VULKAN`. - -- **Shaders**: GLSL 4.60 + `GL_KHR_cooperative_matrix` + `GL_EXT_shader_16bit_storage`/`8bit_storage` + - `GL_EXT_shader_explicit_arithmetic_types_*` + `GL_KHR_shader_subgroup_*`, compiled by `glslc` to SPIR-V at build time - and embedded as `const uint32_t[]` (a generated header), so the binary has no external files. A pipeline cache file is - written next to the model cache and reused. -- **Build**: `-DMOSSINFER_VULKAN=ON` adds the runtime and kernels, and the CLI/server gain `--backend vulkan` (or - `MOSSINFER_BACKEND=vulkan`). Without the option nothing changes, and the package stays as today. A `vulkan` CPack suffix - like the sm75 one can follow once the backend ships. -- **Dependencies**: Vulkan loader and headers only (`vulkan-1.dll`/`libvulkan.so.1` at run time); no Vulkan SDK at run - time. Shader compilation needs `glslc` at build time (vendored or found). - -## Runtime design - -### Device and queues - -- One `VkDevice`, one compute queue (family 0 with `COMPUTE`), one transfer queue if the family differs. -- Single command buffer stream for compute; a second for uploads, synchronized with timeline semaphores. The engine's - CUDA path already separates compute and copy streams (`runtime::StreamSet`), so the abstraction matches. -- `Stream` = a command buffer in recording; `submit(stream)` batches recorded commands; `Event` = timeline semaphore - value; `wait(event)` maps to `vkWaitSemaphores`. - -### Memory - -- A size-class suballocator over a few large `VkDeviceMemory` blocks (256 B alignment, dedicated allocations above - 64 MiB), with `BDA` (buffer device address) so kernels take plain 64-bit device pointers in push constants. -- Heaps: device-local (weights, activations), host-visible coherent (staging ring, 8–16 MiB, mapped), and a - host-cached path for streamed blocks (readback-free upload through a mapped ring). -- Budget: `VK_EXT_memory_budget` (`heap_budget`, `heap_usage`) plus the allocator's own accounting; report - `free_budget_bytes = heap_budget - heap_usage - reserve` to the planner, the same contract `cudaMemGetInfo` provides. -- Weight streaming: the engine's DiT block streaming (5.16 GB INT8 at 256², 9.6 GB bf16) becomes: host-side ring buffer → - transfer queue → device-local block buffers, with timeline semaphores per block and a per-step upload budget. The - planner's existing estimate of copies hidden behind compute applies unchanged. - -### Dispatch and command recording - -- Each kernel is a `VkPipeline` (specialization constants for tile sizes/strides, push constants for pointers/offsets). -- The Vulkan runner records the whole block's kernel sequence into one command buffer (norm → modulation → QKV GEMM → - attention → out GEMM → FFN up/gate GEMM → SiTU → down GEMM → residual), so the CPU cost is one submit per block, not - 20–40. `vkCmdPipelineBarrier2` with compute-to-compute barriers is inserted by the runtime (or explicit per kernel). -- Secondary command buffers (or one pre-recorded buffer per shape) make the per-step recording cost negligible; a - block resident in device memory can replay a pre-recorded buffer, comparable to the CUDA `fused_graphs` path. - -## Kernel plan - -Legend: CM = cooperative matrix, SIMT = plain compute shader. - -| Stage | Operation | Shapes (DiT 3072, its FFN 8192; text encoder 4096, its FFN 12288; VAE 96–384) | Vulkan kernel | -|---|---|---|---| -| DiT block | modulation/norm | [tokens, 3072] | SIMT, vectorized f16 | -| DiT block | QKV / out / FFN up-gate / FFN down / final (N×K, M = tokens) | 9216×3072, 3072×3072, 16384×3072, 3072×8192 | CM f16 16×16×16, f32 acc; epilogues in shader (bias, gated residual) | -| DiT block | joint attention | tokens 4096 (1024²), heads 24, head_dim 128, non-causal varlen | tiled FA: CM QK^T, f32 softmax, SIMT PV or CM PV | -| DiT block | RoPE, SiTU-GLU, residual | elementwise | SIMT | -| Text encoder | RMSNorm, linear, RoPE, causal attention, GDN recurrence | hidden 4096, 32 layers | GEMM CM; attention CM/SIMT; GDN SIMT scan | -| VAE decode | conv2d (NHWC/KRSC), groupnorm, mid attention, upsample | 96/192/384, 64×64→1024² | SIMT f32 conv (like the sm_75 kernels), attention SIMT | -| INT8 path | ConvRot GEMM + dequant | same GEMM shapes | CM s8 16×16×32, s32 acc, dequant epilogue | -| INT4 path | unpack to INT8 then GEMM | same | SIMT unpack + INT8 CM | - -### Devices - -Tile sizes come from the device's limits, not from this document: Vulkan only guarantees 16 KiB of shared memory, and -subgroup sizes differ (the tiles below need 32). A device without a usable cooperative-matrix shape runs a plain compute -GEMM, slower but correct; that path is what makes the backend a fallback rather than a second fast path. - -### GEMM tiling (48 KiB shared memory, f16, f32 accumulate) - -M is free: whole tiles always run, A and C are allocated with M rounded up to 64 rows, and row m of C depends only on row m -of A. N is a multiple of 64 in every engine GEMM. - -- Workgroup 64×64 output tile, 4 warps (2×2), each warp computes 32×32 via four 16×16×16 cooperative-matrix ops. -- K step 32: A tile 64×32 f16 = 4 KiB, B tile 32×64 f16 = 4 KiB; 3-stage pipeline = 24 KiB, leaves room for epilogue - staging and the bias/gate vectors. -- Registers: 32×32 f32 accumulator per warp = 1024 floats = 32 per thread, plus A/B fragments; well within a 255-register - budget at 128 threads/workgroup. -- Split-K for skinny GEMMs (e.g., final layer, modulation) or when M is small (prompt tokens); the engine's CUDA GEMM - already chunks rows, so the Vulkan dispatcher can reuse that policy. -- Epilogues are fused in the shader: bias add, `ffn_down` residual add, gated residual (multiply by gate), SiTU-GLU - (tanh/exp on f32), and the INT8 dequant (`acc * x_scale[row] * w_scale[col] + bias`). - -### Attention kernel - -- One workgroup per (block of queries, head, sequence chunk). Staged f16 tiles must fit 48 KiB, Q + K + V + P: - head_dim 128 (DiT) with 64 queries and 32 keys is 16 + 8 + 8 + 4 = 36 KiB; head_dim 256 (Qwen) with 32 queries and - 16 keys is 16 + 8 + 8 + 1 = 33 KiB. -- QK^T on cooperative matrix (f16 inputs, f32 accumulate), softmax in SIMT f32 with the same online algorithm as the - CUDA sm_75 kernel (`src/kernels/sm75/attention_sm75.cu`), PV on cooperative matrix or SIMT. -- Varlen/GQA: per-sequence offsets and head mapping in push constants; the DiT joint attention is non-causal over - [text + image] tokens, Qwen is causal. -- 4096-token attention at 1024² walks those tiles over 64 query blocks and 24 heads. - -### Quantization formats - -- INT8 ConvRot weights are stored as int8 with per-row (activation) and per-column (weight) f32 scales; the engine's - `Int8Operands` contract carries them. The Vulkan INT8 GEMM uses CM s8 16×16×32 (exact int32 accumulation) and the same - dequant epilogue, so results should be bit-identical to CUTLASS. -- INT4 ConvRot: reuse the CUDA engine's INT4→INT8 unpack (a SIMT kernel) and the INT8 GEMM, as on Hopper/Blackwell and - in the sm_75 port. -- ConvRot (Hadamard rotation) is already folded into the stored weights, so the GEMM kernel is rotation-agnostic. - -## Integration with the engine - -- **Probe**: `src/backend/probe.*` gains a Vulkan path when built with `MOSSINFER_VULKAN`: enumerate physical devices, - report name/vendor/heap/coopmat shapes and, per device, the backend that the fixed rule above gives it. -- **Planner**: the cost model's rates (`bf16_tflops`, PCIe rate) are learned per backend and per device; the tier rules - are backend-neutral (reference/high/turbo), except that on Vulkan the 8-bit attention modes may not exist. -- **Tier/attention**: Vulkan on Turing-class hardware resolves an unset attention to `exact` (same rule as sm_75); on - devices with f16 cooperative matrix a `q8_fp16` port can come later. -- **CLI/server**: `--backend cuda|vulkan` overrides the rule; the report gains `quality.backend = "vulkan"` and the device string; - `moss-server` passes it through `load`'s new optional `backend` parameter (protocol addition). -- **Packaging**: a `-vulkan` package flavour (loader + shaders embedded); the CUDA package stays the default. - -## Testing and validation - -1. **Kernel parity** (`tests/vulkan/`): for each kernel, run the same inputs through CUDA and Vulkan, compare exact - (INT8), cosine ≥ 0.99999 / max|diff| (f16), and report. Reuses the vectors `tests/data` already stores for the CUDA - kernels. -2. **Block parity**: one DiT block (same weights, same input) CUDA vs Vulkan, cosine on the block output. -3. **End-to-end**: generate the same seed/prompt/size on CUDA and Vulkan; compare PSNR/SSIM and the report; expected - differences are accumulation-order only. -4. **Performance gates**: step time at 256²/512²/1024², per-tier, on the 2080 (CUDA baseline in `docs/sm75.md`) and on - at least one AMD device (RX 6000/7000) and one Intel Arc if available. -5. **Stability**: validation layers (`VK_LAYER_KHRONOS_validation`) in debug builds, pipeline cache round-trip, memory - budget accounting, 100 consecutive generations without leaks. - -## Milestones - -| Milestone | Deliverable | Acceptance | Estimate | -|---|---|---|---| -| **V0** | Probe tool + GEMM microbenchmark; done | GEMM correct and measured at the engine's shapes against cuBLASLt | done | -| **V1** | Vulkan runtime + f16 GEMM with shared-memory tiles and a plain compute path; done, except the suballocator and pipeline cache, which wait until the block runner needs them | Checks at engine shapes on a device without the fast path too | done | -| **V2** | Attention + norm/RoPE/SiTU kernels; DiT block runner on real weights; done on the plain compute path | Block output matches the engine's block (cosine ≥ 0.999) | done | -| **V3** | Text encoder (GEMM, RMSNorm, RoPE, causal attention, GDN) and VAE decoder | Encode/decode parity; first full Vulkan image | 2–3 weeks | -| **V4** | Pipeline integration (runtime abstraction, planner, streaming, CLI/server) | `moss-infer --backend vulkan` and `moss-server` produce images on the 2080 | 2–4 weeks | -| **V5** | Quantized paths (INT8/INT4), packaging, CI, docs | `high`/`turbo` tiers on Vulkan; `-vulkan` package; tests in CI | 1–2 weeks | - -The estimates assume one experienced engineer and the existing CUDA code as the specification; the GDN recurrence and -the VAE convolutions are the least predictable parts. - -## Risks and mitigations - -- **Scope explosion.** B-first keeps the engine shippable; V4 (refactor) is gated on V2/V3 parity. -- **Vendor differences.** Cooperative matrix shapes and shared-memory limits vary; the kernel set is selected from a - capability table (NVIDIA first, AMD/Intel fallbacks second). -- **Performance.** As a fallback it never replaces CUDA, so its speed gate is whether a generation is usable on the - fallback device. That can only be measured on one. -- **Maintenance.** Two kernel sets are expensive; the shared parity tests and the runtime abstraction are the - mitigation. If the backend cannot be maintained, V1–V3 remain a standalone tool. - -## Decisions requested - -1. Decided: Vulkan is the fallback for hardware the engine does not otherwise support. -2. Choose the first fallback devices for validation. -3. Decide whether the Vulkan path ships as a separate package flavour or inside the default package (loader dependency). diff --git a/docs/vulkan-backend-research.md b/docs/vulkan-backend-research.md deleted file mode 100644 index 5b515ee..0000000 --- a/docs/vulkan-backend-research.md +++ /dev/null @@ -1,185 +0,0 @@ -# Vulkan backend: research notes - -Status: research only; no Vulkan code is in the engine yet. This document and -[`vulkan-backend-design.md`](vulkan-backend-design.md) are the basis for a decision on whether, and how, MossInfer gets a -Vulkan compute backend. - -## Why a Vulkan backend - -MossInfer is CUDA-only: the pipeline calls cuBLASLt, CUTLASS kernels and CUDA kernels directly, and the probe refuses any -GPU whose compute capability is below 8.0 (7.5 with `MOSSINFER_SM75`, 7.0 with `MOSSINFER_SM70`). That confines the engine -to NVIDIA on Windows and Linux. The Vulkan backend is the fallback for the hardware that leaves out; CUDA stays the path -on every GPU it supports. - -A Vulkan compute backend would give: - -- **Other GPUs.** AMD, Intel and older NVIDIA cards. The engine's model code is portable; only the kernels and the device - runtime are CUDA-specific. -- **No CUDA runtime dependency** in the fallback package. - -The cost is a second implementation of every tensor-core kernel, plus a device/memory/command abstraction the engine -does not have today (see the design document). - -## What this machine exposes (primary evidence) - -Probed on the target machine: Windows 11, RTX 2080 8 GiB, NVIDIA driver 621.01 (Vulkan driver 621.1.0.0). The Vulkan -loader in `C:\Windows\System32\vulkan-1.dll` reports **instance 1.4.357**; the NVIDIA ICD reports **device API 1.4.357** -and one physical device. - -### Limits that shape kernel design - -| Limit | Value | Consequence | -|---|---|---| -| `maxComputeSharedMemorySize` | 48 KiB | 64 KiB is a CUDA opt-in only; Vulkan tiles must fit 48 KiB (Turing's hardware limit is 64 KiB but Vulkan caps at 48 KiB here). | -| `maxComputeWorkGroupInvocations` | 1024 | Up to 32 warps per workgroup. | -| `subgroupSize` | 32 | Cooperative-matrix `gl_ScopeSubgroup` = one warp. | -| `maxStorageBufferRange` | 2^32-1 bytes | 4 GiB per buffer; the largest DiT weight is 96 MiB as f16 (up/gate), so a small number of big buffers is fine. | -| Device-local heap | 7989 MiB | The engine's planner sees the same ~8 GiB. A second device-local heap of 214 MiB exists; host-visible heap is 16 GiB. | - -### Extensions a compute backend needs (all present) - -`VK_KHR_cooperative_matrix` (revision 2), `VK_KHR_16bit_storage`, `VK_KHR_8bit_storage`, `VK_KHR_shader_float16_int8`, -`VK_KHR_shader_integer_dot_product`, `VK_KHR_buffer_device_address`, `VK_EXT_memory_budget`, `VK_KHR_timeline_semaphore`, -`VK_KHR_synchronization2`, `VK_EXT_descriptor_indexing`, `VK_KHR_maintenance4`, `VK_KHR_maintenance5`, -`VK_EXT_subgroup_size_control`, `VK_KHR_pipeline_library`, `VK_EXT_shader_module_identifier`, -`VK_KHR_shader_float_controls` (and `2`), `VK_KHR_vulkan_memory_model`, `VK_KHR_shader_clock`, -`VK_KHR_performance_query`, `VK_NV_cooperative_matrix`, `VK_NV_cooperative_matrix2`, `VK_NV_cooperative_vector`. - -Shader features: `storageBuffer16BitAccess`, `uniformAndStorageBuffer16BitAccess`, `storageBuffer8BitAccess`, -`uniformAndStorageBuffer8BitAccess`, `shaderFloat16`, `shaderInt8`, `shaderBFloat16Type`, -`shaderBFloat16CooperativeMatrix`, `shaderSubgroupClock`, `shaderDeviceClock`, `shaderFmaFloat16`. - -### Cooperative-matrix shapes (queried with `vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR`) - -17 combinations, all `gl_ScopeSubgroup`, non-saturating: - -| Shape (M×N×K) | A | B | C | Result | -|---|---|---|---|---| -| 16×16×16, 16×8×16, 16×8×8 | f16 | f16 | f32 | f32 | -| 16×16×16, 16×8×16, 16×8×8 | f16 | f16 | f16 | f16 | -| 16×16×32, 16×8×32 | s8 | s8 | s32 | s32 | -| 16×16×32, 16×8×32 | s8 | u8 / u8×s8 / u8×u8 | s32 / u32 | s32 / u32 | -| 16×16×16 | bf16 | bf16 | f32 | f32 | - -Notes: - -- **f16 with f32 accumulate is available at 16×16×16 and 16×8×16**, which is the same HMMA shape the CUDA sm_75 port - uses (`m16n8k16`; the 16×8×16 cooperative matrix maps to it). A Vulkan GEMM can therefore reach the same tensor-core - instruction on the same hardware. -- **INT8 at 16×16×32 / 16×8×32 with s32 accumulate** covers the engine's INT8 ConvRot GEMMs (CUTLASS `m8n8k16` on - sm_75 is a different instruction, but the arithmetic is an exact int32 dot product either way). -- **bf16 cooperative matrix is reported by this Turing driver**, but Turing has no bf16 tensor core: the driver must be - emulating it (likely f16 down-conversion or split-f16 math). The engine already converts bf16 weights to fp16 for - sm_75, so the Vulkan backend should keep fp16 as its compute type and treat bf16 as a storage conversion, not rely on - the bf16 cooperative-matrix path. -- `VK_NV_cooperative_matrix` is still exposed (the pre-1.4 NV API); `vkGetPhysicalDeviceCooperativeMatrixPropertiesNV` - returned a table that our probe could not decode reliably (the struct's enums are narrower than the KHR ones). The - KHR query is the authoritative one on this driver; the backend should prefer it and fall back to the NV query only on - older drivers. - -### Memory heaps - -| Heap | Size | Flags | Use | -|---|---|---|---| -| 0 | 7989 MiB | device-local | Model weights, activations | -| 1 | 16347 MiB | host-visible (0x6 = visible + coherent; 0xe adds cached) | Staging ring, streamed blocks, host-cached weights | -| 2 | 214 MiB | device-local | Small warm allocations | - -That matches the CUDA picture the planner already models (8 GiB device + ~16 GiB host), so the existing planner cost -model (bytes per block, PCIe rate, hidden-copy fraction) transfers directly; the Vulkan backend only needs to report the -free device budget from `VK_EXT_memory_budget` instead of `cudaMemGetInfo`. - -## Prior art and what to take from it - -### ggml-vulkan (llama.cpp, whisper.cpp, stable-diffusion.cpp) - -The most complete Vulkan inference backend in the open-source ecosystem. Design points worth copying: - -- **One backend file plus generated shaders.** `ggml-vulkan.cpp` (~20k lines) houses device setup, memory arenas, - pipelines and per-op dispatch; the compute shaders are GLSL templates compiled offline to SPIR-V and embedded (the - build generates a C++ header with the SPIR-V arrays). -- **Specialization constants and push constants.** Workgroup sizes, tile sizes and strides are specialization - constants, so one shader covers many shapes; per-dispatch arguments (pointers via BDA, offsets, strides) go in push - constants, avoiding descriptor-set churn. -- **Cooperative matrix paths.** `coopmat` (KHR 16×16×16) and `coopmat2` (VK_NV_cooperative_matrix2) variants for matrix - multiply and flash attention; a scalar fallback for devices without cooperative matrix. -- **Pipeline cache.** The driver's pipeline cache is serialized to disk and reused; shader compilation is a large part - of a cold start otherwise. -- **Memory arenas.** Buffers are suballocated from a few large `VkDeviceMemory` blocks with size classes and a - buffer-device-address allocator; the engine's streaming design can use the same idea for its block buffers. -- **Flash attention shaders.** Tiled FA with online softmax, GQA support (`coopmat2` shader in llama.cpp #12559) and - separate paths for quantized KV; attention is the piece most likely to be slower than CUDA. - -### stable-diffusion.cpp - -Runs SD1.5/SDXL/Flux families through ggml, with the same Vulkan backend. It demonstrates diffusion inference on Vulkan -end-to-end, including VAE decode and img2img, on consumer GPUs. Its known weaknesses map to ours: performance depends -strongly on the GPU vendor's Vulkan driver, AMD RDNA3/4 is often close to CUDA, small Intel Arc parts are slower, and -attention remains the bottleneck. It is the closest prior art to what MossInfer needs. - -### Others - -- **ncnn Vulkan**: mobile/desktop CNN inference; mature Vulkan kernels (convolution, pooling), good evidence that - Vulkan convolution is practical (relevant for the VAE decoder). -- **Kompute / vuh**: C++ Vulkan compute wrappers; convenient for prototyping but no tensor-core help. -- **IREE HAL / ONNX Runtime**: IREE has a Vulkan HAL; ONNX Runtime's GPU EP is DirectML on Windows, not Vulkan. Neither - offers a drop-in path for a custom engine, but their Vulkan memory/scheduling designs are useful references. -- **wgpu / WebGPU**: a higher-level portability layer with a compute language (WGSL) and no cooperative matrix in the - standard; shader-f16 and subgroup support are optional. It would be easier to port but cannot reach tensor-core - throughput on Turing; not a candidate for the "fast" backend, only for a "works everywhere, slowly" fallback. - -## Performance model - -Measured on an RTX 4090 Laptop against cuBLASLt on the same GPU. See [`vulkan/README.md`](../vulkan/README.md) for the table: -the cooperative-matrix GEMM reaches 35-58 % of cuBLASLt at the engine's shapes, and the plain compute path 1-1.6 TFLOPS. -These are NVIDIA numbers: they show the kernels work and where they lose time, not how a fallback device will perform. - -For scale, `docs/sm75.md` measures CUDA on an RTX 2080 Ti at 256x256: 0.175 s per step in `reference`, 0.091 s in `high` -(INT8) and 0.103 s in `turbo` (INT4). - -The critical unknowns to measure in a prototype (M0): - -1. Cooperative-matrix GEMM throughput at the engine's shapes: measured above. -2. Dispatch overhead per kernel (engine runs ~20–40 dispatches per block; command-buffer batching must keep the CPU - cost per step below the GPU cost). -3. Attention throughput (fp16 QK^T + softmax + PV, 4096-token joint attention at 1024²). -4. PCIe upload rate for streamed DiT blocks (the CUDA path measures 2.9–4.6 GB/s from file/page cache). -5. fp16 vs bf16 numerics: the engine's reference tier is bf16 storage with fp32 accumulate; converting to fp16 for - Vulkan changes rounding in the same way the sm_75 CUDA path does (documented in `docs/sm75.md`). - -## Numerics and formats - -- **Storage:** the engine's official bundle is bf16; the sm_75 port converts it to fp16 in place, and the Vulkan backend - would do the same (bf16→fp16 is exact inside fp16's range, and the reverse conversion recovers the bf16 bits). - Community INT8/INT4 ConvRot files are already quantized; no conversion is needed for weights. -- **Compute:** f16 inputs with f32 accumulate via cooperative matrix (matching cuBLASLt's `CUBLAS_COMPUTE_32F` use in the - engine's fp16 path); INT8 weights use s32 accumulate (exact, same result as CUTLASS on CUDA); the VAE decoder stays - fp32 (SIMT), since Vulkan exposes no fp32 tensor core. -- **No FP8 on Turing** (`VK_EXT_shader_float8` is not exposed), so the engine's `q8_fp8` attention path cannot be ported - to this GPU; `q8_fp16` would be the 8-bit attention variant on Vulkan, and `exact` remains the default on sm_75-class - hardware. -- **Denormals:** f16 cooperative matrix on NVIDIA flushes subnormal inputs/outputs? The CUDA sm_75 port relies on fp16 - exactly matching bf16 inside the normal range; the prototype must compare bit-for-bit or cosine against the CUDA path - and document the deviation. - -## Risks - -| Risk | Impact | Mitigation | -|---|---|---| -| Cooperative-matrix shapes, subgroup sizes and shared memory differ per device, or there is no cooperative matrix at all | The fast kernel cannot run everywhere | Query properties at startup and select kernels; a plain compute path is required, not optional | -| Shader compilation time and driver bugs | Cold start and hangs | Precompile SPIR-V at build time, ship pipeline cache, validation layers in CI | -| 48 KiB shared memory | Smaller tiles than CUDA, lower GEMM efficiency | 3-stage 64×64×32 f16 tiles; split-K for skinny shapes | -| Dispatch overhead (CPU-bound at small sizes) | Step time worse than CUDA | Batch dispatches into one command buffer; reuse command buffers; push constants + BDA | -| `VK_EXT_memory_budget` imprecision | Planner over/under-estimates free VRAM | Use budget + own allocation accounting; keep the CUDA planner model | -| bf16 coopmat is emulated on Turing | Wrong performance/numerics assumptions | Compute in f16, store bf16 as f16, never use bf16 coopmat | -| Windows WDDM TDR (2 s) on long kernels | Device reset on big attention/GEMM | Keep dispatches short; split work (the engine already tiles) | -| Two backends to test | Maintenance cost | Shared test vectors and CI matrix; backend-agnostic golden images | - -## Open questions for the design phase - -1. Decided: Vulkan is the fallback for hardware the engine does not otherwise support. It therefore needs a path for devices - without cooperative matrices, not only the fast one. -2. Should the Vulkan path reuse the engine's existing planner/streaming code through a device abstraction, or start as a - separate `moss-vulkan` runner and merge later? (The design document recommends the latter first.) -3. Which attention variant is acceptable on Vulkan: `exact` only (like sm_75), or a `q8_fp16` port? -4. Which fallback devices are the first validation targets? No non-NVIDIA device has run any of this. From 398b529a16b0aafb03ed1ab6a1a940e5570b0299 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 04:25:29 -0500 Subject: [PATCH 37/41] probe.py: no reference to the dropped research notes --- vulkan/tools/probe.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vulkan/tools/probe.py b/vulkan/tools/probe.py index 1e3e5f9..a44c3ee 100644 --- a/vulkan/tools/probe.py +++ b/vulkan/tools/probe.py @@ -1,8 +1,7 @@ """Lists a Vulkan device's cooperative-matrix shapes and memory heaps, using only the Vulkan loader. -The C++ tool needs Vulkan headers and an import library; this script needs neither (ctypes only), so it works on a -machine with just the GPU driver installed. It is how the capability table in docs/vulkan-backend-research.md was -produced. +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 """ From 21596f5fe0706837957c4cb1ee901a3e3ae19c56 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 04:29:31 -0500 Subject: [PATCH 38/41] vulkan: build against the Vulkan SDK 1.4.357 the CI installs - 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. --- vulkan/CMakeLists.txt | 12 ++++++++++-- vulkan/shaders/te_attention.comp | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 53cfab3..26112ee 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -10,8 +10,16 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) set(CMAKE_BUILD_TYPE Release) # the CPU reference is unusably slow unoptimized endif() -find_package(Vulkan REQUIRED COMPONENTS glslc) +# 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) @@ -43,7 +51,7 @@ target_include_directories(mossinfer_vulkan PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} P 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 Vulkan::Headers Threads::Threads ${CMAKE_DL_LIBS}) +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 "") diff --git a/vulkan/shaders/te_attention.comp b/vulkan/shaders/te_attention.comp index 058ebe2..e05d966 100644 --- a/vulkan/shaders/te_attention.comp +++ b/vulkan/shaders/te_attention.comp @@ -48,7 +48,7 @@ shared float s_tile[BR * SP]; shared float rescale_shared[BR]; shared float inv_sum_shared[BR]; -const float NEG_INF = uintBitsToFloat(0xff800000u); +#define NEG_INF uintBitsToFloat(0xff800000u) void main() { const uint t = gl_LocalInvocationIndex; From 4b9ebc784503da34d3981a7d368c6f38f99b80e1 Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 04:34:20 -0500 Subject: [PATCH 39/41] vulkan_te_kernel_check: the g gate's limit allows an unfused fma; CI 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. --- .github/workflows/build-vulkan.yml | 9 ++++++--- vulkan/tools/vulkan_te_kernel_check.cpp | 4 +++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build-vulkan.yml b/.github/workflows/build-vulkan.yml index 5178289..7b6e87c 100644 --- a/.github/workflows/build-vulkan.yml +++ b/.github/workflows/build-vulkan.yml @@ -96,9 +96,12 @@ jobs: export LD_LIBRARY_PATH="$VULKAN_SDK/lib" export VK_INSTANCE_LAYERS=VK_LAYER_KHRONOS_validation build-vulkan/moss-infer-vulkan --list-devices - build-vulkan/vulkan_block_check --no-block - build-vulkan/vulkan_te_kernel_check - build-vulkan/vulkan_vae_check --kat-only + # 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 diff --git a/vulkan/tools/vulkan_te_kernel_check.cpp b/vulkan/tools/vulkan_te_kernel_check.cpp index d443b61..bd47e3f 100644 --- a/vulkan/tools/vulkan_te_kernel_check.cpp +++ b/vulkan/tools/vulkan_te_kernel_check.cpp @@ -365,7 +365,9 @@ bool check_conv_gates(Device& device, const TeKernels& k) { } } ok = report("gates beta", read(beta, beta_want.size()), beta_want, {1.0, 0.0}) && ok; - ok = report_f32("gates g", read(g, g_want.size()), g_want, 8.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. From 4c57001255babf80e91ea258f58cefc8c3f43e8e Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 05:11:15 -0500 Subject: [PATCH 40/41] build-vulkan.yml: cache the Vulkan SDK, download it with retries 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. --- .github/workflows/build-vulkan.yml | 51 ++++++++++++++++++++++-------- 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/.github/workflows/build-vulkan.yml b/.github/workflows/build-vulkan.yml index 7b6e87c..c7f6cb7 100644 --- a/.github/workflows/build-vulkan.yml +++ b/.github/workflows/build-vulkan.yml @@ -49,33 +49,56 @@ jobs: sudo apt-get clean echo "CXX=g++-12" >> "$GITHUB_ENV" - - name: Vulkan SDK (Linux) - if: matrix.platform == 'linux' + # 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 3 -o "$RUNNER_TEMP/$f" "https://sdk.lunarg.com/sdk/download/$VULKAN_SDK_VERSION/linux/$f" + 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" - sdk="$RUNNER_TEMP/vulkan-sdk/$VULKAN_SDK_VERSION/x86_64" - echo "VULKAN_SDK=$sdk" >> "$GITHUB_ENV" - echo "$sdk/bin" >> "$GITHUB_PATH" - "$sdk/bin/glslc" --version + rm "$RUNNER_TEMP/$f" - - name: Vulkan SDK (Windows) - if: matrix.platform == 'windows' + - 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 - Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/$env:VULKAN_SDK_VERSION/windows/$f" -OutFile $exe + 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 "VulkanSDK" + $root = Join-Path $env:RUNNER_TEMP "vulkan-sdk" Start-Process -FilePath $exe -Wait -NoNewWindow -ArgumentList @("--root", $root, "--accept-licenses", "--default-answer", "--confirm-command", "install") - "VULKAN_SDK=$root" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 - "$root\Bin" | Out-File -FilePath $env:GITHUB_PATH -Append -Encoding utf8 - & "$root\Bin\glslc.exe" --version + 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 From 8d403b1f9b754474bd8c32d723fe78fd51a97aaa Mon Sep 17 00:00:00 2001 From: xycld Date: Tue, 3 Nov 2026 10:26:28 -0500 Subject: [PATCH 41/41] vulkan: coop-f16, the cooperative-matrix kernels on f16 fragments for 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. --- docs/vulkan.md | 11 ++- vulkan/CMakeLists.txt | 17 ++-- vulkan/README.md | 21 +++-- vulkan/block.cpp | 119 ++++++++++++++++++++++------ vulkan/block.h | 29 +++++-- vulkan/dit.cpp | 2 +- vulkan/generate.h | 3 +- vulkan/runtime.cpp | 14 +++- vulkan/runtime.h | 4 + vulkan/shaders/attention_coop.comp | 30 +++++-- vulkan/shaders/common.glsl | 19 ++++- vulkan/shaders/gemm_bf16_coop.comp | 76 +++++++++++++++++- vulkan/shaders/row_scale.comp | 57 +++++++++++++ vulkan/text_encoder.cpp | 4 +- vulkan/tools/moss_infer_vulkan.cpp | 34 +++++--- vulkan/tools/vulkan_block_check.cpp | 67 +++++++++++----- vulkan/tools/vulkan_block_real.cpp | 4 +- vulkan/tools/vulkan_check.cpp | 5 ++ vulkan/tools/vulkan_dit_run.cpp | 6 +- vulkan/tools/vulkan_te_check.cpp | 12 +-- vulkan/vae.h | 2 +- 21 files changed, 429 insertions(+), 107 deletions(-) create mode 100644 vulkan/shaders/row_scale.comp diff --git a/docs/vulkan.md b/docs/vulkan.md index c3a121b..49429f2 100644 --- a/docs/vulkan.md +++ b/docs/vulkan.md @@ -19,8 +19,8 @@ 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|plain` the kernels (default: `coop` where the device has them); both are printed -at the start of every run. Progress goes to stderr one line per event: `text encoder`, `denoise / +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. @@ -29,8 +29,15 @@ request, 1 for any other failure. | 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 diff --git a/vulkan/CMakeLists.txt b/vulkan/CMakeLists.txt index 26112ee..3b2248e 100644 --- a/vulkan/CMakeLists.txt +++ b/vulkan/CMakeLists.txt @@ -26,21 +26,28 @@ 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. -function(_mi_add_shader shader target_env) +# _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}/${shader}.spv.inc) + 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 + 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") + 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) -foreach(shader IN ITEMS gemm_plain norm_modulate qk_norm_rope attention_plain elementwise) +_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() diff --git a/vulkan/README.md b/vulkan/README.md index 0dab00c..1e27f48 100644 --- a/vulkan/README.md +++ b/vulkan/README.md @@ -5,11 +5,12 @@ directory without CUDA (user guide: [docs/vulkan.md](../docs/vulkan.md)). CUDA s 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 two kernel paths, chosen per device and reported: +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. @@ -20,7 +21,7 @@ The Vulkan loader is opened at run time (`vk_api.h`), so nothing links against i 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 plain] # the DiT block's kernels against host replays, then one block's timing +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 @@ -35,7 +36,7 @@ validation layer and on llvmpipe (Mesa's CPU driver: 8-wide subgroups, 32 KiB of ## DiT -`block.h` holds one DiT block's kernels on both paths, with the engine's rounding points: RMSNorm with modulation, the +`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 @@ -57,6 +58,16 @@ Against the engine's own output on real blocks of the 1024x1024 golden case (cap 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 @@ -69,7 +80,7 @@ bytes identical and none off by more than 1, as the CUDA engine's reference deco `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 both paths share them: the Qwen RMSNorms (with +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 @@ -116,7 +127,7 @@ At 4112 tokens on the cooperative-matrix path the GEMMs take 5.6 s of the 7.2 s, |---|---| | `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 both paths | +| `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 | diff --git a/vulkan/block.cpp b/vulkan/block.cpp index 1bac856..ae4d04c 100644 --- a/vulkan/block.cpp +++ b/vulkan/block.cpp @@ -29,11 +29,28 @@ const std::uint32_t kAttentionSpirv[] = { 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; @@ -68,6 +85,10 @@ std::unique_ptr gemm_pipeline(Device& device, BlockPath path, GemmMode 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}); } @@ -78,34 +99,79 @@ void require(bool condition, const std::string& message) { constexpr std::uint32_t kGemmTileRows = 128, kGemmTileCols = 128; -void dispatch_gemm(const Pipeline& pipeline, BlockPath path, GemmMode mode, VkCommandBuffer cb, const GemmPush& push) { - const std::uint32_t align = path == BlockPath::kCoop ? 8 : 4; - require(push.m > 0 && push.n > 0 && push.k > 0, "bad GEMM shape"); - require(push.n % align == 0 && push.k % align == 0 && push.lda % align == 0 && push.ldb % align == 0 && - push.ldo % align == 0, +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)); - if (path == BlockPath::kCoop) { + 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, (push.n + cols - 1) / cols, (push.m + kCoopTileRows - 1) / kCoopTileRows); + 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, (push.n + cols - 1) / cols, (push.m + kGemmTileRows - 1) / kGemmTileRows); + pipeline.dispatch(cb, &push, (n + cols - 1) / cols, (m + kGemmTileRows - 1) / kGemmTileRows); } compute_barrier(cb); } -} // namespace - -const char* block_path_name(BlockPath path) { return path == BlockPath::kCoop ? "coop" : "plain"; } - -BlockKernels::BlockKernels(Device& device) - : BlockKernels(device, device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) {} +BlockKernels::BlockKernels(Device& device) : BlockKernels(device, default_block_path(device.info().caps)) {} BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { - if (path == BlockPath::kCoop && !device.info().caps.coop_bf16) { - throw std::runtime_error(device.info().name + " cannot run the cooperative-matrix DiT block: " + - device.info().caps.coop_bf16_reason); + 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); @@ -121,6 +187,11 @@ BlockKernels::BlockKernels(Device& device, BlockPath path) : path_(path) { 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)); } @@ -143,33 +214,33 @@ void BlockKernels::record_norm_modulate(VkCommandBuffer cb, VkDeviceAddress x, s 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_, path_, kBF16, cb, GemmPush{x, w, out, 0, m, n, k, ldx, ldw, ldo}); + 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_, path_, kGated, cb, GemmPush{x, w, residual, gate, m, n, k, ldx, ldw, ldr}); + 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_, path_, kSitu, cb, GemmPush{x, w_up_gate, out, 0, m, f, k, ldx, ldw, ldo}); + 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_, path_, kBias, cb, GemmPush{x, w, out, bias, m, n, k, ldx, ldw, ldo}); + 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_, path_, kBiasResidual, cb, GemmPush{x, w, residual, bias, m, n, k, ldx, ldw, ldr}); + 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, @@ -213,7 +284,7 @@ void BlockKernels::record_attention(VkCommandBuffer cb, const AttentionArgs& a) 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 (path_ == BlockPath::kCoop) { + 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"); @@ -225,7 +296,7 @@ void BlockKernels::record_attention(VkCommandBuffer cb, const AttentionArgs& a) 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 = path_ == BlockPath::kCoop ? kCoopAttentionRows : kPlainAttentionRows; + 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); diff --git a/vulkan/block.h b/vulkan/block.h index 0369892..2f6e820 100644 --- a/vulkan/block.h +++ b/vulkan/block.h @@ -7,13 +7,18 @@ // 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) -// Two 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 on bf16 cooperative -// matrices (Capabilities::coop_bf16) and needs strides, N and K that are multiples of 8. Each record_* adds its own barrier. +// 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" @@ -34,8 +39,14 @@ struct AttentionArgs { float scale; // SDPA's 1 / sqrt(head_dim) for the DiT }; -enum class BlockPath { kCoop, kPlain }; -const char* block_path_name(BlockPath path); +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: @@ -77,9 +88,17 @@ class BlockKernels { 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 { diff --git a/vulkan/dit.cpp b/vulkan/dit.cpp index 7b0cb93..5932a62 100644 --- a/vulkan/dit.cpp +++ b/vulkan/dit.cpp @@ -65,7 +65,7 @@ std::vector rope_table(std::uint32_t height, std::uint32_t width, 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::kCoop ? 200.0 : 20.0; + 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 diff --git a/vulkan/generate.h b/vulkan/generate.h index c34750e..d770166 100644 --- a/vulkan/generate.h +++ b/vulkan/generate.h @@ -42,7 +42,8 @@ struct GenerateResult { 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); the VAE follows it. +// 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/runtime.cpp b/vulkan/runtime.cpp index 5b09453..4c7007c 100644 --- a/vulkan/runtime.cpp +++ b/vulkan/runtime.cpp @@ -20,7 +20,8 @@ bool pipeline_stats_requested() { 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 bf16 DiT kernels' largest shared memory: gemm_bf16_coop.comp's two stages of 128x32 A and B tiles, rows padded to 80 B. +// 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) { @@ -151,6 +152,15 @@ void describe_coop(VkInstance instance, VkPhysicalDevice physical, const VkPhysi 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) { @@ -262,7 +272,7 @@ Device::Device(int index) { coop.pNext = &bf16; extensions.push_back(VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); } - if (info_.caps.coop || info_.caps.coop_bf16) { + if (info_.caps.coop || info_.caps.coop_bf16 || info_.caps.coop_f16) { v12.pNext = &v13; v12.shaderFloat16 = VK_TRUE; v12.vulkanMemoryModel = VK_TRUE; diff --git a/vulkan/runtime.h b/vulkan/runtime.h index 39f407b..1d4dd73 100644 --- a/vulkan/runtime.h +++ b/vulkan/runtime.h @@ -23,6 +23,10 @@ struct Capabilities { // 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). diff --git a/vulkan/shaders/attention_coop.comp b/vulkan/shaders/attention_coop.comp index 2be652b..03e96c2 100644 --- a/vulkan/shaders/attention_coop.comp +++ b/vulkan/shaders/attention_coop.comp @@ -10,10 +10,23 @@ // 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 @@ -60,8 +73,8 @@ 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 FRAG_A coopmat +#define FRAG_B coopmat #define ACC coopmat void sync_subgroup() { @@ -87,8 +100,8 @@ 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] = rk[j]; - smem[SV + at] = rv[j]; + smem[SK + at] = TO_FRAG(rk[j], 1.0); + smem[SV + at] = TO_FRAG(rv[j], 1.0); } } @@ -108,7 +121,7 @@ void main() { 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] = v; + smem[SK + r * KV_ROW + i % 16u] = TO_FRAG(v, 1.0); } barrier(); FRAG_A fq[8]; @@ -185,8 +198,8 @@ void main() { } 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] = uvec4(p_bits[0], p_bits[1], p_bits[2], p_bits[3]); - smem[sp + frow * P_ROW + half_ * 2u + 1u] = uvec4(p_bits[4], p_bits[5], p_bits[6], p_bits[7]); + 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(); @@ -229,7 +242,8 @@ void main() { } // out = bf16(acc_o * (1 / row_sum)), through the staging two fragments at a time. - const float inv_sum = (row_sum == 0.0 || isnan(row_sum)) ? 1.0 : 1.0 / row_sum; + // 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; diff --git a/vulkan/shaders/common.glsl b/vulkan/shaders/common.glsl index ccac00b..6b00802 100644 --- a/vulkan/shaders/common.glsl +++ b/vulkan/shaders/common.glsl @@ -23,11 +23,28 @@ uint bf(float v) { 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); - return sign(x) * t; + 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) { diff --git a/vulkan/shaders/gemm_bf16_coop.comp b/vulkan/shaders/gemm_bf16_coop.comp index 8b8e95b..7ca452e 100644 --- a/vulkan/shaders/gemm_bf16_coop.comp +++ b/vulkan/shaders/gemm_bf16_coop.comp @@ -13,10 +13,22 @@ // 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 @@ -26,6 +38,8 @@ 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; @@ -38,6 +52,8 @@ layout(push_constant) uniform PC { uint lda; uint ldb; uint ldo; + uint64_t a_exponents; // F16 only + uint64_t b_exponents; // F16 only } pc; const uint BM = 128u; @@ -51,8 +67,8 @@ 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 FRAG_A coopmat +#define FRAG_B coopmat #define ACC coopmat uint b_row(uint n_base, uint r) { @@ -61,6 +77,9 @@ uint b_row(uint n_base, uint r) { } 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; @@ -74,9 +93,43 @@ void load_tiles(uint m_base, uint n_base, uint k0, out uvec4 ra[2], out uvec4 rb 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) { @@ -115,6 +168,18 @@ void main() { [[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); @@ -148,13 +213,15 @@ void main() { 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 vec4 lo = staged(base0, frow, fcol / 4u), hi = staged(base0, frow, fcol / 4u + 1u); 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); @@ -169,7 +236,8 @@ void main() { 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 = staged(base1, frow, fcol / 4u), g_hi = staged(base1, frow, fcol / 4u + 1u); + 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 { 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/text_encoder.cpp b/vulkan/text_encoder.cpp index a1b4e5d..68dadec 100644 --- a/vulkan/text_encoder.cpp +++ b/vulkan/text_encoder.cpp @@ -246,7 +246,7 @@ struct TextEncoder::Impl { } double budget() const { - const double gflop = options.submit_gflop > 0 ? options.submit_gflop : (path == BlockPath::kCoop ? 200.0 : 20.0); + const double gflop = options.submit_gflop > 0 ? options.submit_gflop : (path != BlockPath::kPlain ? 200.0 : 20.0); return gflop * 1e9; } @@ -291,7 +291,7 @@ struct TextEncoder::Impl { }; TextEncoder::TextEncoder(Device& device, const std::filesystem::path& dir) - : TextEncoder(device, dir, device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) {} + : 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)) {} diff --git a/vulkan/tools/moss_infer_vulkan.cpp b/vulkan/tools/moss_infer_vulkan.cpp index 1c8bdaf..6a1dc74 100644 --- a/vulkan/tools/moss_infer_vulkan.cpp +++ b/vulkan/tools/moss_infer_vulkan.cpp @@ -37,15 +37,16 @@ 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|plain] [--no-metadata] [--quiet]\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 bf16 cooperative-matrix kernels or the\n" - "plain ones (default: coop where the device has them). The device and kernels are printed at the start. Progress goes\n" - "to stderr one line per event: `text encoder`, `denoise / `, `vae`, then a summary.\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; @@ -166,7 +167,9 @@ Args parse(const std::vector& argv) { 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 != "plain") usage("--kernels takes coop or plain"); + if (!a.kernels.empty() && a.kernels != "coop" && a.kernels != "coop-f16" && a.kernels != "plain") { + usage("--kernels takes coop, coop-f16 or plain"); + } return a; } @@ -178,10 +181,10 @@ int list_devices() { d.discrete ? " discrete" : ""); if (!d.unusable_reason.empty()) { std::printf(" unusable: %s\n", d.unusable_reason.c_str()); - } else if (d.caps.coop_bf16) { - std::printf(" kernels: coop\n"); } else { - std::printf(" kernels: plain (no coop: %s)\n", d.caps.coop_bf16_reason.c_str()); + 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; @@ -190,13 +193,18 @@ int list_devices() { int run(const Args& a) { vk::Device device(a.device); const vk::Capabilities& caps = device.info().caps; - const vk::BlockPath path = a.kernels.empty() ? (caps.coop_bf16 ? vk::BlockPath::kCoop : vk::BlockPath::kPlain) - : a.kernels == "coop" ? vk::BlockPath::kCoop - : vk::BlockPath::kPlain; + 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), - path == vk::BlockPath::kPlain && !caps.coop_bf16 ? (" (no coop: " + caps.coop_bf16_reason + ")").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; diff --git a/vulkan/tools/vulkan_block_check.cpp b/vulkan/tools/vulkan_block_check.cpp index 6c50fdf..e6d7cba 100644 --- a/vulkan/tools/vulkan_block_check.cpp +++ b/vulkan/tools/vulkan_block_check.cpp @@ -195,15 +195,24 @@ struct Limit { double floor_of_rms; }; -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; +// 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 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)); + 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; @@ -288,10 +297,22 @@ bool check_norms(Device& device, const BlockKernels& kernels) { } // m, n and k: tails on every axis of both paths' tiles; the larger shape also crosses tiles and bands of tiles. -bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int k, int f) { +// 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; - const auto x_h = synth_bf16(21, static_cast(m) * ld, 0), 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); + 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); @@ -308,14 +329,15 @@ bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int situ[static_cast(r) * f + c] = host_situ(h, g); } } - std::printf(" GEMMs at M %d, N %d, K %d, SiTU F %d\n", m, n, k, f); + 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}); + 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}) && ok; + 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) { @@ -328,16 +350,16 @@ bool check_gemms(Device& device, const BlockKernels& kernels, int m, int n, int 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}) && ok; + 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}) && ok; + 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}) && ok; + 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) { @@ -438,24 +460,25 @@ int main(int argc, char** argv) { 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|plain] [--rows L1,L2,...] [--iters N] [--no-block]\n", argv[0]); + 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 != "plain") { - std::printf("--path takes coop or plain\n"); + 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, path_name == "coop" ? BlockPath::kCoop : BlockPath::kPlain); - std::printf("path: %s%s%s\n", block_path_name(kernels.path()), device.info().caps.coop_bf16 ? "" : " (no coop: ", - device.info().caps.coop_bf16 ? "" : (device.info().caps.coop_bf16_reason + ")").c_str()); + : 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"); diff --git a/vulkan/tools/vulkan_block_real.cpp b/vulkan/tools/vulkan_block_real.cpp index 17ebce7..015e22c 100644 --- a/vulkan/tools/vulkan_block_real.cpp +++ b/vulkan/tools/vulkan_block_real.cpp @@ -78,7 +78,7 @@ bool compare(const std::vector& got, const std::vector& engine) { int main(int argc, char** argv) { if (argc < 2) { - std::printf("usage: %s DIR [--device N] [--path coop|plain]\n", argv[0]); + std::printf("usage: %s DIR [--device N] [--path coop|coop-f16|plain]\n", argv[0]); return 2; } const fs::path dir = argv[1]; @@ -92,7 +92,7 @@ int main(int argc, char** argv) { try { Device device(device_index); const BlockKernels kernels = path_name.empty() ? BlockKernels(device) - : BlockKernels(device, path_name == "plain" ? BlockPath::kPlain : BlockPath::kCoop); + : 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; diff --git a/vulkan/tools/vulkan_check.cpp b/vulkan/tools/vulkan_check.cpp index 937ea33..bd5fcc4 100644 --- a/vulkan/tools/vulkan_check.cpp +++ b/vulkan/tools/vulkan_check.cpp @@ -13,6 +13,7 @@ #include #include +#include "block.h" #include "gemm.h" #include "half.h" #include "runtime.h" @@ -127,6 +128,10 @@ void print_devices() { 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()); } } diff --git a/vulkan/tools/vulkan_dit_run.cpp b/vulkan/tools/vulkan_dit_run.cpp index a422720..75011b7 100644 --- a/vulkan/tools/vulkan_dit_run.cpp +++ b/vulkan/tools/vulkan_dit_run.cpp @@ -1,7 +1,7 @@ // 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|plain] [--device N] +// 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}. @@ -68,7 +68,7 @@ double now_s() { 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|plain] [--device N]\n", + 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; } @@ -106,7 +106,7 @@ int main(int argc, char** argv) { Device device(device_index); const BlockKernels kernels = path_name.empty() ? BlockKernels(device) - : BlockKernels(device, path_name == "plain" ? BlockPath::kPlain : BlockPath::kCoop); + : 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; diff --git a/vulkan/tools/vulkan_te_check.cpp b/vulkan/tools/vulkan_te_check.cpp index 71bd561..cafed55 100644 --- a/vulkan/tools/vulkan_te_check.cpp +++ b/vulkan/tools/vulkan_te_check.cpp @@ -3,7 +3,7 @@ // with them: cosine over the sequence, worst row cosine, row relative L2, share of identical bf16 bits, worst bf16 ULPs // scaled to max(|golden|, rms(golden) / 64), NaN count. The cond and uncond prompts of a case are encoded in one call; // --iters repeats the call for timing; --profile submits each op alone and prints seconds per kind of op. Exit status: -// 0 ran without NaN, 1 otherwise. Usage: vulkan_te_check [--device N] [--path coop|plain] [--weights DIR] [--golden +// 0 ran without NaN, 1 otherwise. Usage: vulkan_te_check [--device N] [--path coop|coop-f16|plain] [--weights DIR] [--golden // DIR] [--case ID]... [--iters N] [--profile] #include #include @@ -102,21 +102,21 @@ int main(int argc, char** argv) { iters = std::max(1, std::atoi(argv[++i])); } else { std::printf( - "usage: %s [--device N] [--path coop|plain] [--weights DIR] [--golden DIR] [--case ID]... [--iters N] " + "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 != "plain") { - std::printf("--path takes coop or plain\n"); + 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() ? (device.info().caps.coop_bf16 ? BlockPath::kCoop : BlockPath::kPlain) - : (path_name == "coop" ? BlockPath::kCoop : BlockPath::kPlain); + 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())); diff --git a/vulkan/vae.h b/vulkan/vae.h index 6f3ee38..eebf469 100644 --- a/vulkan/vae.h +++ b/vulkan/vae.h @@ -10,7 +10,7 @@ // 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, chosen as BlockKernels does: kCoop runs the convolutions and GEMMs on bf16 cooperative matrices +// 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. //