LiSsHhUuAaIi opened a new issue, #18440:
URL: https://github.com/apache/tvm/issues/18440
### Expected behavior
When compiling a PyTorch model containing `nn.MultiheadAttention` with TVM
via `torch.export`, an InternalError occurs during the legalize_ops pass. The
error indicates that TVM is attempting to convert negative infinity (-inf) to
an unsigned integer type, which is mathematically impossible.
The PyTorch model with MultiheadAttention should be successfully compiled by
TVM without internal errors, as it runs correctly in native PyTorch.
### Actual behavior
An InternalError occurs during `relax.build` with the message `cannot make
uint from negative value -inf`, indicating that TVM's legalization pass
encounters a -inf value and attempts to convert it to an unsigned integer type.
### Environment
* **OS:** Ubuntu 20.04.6 LTS
* **TVM version:** 0.23.dev0
* **Python version:** 3.11.14
### Steps to reproduce
```python
import torch
import torch.nn as nn
import tvm
from tvm import relax
class TestModel(nn.Module):
def __init__(self):
super().__init__()
self.attention = nn.MultiheadAttention(64, 4, dropout=0.1,
batch_first=True)
def forward(self, x, mask):
attn_output, _ = self.attention(x, x, x, key_padding_mask=mask)
return attn_output
model = TestModel()
model.eval()
x = torch.randn(8, 16, 64)
mask = torch.randint(0, 2, (8, 16)).bool()
# PyTorch execution works
with torch.no_grad():
output = model(x, mask)
# PyTorch export works
exported_program = torch.export.export(model, (x, mask))
# TVM conversion works
from tvm.relax.frontend.torch import from_exported_program
mod = from_exported_program(exported_program)
# TVM compilation fails
target = tvm.target.Target("llvm")
with tvm.transform.PassContext(opt_level=3):
ex = relax.build(mod, target) # InternalError here
```
### Error logs
```
Traceback (most recent call last):
...
tvm.error.InternalError: cannot make uint from negative value -inf
```
### Triage
* needs-triage
* bug
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
