Add a new SdpaFwdOp IR node for Flash Attention

[Priya2698]

Issue #2278.
This PR adds the new node with same functionality as torch.nn.functional.scaled_dot_product_attention, and enables scheduling it through ExprEvalScheduler.

Based on the PR discussions, this PR is repurposed to introduce a new IR node SdpaFwdOp for scaled dot product flash attention forward (see #2278 for details).
This PR does not include changes to the scheduler.

The next PRs will:

1. Add mapping for producer-consumer in root_domain_map and enable this op in ExprEvalScheduler.
2. Python API
3. Add a node for backward pass similar to this forward node.

After the completion of these tasks, we also aim at introducing Memory Efficient Attention.

[Priya2698]

!build

[jacobhinkle]

Is it correct to say that before accepting this in ExprEval scheduler, we need to handle the other execution modes?

[Priya2698]

Is it correct to say that before accepting this in ExprEval scheduler, we need to handle the other execution modes?

No. I mainly separated them to handle the mapping, and the ID graphs workarounds separately than the node to reduce the scope of this PR.

As we discussed in today's meeting, at the moment we only plan on supporting Flash Attention to support multi-GPU development. Once we support Flash Attention, we can revisit, if we need to add Memory Efficient Attention as well. There could be a few ways:

1. Plumbing down the backend info from Thunder and using that within our nodes: While the two implementations have different function signatures, there are overlaps and hence, one possibility is to use a superset of the inputs and outputs. The other design here would be distinct nodes for each implementation.
2. We make the decision about the backend within nvFuser using the same logic as ATen/Thunder. See: https://github.com/Lightning-AI/lightning-thunder/blob/9f0c50cc6df187cf5fd2e31240690fe2b5e9ccc1/thunder/executors/sdpaex.py#L618-L680

[jacobhinkle]

Is it correct to say that before accepting this in ExprEval scheduler, we need to handle the other execution modes?

No. I mainly separated them to handle the mapping, and the ID graphs workarounds separately than the node to reduce the scope of this PR.

As we discussed in today's meeting, at the moment we only plan on supporting Flash Attention to support multi-GPU development. Once we support Flash Attention, we can revisit, if we need to add Memory Efficient Attention as well. There could be a few ways:

1. Plumbing down the backend info from Thunder and using that within our nodes: While the two implementations have different function signatures, there are overlaps and hence, one possibility is to use a superset of the inputs and outputs. The other design here would be distinct nodes for each implementation.
2. We make the decision about the backend within nvFuser using the same logic as ATen/Thunder. See: https://github.com/Lightning-AI/lightning-thunder/blob/9f0c50cc6df187cf5fd2e31240690fe2b5e9ccc1/thunder/executors/sdpaex.py#L618-L680

I dont understand how to do partial support. If we are given inputs that we cannot evaluate with flashattention, what will happen?

[Priya2698]

Is it correct to say that before accepting this in ExprEval scheduler, we need to handle the other execution modes?

No. I mainly separated them to handle the mapping, and the ID graphs workarounds separately than the node to reduce the scope of this PR.
As we discussed in today's meeting, at the moment we only plan on supporting Flash Attention to support multi-GPU development. Once we support Flash Attention, we can revisit, if we need to add Memory Efficient Attention as well. There could be a few ways:

1. Plumbing down the backend info from Thunder and using that within our nodes: While the two implementations have different function signatures, there are overlaps and hence, one possibility is to use a superset of the inputs and outputs. The other design here would be distinct nodes for each implementation.
2. We make the decision about the backend within nvFuser using the same logic as ATen/Thunder. See: https://github.com/Lightning-AI/lightning-thunder/blob/9f0c50cc6df187cf5fd2e31240690fe2b5e9ccc1/thunder/executors/sdpaex.py#L618-L680

I dont understand how to do partial support. If we are given inputs that we cannot evaluate with flashattention, what will happen?

We will only accept the op if the backend identified is Flash Attention in Thunder: https://github.com/Lightning-AI/lightning-thunder/blob/9f0c50cc6df187cf5fd2e31240690fe2b5e9ccc1/thunder/executors/sdpaex.py#L618-L680.

Do you think this will not be sufficient?

[jacobhinkle]

Do you think this will not be sufficient?

Makes sense to me. The logic is taking place in thunder to dispatch to flash attention, which seems fine.

[naoyam]

Is anyone still reviewing this PR? @jacobhinkle?

[jacobhinkle]

LGTM once this is addressed: #2294 (comment)

[Priya2698]

!build

[Priya2698]

!build

[naoyam]

!build

[Priya2698]

!build

[Priya2698]

!build

[Priya2698]

The failing tests look unrelated.