If you uncomment the prints, the script works without errors (even though the outputs are incorrect)

weird_shape_errors_without_prints.py

from collections import OrderedDict
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.nn.attention.flex_attention import flex_attention, BlockMask


def norm(x):
    return F.rms_norm(x, (x.size(-1),))


class NoBiasLinear(nn.Linear):
    def __init__(self, in_features: int, out_features: int, device: torch.device | None, dtype: torch.dtype | None):
        super().__init__(in_features, out_features, bias=False, device=device, dtype=dtype)

    def forward(self, x):
        return super().forward(x)


class Rotary(torch.nn.Module):

    def __init__(self, dims: int, base: float, device: torch.device, dtype: torch.dtype):
        super().__init__()
        self.dims = dims
        self.base = base
        self.device = device
        self.dtype = dtype
        self.inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dims, 2, device=self.device, dtype=torch.float) / self.dims))

    def forward(
            self,
            x: torch.Tensor,  # (B, L, H)
            x_pos: torch.Tensor,  # (B, L)
    ) -> torch.Tensor:
        
        if x.shape[0] != 1:
            # FIXME: torch.outer will not work otherwise
            raise NotImplementedError() 
        
        x_dtype = x.dtype

        t = x_pos.type_as(self.inv_freq).squeeze(0)
        freqs = torch.outer(t, self.inv_freq)

        cos_freq = freqs.cos().to(x_dtype)[None, :, None, :]
        sin_freq = freqs.sin().to(x_dtype)[None, :, None, :]

        # apply_rotary_emb(x, cos, sin)
        assert x.ndim == 4  # multihead attention

        d = x.shape[3] // 2
        x1 = x[..., :d]
        x2 = x[..., d:]
        y1 = x1 * cos_freq + x2 * sin_freq
        y2 = x1 * (-sin_freq) + x2 * cos_freq

        return torch.cat([y1, y2], 3).to(x_dtype)



def lambda_init_fn(depth):
    return 0.8 - 0.6 * math.exp(-0.3 * depth)




class DiffFlexSelfAttention(nn.Module):
    def __init__(
            self, 
            layer_idx: int,
            n_q_heads: int,
            n_kv_heads: int,
            head_dims: int,
            rotary_inv_freq_base: float,
            device: torch.device,
            dtype: torch.dtype
        ):
        super().__init__()

        extras = {'device': device, 'dtype': dtype}

        self.layer_idx = layer_idx
        self.gqa_enabled = (n_q_heads != n_kv_heads)

        # for keys and queries we split each head in 2
        self.n_q_heads = n_q_heads * 2
        self.n_k_heads = n_kv_heads * 2

        self.q_head_dims = head_dims // 2
        self.k_head_dims = head_dims // 2

        self.n_v_heads = n_kv_heads
        self.v_head_dims = head_dims
        self.hidden_dims = head_dims * n_q_heads

        self.q_dims = self.n_q_heads * self.q_head_dims
        self.k_dims = self.n_k_heads * self.k_head_dims
        self.v_dims = self.n_v_heads * self.v_head_dims

        self.q_proj = NoBiasLinear(self.hidden_dims, self.q_dims, **extras)
        self.k_proj = NoBiasLinear(self.hidden_dims, self.k_dims, **extras)
        self.v_proj = NoBiasLinear(self.hidden_dims, self.v_dims, **extras)

        self.rotary = Rotary(self.q_head_dims, rotary_inv_freq_base, **extras)

        self.out_proj = NoBiasLinear(self.hidden_dims, self.hidden_dims, **extras)
        self.out_proj.weight.data.zero_()

        # see https://github.com/microsoft/unilm/blob/master/Diff-Transformer/multihead_flashdiff_1.py
        self.lambda_init = lambda_init_fn(self.layer_idx)

        half_head = self.q_head_dims  # same as k_head_dims
        self.lambda_q1 = nn.Parameter(torch.zeros(half_head, dtype=torch.float32).normal_(mean=0,std=0.1))
        self.lambda_k1 = nn.Parameter(torch.zeros(half_head, dtype=torch.float32).normal_(mean=0,std=0.1))
        self.lambda_q2 = nn.Parameter(torch.zeros(half_head, dtype=torch.float32).normal_(mean=0,std=0.1))
        self.lambda_k2 = nn.Parameter(torch.zeros(half_head, dtype=torch.float32).normal_(mean=0,std=0.1))

    
    def forward(
            self,
            x: torch.Tensor,  # (B, n_mem + L, H)
            # x_pos: torch.Tensor,  # (B, n_mem + L)
            block_mask: BlockMask | None = None,
    ) -> torch.Tensor:
        
        batch_size, seq_length, _ = x.shape
        x_pos = torch.arange(seq_length, device='cuda').unsqueeze(0)  # for simplicity

        k = self.k_proj(x).view(batch_size, seq_length, self.n_k_heads, self.k_head_dims)
        v = self.v_proj(x).view(batch_size, seq_length, self.n_v_heads, self.v_head_dims)
        q = self.q_proj(x).view(batch_size, seq_length, self.n_q_heads, self.q_head_dims)

        # QK norm
        k = self.rotary(norm(k), x_pos)
        q = self.rotary(norm(q), x_pos)

        # see https://github.com/microsoft/unilm/blob/master/Diff-Transformer/multihead_flashdiff_1.py
        q = q.reshape(batch_size, seq_length, self.n_q_heads // 2, 2, self.q_head_dims)
        k = k.reshape(batch_size, seq_length, self.n_k_heads // 2, 2, self.k_head_dims)

        q1, q2 = q[:, :, :, 0, :], q[:, :, :, 1, :]
        k1, k2 = k[:, :, :, 0, :], k[:, :, :, 1, :]

        attn1 = flex_attention(
            q1.transpose(1, 2).contiguous(),
            k1.transpose(1, 2).contiguous(),
            v.transpose(1, 2).contiguous(),
            block_mask=block_mask,
            enable_gqa=self.gqa_enabled
        )
        # print(q1.shape, k1.shape, v.shape, '->', attn1.shape)

        attn2 = flex_attention(
            q2.transpose(1, 2).contiguous(),
            k2.transpose(1, 2).contiguous(),
            v.transpose(1, 2).contiguous(),
            block_mask=block_mask,
            enable_gqa=self.gqa_enabled
        )
        # print(q2.shape, k2.shape, v.shape, '->', attn2.shape)
        
        lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(q1)
        lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(q2)
        lambda_full = lambda_1 - lambda_2 + self.lambda_init

        attn = norm(attn1 - lambda_full * attn2) * (1 - self.lambda_init)
        attn = attn.transpose(1, 2).contiguous().view(batch_size, seq_length, self.hidden_dims)

        return self.out_proj(attn)


torch.set_float32_matmul_precision('medium')
    

model = nn.Sequential(OrderedDict({
    f'layer_{idx}': DiffFlexSelfAttention(
        layer_idx=idx,
        n_q_heads=8,
        n_kv_heads=4,
        head_dims=128,
        rotary_inv_freq_base=500.0,
        device='cuda',
        dtype=torch.float
    ) for idx in range(12)
}))

model = torch.compile(model)

model(
    torch.randn((1, 10000, 128 * 8), device='cuda', dtype=torch.float),
).sum().backward()