Lightning-AI / Lightning-AI/lightning-thunder

Compiling torch.nn.MultiheadAttention with key_padding_mask fails in F. _canonical_mask

Open
#2,307 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

When compiling torch.nn.MultiheadAttention with a key_padding_mask, the compilation fails because the mask tensor gets converted from dtype=torch.bool to thunder.dtypes.bool8. F._canonical_mask then throws an exception.

From my explain output, it looks like multi_head_attention_forward has a registered thunder module, but the _canonical_mask check happens outside of that function.

Code snippet:

```
self.attn = torch.nn.MultiheadAttention(
embed_dim=latent_dim, num_heads=latent_dim // kv_dim, batch_first=True, dropout=dropout
)

...

key_padding_mask = ~attention_mask.bool() if attention_mask is not None else None
attn_out = self.attn(q, kv, kv, key_padding_mask=key_padding_mask)[0]

...

model = thunder.compile(model) # or model = thunder.jit(model)
```

Runtime messages:

Compiling model
Backend tkagg is interactive backend. Turning interactive mode on.
Found 22 distinct operations, of which 22 (100.0%) are supported
Note 3 operators are automatically registered:
_VariableFunctionsClass.conv_transpose1d of torch
multi_head_attention_forward of torch.nn.functional
TensorBase.bool

Stack Trace:

```
File "/home/wags/dev/lcm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/home/wags/dev/lcm/.venv/lib/python3.12/site-packages/thunder/core/interpreter.py", line 6840, in _impl
return fn.__func__(fn.__self__, *args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/home/wags/dev/lcm/.venv/lib/python3.12/site-packages/torch/nn/modules/activation.py", line 1220, in forward
key_padding_mask = F._canonical_mask(
^^^^^^^^^^^^^^^^^^
File "/home/wags/dev/lcm/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 5957, in _canonical_mask
raise AssertionError(
^^^^^^^^^^^^^^^^^^
AssertionError: only bool and floating types of key_padding_mask are supported.
```

type is thunder.dtypes.bool8 at the time of the exception.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing torch.nn.MultiheadAttention.forward into F._canonical_mask and the registered multi_head_attention_forward path, using the provided key_padding_mask example with thunder.compile or thunder.jit. Determine where bool becomes thunder.dtypes.bool8 and verify that compilation completes without the reported assertion for this reproducer.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.