Skip to content

test: align test_fmha_v2_prefill SM gating with is_sm12x_supported #3182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kahyunnam merged 2 commits into from
May 28, 2026
Merged

test: align test_fmha_v2_prefill SM gating with is_sm12x_supported #3182

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e170b97
test: align test_fmha_v2_prefill SM gating with is_sm12x_supported
leonardHONG Apr 26, 2026
c9e2c3f
Merge branch 'main' into fix/3170-fmha-v2-test-sm12x
kahyunnam May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 9 additions & 13 deletions tests/attention/test_fmha_v2_prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from flashinfer.prefill import fmha_v2_prefill_deepseek
from tests.utils_fp8 import to_float8
from flashinfer.utils import is_sm12x_supported, is_sm120a_supported
from flashinfer.utils import is_sm12x_supported

_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
_workspace_buffer: Optional[torch.Tensor] = None
Expand Down Expand Up @@ -484,27 +484,23 @@ def run_trtllm_fmha_v2_prefill_case(
from flashinfer.prefill import trtllm_fmha_v2_prefill
from flashinfer.utils import is_sm90a_supported

if not is_sm90a_supported(torch.device("cuda")) and not is_sm120a_supported(
if not is_sm90a_supported(torch.device("cuda")) and not is_sm12x_supported(
torch.device("cuda")
):
pytest.skip("FMHA v2 requires SM90+ (Hopper) or SM12x GPUs.")

# Skip invalid combinations
is_sm120_plus = is_sm120a_supported(torch.device("cuda"))
if dtype == torch.float8_e4m3fn and is_sm120_plus:
pytest.skip("FP8 FMHAv2 not yet supported on SM120+")
is_sm12x = is_sm12x_supported(torch.device("cuda"))
if dtype == torch.float8_e4m3fn and is_sm12x:
pytest.skip("FP8 FMHAv2 not yet supported on SM12x")
if input_layout == "SEPARATE_Q_K_V" and dtype == torch.float8_e4m3fn:
pytest.skip("FP8 not supported for SEPARATE_Q_K_V layout")
if input_layout == "SEPARATE_Q_K_V" and is_sm120_plus:
if input_layout == "SEPARATE_Q_K_V" and is_sm12x:
pytest.skip(
"SEPARATE_Q_K_V requires SM90 warp-specialization, not available on SM120+"
"SEPARATE_Q_K_V requires SM90 warp-specialization, not available on SM12x"
)
if (
is_sm120_plus
and mask_mode is not None
and mask_mode.upper() == "SLIDING_WINDOW"
):
pytest.skip("SLIDING_WINDOW mask not yet supported on SM120+ (only causal)")
if is_sm12x and mask_mode is not None and mask_mode.upper() == "SLIDING_WINDOW":
pytest.skip("SLIDING_WINDOW mask not yet supported on SM12x (only causal)")
if input_layout == "SEPARATE_Q_K_V" and logits_soft_cap > 0:
pytest.skip("Logits soft capping not supported for SEPARATE_Q_K_V layout")
# save_softmax_stats only supported for CONTIGUOUS_Q_KV (normal attention)
Expand Down
Loading