Try lowering aten._scaled_dot_product_flash_attention

[jdh8]

Ticket

• Resolves aten._scaled_dot_product_flash_attention.default #541

Problem description

Convert aten._scaled_dot_product_flash_attention to a series of ops. Future goal might be implementing it as a composite kernel op instead.

The source op is functionally equivalent to its high-level counterpart:
https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

What's changed

[jerrysky3]

ttnn seems to already have the corresponding API https://docs.tenstorrent.com/ttnn/latest/ttnn/api/ttnn.transformer.scaled_dot_product_attention.html, can it be used here?

[swimdi]

There's its coressponding test

https://github.com/tenstorrent/tt-metal/blob/main/tests/tt_eager/python_api_testing/unit_testing/misc/test_scaled_dot_product_attention.py

[ayerofieiev-tt]

@jdh8 please let us know whatever question you got, I will help to connect with the right stakeholder to resolve this fast.

[jdh8]

Inferred 0 batch size 🤔

FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape0-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 16, 197, 64] vs list(actual_pytorch_result.shape)=[0, 16, 197, 64]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape1-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 12, 197, 64] vs list(actual_pytorch_result.shape)=[0, 12, 197, 64]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape2-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 16, 50, 64] vs list(actual_pytorch_result.shape)=[0, 16, 50, 64]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape3-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 8, 4096, 40] vs list(actual_pytorch_result.shape)=[0, 8, 4096, 40]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape4-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 8, 1024, 80] vs list(actual_pytorch_result.shape)=[0, 8, 1024, 80]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape5-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 8, 256, 160] vs list(actual_pytorch_result.shape)=[0, 8, 256, 160]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape6-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 8, 64, 160] vs list(actual_pytorch_result.shape)=[0, 8, 64, 160]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape7-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 12, 50, 64] vs list(actual_pytorch_result.shape)=[0, 12, 50, 64]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape8-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 16, 1370, 80] vs list(actual_pytorch_result.shape)=[0, 16, 1370, 80]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape9-False] - AssertionError: list(expected_pytorch_result.shape)=[1, 12, 1, 64] vs list(actual_pytorch_result.shape)=[0, 12, 1, 64]
FAILED tests/lowering/misc/test_scaled_dot_product_attention.py::test_sdpa[input_shape10-True] - AssertionError: list(expected_pytorch_result.shape)=[1, 12, 4, 64] vs list(actual_pytorch_result.shape)=[0, 12, 4, 64]
=================================================================== 11 failed in 14.22s ====================================================================
                 Device | INFO     | Closing user mode device drivers

    def assert_with_pcc(expected_pytorch_result, actual_pytorch_result, pcc=0.999):
>       assert list(expected_pytorch_result.shape) == list(
            actual_pytorch_result.shape
        ), f"list(expected_pytorch_result.shape)={list(expected_pytorch_result.shape)} vs list(actual_pytorch_result.shape)={list(actual_pytorch_result.shape)}"
E       AssertionError: list(expected_pytorch_result.shape)=[1, 12, 4, 64] vs list(actual_pytorch_result.shape)=[0, 12, 4, 64]

[ayerofieiev-tt]

Is this still an issue?

[jdh8]

I just confirmed that this issue still lingers. I filed tenstorrent/tt-metal#16021 to keep track on this.

[ayerofieiev-tt]

need to fire a ticket for unsupported cases?

[jdh8]

I'm marking this issue as a feature request tenstorrent/tt-metal#16022. No input variation now has nonzero dropout_p yet. It's still good to keep an eye.

[cglagovichTT]

Doesn't this logic drop the attention mask, which must be provided if is_causal == False

[jdh8]

aten._scaled_dot_product_flash_attention does not provide attention mask as far as I know. I have not yet found better documentation. Please correct me if I am wrong.

[cglagovichTT]

I see. I am unfamiliar with the aten API. My understanding of the op comes from the functional API https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html