[Vulkan] Reduce 4w decode overhead with specialization

[mergennachin]

Summary

Related to #22968.

The 4w decode kernel keeps a scale-cache check even when each lane advances by at least one quantization group and cannot reuse the previous scales. Use the existing group-size specialization constant to eliminate that check and its loop-carried state in those cases. The reuse path is retained when a lane can share scales across iterations.

For FP16 decode, specialize logical activation K and weight N to give the pipeline compiler constant loop bounds and indexing strides. K comes from the input because export can pad packed weights beyond the activation width. Forced cooperative-kernel selectors use the same specialization constants as production. FP32 keeps runtime dimensions because its controls showed a small slowdown with shape specialization. Weight packing, workgroup selection, and prefill shaders are unchanged.

On an M1 Pro / MoltenVK 1.4.1, the combined change gave the following Qwen3-0.6B 4w results:

Median of six runs	Original baseline	Combined change
Decode, tokens/s	34.61	35.76 (+3.32%)
Prefill, tokens/s	369.49	371.05

Six alternating baseline/candidate pairs, FP16 activations, group size 32, a 41-token prompt, 214 timed decode steps per run, temperature zero, eight CPU threads, and warmup enabled. Both versions used the same exported PTE and runner dependencies, with only the Vulkan backend archive replaced. Five of six paired decode comparisons improved. All 12 runs generated identical text. These are Mac measurements; Adreno performance still needs testing.

The prefill difference is within observed run-to-run variation. Performance runs had tracing and timestamp queries disabled. Timings were collected before the padded-K fix; the compiled shaders and specialization values are unchanged for these unpadded model shapes, and timings were not rerun.

Specializing N can add pipeline-cache entries for layers with the same K and different N. Startup compilation cost was not isolated.

Test plan

The Release Vulkan backend build and all four native decode tests passed. The decode/prefill tests cover 480 graph executions across FP16/FP32, buffer/Texture3D, group sizes 32/128, optional bias, aligned and partial output tiles, and M transitions 1 → 8 → 1 → 3 → 1. They compare against an independent CPU reference and check that GEMV/GEMM dispatch follows M.

Four additional cases use logical K=12, weights padded to K=16, and NaNs in unused activation capacity. These pass for both dtypes and storage types, bringing the total to 484 graph executions. The FP16 regression test fails if decode uses padded weight K as its loop bound. A separate local probe also passed all 16 combinations of forced selectors 13–16, FP16/FP32, and buffer/Texture3D.

cmake --build build --target vulkan_backend --parallel 6
VK_DRIVER_FILES=/usr/local/share/vulkan/icd.d/MoltenVK_icd.json \
  test-build/vulkan-test/vulkan_compute_api_test \
  '--gtest_filter=VulkanComputeGraphTest.q4gsw_decode_*'

Targeted lintrunner checks and git diff --check passed. Comparing all 1,151 production SPIR-V binaries against the original baseline, exactly the 12 4w decode variants change; the other 1,139, including all prefill shaders, are byte-identical.

Authored with AI assistance from OpenAI Codex.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22973

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• linux.dgx.b200 runner queuing exceeds 4 hours and dgxb200-05 "CUDA not available"

❌ 1 New Failure

As of commit 23d7924 with merge base 3b065a6 ():

NEW FAILURE - The following job has failed:

• pull / test-models-linux (add, xnnpack-quantization-delegation, mt-l-x86iavx512-8-64) / linux-job (gh)
  ##[error]fatal: Fetched in submodule path 'extension/llm/tokenizers/third-party/json', but it did not contain e6cafa573aac6ed9227f752a5371c0b3f436307d. Direct fetching of that commit failed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mergennachin]

Looks like it also showed improvement on Adreno too.

See #22968 (comment)

[Copilot]

Copilot review overview

🔵 Needs a closer look

Add explicit <random> and <cmath> includes to keep the test translation unit self-contained.

Review effort: Lite
Findings: None

What changed in this PR

This PR optimizes Vulkan 4-bit decode by specializing FP16 dimensions and removing unnecessary scale-cache checks, with expanded correctness coverage.

Changes:

• Specializes FP16 logical K/N and forced selector constants.
• Simplifies scale loading when scale reuse is impossible.
• Adds decode, prefill, padding, storage, dtype, and selector tests.

File	Summary
backends/​vulkan/​test/​vulkan_compute_api_test.cpp	Adds comprehensive decode, prefill, dispatch, and padded-weight tests.
backends/​vulkan/​test/​custom_ops/​impl/​TestFpaQ4gswLinear.cpp	Updates forced shader specialization values.
backends/​vulkan/​runtime/​graph/​ops/​impl/​Q4gswLinear.cpp	Supplies dtype-specific logical dimension specializations.
backends/​vulkan/​runtime/​graph/​ops/​glsl/​q4gsw_linear_gemv_coop__w_4x8.glsl	Specializes dimensions and optimizes scale-cache logic.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[meta-codesync]

@mergennachin has imported this pull request. If you are a Meta employee, you can view this in D121197981.