Add Strided Input test for flex attention #128915

joydddd · 2024-06-18T00:38:04Z

Test strided inputs to the flex_attention HOP. Similar to how inputs are generated in
https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/score_mod.py.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

pytorch-bot · 2024-06-18T00:38:07Z

🔗 Helpful Links

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

📄 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 or our office hours

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

✅ You can merge normally! (6 Unrelated Failures)

As of commit 72c63ff with merge base b24787b ():

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

trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (similar failure)
inductor/test_torchinductor_dynamic_shapes.py::DynamicShapesCpuTests::test__unsafe_masked_index_dynamic_shapes_cpu

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

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

trunk / macos-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh) (trunk failure)
test_fx.py::TestSourceMatcher::test_legalize_slice
trunk / macos-py3-arm64 / test (default, 3, 3, macos-m1-stable) (gh) (trunk failure)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_H_cpu_float32
trunk / macos-py3-arm64-mps / test (mps, 1, 1, macos-m1-13) (gh) (trunk failure)
test_mps.py::TestMPS::test_mps_allocator_module
trunk / macos-py3-arm64-mps / test (mps, 1, 1, macos-m1-14) (gh) (trunk failure)
test_mps.py::TestMPS::test_mps_allocator_module

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

inductor / rocm6.1-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2, unstable) (gh) (#128871)
'test/inductor/test_max_autotune.py::TestTuningProcess::test_tuning_pool_multiple_devices'

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

joydddd · 2024-06-19T02:32:13Z

Fix flex attention forward for strided inputs.

test/inductor/test_flex_attention.py

drisspg · 2024-06-19T16:37:48Z

torch/_inductor/kernel/flex_attention.py

@@ -194,10 +194,12 @@ def build_subgraph_buffer(
 MATMUL_PRECISION = Q.dtype.element_ty

 start_m = tl.program_id(0)
- off_hz = tl.program_id(1)
+ off_z = tl.program_id(1) // H
+ off_h = tl.program_id(1) % H


great catch though!

test/inductor/test_flex_attention.py

joydddd · 2024-06-24T22:41:54Z

@pytorchbot merge

pytorchmergebot · 2024-06-24T22:44:08Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

pytorchmergebot · 2024-06-24T22:44:19Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-focal-py3.8-clang10-onnx / test (default, 1, 2, linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

joydddd · 2024-06-24T22:53:58Z

@pytorchbot merge -f 'Unrelated failures'

pytorchmergebot · 2024-06-24T22:56:25Z

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

huydhn · 2024-06-25T00:42:04Z

@pytorchbot revert -m 'Sorry for reverting your change but its tests are failing in trunk, i.e. https://hud.pytorch.org/pytorch/pytorch/commit/41bb81b58279f492e72bd270b3b071dd2953ed8c#26627138290' -c weird

I think this is a land race because your base commit is too old (20 days). It would also explain why there are many unrelated failures in the PR

pytorchmergebot · 2024-06-25T00:43:29Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2024-06-25T00:43:37Z

@joydddd your PR has been successfully reverted.

This reverts commit 41bb81b. Reverted #128915 on behalf of https://github.com/huydhn due to Sorry for reverting your change but its tests are failing in trunk, i.e. https://hud.pytorch.org/pytorch/pytorch/commit/41bb81b58279f492e72bd270b3b071dd2953ed8c#26627138290 ([comment](#128915 (comment)))

joydddd · 2024-06-25T22:34:38Z

@pytorchmergebot merge

pytorchmergebot · 2024-06-25T22:36:07Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

pytorchmergebot · 2024-06-25T22:36:19Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable)

Details for Dev Infra team

Raised by workflow job

joydddd · 2024-06-25T23:24:32Z

@pytorchmergebot merge -f "Unrelated failure"

pytorchmergebot · 2024-06-25T23:26:25Z

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

pytorch-bot bot added module: inductor topic: not user facing topic category labels Jun 18, 2024

joydddd requested review from yanboliang, Chillee, drisspg and BoyuanFeng June 18, 2024 00:38

pytorch-bot bot added the ciflow/inductor label Jun 19, 2024

Chillee reviewed Jun 19, 2024

View reviewed changes

test/inductor/test_flex_attention.py Outdated Show resolved Hide resolved

drisspg reviewed Jun 19, 2024

View reviewed changes

joydddd marked this pull request as ready for review June 21, 2024 02:25

Chillee approved these changes Jun 21, 2024

View reviewed changes

test/inductor/test_flex_attention.py Outdated Show resolved Hide resolved

drisspg approved these changes Jun 24, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 24, 2024

pytorchmergebot added the merging label Jun 24, 2024

pytorchmergebot removed the merging label Jun 24, 2024

pytorchmergebot added the merging label Jun 24, 2024

pytorchmergebot closed this in 41bb81b Jun 24, 2024

pytorchmergebot added Merged and removed merging labels Jun 24, 2024

pytorchmergebot added the Reverted label Jun 25, 2024

pytorchmergebot reopened this Jun 25, 2024

joydddd added 6 commits June 24, 2024 20:18

Add Strided Input test for flex attention

84afacb

Fix flex attention forward for strided input

09fa591

Add more strided input tests

744817e

Fix code style

c1c11a1

Lint test_flex_attention

cf8c279

Fix strided_input test for flex_decoding

e17f87c

joydddd force-pushed the strided_input_test branch from ddeb103 to e17f87c Compare June 25, 2024 03:32

Skip query braodcasting layout

afe8891

joydddd force-pushed the strided_input_test branch from 01b7f30 to afe8891 Compare June 25, 2024 18:14

Apply Lint

72c63ff

pytorchmergebot added the merging label Jun 25, 2024

pytorchmergebot removed the merging label Jun 25, 2024

pytorchmergebot added the merging label Jun 25, 2024

pytorchmergebot closed this in f389541 Jun 25, 2024

pytorchmergebot removed the merging label Jun 25, 2024

joydddd removed the Reverted label Jun 26, 2024

joydddd deleted the strided_input_test branch June 26, 2024 17:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Strided Input test for flex attention #128915

Add Strided Input test for flex attention #128915

joydddd commented Jun 18, 2024 •

edited

Loading

pytorch-bot bot commented Jun 18, 2024 •

edited

Loading

joydddd commented Jun 19, 2024

drisspg Jun 19, 2024

joydddd commented Jun 24, 2024

pytorchmergebot commented Jun 24, 2024

pytorchmergebot commented Jun 24, 2024

joydddd commented Jun 24, 2024

pytorchmergebot commented Jun 24, 2024

huydhn commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

joydddd commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

joydddd commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

Add Strided Input test for flex attention #128915

Add Strided Input test for flex attention #128915

Conversation

joydddd commented Jun 18, 2024 • edited Loading

pytorch-bot bot commented Jun 18, 2024 • edited Loading

🔗 Helpful Links

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

✅ You can merge normally! (6 Unrelated Failures)

joydddd commented Jun 19, 2024

drisspg Jun 19, 2024

Choose a reason for hiding this comment

joydddd commented Jun 24, 2024

pytorchmergebot commented Jun 24, 2024

Merge started

pytorchmergebot commented Jun 24, 2024

Merge failed

joydddd commented Jun 24, 2024

pytorchmergebot commented Jun 24, 2024

Merge started

huydhn commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

joydddd commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

Merge started

pytorchmergebot commented Jun 25, 2024

Merge failed

joydddd commented Jun 25, 2024

pytorchmergebot commented Jun 25, 2024

Merge started

joydddd commented Jun 18, 2024 •

edited

Loading

pytorch-bot bot commented Jun 18, 2024 •

edited

Loading