Skip to content

[ET-VK][sdpa] Reuse shared V cache across GQA query heads in AV coop-GEMV#21063

Merged
meta-codesync[bot] merged 2 commits into
gh/SS-JIA/576/basefrom
gh/SS-JIA/576/head
Jul 22, 2026
Merged

[ET-VK][sdpa] Reuse shared V cache across GQA query heads in AV coop-GEMV#21063
meta-codesync[bot] merged 2 commits into
gh/SS-JIA/576/basefrom
gh/SS-JIA/576/head

Conversation

@SS-JIA

@SS-JIA SS-JIA commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

The LLM decode AV coop-GEMV reloads the shared V cache once per query head. In
grouped-query attention Hq = G * Hkv query heads share each KV head (Llama G=4,
Phi G=3, Qwen G=2), and out[q_h, d] = sum_c attn[c, q_h] * V[c, kv_h, d] reads
the SAME V texel for every query head in a group. The per-query-head coop shader
gives each of the Hq heads its own workgroup, so V -- the dominant traffic
(head_dim-wide per context texel, vs a scalar attn weight) -- is read G times.

This adds a GQA-reuse AV variant that assigns ONE workgroup per (d4, kv_h): it
loads each V texel once and reuses it across all G query heads in the group,
producing G output texels. For this bandwidth-bound kernel that cuts V-cache
traffic ~Gx.

Implementation:

  • The variant is a codegen flag (GQA) on the existing
    sdpa_compute_out_coop.glsl template, not a separate file: one shared header
    plus two #ifdef GQA main()s (per-head and GQA-reuse), so the shared setup
    lives in one place while each algorithm reads end-to-end. It emits the shader
    sdpa_compute_out_gqa_coop.
  • Reduction reuses the per-head coop shader's shared-memory tree reduction (no
    subgroup arithmetic), so the variant runs on any Vulkan device -- Adreno and
    Mali alike -- with no capability gate.
  • Each thread holds G output accumulators; the array is sized to a compile-time
    MAX_GROUP_SIZE = 8 and the group loop is bounded by the group_size =
    Hq/Hkv spec constant, so the driver fully unrolls it at pipeline creation.
  • Dispatch (pick_sdpa_av_shader + global-wg picker + spec-const wiring in
    add_sdpa_compute_out_node): the GQA variant is selected on the LLM decode
    coop path when Hq > Hkv, evenly divisible, and G <= 8 (use_gqa_av_coop); it
    sets group_size and changes the global workgroup z-dim from Hq to Hkv.
    Everything else -- MHA (Hq == Hkv), groups exceeding the cap (G > 8, e.g. MQA
    with Hq > 8), and non-divisible shapes -- falls back to the unchanged per-head
    sdpa_compute_out_coop. (Low-ratio MQA -- Hkv == 1 with Hq <= 8 -- is
    eligible and takes the GQA path.)
  • A test-only gqa_override knob is threaded through add_sdpa_compute_out_node
    (declared in the new SDPA.h): -1 auto-select, 0 force per-head, 1 force GQA,
    so a benchmark can exercise both AV shaders on the same shape; forcing GQA is
    VK_CHECK'd against shape eligibility.

Differential Revision: D112906311

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 21, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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:

⏳ No Failures, 2 Pending

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

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 21, 2026
@github-actions

Copy link
Copy Markdown

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 digantdesai left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review automatically exported from Phabricator review in Meta.

[ghstack-poisoned]
@meta-codesync
meta-codesync Bot merged commit ec47c4b into gh/SS-JIA/576/base Jul 22, 2026
188 of 189 checks passed
@meta-codesync
meta-codesync Bot deleted the gh/SS-JIA/576/head branch July 22, 2026 02:34
@meta-codesync
meta-codesync Bot temporarily deployed to cherry-pick-bot July 22, 2026 02:34 Inactive
SS-JIA added a commit that referenced this pull request Jul 22, 2026
…GEMV

Pull Request resolved: #21063

The LLM decode AV coop-GEMV reloads the shared V cache once per query head. In
grouped-query attention Hq = G * Hkv query heads share each KV head (Llama G=4,
Phi G=3, Qwen G=2), and out[q_h, d] = sum_c attn[c, q_h] * V[c, kv_h, d] reads
the SAME V texel for every query head in a group. The per-query-head coop shader
gives each of the Hq heads its own workgroup, so V -- the dominant traffic
(head_dim-wide per context texel, vs a scalar attn weight) -- is read G times.

This adds a GQA-reuse AV variant that assigns ONE workgroup per (d4, kv_h): it
loads each V texel once and reuses it across all G query heads in the group,
producing G output texels. For this bandwidth-bound kernel that cuts V-cache
traffic ~Gx.

