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

Code before:
        ("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,
                dtype=torch.float16,
                kv_cache_dtype="fp8",
                block_size=64,
                use_per_head_quant_scales=True,
            )
            assert backend.get_name() == backend_name
        else:
            with pytest.raises(ValueError) as exc_info:
                get_attn_backend(
                    head_size=128,
                    dtype=torch.float16,
                    kv_cache_dtype="fp8",
                    block_size=64,
                    use_per_head_quant_scales=True,
                )
            assert backend_name in str(exc_info.value)


Review comment (suggestion):
By patching this I think the test isn't actually exercising the backend support, it'll just always pass as long as `supports_per_head` matches `should_succeed` in the test case. Can you get rid of this patch?
```suggestion
```

Code after:
        ("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,
                dtype=torch.float16,
                kv_cache_dtype="fp8",
                block_size=64,
                use_per_head_quant_scales=True,
            )
            assert backend.get_name() == backend_name
        else:
            with pytest.raises(ValueError) as exc_info:
                get_attn_backend(
                    head_size=128,
                    dtype=torch.float16,
                    kv_cache_dtype="fp8",
                    block_size=64,
                    use_per_head_quant_scales=True,
                )
            assert backend_name in str(exc_info.value)