Skip to content

[Executorch][llama] Allow custom sdpa op replacement pass to leverage attention mask #10285

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

Merged
merged 4 commits into from
Apr 21, 2025
Merged

[Executorch][llama] Allow custom sdpa op replacement pass to leverage attention mask #10285

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
013874c
[Executorch][llama] Allow custom sdpa op replacement pass to leverage…
kimishpatel Apr 17, 2025
e7109f4
Update on "[Executorch][llama] Allow custom sdpa op replacement pass …
kimishpatel Apr 18, 2025
2d9ba14
Update on "[Executorch][llama] Allow custom sdpa op replacement pass …
kimishpatel Apr 18, 2025
06831fd
Update on "[Executorch][llama] Allow custom sdpa op replacement pass …
kimishpatel Apr 21, 2025
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
64 changes: 50 additions & 14 deletions examples/models/llama/source_transformation/sdpa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,15 @@ class SDPACustom(torch.nn.Module):
def __init__(
self,
dim: int,
max_context_len,
enable_dynamic_shape,
use_attention_mask: bool = False,
):
super().__init__()
self.dim = dim
self.max_context_len = max_context_len
self.use_attention_mask = use_attention_mask
self.enable_dynamic_shape = enable_dynamic_shape

def forward(
self,
Expand All @@ -36,6 +42,16 @@ def forward(
seqlen,
mask,
):
if self.enable_dynamic_shape:
start_pos = input_pos[-1].item()
torch._check_is_size(start_pos)
torch._check(start_pos < self.max_context_len)
seq_length = q.size(2)
# pyre-ignore: Incompatible parameter type [6]
mask = mask.narrow(0, start_pos, seq_length)
else:
mask = mask[input_pos]

q = q.transpose(1, 2) # (bs, seqlen, n_local_heads, head_dim)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
Expand All @@ -47,34 +63,54 @@ def forward(
k = k.to(dtype=torch.float)
v = v.to(dtype=torch.float)

output = torch.ops.llama.custom_sdpa(
q,
k,
v,
input_pos[0].item(),
None, # Attention mask
0, # dropout probability. Ignored by the code
True, # is_causal
)
if self.use_attention_mask:
Copy link
Contributor

@lucylq lucylq Apr 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

nit, maybe can move this up to where we handle if self.enable_dynamic_shape, and do something similar:

if not self.use_attention_mask:
    mask = None

And then have one call to custom_sdpa

output = torch.ops.llama.custom_sdpa(
                q,
                k,
                v,
                input_pos[0].item(),
                mask,  # Attention mask
                0,  # dropout probability. Ignored by the code
                False,  # is_causal
            )

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Note that is_causal is True in one case, while not true in the other

output = torch.ops.llama.custom_sdpa(
q,
k,
v,
input_pos[0].item(),
mask, # Attention mask
0, # dropout probability. Ignored by the code
False, # is_causal
)
else:
output = torch.ops.llama.custom_sdpa(
q,
k,
v,
input_pos[0].item(),
None, # Attention mask
0, # dropout probability. Ignored by the code
True, # is_causal
)
return output.view(bsz, seqlen, self.dim).to(dtype=input_dtype)


def _replace_sdpa_with_custom_op(module: torch.nn.Module):
def _replace_sdpa_with_custom_op(
module: torch.nn.Module, use_attention_mask: bool = False
):
for name, child in module.named_children():
if isinstance(child, SDPA):
setattr(
module,
name,
SDPACustom(child.dim),
SDPACustom(
child.dim,
child.max_context_len,
child.enable_dynamic_shape,
use_attention_mask=use_attention_mask,
),
)
else:
_replace_sdpa_with_custom_op(child)
_replace_sdpa_with_custom_op(child, use_attention_mask=use_attention_mask)


def replace_sdpa_with_custom_op(module: torch.nn.Module) -> torch.nn.Module:
def replace_sdpa_with_custom_op(
module: torch.nn.Module, use_attention_mask: bool = False
) -> torch.nn.Module:
from executorch.extension.llm.custom_ops import custom_ops # noqa

_replace_sdpa_with_custom_op(module)
_replace_sdpa_with_custom_op(module, use_attention_mask=use_attention_mask)
return module


Expand Down
Loading