You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit 266e0dc
Browse filesBrowse the repository at this point in the historyBrowse files
[ET-VK][sdpa] Vendor-adaptive head_dim output-tiling (TILE_N4=2) in GQA AV coop-GEMV
Pull Request resolved: pytorch#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/)
0 commit comments