Repository: vllm-project/vllm
File: tests/compile/fusions_e2e/test_tp1_quant.py (Python)

Code before:


@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
    [
        (*llama3_8b_fp8, False),
        (*qwen3_a3b_fp8, False),
        pytest.param(
            *llama4_scout_fp8,
            False,
            marks=pytest.mark.skipif(
                not current_platform.is_cuda(),
                reason="Llama4 Scout FP8 only supported on CUDA",
            ),
        ),
        pytest.param(
            *qwen3_a3b_fp8,
            True,
            marks=pytest.mark.skipif(
                not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
            ),
        ),
    ],
)
@pytest.mark.parametrize(
    "attn_backend",
    [
        TRITON_ATTN,
        FLASHINFER_ATTN,
        ROCM_ATTN,
        ROCM_AITER_UNIFIED_ATTN,
    ],
)
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
def test_tp1_fp8_fusions(
    model_name: str,
    matches_fn: Callable[[int], Matches],
    model_kwargs: dict,
    hf_overrides: Callable[[int], dict],
    attn_backend: AttentionBackendCase,
    n_layers: int,
    custom_ops: str,
    inductor_graph_partition: bool,
    use_deepgemm: bool,
    run_e2e_fusion_test,
    monkeypatch,
):
    if use_deepgemm:
        # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported

Review comment (suggestion):
you can add `attn_backend` to the parametrize list for the model, or have the model list and backends be set according to the platform e.g. https://github.com/vllm-project/vllm/blob/main/tests/compile/passes/test_fusion_attn.py#L223

Code after:
)


@pytest.mark.parametrize(
    "model_name, matches_fn, model_kwargs, hf_overrides, use_deepgemm",
    [
        (*llama3_8b_fp8, False),
        (*qwen3_a3b_fp8, False),
        pytest.param(
            *llama4_scout_fp8,
            False,
            marks=pytest.mark.skipif(
                not current_platform.is_cuda(),
                reason="Llama4 Scout FP8 only supported on CUDA",
            ),
        ),
        pytest.param(
            *qwen3_a3b_fp8,
            True,
            marks=pytest.mark.skipif(
                not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
            ),
        ),
    ],
)
@pytest.mark.parametrize(
    "attn_backend",
    [
        TRITON_ATTN,
        FLASHINFER_ATTN,
        ROCM_ATTN,
        ROCM_AITER_UNIFIED_ATTN,
    ],
)
@pytest.mark.parametrize("n_layers", [6])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
def test_tp1_fp8_fusions(
    model_name: str,
    matches_fn: Callable[[int], Matches],
    model_kwargs: dict,
    hf_overrides: Callable[[int], dict],
    attn_backend: AttentionBackendCase,
    n_layers: int,
    custom_ops: str,
    inductor_graph_partition: bool,
    use_deepgemm: bool,
    run_e2e_fusion_test,
    monkeypatch,
):
    if use_deepgemm and is_flashinfer_fp8_blockscale_gemm_supported():