def test_reuse_prev_layer_kv_cache(
    pos_emb_config: dict,
    attn_impl: str,
    device: str = 'cuda',
):
    """Checks reusing previous layer's kv cache."""
    alibi = pos_emb_config['alibi']
    rope = pos_emb_config['rope']

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

    n, s, f = 2, 4, cfg.d_model
    assert cfg.d_model % cfg.n_heads == 0
    cfg.kv_n_heads = 2

    sequence_id = torch.LongTensor([
        [0] * 2 + [1] * (s - 2),
        [0] * 4 + [1] * (s - 4),
    ]).to(device=device)

    # Computes its own kv cache
    cfg.reuse_kv_layer_idx = None
    attn0 = build_attention_layer(
        name='grouped_query_attention',
        attn_kwargs=dict(cfg),  # type: ignore
    ).to(device)

    # Reuses layer 0's kv cache
    cfg.reuse_kv_layer_idx = 0
    attn1 = build_attention_layer(
        name='grouped_query_attention',
        attn_kwargs=cfg,  # type: ignore
    ).to(device)
    attn0_sd = attn0.state_dict()
    attn0_sd['Wq.weight'] = attn0_sd['Wqkv.weight'][:cfg.d_model]
    attn0_sd['Wq.bias'] = attn0_sd['Wqkv.bias'][:cfg.d_model]
    del attn0_sd['Wqkv.weight']
    del attn0_sd['Wqkv.bias']
    attn1.load_state_dict(attn0_sd)

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

    attention_mask_in_length = gen_attention_mask_in_length(
        sequence_id=sequence_id,
        S=s,
        attn_uses_sequence_id=True,
        attn_impl=attn_impl,
        attention_mask=attention_mask,
    )

    flash_attn_padding_info = gen_flash_attn_padding_info(
        n,
        s,
        0,
        torch.device(device),
        attention_mask_in_length,
        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,
            cfg,
            s,
            alibi,
            True,
            device,
            sequence_id,
        )
        alibi_slopes_0 = None
        if alibi:
            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, _, prev_layer_key_value = attn0(
            x0,
            past_key_value=(),
            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,
            alibi_slopes=alibi_slopes_0,
        )
        attn_bias_1 = gen_bias(
            attn_impl,
            cfg,
            s,
            alibi,
            True,
            device,
            sequence_id,
        )
        alibi_slopes_1 = None
        if alibi:
            alibi_slopes_1 = gen_slopes(
                n_heads=cfg.n_heads,
                alibi_bias_max=8,
                device=torch.device(device),
                return_1d=True,
            )

        prev_layer_key_value = [
            t.clone().detach() for t in prev_layer_key_value
        ]
        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,
            alibi_slopes=alibi_slopes_1,
            prev_layer_key_value=prev_layer_key_value,
        )
        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():
        if 'Wq' in n:
            tp = torch_name_param_map[n.replace('Wqkv', 'Wq')]
            assert p.grad is not None
            assert tp.grad is not None
            assert allclose_helper(p[:cfg.d_model], tp)
            assert allclose_helper(p.grad[:cfg.d_model], tp.grad)
        else:
            tp = torch_name_param_map[n]
            assert p.grad is not None
            assert tp.grad is not None
            assert allclose_helper(p, tp)
            assert allclose_helper(p.grad, tp.grad)
Checks reusing previous layer's kv cache.
tests/models/layers/test_flash_torch.py