def test_vs_mha(attn_impl: str, device: str = 'cuda'):
    """Compare diff attn_impl to torch.nn.MultiheadAttention."""
    from llmfoundry.models.layers import attention

    cfg = om.create({
        'attn_impl': attn_impl,
        'd_model': 64,
        'n_heads': 2,
        'attn_pdrop': 0,
        'clip_qkv': False,
        'qk_ln': False,
    })

    n, s, f = 2, 16, cfg.d_model

    mmhsa = attention.MultiheadAttention(**cfg).to(device)
    tmhsa = torch.nn.MultiheadAttention(
        embed_dim=cfg.d_model,
        num_heads=cfg.n_heads,
        dropout=cfg.attn_pdrop,
        bias=True,
        batch_first=True,
        device=device,
    )

    def gen_tca_mask():
        # generate causal mask for torch attn
        ms = (s, s)
        attn_mask = torch.empty(*ms).to(device)
        attn_mask.fill_(float('-inf'))
        attn_mask.masked_fill_(attn_mask.to(torch.bool).fill_(1).tril_(), 0.)
        return attn_mask

    # clone weights
    tmhsa.in_proj_weight.data = mmhsa.Wqkv.weight.data.clone().detach()
    tmhsa.in_proj_bias.data = mmhsa.Wqkv.bias.data.clone().detach()
    tmhsa.out_proj.weight.data = mmhsa.out_proj.weight.data.clone().detach()
    tmhsa.out_proj.bias.data = mmhsa.out_proj.bias.data.clone().detach()

    attention_mask = torch.ones(n, s).to(device).bool()
    x0 = torch.randn(n, s, f).to(device)
    x1 = x0.clone().detach()
    x0.requires_grad = True
    x1.requires_grad = True

    with torch.autocast(x0.device.type):
        flash_attn_padding_info = None
        if attn_impl == 'flash':
            flash_attn_padding_info = gen_flash_attn_padding_info(
                n,
                s,
                0,
                torch.device(device),
                None,
                attention_mask,
            )
        y0, _, _ = mmhsa(
            x0,
            past_key_value=None,
            attn_bias=None,
            attention_mask=attention_mask,
            is_causal=True,
            flash_attn_padding_info=flash_attn_padding_info,
        )
        y1, _ = tmhsa(
            x1,
            x1,
            x1,
            attn_mask=gen_tca_mask(),
            key_padding_mask=~attention_mask,
            need_weights=True,
        )
        y0 *= attention_mask.unsqueeze(-1)
        y1 *= attention_mask.unsqueeze(-1)

        loss0 = y0.sum()
        loss1 = y1.sum()

    loss0.backward()
    loss1.backward()

    assert y0 is not None
    assert y1 is not None
    assert tmhsa.out_proj.bias.grad is not None
    assert mmhsa.out_proj.bias.grad is not None
    assert tmhsa.out_proj.weight.grad is not None
    assert mmhsa.out_proj.weight.grad is not None
    assert tmhsa.in_proj_bias.grad is not None
    assert mmhsa.Wqkv.bias.grad is not None
    assert tmhsa.in_proj_weight.grad is not None
    assert mmhsa.Wqkv.weight.grad is not None
    assert x0.grad is not None
    assert x1.grad is not None

    assert allclose_helper(y0, y1)

    assert allclose_helper(tmhsa.out_proj.bias.grad, mmhsa.out_proj.bias.grad)
    assert allclose_helper(
        tmhsa.out_proj.weight.grad,
        mmhsa.out_proj.weight.grad,
    )
    assert allclose_helper(tmhsa.in_proj_bias.grad, mmhsa.Wqkv.bias.grad)
    assert allclose_helper(tmhsa.in_proj_weight.grad, mmhsa.Wqkv.weight.grad)

    assert allclose_helper(x0.grad, x1.grad)
Compare diff attn_impl to torch.nn.MultiheadAttention.
tests/models/layers/test_flash_torch.py