Skip to content

Add TLX attention (WS pipelined pingpong hopper) #320

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

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

Add TLX attention (WS pipelined pingpong hopper) #320

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions tritonbench/operators/flash_attention/operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,18 @@
except (ImportError, IOError, AttributeError, TypeError):
HAS_XFORMERS = False

# [Optional] TLX backend
try:
Copy link
Contributor

@xuzhao9 xuzhao9 Sep 5, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

We introduced a cleaner version of the guards in #387. This can be simplified to

Suggested change
try:
from tritonbench.utils.python_utils import try_import
with try_import("HAS_TLX"):
import triton.language.extra.tlx as tlx
from .tlx_attn_ws_pipelined_pingpong_hopper import (
attention as tlx_attn_ws_pipelined_pingpong_hopper,
)

import triton.language.extra.tlx as tlx

from .tlx_attn_ws_pipelined_pingpong_hopper import (
attention as tlx_attn_ws_pipelined_pingpong_hopper,
)

HAS_TLX = True
except (ImportError, IOError, AttributeError):
HAS_TLX = False

from typing import Any, Generator, List

from tritonbench.utils.input import input_filter
Expand Down Expand Up @@ -299,6 +311,16 @@ def triton_tutorial_flash_v2_tma(
q, k, v, self.causal, self.sm_scale, "tma"
)

@register_benchmark(enabled=HAS_TLX)
Copy link
Contributor

Choose a reason for hiding this comment

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

Is the intention to run on CI? If you just want to trigger this test directly you can do with enabled=False and without any risk of it running on hardware you don't expect (e.g. MI300), which it will by default right now.

def tlx_attn_ws_pipelined_pingpong_hopper(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> Callable:
# TLX flash attention with Hopper optimizations
return lambda: tlx_attn_ws_pipelined_pingpong_hopper(q, k, v, self.sm_scale)

def xformers_preprocess(
self,
q: torch.Tensor,
Expand Down
Loading