[ET-VK][sdpa] Vendor-adaptive head_dim output-tiling (TILE_N4=2) in GQA AV coop-GEMV

[SS-JIA]

Stack from ghstack (oldest at bottom):

• -> [ET-VK][sdpa] Vendor-adaptive head_dim output-tiling (TILE_N4=2) in GQA AV coop-GEMV #21064
• [ET-VK][sdpa] Reuse shared V cache across GQA query heads in AV coop-GEMV #21063
• [ET-VK][sdpa] Add SDPA operator perf benchmark binary (test_sdpa) #21062

Builds on the GQA-reuse AV coop-GEMV (parent commit). The base
sdpa_compute_out_gqa_coop shader assigns one head_dim texel (TILE_N4=1) per
workgroup x-slot. This adds a head_dim output-tiled variant,
sdpa_compute_out_gqa_coop_tile2 (TILE_N4=2): each workgroup owns 2 head_dim
texels and emits G x TILE_N4 outputs, so the per-context-texel attn-weight loads
and the shared-memory tree reduction are amortized over twice as many outputs.
The AV coop GEMV GLSL is already a TILE_N4-parameterized template (the
partial_n_tile uniform branch handles the D4 % TILE_N4 != 0 tail), so this
change is purely a new codegen variant plus dispatch wiring — no shader-body
change.

Selection is vendor-adaptive: pick_sdpa_av_shader appends the _tile2 suffix
only on Adreno (graph->device_is_adreno()). Tiling is a consistent win on
Adreno (AV ~1.14-1.63x over the base GQA variant) but a regression on Mali at
common decode contexts (~0.67-0.86x, interleaved median-of-N), so Mali and other
vendors keep the base sdpa_compute_out_gqa_coop.
pick_sdpa_av_global_wg_size keys off the same _tile2 suffix: the tiled
variant's x-dim collapses from D4 to div_up(D4, 2) since each workgroup now
covers 2 head_dim texels.

Because the variant is Adreno-only in production, the test-only gqa_override
knob (see SDPA.h) is extended so tests can pin the variant on any device:
kGqaOverrideForceTile2 / kGqaOverrideForceBase force the tiled / base
variant regardless of vendor (via resolve_use_tile2), giving the tiled shader
deterministic coverage on Mali / SwiftShader. resolve_use_gqa is unchanged
(any non-ForceNonGqa value still forces the GQA family, VK_CHECK'd for
eligibility).

Differential Revision: D112906313

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• Unrelated MacOS MPS test jobs time out (30 min) due to stale cpp_extension lock file poisoning self-hosted runners

⏳ No Failures, 2 Pending

As of commit 8bea766 with merge base 37400d9 ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

Review automatically exported from Phabricator review in Meta.