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

Code before:
                            head_size=head_size,
                            causal=causal,
                            past_sequence_length=0,
                            kv_sequence_length=sequence_length,
                            max_cache_sequence_length=None,
                            provider=provider,
                            device=device,
                            dtype=dtype,
                            use_kv_cache=False,
                            share_past_present_buffer=False,
                            input_format=format,
                            has_bias=has_bias,
                            mask_format=mask_format,
                            has_attn_bias=has_attn_bias,
                            broadcast_attn_bias_dim_0=broadcast_attn_bias_dim_0,
                            broadcast_attn_bias_dim_1=broadcast_attn_bias_dim_1,
                        )
                        yield config


def 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, 15, 16, 255, 256, 384, 512]
    heads = [1, 2, 3, 4, 16]
    head_sizes = [8, 16, 32, 40, 64, 80, 96, 128, 160, 224, 256]

    device, dtypes, formats = get_provider_support_info(provider, True)
    mask_formats = [
        AttentionMaskFormat.Mask_None,
        AttentionMaskFormat.Mask_1D_Key_SeqLen,
        AttentionMaskFormat.Mask_2D_Key_PaddingMask,
    ]

    atten_bias_options = get_atten_bias_support()

    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 past_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 has_past_input in [True, False]:
                                    for mask_format in mask_formats:
                                        for has_bias in get_bias_support(format):

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

Code after:
                            head_size=head_size,
                            causal=causal,
                            past_sequence_length=0,
                            kv_sequence_length=sequence_length,
                            max_cache_sequence_length=None,
                            provider=provider,
                            device=device,
                            dtype=dtype,
                            use_kv_cache=False,
                            share_past_present_buffer=False,
                            input_format=format,
                            has_bias=has_bias,
                            mask_format=mask_format,
                            has_attn_bias=has_attn_bias,
                            broadcast_attn_bias_dim_0=broadcast_attn_bias_dim_0,
                            broadcast_attn_bias_dim_1=broadcast_attn_bias_dim_1,
                        )
                        yield config


def 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, 15, 16, 255, 256, 384, 512]
    heads = [1, 2, 3, 4, 16]
    head_sizes = [8, 16, 32, 40, 64, 80, 96, 128, 160, 224, 256]

    device, dtypes, formats = get_provider_support_info(provider, True)
    mask_formats = [
        AttentionMaskFormat.Mask_None,
        AttentionMaskFormat.Mask_1D_Key_SeqLen,
        AttentionMaskFormat.Mask_2D_Key_PaddingMask,
    ]

    atten_bias_options = get_atten_bias_support()

    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 past_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 has_past_input in [True, False]:
                                    for mask_format in mask_formats:
                                        for has_bias in get_bias_support(format):