SIGN IN SIGN UP

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

Pull Request resolved: https://github.com/pytorch/executorch/pull/21064

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).
ghstack-source-id: 405400514
@exported-using-ghexport

Differential Revision: [D112906313](https://our.internmc.facebook.com/intern/diff/D112906313/)
S
Stephen Jia committed
266e0dc4556d884cb8548455ff9d8908720e9859
Parent: bd7991f