Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[MPS] Support captured buffers in flex attention#186215

Open
Isalia20 wants to merge 1 commit into
pytorch:mainfrom
Isalia20:mps-fa-captured-buffers
Open

[MPS] Support captured buffers in flex attention#186215
Isalia20 wants to merge 1 commit into
pytorch:mainfrom
Isalia20:mps-fa-captured-buffers

Conversation

@Isalia20
Copy link
Copy Markdown
Collaborator

@Isalia20 Isalia20 commented Jun 4, 2026 •
edited by pytorch-bot Bot
Loading

Support captured buffers in flex attention mps. Before this wasn't possible:

import torch
from torch.nn.attention.flex_attention import flex_attention

q, k, v = (torch.randn(1, 4, 256, 64, device="mps") for _ in range(3))
bias = torch.randn(256, 256, device="mps")

out = torch.compile(flex_attention)(q, k, v, score_mod=lambda s, b, h, i, j: s + bias[i, j])
print(out.shape)

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

@Isalia20 Isalia20 requested a review from malfet June 4, 2026 09:19
@Isalia20 Isalia20 requested a review from drisspg as a code owner June 4, 2026 09:19
@Isalia20 Isalia20 added ciflow/trunk Trigger trunk jobs on your pull request topic: improvements topic category labels Jun 4, 2026
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Jun 4, 2026 •
edited
Loading

🔗 Helpful Links

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

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

❌ 1 New Failure, 82 Pending, 2 Unrelated Failures

As of commit 1fb65ae with merge base 2b2898d (image):

NEW FAILURE - The following job has failed:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

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

@Isalia20 Isalia20 added the release notes: mps Release notes category label Jun 4, 2026
@Isalia20 Isalia20 changed the title [MPS] Support captured buffers in flex attention mps [MPS] Support captured buffers in flex attention Jun 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@malfet malfet Awaiting requested review from malfet

@drisspg drisspg Awaiting requested review from drisspg drisspg is a code owner

Assignees

No one assigned

Labels

ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests ciflow/trunk Trigger trunk jobs on your pull request module: inductor open source release notes: mps Release notes category topic: improvements topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Isalia20 @pytorchbot