def test_attn_impl(
    attn_impl_0: str,
    attn_impl_1: str,
    clip_qkv: bool,
    qk_ln: bool,
    qk_gn: bool,
    pos_emb_config: dict,
    attn_type: str,
    attn_uses_sequence_id: bool,
    pad_attention_mask: bool,
    sliding_window_size: int,
    device: str = 'cuda',
):
    """Compare all attn impl with each other.

    Includes testing with and without attn_clip_qkv, attn_qk_ln, attn_qk_gn,
    alibi, and rope.
    """
    alibi = pos_emb_config['alibi']
    rope = pos_emb_config['rope']
    if alibi and not (
        check_alibi_support(attn_impl_0) and check_alibi_support(attn_impl_1)
    ):
        pytest.skip('flash attention below v2.4.2 does not support alibi.')
    if rope and (pos_emb_config['rope_impl']
                 == 'dail') and (not is_flash_v2_installed()):
        pytest.skip('dail implementation of rope requires flash attention 2.')

    if attn_uses_sequence_id and (
        attn_impl_0 == 'flash' or attn_impl_1 == 'flash'
    ) and (not is_flash_v2_installed(v2_version='v2.1.2')):
        pytest.skip(
            'Using sequence id with flash attention requires flash attention v2.1.2 or higher.',
        )

    if not (alibi or rope) and attn_uses_sequence_id:
        pytest.skip('attn_uses_sequence_id requires alibi or rope.')

    cfg = om.create({
        'attn_impl': 'flash',
        'd_model': 64,
        'n_heads': 4,
        'attn_pdrop': 0,
        'clip_qkv': clip_qkv,
        'qk_ln': qk_ln,
        'qk_gn': qk_gn,
        'sliding_window_size': sliding_window_size,
    })

    n, s, f = 2, 4, cfg.d_model
    assert cfg.d_model % cfg.n_heads == 0
    if attn_type == 'grouped_query_attention':
        cfg.kv_n_heads = 2

    sequence_id = None
    if attn_uses_sequence_id:
        assert n == 2
        assert s >= 4
        sequence_id = torch.LongTensor([
            [0] * 2 + [1] * (s - 2),
            [0] * 4 + [1] * (s - 4),
        ]).to(device=device)

    cfg.attn_impl = attn_impl_0
    attn0 = build_attention_layer(
        name=attn_type,
        attn_kwargs=om.to_container(cfg),  # type: ignore
    ).to(device)
    cfg.attn_impl = attn_impl_1
    attn1 = build_attention_layer(
        name=attn_type,
        attn_kwargs=om.to_container(cfg),  # type: ignore
    ).to(device)

    attn1.load_state_dict(attn0.state_dict())

    attention_mask = torch.ones(n, s).to(device).bool()

    if pad_attention_mask:
        # zero out the last third of the attention mask
        # to simulate padding
        attention_mask[:, -s // 3:] = 0
        if sequence_id is not None:
            sequence_id = sequence_id.masked_fill(
                ~attention_mask,
                -1,
            )  # Similar to how we set sequence id for padded tokens: https://github.com/mosaicml/llm-foundry/blob/706ea7dd40ba60a98dea5f37695d143d91c98b6c/llmfoundry/data/packing.py#L249

    attention_mask_in_length_0 = gen_attention_mask_in_length(
        sequence_id=sequence_id,
        S=s,
        attn_uses_sequence_id=attn_uses_sequence_id,
        attn_impl=attn_impl_0,
        attention_mask=attention_mask,
    )

    flash_attn_padding_info_0 = {}
    if attn_impl_0 == 'flash':
        flash_attn_padding_info_0 = gen_flash_attn_padding_info(
            n,
            s,
            0,
            torch.device(device),
            attention_mask_in_length_0,
            attention_mask,
        )

    attention_mask_in_length_1 = gen_attention_mask_in_length(
        sequence_id=sequence_id,
        S=s,
        attn_uses_sequence_id=attn_uses_sequence_id,
        attn_impl=attn_impl_1,
        attention_mask=attention_mask,
    )

    flash_attn_padding_info_1 = {}
    if attn_impl_1 == 'flash':
        flash_attn_padding_info_1 = gen_flash_attn_padding_info(
            n,
            s,
            0,
            torch.device(device),
            attention_mask_in_length_1,
            attention_mask,
        )

    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):
        attn_bias_0 = gen_bias(
            attn_impl_0,
            cfg,
            s,
            alibi,
            attn_uses_sequence_id,
            device,
            sequence_id,
        )
        alibi_slopes_0 = None
        if alibi and attn_impl_0 == 'flash':
            alibi_slopes_0 = gen_slopes(
                n_heads=cfg.n_heads,
                alibi_bias_max=8,
                device=torch.device(device),
                return_1d=True,
            )
        rotary_emb_w_meta_info = None
        if rope:
            rotary_embedding = gen_rotary_embedding(
                rope_impl=pos_emb_config['rope_impl'],
                rope_theta=pos_emb_config['rope_theta'],
                rope_dail_config=pos_emb_config.get('rope_dail_config', {}),
                rope_hf_config=pos_emb_config.get('rope_hf_config', {}),
                max_seq_len=s,
                d_model=cfg.d_model,
                n_heads=cfg.n_heads,
            ).to(device)
            pos = torch.arange(s).unsqueeze(0).to(device=device)
            # adjust the position indices to account for padding tokens
            pos = torch.clamp(
                pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1),
                min=0,
            )
            rotary_emb_w_meta_info = {
                'impl':
                    pos_emb_config['rope_impl'],
                'rotary_emb':
                    rotary_embedding,
                'offset_info':
                    pos if (pos_emb_config['rope_impl'] == 'hf') else 0,
                'seq_len':
                    s,
            }

        y0, _, _ = attn0(
            x0,
            past_key_value=None,
            attn_bias=attn_bias_0,
            attention_mask=attention_mask,
            rotary_emb_w_meta_info=rotary_emb_w_meta_info,
            is_causal=True,
            flash_attn_padding_info=flash_attn_padding_info_0,
            alibi_slopes=alibi_slopes_0,
        )
        attn_bias_1 = gen_bias(
            attn_impl_1,
            cfg,
            s,
            alibi,
            attn_uses_sequence_id,
            device,
            sequence_id,
        )
        alibi_slopes_1 = None
        if alibi and attn_impl_1 == 'flash':
            alibi_slopes_1 = gen_slopes(
                n_heads=cfg.n_heads,
                alibi_bias_max=8,
                device=torch.device(device),
                return_1d=True,
            )
        y1, _, _ = attn1(
            x1,
            past_key_value=None,
            attn_bias=attn_bias_1,
            attention_mask=attention_mask,
            rotary_emb_w_meta_info=rotary_emb_w_meta_info,
            is_causal=True,
            flash_attn_padding_info=flash_attn_padding_info_1,
            alibi_slopes=alibi_slopes_1,
        )
        y0 *= attention_mask.unsqueeze(-1)
        y1 *= attention_mask.unsqueeze(-1)

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

    loss0.backward()
    loss1.backward()

    assert allclose_helper(y0, y1)

    torch_name_param_map = dict(attn1.named_parameters())
    for n, p in attn0.named_parameters():
        tp = torch_name_param_map[n]
        assert p.grad is not None
        assert tp.grad is not None
        assert allclose_helper(p, tp)

        using_hf_rope = pos_emb_config['rope'] and pos_emb_config['rope_impl'
                                                                 ] == 'hf'

        # special case that (likely) fails due to numerics
        if (
            clip_qkv and (qk_ln or qk_gn) and using_hf_rope and
            attn_type == 'grouped_query_attention'
        ):
            assert allclose_helper(p.grad, tp.grad, atol=2.e-2, rtol=2.e-2)
        else:
            assert allclose_helper(p.grad, tp.grad)

    assert x0.grad is not None
    assert x1.grad is not None
    assert allclose_helper(x0.grad, x1.grad)
Compare all attn impl with each other.

    Includes testing with and without attn_clip_qkv, attn_qk_ln, attn_qk_gn,
    alibi, and rope.
    
tests/models/layers/test_flash_torch.py