Debug a Transformer bug in an unfamiliar repo

The Problem

You are working on a new generative Transformer architecture. During training, the loss is not converging, and the model produces NaNs after a few hundred iterations. You have isolated the issue to the SelfAttention module. Below is the provided implementation snippet:

import torch
import torch.nn.functional as F

class SelfAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv_proj = torch.nn.Linear(d_model, 3 * d_model)

    def forward(self, x):
        # x shape: [batch_size, seq_len, d_model]
        batch_size, seq_len, _ = x.shape
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)
        
        # Reshape for multi-head: [batch, heads, seq_len, head_dim]
        q = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)

        # Attention scores
        attn_scores = torch.matmul(q, k.transpose(-2, -1))
        
        # Causal mask
        mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len)
        attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax
        attn_weights = F.softmax(attn_scores, dim=-1)
        return torch.matmul(attn_weights, v)

Your Tasks

1. Identify the logical flaw(s) in the implementation that cause numerical instability and incorrect causal behavior.
2. Provide the corrected implementation for the forward pass.
3. Briefly explain why float('-inf') might be problematic if not handled alongside the scaling factor.

Constraints

• Assume d_model is always divisible by n_heads.
• The input x uses standard float32 precision.
• The model is intended for autoregressive sequence generation.