Skip to content

Patch validate_device in torch.nn.attention.flex_attention #33

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
wenqf11 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Patch validate_device in torch.nn.attention.flex_attention #33

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
17 changes: 17 additions & 0 deletions src/torchada/_patch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1176,6 +1176,23 @@ def _patch_autotune_process():
autotune_process.CUDA_VISIBLE_DEVICES = "MUSA_VISIBLE_DEVICES"


@patch_function
@requires_import("torch_musa")
def _patch_validate_device():
import torch.nn.attention.flex_attention
from torch import Tensor
def _validate_device(query: Tensor, key: Tensor, value: Tensor):
"""TODO: Remove once non cuda/cpu devices support is added
We only need to check query since we have already that q,k,v are on the same device
"""
if query.device.type != "cuda" and query.device.type != "cpu" and query.device.type != "musa":
raise ValueError(
"FlexAttention is only supported on CUDA or CPU or MUSA devices. "
f"Found input tensors on {query.device.type} device."
)
torch.nn.attention.flex_attention._validate_device = _validate_device


class _CDLLWrapper:
"""
Wrapper for ctypes.CDLL that automatically translates CUDA/NCCL function names
Expand Down