def _test_correctness_prefix(
    B=2,
    H=8,
    P=512,
    C=256,
    R=256,
    D=32,
    dtype=torch.float32,
    atol=1e-3,
    rtol=5e-4,
    device=infer_device(),
):
    """
    Test that prefix sharing attention matches separate computations (i.e. two separate casual masked attention, prefix+chosen and prefix+rejected).
    The mental model is:

    A. prefix + chosen
    P
    P P
    P P P
    P P P C
    P P P C C
    P P P C C C

    B. prefix + rejected
    P
    P P
    P P P
    P P P R
    P P P R R
    P P P R R R

    C. shared prefix + chosen + rejected
    P
    P P
    P P P
    P P P C
    P P P C C
    P P P C C C
    P P P       R
    P P P       R R
    P P P       R R R


    We test them as below to ensure attention value equivalence:
    1. prefix of shared attn (upper of C.) == prefix of chosen attn (upper of A.)
    2. prefix of shared attn (upper of C.) == prefix of rejected attn (upper of B.)
    P       P
    P P   = P P
    P P P   P P P

    3. prefix of shared attn (middle right of C.) == prefix of chosen attn (lower right of A.)
    C       C
    C C   = C C
    C C C   C C C

    4. prefix of shared attn (lower right of C.) == prefix of rejected attn (lower right of B.)
    R       R
    R R   = R R
    R R R   R R R

    Args:
        B: batch size
        H: number of heads
        P: prefix length
        C: chosen response length
        R: rejected response length
        D: hidden dimension per head
    """
    torch.manual_seed(0)

    # Total sequence length for shared version
    S = P + C + R

    # Initialize input tensors, i.e. the tensors after q, k, and v projections of hidden states (attention head input)
    query = torch.randn(B, H, S, D, device=device, dtype=dtype)
    key = torch.randn(B, H, S, D, device=device, dtype=dtype)
    value = torch.randn(B, H, S, D, device=device, dtype=dtype)

    # Split tensors for separate computation
    query_prefix = query[:, :, :P, :]
    key_prefix = key[:, :, :P, :]
    value_prefix = value[:, :, :P, :]

    query_chosen = query[:, :, P : P + C, :]
    key_chosen = key[:, :, P : P + C, :]
    value_chosen = value[:, :, P : P + C, :]

    query_rejected = query[:, :, P + C :, :]
    key_rejected = key[:, :, P + C :, :]
    value_rejected = value[:, :, P + C :, :]

    chosen_index = torch.full((B,), P + C, device=device)
    rejected_index = torch.full((B,), S, device=device)

    def wrapped_prefix_mask(b, h, q_idx, kv_idx):
        return prefix_mask(b, h, q_idx, kv_idx, rejected_index, chosen_index)

    block_mask = create_block_mask(wrapped_prefix_mask, B, H, S, S, device=device)
    shared_out = flex_attention(query, key, value, block_mask=block_mask)

    # Compute attention for prefix + chosen separately
    PC = P + C
    query_pc = torch.cat([query_prefix, query_chosen], dim=2)
    key_pc = torch.cat([key_prefix, key_chosen], dim=2)
    value_pc = torch.cat([value_prefix, value_chosen], dim=2)

    def causal_mask(b, h, q_idx, kv_idx):
        return q_idx >= kv_idx

    pc_block_mask = create_block_mask(causal_mask, B, H, PC, PC, device=device)
    pc_out = flex_attention(query_pc, key_pc, value_pc, block_mask=pc_block_mask)

    # Compute attention for prefix + rejected separately
    PR = P + R
    query_pr = torch.cat([query_prefix, query_rejected], dim=2)
    key_pr = torch.cat([key_prefix, key_rejected], dim=2)
    value_pr = torch.cat([value_prefix, value_rejected], dim=2)

    pr_block_mask = create_block_mask(causal_mask, B, H, PR, PR, device=device)
    pr_out = flex_attention(query_pr, key_pr, value_pr, block_mask=pr_block_mask)

    shared_prefix = shared_out[:, :, :P, :P]
    shared_chosen = shared_out[:, :, P : P + C, P : P + C]
    shared_rejected = shared_out[:, :, P + C :, P + C :]

    separate_prefix_c = pc_out[:, :, :P, :P]
    separate_chosen = pc_out[:, :, P:, P:]
    separate_prefix_r = pr_out[:, :, :P, :P]
    separate_rejected = pr_out[:, :, P:, P:]

    # Verify prefix outputs are identical
    assert torch.allclose(shared_prefix, separate_prefix_c, atol=atol, rtol=rtol), (
        "Prefix attention from shared computation doesn't match prefix+chosen computation"
    )
    assert torch.allclose(shared_prefix, separate_prefix_r, atol=atol, rtol=rtol), (
        "Prefix attention from shared computation doesn't match prefix+rejected computation"
    )

    # Verify chosen and rejected outputs
    assert torch.allclose(shared_chosen, separate_chosen, atol=atol, rtol=rtol), (
        "Chosen response attention doesn't match between shared and separate computation"
    )
    assert torch.allclose(shared_rejected, separate_rejected, atol=atol, rtol=rtol), (
        "Rejected response attention doesn't match between shared and separate computation"
    )

    print("All attention values match between shared and separate computations!")

    Test that prefix sharing attention matches separate computations (i.e. two separate casual masked attention, prefix+chosen and prefix+rejected).
    The mental model is:

    A. prefix + chosen
    P
    P P
    P P P
    P P P C
    P P P C C
    P P P C C C

    B. prefix + rejected
    P
    P P
    P P P
    P P P R
    P P P R R
    P P P R R R

    C. shared prefix + chosen + rejected
    P
    P P
    P P P
    P P P C
    P P P C C
    P P P C C C
    P P P       R
    P P P       R R
    P P P       R R R


    We test them as below to ensure attention value equivalence:
    1. prefix of shared attn (upper of C.) == prefix of chosen attn (upper of A.)
    2. prefix of shared attn (upper of C.) == prefix of rejected attn (upper of B.)
    P       P
    P P   = P P
    P P P   P P P

    3. prefix of shared attn (middle right of C.) == prefix of chosen attn (lower right of A.)
    C       C
    C C   = C C
    C C C   C C C

    4. prefix of shared attn (lower right of C.) == prefix of rejected attn (lower right of B.)
    R       R
    R R   = R R
    R R R   R R R

    Args:
        B: batch size
        H: number of heads
        P: prefix length
        C: chosen response length
        R: rejected response length
        D: hidden dimension per head
    
test/transformers/test_flex_attention.py