[danielhanchen]

Stacked on the upstream sync branch (#35, sync-upstream-a77d362, now at upstream master-995-fd21ecf, commit 344a3b66). Adds scripts/unsloth/ggml-patches/0015-ggml-vulkan-rms-norm-short-rows.patch (numbered 0015 because #36 adds 0014; the CUDA and CPU branches use 0016 to 0018, and all of them apply together in order on ggml c1f80e5).

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):

Op	Shape	Calls	Time
RMS_NORM_MUL	(128, 30, 4128)	60	173.5 ms (2.89 ms each)
RMS_NORM_MUL	(3840, 4128), same number of bytes	120	5.5 ms (0.046 ms each)

The 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.comp with BLOCK_SIZE = 128 (rms_norm_small_f32, plus the fused-multiply variant through the same specialisation constant as today). ggml_vk_op_get_pipeline picks it for GGML_OP_RMS_NORM when ne00 <= 128 and 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 t holds x[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 on here 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:

Model, size	RMS_NORM GPU time per step, base	head	ratio
Z-Image Turbo Q4_K_M, 1024, 8 steps	74.5 / 49.9 ms	28.1 / 18.4 ms	2.65x / 2.71x
FLUX.2 klein 4B Q4_K_M, 1024, 4 steps	31.8 / 31.9 ms	14.9 / 12.8 ms	2.13x / 2.49x
Z-Image Turbo Q4_K_M, 512, 8 steps	13.1 / 33.0 ms (second run on a busier card)	5.4 / 5.4 ms	2.4x / -

At 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, then scripts/sdcpp_speed/vk_perf_parse.py <log>.

Accuracy

• Renders are bit-identical (sha256 of the PNG) between base and head for Z-Image 512 and 1024 and FLUX.2 klein 1024, in every run, with and without the perf logger. Also identical for Wan 2.1 T2V 1.3B Q4_K_M (480x320, 17 frames, 10 steps), whose norms are over the full 1536-wide rows and so keep the old pipeline.
• test-backend-ops test -o RMS_NORM -b Vulkan0 on the head: 55/55 OK (includes 64-column rows, which take the new pipeline).

Platforms

• Measured: Linux NVIDIA B200 Vulkan.
• Not measured: AMD RADV (gfx1151), Intel ANV, Windows, MoltenVK. The change only lowers the workgroup size for short rows, so it does not depend on vendor extensions; the gfx1151 runner only ran the published bundle in this pass, since timing this branch there needs the branch pushed somewhere the runner can build it.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-23T04:52:25.093366Z	—	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[danielhanchen]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. You're on a roll.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

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".