alibaba / alibaba/TinyNeuralNetwork

[converter] implement torch's `aten::scaled_dot_product_attention` operator

Open
#265 2 comments 0 reactions 0 assignees View on GitHub
enhancement work/medium
Dominant language
Python
Stars
879
Forks
134
PR merge metrics
No merged PRs in 30d

Description

Is there any chance to implement torch `aten::scaled_dot_product_attention`?

https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html says it could be done as
```
# Efficient implementation equivalent to the following:
scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale
attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask
attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask
attn_weight = torch.softmax((Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1)
attn_weight = torch.dropout(attn_weight, dropout_p)
return attn_weight @ V
```

Contributor guide

Open the contributing guide

Research direction

Locate the converter's existing PyTorch operator handling and compare its conventions with the linked PyTorch scaled_dot_product_attention documentation. Verify completion against the documented behavior for scale, causal and explicit masks, softmax, dropout, and the final value projection, using any nearby operator tests that the repository exposes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, tooling
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.