pytorch - 💡(How to fix) Fix `torch.compile` produces different output for single-head attention with `matmul → div → softmax → dropout → matmul` pattern (SFDP pattern 11) compared to eager mode [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
import math

class SingleHeadAttention(nn.Module):
    def __init__(self, embed_dim=256, dropout_p=0.0):
        super().__init__()
        self.embed_dim = embed_dim
        self.scale = math.sqrt(embed_dim)
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        query = self.q_proj(x)
        key = self.k_proj(x)
        value = self.v_proj(x)

        # Single-head: unsqueeze head dim + permute to [B, 1, S, D]
        query = query.unsqueeze(1).permute(0, 1, 2, 3)  # [B, 1, S, D]
        key = key.unsqueeze(1).permute(0, 1, 2, 3)
        value = value.unsqueeze(1).permute(0, 1, 2, 3)

        # SFDP pattern 11: matmul → div(scale) → softmax → dropout → matmul
        attn_scores = torch.matmul(query, key.transpose(-2, -1))
        attn_scores = attn_scores / self.scale
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        context = torch.matmul(attn_weights, value)

        # Remove head dim and project
        context = context.squeeze(1)  # [B, S, D]
        return self.out_proj(context)


class SingleHeadTransformer(nn.Module):
    def __init__(self, vocab_size=10000, embed_dim=256, max_seq_len=64, dropout_p=0.0):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb = nn.Embedding(max_seq_len, embed_dim)
        self.attn = SingleHeadAttention(embed_dim, dropout_p=dropout_p)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim),
        )

    def forward(self, token_ids):
        batch_size, seq_len = token_ids.shape
        positions = torch.arange(seq_len, device=token_ids.device).unsqueeze(0).expand(batch_size, -1)
        x = self.token_emb(token_ids) + self.pos_emb(positions)

        # Pre-norm residual attention
        x_norm = self.norm1(x)
        attn_out = self.attn(x_norm)
        x = x + attn_out

        # Pre-norm residual FFN
        x = x + self.ffn(self.norm2(x))
        return x


device = "cuda"
torch.manual_seed(42)
model = SingleHeadTransformer(
    vocab_size=10000, embed_dim=256, max_seq_len=64, dropout_p=0.0
).to(device).eval()
token_ids = torch.randint(0, 10000, (4, 16), device=device)

# Eager: deterministic
with torch.no_grad():
    ref = model(token_ids)
    ref2 = model(token_ids)
print(f"Eager deterministic: {(ref - ref2).abs().max().item():.6e}")

# Compiled
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp = compiled(token_ids)

diff = (ref.float() - comp.float()).abs()
print(f"Max diff: {diff.max().item():.6e}")
print(f"Mean diff: {diff.mean().item():.6e}")
print(f"Match (atol=1e-5): {torch.allclose(ref, comp, atol=1e-5, rtol=1e-4)}")

---

Eager deterministic: 0.000000e+00
Max diff and mean diff show systematic numerical divergence from SFDP fusion

---

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