Repository navigation
Vulkan: RMS norm for short rows (per-head Q/K norm 2.4-2.7x faster, bit-identical) - #37
danielhanchen wants to merge 5 commits into
Conversation
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…subgroup kernel ggml 135ce1e0 (e995244) adds rms_norm_128 for rows of exactly 128 on devices with full 32-wide subgroups. Keep that kernel first; the 128-thread variant here now serves every other row of at most 128 elements, and rows of 128 on devices without the subgroup path.
💡 Codex Reviewstable-diffusion.cpp/include/stable-diffusion.h Lines 556 to 558 in 725579a  Preserve the existing C ABI when adding generation results
For applications or language bindings built against the previous public API, retaining the exported AGENTS.md reference: AGENTS.md:L74-L76 stable-diffusion.cpp/examples/server/routes_openai.cpp Lines 207 to 209 in 725579a Restore uploaded-image initialization for non-reference edits
For ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Based on master (ggml
135ce1e0, merged in on Nov 7). Addsscripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch.Update (Nov 7): ggml now has its own 128-column kernel
ggml
135ce1e0(e995244, "fuse attention normalization and layout conversion") addsrms_norm_128.comp, a subgroup-shuffle RMS norm for rows of exactly 128, created only on devices with full 32-wide subgroups. The old patch no longer applied, so it was ported:ggml_vk_op_get_pipelinetries upstream'srms_norm_128first (rows of 128, subgroup size 32), then this patch's 128-threadrms_norm_smallfor every other row of at most 128 elements, and for rows of 128 on devices without the subgroup path (for example wave64-only AMD GCN, or drivers that cannot require full subgroups).The timings below were measured before upstream's kernel existed. On devices that create
rms_norm_128, head size 128 (Z-Image, FLUX.2 klein) now takes upstream's kernel on master already, so the win measured below only remains for head sizes under 128 and for devices without full 32-wide subgroups. Not re-timed on a real GPU after the port (no GPU on the CI runners; lavapipe only checks correctness).Problem
ggml-vulkan's RMS norm runs one 512-thread workgroup per row. The DiTs Studio runs on sd.cpp normalise Q and K per head, so the rows are 128 elements long and there are many of them: Z-Image at 1024x1024 has 30 heads x 4128 tokens = 123,840 rows per call. Three quarters of every workgroup only adds zeros and waits on barriers.
Per-op GPU time (
GGML_VK_PERF_LOGGER=1, B200, Vulkan, one steady denoise step of Z-Image Turbo Q4_K_M at 1024x1024, 1198 ms in total):RMS_NORM_MULRMS_NORM_MULThe per-head norm was 15% of the step, behind only the quantised matmuls (48%) and flash attention (18%).
Fix
A second build of
rms_norm.compwithBLOCK_SIZE = 128(rms_norm_small_f32, plus the fused-multiply variant through the same specialisation constant as today).ggml_vk_op_get_pipelinepicks it forGGML_OP_RMS_NORMwhenne00 <= 128and the add-partials path is not in use. Everything else keeps the 512-thread pipelines.The result is bit-identical by construction. With at most 128 columns, thread
tholdsx[t]^2(or 0) in both versions, and the 512-thread tree only adds exact zeros in its first two halving steps (strides 256 and 128) before running the same strides 64 down to 1 as the 128-thread tree. The scale and the multiply loop are unchanged.Results
B200 (Vulkan,
NV_coopmat2), base =sync-upstream-fd21ecf, head = this branch, same build flags (the prebuilt Vulkan leg's). Studio's Vulkan argv (--diffusion-fa, with--conv-direct onhere because the perf logger aborts in ggml-vulkan on the im2col VAE graph at 1024). The card is shared with other jobs, so totals move between runs; base and head ran back to back, interleaved. Medians over the steady denoise steps, two runs per arm:RMS_NORMGPU time per step, baseAt Z-Image 1024 that is about 30 to 55 ms of a 390 to 570 ms step in these runs, roughly 8%. Wall-clock sampling did not resolve it on this host (host load moves Vulkan wall time by tens of percent between identical runs), and the exclusive-card timing was not run, so the end-to-end win is the GPU-time estimate above, not a measured wall-clock number.
Reproduce:
scripts/sdcpp_speed/vk_bench.py --cases zimage1024,klein1024,zimage512 --arms base=<base sd-cli>,rms=<head sd-cli> --variants studio_cd_on,studio_cd_on+perf --env studio_cd_on+perf:GGML_VK_PERF_LOGGER=1 --reps 2, thenscripts/sdcpp_speed/vk_perf_parse.py <log>.Accuracy
test-backend-ops test -o RMS_NORM -b Vulkan0on the head: 55/55 OK (includes 64-column rows, which take the new pipeline).Platforms