Repository: vllm-project/vllm
File: tests/kernels/attention/test_attention_selector.py (Python)

Code before:

        # Unsupported head size
        backend = get_attn_backend(17, torch.float16, None, 16)
        assert backend.get_name() != "FLASH_ATTN"


def test_invalid_backend():
    """Test that invalid attention backend names raise ValueError."""
    with (
        pytest.raises(ValueError),
    ):
        # Invalid backend name should raise ValueError when creating enum
        AttentionConfig(backend=AttentionBackendEnum["INVALID"])


@pytest.mark.parametrize(
    "backend_name,supports_per_head,should_succeed",
    [
        ("FLASH_ATTN", True, True),  # FA3 supports per-head quant scales
        ("FLASH_ATTN", False, False),  # FA2 does not support per-head quant scales
        ("FLASHINFER", False, False),  # FlashInfer does not support
        ("FLEX_ATTENTION", False, False),  # Flex does not support
    ],
)
def test_per_head_quant_scales_backend_selection(
    backend_name: str, supports_per_head: bool, should_succeed: bool
):
    """Test backend selection when use_per_head_quant_scales=True."""
    # Clear cache to ensure fresh backend selection
    _cached_get_attn_backend.cache_clear()

    attention_config = AttentionConfig(backend=AttentionBackendEnum[backend_name])
    vllm_config = VllmConfig(attention_config=attention_config)

    # Get the backend class to patch
    prefix = "vllm.v1.attention.backends"
    backend_module = {
        "FLASH_ATTN": f"{prefix}.flash_attn.FlashAttentionBackend",
        "FLASHINFER": f"{prefix}.flashinfer.FlashInferBackend",
        "FLEX_ATTENTION": f"{prefix}.flex_attention.FlexAttentionBackend",
    }[backend_name]

    supports_attr = f"{backend_module}.supports_per_head_quant_scales"
    with (
        set_current_vllm_config(vllm_config),
        patch("vllm.platforms.current_platform", CudaPlatform()),
        patch(supports_attr, return_value=supports_per_head),
    ):
        if should_succeed:
            backend = get_attn_backend(
                head_size=128,

Review comment (suggestion):
see below
```suggestion
    backend_name: str, should_succeed: bool
```

Code after:

        # Unsupported head size
        backend = get_attn_backend(17, torch.float16, None, 16)
        assert backend.get_name() != "FLASH_ATTN"


def test_invalid_backend():
    """Test that invalid attention backend names raise ValueError."""
    with (
        pytest.raises(ValueError),
    ):
        # Invalid backend name should raise ValueError when creating enum
        AttentionConfig(backend=AttentionBackendEnum["INVALID"])


@pytest.mark.parametrize(
    "backend_name,flash_attn_version,should_succeed",
    [
        ("FLASH_ATTN", 3, True),  # FA3 supports per-head quant scales
        ("FLASH_ATTN", 2, False),  # FA2 does not support per-head quant scales
        ("FLASHINFER", None, False),  # FlashInfer does not support
        ("FLEX_ATTENTION", None, False),  # Flex does not support
    ],
)
def test_per_head_quant_scales_backend_selection(
    backend_name: str, flash_attn_version: int | None, should_succeed: bool
):
    """Test backend selection when use_per_head_quant_scales=True."""
    # Clear cache to ensure fresh backend selection
    _cached_get_attn_backend.cache_clear()

    attention_config = AttentionConfig(
        backend=AttentionBackendEnum[backend_name],
        flash_attn_version=flash_attn_version,
    )
    vllm_config = VllmConfig(attention_config=attention_config)

    with (
        set_current_vllm_config(vllm_config),
        patch("vllm.platforms.current_platform", CudaPlatform()),
    ):
        if backend_name == "FLASH_ATTN" and flash_attn_version == 3:
            if not torch.cuda.is_available():
                pytest.skip("FA3 requires CUDA")
            capability = torch.cuda.get_device_capability()
            if capability[0] != 9:
                pytest.skip("FA3 is only supported on Hopper (SM 9.x) GPUs")

        if should_succeed:
            backend = get_attn_backend(
                head_size=128,