Implementation:
- The variant is a codegen flag (`GQA`) on the existing
  `sdpa_compute_out_coop.glsl` template, not a separate file: one shared header
  plus two `#ifdef GQA` `main()`s (per-head and GQA-reuse), so the shared setup
  lives in one place while each algorithm reads end-to-end. It emits the shader
  `sdpa_compute_out_gqa_coop`.
- Reduction reuses the per-head coop shader's shared-memory tree reduction (no
  subgroup arithmetic), so the variant runs on any Vulkan device -- Adreno and
  Mali alike -- with no capability gate.
- Each thread holds G output accumulators; the array is sized to a compile-time
  `MAX_GROUP_SIZE` = 8 and the group loop is bounded by the `group_size` =
  Hq/Hkv spec constant, so the driver fully unrolls it at pipeline creation.
- Dispatch (`pick_sdpa_av_shader` + global-wg picker + spec-const wiring in
  `add_sdpa_compute_out_node`): the GQA variant is selected on the LLM decode
  coop path when Hq > Hkv, evenly divisible, and G <= 8 (`use_gqa_av_coop`); it
  sets `group_size` and changes the global workgroup z-dim from Hq to Hkv.
  Everything else -- MHA (Hq == Hkv), groups exceeding the cap (G > 8, e.g. MQA
  with Hq > 8), and non-divisible shapes -- falls back to the unchanged per-head
  `sdpa_compute_out_coop`. (Low-ratio MQA -- Hkv == 1 with Hq <= 8 -- is
  eligible and takes the GQA path.)
- A test-only `gqa_override` knob is threaded through `add_sdpa_compute_out_node`
  (declared in the new `SDPA.h`): -1 auto-select, 0 force per-head, 1 force GQA,
  so a benchmark can exercise both AV shaders on the same shape; forcing GQA is
  VK_CHECK'd against shape eligibility.
ghstack-source-id: 405400503
@exported-using-ghexport

Differential Revision: [D112906311](https://our.internmc.facebook.com/intern/diff/D112906311/)
SS-JIA added a commit that referenced this pull request Jul 22, 2026
…GEMV

Pull Request resolved: #21063

The LLM decode AV coop-GEMV reloads the shared V cache once per query head. In
grouped-query attention Hq = G * Hkv query heads share each KV head (Llama G=4,
Phi G=3, Qwen G=2), and out[q_h, d] = sum_c attn[c, q_h] * V[c, kv_h, d] reads
the SAME V texel for every query head in a group. The per-query-head coop shader
gives each of the Hq heads its own workgroup, so V -- the dominant traffic
(head_dim-wide per context texel, vs a scalar attn weight) -- is read G times.

This adds a GQA-reuse AV variant that assigns ONE workgroup per (d4, kv_h): it
loads each V texel once and reuses it across all G query heads in the group,
producing G output texels. For this bandwidth-bound kernel that cuts V-cache
traffic ~Gx.

Implementation:
- The variant is a codegen flag (`GQA`) on the existing
  `sdpa_compute_out_coop.glsl` template, not a separate file: one shared header
  plus two `#ifdef GQA` `main()`s (per-head and GQA-reuse), so the shared setup
  lives in one place while each algorithm reads end-to-end. It emits the shader
  `sdpa_compute_out_gqa_coop`.
- Reduction reuses the per-head coop shader's shared-memory tree reduction (no
  subgroup arithmetic), so the variant runs on any Vulkan device -- Adreno and
  Mali alike -- with no capability gate.
- Each thread holds G output accumulators; the array is sized to a compile-time
  `MAX_GROUP_SIZE` = 8 and the group loop is bounded by the `group_size` =
  Hq/Hkv spec constant, so the driver fully unrolls it at pipeline creation.
- Dispatch (`pick_sdpa_av_shader` + global-wg picker + spec-const wiring in
  `add_sdpa_compute_out_node`): the GQA variant is selected on the LLM decode
  coop path when Hq > Hkv, evenly divisible, and G <= 8 (`use_gqa_av_coop`); it
  sets `group_size` and changes the global workgroup z-dim from Hq to Hkv.
  Everything else -- MHA (Hq == Hkv), groups exceeding the cap (G > 8, e.g. MQA
  with Hq > 8), and non-divisible shapes -- falls back to the unchanged per-head
  `sdpa_compute_out_coop`. (Low-ratio MQA -- Hkv == 1 with Hq <= 8 -- is
  eligible and takes the GQA path.)
- A test-only `gqa_override` knob is threaded through `add_sdpa_compute_out_node`
  (declared in the new `SDPA.h`): -1 auto-select, 0 force per-head, 1 force GQA,
  so a benchmark can exercise both AV shaders on the same shape; forcing GQA is
  VK_CHECK'd against shape eligibility.
ghstack-source-id: 405400503
@exported-using-ghexport

Differential Revision: [D112906311](https://our.internmc.facebook.com/intern/diff/D112906311/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants