[Isalia20]

Support returning lse in flex attention, setup for backward implementation of FA on MPS

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

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:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 2 New Failures, 4 Unrelated Failures

As of commit 238f71d with merge base fd03d74 ():

NEW FAILURES - The following jobs have failed:

• inductor / inductor-test / test-osdc (inductor_timm, 1, 2, mt-l-x86aavx2-29-113-a10g) (gh)
  mobilenetv2_100
• inductor / unit-test / inductor-test / test-osdc (inductor, 1, 2, mt-l-x86aavx2-29-113-a10g) (gh)
  test/inductor/test_aot_inductor.py::AOTInductorLoggingTest::test_shape_env_reuse

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

• inductor / inductor-test / test-osdc (inductor_torchbench, 2, 2, mt-l-x86aavx2-29-113-a10g) (gh) (similar failure)
  shufflenet_v2_x1_0

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

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

• inductor / inductor-test / test-osdc (inductor_timm, 2, 2, mt-l-x86aavx2-29-113-a10g) (gh) (trunk failure)
  tf_efficientnet_b0
• torchtitan-test / torchtitan-x-pytorch-test / test-osdc (torchtitan_features_integration, 1, 1, mt-l-x86aavx2-189-704-a10g-8) (gh) (trunk failure)
  RuntimeError: 1 test steps failed: ['scripts/ci/pytorch_ci_test_runner.sh feature_tests']
• torchtitan-test / torchtitan-x-pytorch-test / test-osdc (torchtitan_models_integration, 1, 1, mt-l-x86aavx2-189-704-a10g-8) (gh) (trunk failure)
  RuntimeError: 1 test steps failed: ['scripts/ci/pytorch_ci_test_runner.sh model_tests']

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

[jansel]

@claude explain those CI failures, check if they are related to this PR, and review the changes

[claude]

Claude finished @jansel's task in 3m 36s —— View job

---

Task

• Analyze the CI failures and check if related to this PR
• Review the code changes

---

CI failures

None of the reported failures touch the code paths this PR changes. This PR only modifies the MPS flex-attention codegen (metal_flex_attention_template.py, flex_mps.py) and one test decorator. The Metal shader path only executes on Apple Silicon, while every failing job below runs on x86/a10g (CPU/CUDA inductor). So they cannot be caused by these changes.

New failures (unrelated):

• inductor_timm → mobilenetv2_100 and inductor_torchbench → shufflenet_v2_x1_0 — x86 inductor model benchmarks, no MPS involvement.
• test_aot_inductor.py::AOTInductorLoggingTest::test_shape_env_reuse — AOTInductor logging test, unrelated to flex attention; this one is a known-flaky/logging test.

Broken trunk (explicitly pre-existing): the inductor_timm (tf_efficientnet_b0) and two torchtitan jobs are flagged by Dr. CI as already failing on the merge base fd03d74. Rebasing onto viable/strict should clear them.

Verdict: all 6 failures are unrelated to this PR. A rebase onto viable/strict is the right next step; the 3 "new" ones look like flake/infra and are worth a re-run.

---

Code review

The change is clean and follows the existing structure well. The wiring of the new lse buffer through both the MMA and fallback shaders, the input list, and the mutation tracking all line up correctly.

Correctness checks that pass:

• Buffer ordering matches. Shader order is out, Q, K, V, kv_num_blocks, kv_indices, [full_kv...], captured(score, mask), lse, scalar_params. In codegen, arg_parts is out + self.inputs + params, and node_inputs appends logsumexp after the mask tensors and before the params buffer — so the positions agree. lse_param correctly consumes a buf_idx before scalar_params_str (metal_flex_attention_template.py:638-649).
• Log-base convention is correct. The kernel accumulates in natural-log space (metal::precise::exp), so lse_natural = row_max + ln(row_sum). Multiplying by 1.44269504... (= 1/ln2) stores it in log2 space, matching the Triton template (flex_attention.py.jinja:212, lse = m_i + tl.math.log2(l_i)). The top-level op converts back with * math.log(2) (_higher_order_ops/flex_attention.py:1053), so the round-trip is consistent across backends.
• Read-before-clamp ordering is correct. lse_write_code is emitted before if (row_sum == 0.0f) row_sum = 1.0f; in both shaders (:505, :838), so the LSE reads the true row_sum. The fully-masked case (row_sum == 0) is handled explicitly by writing -INFINITY.
• Mutation tracking via MutationOutput on the last input (lse_buf) plus logsumexp.realize() is the standard inductor pattern for an in-place output buffer, and the returned logsumexp IR node is the same buffer that is marked mutated.
• The removed NotImplementedError and the dropped @expected_not_implemented_on_mps on test_force_write_lse are consistent — the feature is now implemented.

Minor observations (non-blocking):

1. max_scores is returned uninitialized. lower_mps still returns max_scores from an empty_strided that nothing writes to (OUTPUT_MAX raises). This is pre-existing behavior (the old code did the same), and test_lse_masked_output keeps @expected_not_implemented_on_mps, so it's fine for now — just flagging that the third tuple element is garbage memory if anything ever reads it without the guard.
2. LSE index math is hardcoded contiguous. lse[b_idx * (Hq * N_Q) + h_idx * N_Q + m_idx] assumes a fully contiguous [B, Hq, N_Q] layout, which is guaranteed here because logsumexp = empty_strided(lse_shape, None, ...) produces a contiguous tensor. That's a safe coupling today, but if the LSE layout ever becomes non-contiguous (e.g. a strided/sliced output) this silently breaks. Mirrors the Triton kernel's off_hz * Q_LEN + offs_m assumption, so consistent with the rest of the codebase — fine to leave.
3. Since the PR description says this is setup for the backward pass, it's worth confirming test_force_write_lse actually exercises the MPS path on a runner with MPS available (the comparison is lse_eager vs lse_compiled); on CI without MPS hardware this won't validate the new shader. Presumably covered by the MacOS MPS CI lane.

Overall this looks correct and well-scoped. The CI is not a blocker for this change.

[jansel]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Tried to rebase and push PR #187768, but it was already up to date. Try rebasing against main by issuing:
@pytorchbot rebase -b main

[Isalia20]

@pytorchbot merge -f "Lint and MPS are green, failures unrelated"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here