pytorch - 💡(How to fix) Fix `torch.compile` produces different output for model using `torch.einsum` projections with element-wise addition pattern [1 comments, 1 participants]

import os
os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, batch_size=4, seq_len=16):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Batch-specific projection parameters
        self.q_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model))
        self.k_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model))
        self.v_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model))
        self.out_proj = nn.Parameter(torch.randn(batch_size, d_model, d_model))

        self.q_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model))
        self.k_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model))
        self.v_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model))
        self.out_bias = nn.Parameter(torch.randn(batch_size, seq_len, d_model))

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.ReLU(), nn.Linear(d_model * 4, d_model)
        )

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_norm = self.norm1(x)

        # einsum projections + element-wise add (the target pattern)
        q = torch.einsum("b i d, b d h -> b i h", x_norm, self.q_proj) + self.q_bias[:, :seq_len, :]
        k = torch.einsum("b i d, b d h -> b i h", x_norm, self.k_proj) + self.k_bias[:, :seq_len, :]
        v = torch.einsum("b i d, b d h -> b i h", x_norm, self.v_proj) + self.v_bias[:, :seq_len, :]

        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        attn = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)

        out = torch.einsum("b i d, b d h -> b i h", out, self.out_proj) + self.out_bias[:, :seq_len, :]
        out = x + out
        out = out + self.ffn(self.norm2(out))
        return out


device = "cuda"
torch.manual_seed(42)
model = MultiHeadAttention(d_model=512, num_heads=8, batch_size=4, seq_len=16).to(device).eval()
x = torch.randn(4, 16, 512, device=device)

with torch.no_grad():
    eager_out = model(x)

torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp_out = compiled(x)

diff = (eager_out.float() - comp_out.float()).abs().max().item()
print(f"Max diff: {diff}")
print(f"Match: {torch.allclose(eager_out, comp_out, atol=1e-5, rtol=1e-4)}")

---

PyTorch version: 2.12.0.dev20260327+cu126
CUDA used to build PyTorch: 12.6
OS: Ubuntu 22.04.5 LTS (x86_64) — WSL2
Python version: 3.10.12
GPU: NVIDIA GeForce RTX 3060 Laptop GPU