8000 [pytorch][triton] Enabling TMA for flex-attention for supported device types by mandroid6 · Pull Request #153662 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[pytorch][triton] Enabling TMA for flex-attention for supported device types #153662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[pytorch][triton] Enabling TMA for flex-attention for supported device types #153662

wants to merge 1 commit into from
+14 −8

Conversation

mandroid6
Copy link
Contributor
@mandroid6 mandroid6 commented May 15, 2025 •
edited by pytorch-bot bot
Loading

Summary:
Currently flex-attention defaults to USE_TMA=False.

We can enable TMA on devices which support it based on has_triton_tma_device.

Test Plan:

Tritonbench results

buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322

Differential Revision: D74841543

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

@pytorch-bot PyTorch Bot
Copy link
pytorch-bot bot commented May 15, 2025 •
edited
Loading

🔗 Helpful Links

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

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:

✅ You can merge normally! (1 Unrelated Failure)

As of commit 3317a9e with merge base 9fe2d15 (image):

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.

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

@mandroid6 mandroid6 requested a review from drisspg May 15, 2025 21:43
@mandroid6 mandroid6 added the topic: not user facing topic category label May 15, 2025
@drisspg
Copy link
Contributor
drisspg commented May 16, 2025

Can you also do a perf bench w/ larger sequen lengths, I am curious

Copying over comments:
Looks good, the only thing is that we dont have CI testing for this, I am going to run the tests on the devvm

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 16, 2025
@drisspg
Copy link
Contributor
drisspg commented May 16, 2025

https://www.internalfb.com/intern/paste/P1813676625/

Some failing tests

