Repository: microsoft/onnxruntime
File: onnxruntime/test/python/transformers/test_mha.py (Python)

Code before:
    if past_v is not None:
        assert past_v.dim() == 4 and v.dim() == 4 and past_v.size(1) == v.size(1), (
            f"expect BNSH format: {past_v.shape=} {v.shape=}"
        )

    present_k = torch.cat((past_k, k), dim=2) if past_k is not None else k
    present_v = torch.cat((past_v, v), dim=2) if past_v is not None else v
    out = attention_reference(config.head_size, q, present_k, present_v, scale=scale, attn_bias=attn_bias, mask=mask)

    return out, present_k, present_v


def get_compute_capability():
    if torch.cuda.is_available() and "CUDAExecutionProvider" in onnxruntime.get_available_providers():
        major, minor = torch.cuda.get_device_capability()
        sm = major * 10 + minor
        return sm
    return 0


def no_kv_cache_test_cases(provider: str, comprehensive: bool):
    if provider == "CUDAExecutionProvider" and get_compute_capability() < 60:
        return
        yield

    # Lengths of arraies are prime numbers since modulo (% length) is used in non comprehensive mode.
    batch_sizes = [1, 2, 3]
    sequence_lengths = [1, 16, 127, 128, 256, 384, 512]
    heads = [1, 2, 3, 4, 16]
    head_sizes = [8, 16, 32, 40, 64, 80, 96, 128, 160, 192, 256]

    mask_formats = [
        AttentionMaskFormat.Mask_None,
        AttentionMaskFormat.Mask_1D_Key_SeqLen,
        AttentionMaskFormat.Mask_2D_Key_PaddingMask,
    ]
    atten_bias_options = get_atten_bias_support()

    device, dtypes, formats = get_provider_support_info(provider, False)
    if comprehensive:
        sequence_lengths = [*sequence_lengths, 2048]  # Large sequence length is slow and need a lot of memory
        for batch_size in batch_sizes:
            for sequence_length in sequence_lengths:
                for num_heads in heads:
                    for head_size in head_sizes:
                        for format in formats:
                            for causal in get_causal_support(format):
                                for mask_format in mask_formats:
                                    for has_bias in get_bias_support(format):
                                        for (
                                            has_attn_bias,

Review comment (suggestion):
```suggestion
    # Lengths of arrays are prime numbers since modulo (% length) is used in non comprehensive mode.
```

Code after:
    if past_v is not None:
        assert past_v.dim() == 4 and v.dim() == 4 and past_v.size(1) == v.size(1), (
            f"expect BNSH format: {past_v.shape=} {v.shape=}"
        )

    present_k = torch.cat((past_k, k), dim=2) if past_k is not None else k
    present_v = torch.cat((past_v, v), dim=2) if past_v is not None else v
    out = attention_reference(config.head_size, q, present_k, present_v, scale=scale, attn_bias=attn_bias, mask=mask)

    return out, present_k, present_v


def get_compute_capability():
    if torch.cuda.is_available() and "CUDAExecutionProvider" in onnxruntime.get_available_providers():
        major, minor = torch.cuda.get_device_capability()
        sm = major * 10 + minor
        return sm
    return 0


def no_kv_cache_test_cases(provider: str, comprehensive: bool):
    if provider == "CUDAExecutionProvider" and get_compute_capability() < 60:
        return
        yield

    # Lengths of arrays are prime numbers since modulo (% length) is used in non comprehensive mode.
    batch_sizes = [1, 2, 3]
    sequence_lengths = [1, 16, 127, 128, 256, 384, 512]
    heads = [1, 2, 3, 4, 16]
    head_sizes = [8, 16, 32, 40, 64, 80, 96, 128, 160, 192, 256]

    mask_formats = [
        AttentionMaskFormat.Mask_None,
        AttentionMaskFormat.Mask_1D_Key_SeqLen,
        AttentionMaskFormat.Mask_2D_Key_PaddingMask,
    ]
    atten_bias_options = get_atten_bias_support()

    device, dtypes, formats = get_provider_support_info(provider, False)
    if comprehensive:
        sequence_lengths = [*sequence_lengths, 2048]  # Large sequence length is slow and need a lot of memory
        for batch_size in batch_sizes:
            for sequence_length in sequence_lengths:
                for num_heads in heads:
                    for head_size in head_sizes:
                        for format in formats:
                            for causal in get_causal_support(format):
                                for mask_format in mask_formats:
                                    for has_bias in get_bias_support(format):
                                        for (
                                            has_attn_bias,