pytorch-bot bot pushed a commit that referenced this pull request May 27, 2025
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
9972d12 
…e types (#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
https://www.internalfb.com/intern/testinfra/testrun/3096224974340688
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

mandroid6 added a commit to mandroid6/pytorch that referenced this pull request May 28, 2025 
…e types (pytorch#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
https://www.internalfb.com/intern/testinfra/testrun/3096224974340688
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

mandroid6 added a commit to mandroid6/pytorch that referenced this pull request May 28, 2025
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
2e14cda 
…e types (pytorch#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
https://www.internalfb.com/intern/testinfra/testrun/3096224974340688
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

pytorch-bot bot pushed a commit that referenced this pull request May 30, 2025
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
db2c9b6 
…e types (#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
 https://www.internalfb.com/intern/testinfra/testrun/7318349664976675
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
pytorch-bot bot pushed a commit that referenced this pull request Jun 10, 2025
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
fec2acc 
…e types (#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

mandroid6 added a commit to mandroid6/pytorch that referenced this pull request Jun 10, 2025
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
f45bd01 
…e types (pytorch#153662)

Summary:
Pull Request resolved: pytorch#153662

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Differential Revision: D74841543
@davidberard98 davidberard98 mentioned this pull request Jun 12, 2025
Closed
davidberard98
davidberard98 reviewed Jun 12, 2025
View reviewed changes
torch/_inductor/kernel/flex_attention.py Outdated
@@ -1613,7 +1614,7 @@ def flex_attention(
original_kernel_options = kernel_options.copy()
# Default config for warp specialization
num_consumer_groups, num_buffers_warp_spec = 0, 0

USE_TMA = has_triton_tma_device()
Copy link
Contributor
@davidberard98 davidberard98 Jun 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

@mandroid6 once #155771 lands, can you change this to has_triton_stable_tma_api() ? This is because #155771 switches to using the stable TMA API (whereas has_triton_tma_device() just checks if there's support for any TMA API)

davidberard98 added a commit that referenced this pull request Jun 12, 2025
@davidberard98
Update base for Update on "[flex attention][triton pin] use new TMA API"
81774fb 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100, [TODO confirm results]

TODO: When #153662 lands, turning on TMA support by default, it should be modified slightly to check specifically for stable TMA API support.

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

[ghstack-poisoned]
davidberard98 added a commit that referenced this pull request Jun 12, 2025
@davidberard98
Update on "[flex attention][triton pin] use new TMA API"
cec0b53 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100, [TODO confirm results]

TODO: When #153662 lands, turning on TMA support by default, it should be modified slightly to check specifically for stable TMA API support.

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

[ghstack-poisoned]
davidberard98 added a commit that referenced this pull request Jun 12, 2025
@davidberard98
Update base for Update on "[flex attention][triton pin] use new TMA API"
037b6ab 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100, [TODO confirm results]

Note: When #153662 lands, turning on TMA support by default, it should be checking specifically for stable TMA API support (commented on PR)

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

[ghstack-poisoned]
davidberard98 added a commit that referenced this pull request Jun 12, 2025
@davidberard98
Update on "[flex attention][triton pin] use new TMA API"
9c7b915 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100, [TODO confirm results]

Note: When #153662 lands, turning on TMA support by default, it should be checking specifically for stable TMA API support (commented on PR)

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

[ghstack-poisoned]
pytorchmergebot pushed a commit that referenced this pull request Jun 14, 2025
@davidberard98 @pytorchmergebot
[flex attention][triton pin] use new TMA API (#155771)
c843909 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100 with @mandroid6's PR #153662 patched in to turn on TMA [TODO: confirm results once all the local tests pass, but from the first 100 tests I ran locally, all the failing tests were also failing on #153662 alone]

Note: When #153662 lands, turning on TMA support by default, it should be checking specifically for stable TMA API support (commented on PR)

Pull Request resolved: #155771
Approved by: https://github.com/mandroid6, https://github.com/nmacchioni
thatgeeman pushed a commit to thatgeeman/pytorch-docathon that referenced this pull request Jun 15, 2025
@davidberard98 @thatgeeman
[flex attention][triton pin] use new TMA API (pytorch#155771)
397045f 
Triton 3.4 will remove the experimental TMA APIs: triton-lang/triton#6488. Ahead of this, we are **replacing the experimental TMA API usage with the stable TMA API** in flex attention. This means that **flex attention TMA will stop working with Triton 3.2 or Triton 3.3/3.3.1** for now (but it should work for Triton 3.4 in the PyTorch 2.8 release, and Meta-internal triton 3.3.1fb, which have the new TMA API).

This PR does the following:
* replace the experimental TMA APIs with the stable TMA APIs
* remove the workspace args.

Testing: I ran test/inductor/test_flex_attention.py on a H100 with @mandroid6's PR pytorch#153662 patched in to turn on TMA [TODO: confirm results once all the local tests pass, but from the first 100 tests I ran locally, all the failing tests were also failing on pytorch#153662 alone]

Note: When pytorch#153662 lands, turning on TMA support by default, it should be checking specifically for stable TMA API support (commented on PR)

Pull Request resolved: pytorch#155771
Approved by: https://github.com/mandroid6, https://github.com/nmacchioni
@mandroid6 @facebook-github-bot
[pytorch][triton] Enabling TMA for flex-attention for supported devic…
3317a9e 
…e types (pytorch#153662)

Summary:

Currently flex-attention defaults to `USE_TMA=False`.

We can enable TMA on devices which support it based on `has_triton_tma_device`.

Test Plan:
## Unit tests on H100
```
buck test 'fbcode//mode/opt' fbcode//caffe2/test/inductor:flex_attention
```
# Tritonbench results

```
buck2 run mode/opt //pytorch/tritonbench:run -- --op flex_attention --use-tma --mod-type all
.
.
.
(B, Hq, M, Hkv, N, D)     |       Mask Type    compiled-latency    compiled-tflops
- USE_TMA =  False
(8, 16, 128, 16, 128, 128) |            noop   0.027936 (±5.61%)            38.5859
(8, 16, 128, 16, 128, 128) |          causal   0.017760 (±3.42%)            60.6946
(8, 16, 128, 16, 128, 128) |             rel   0.028384 (±5.07%)            37.9769
(8, 16, 128, 16, 128, 128) |       head_bias   0.027712 (±4.27%)            38.8978
(8, 16, 128, 16, 128, 128) |           alibi   0.017920 (±3.21%)            60.1527
- USE_TMA = True
(8, 16, 128, 16, 128, 128) |            noop   0.025632 (±5.74%)            42.0543
(8, 16, 128, 16, 128, 128) |          causal   0.015328 (±3.97%)            70.3246
(8, 16, 128, 16, 128, 128) |             rel   0.025824 (±4.96%)            41.7416
(8, 16, 128, 16, 128, 128) |       head_bias   0.025472 (±4.90%)            42.3185
(8, 16, 128, 16, 128, 128) |           alibi   0.015392 (±3.74%)            70.0322
```

Rollback Plan:

Differential Revision: D74841543
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D74841543

@drisspg
Copy link
Contributor
drisspg commented Jul 3, 2025

@mandroid6 ping me when this is re-ready for review

3D16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@davidberard98 davidberard98 davidberard98 left review comments

@Skylion007 Skylion007 Skylion007 approved these changes

@drisspg drisspg Awaiting requested review from drisspg

Assignees
No one assigned
Labels
ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request fb-exported module: inductor topic: not user facing topic category
Projects
None yet
Milestone
No milestone
< 3DF0 input type="hidden" name="_method" value="put" autocomplete="off" />
Development

Successfully merging this pull request may close these issues.
5 participants
@mandroid6 @facebook-github-bot @drisspg @Skylion007 @davidberard98
0