Skip to content

[fmhav2] skip fp8 tests and add warning #3050

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
jimmyzho merged 4 commits into from
Apr 21, 2026
+6 −21
Merged

[fmhav2] skip fp8 tests and add warning #3050

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c4c4559
skip fp8 tests and add warning
jimmyzho Apr 13, 2026
acb0192
pytest.skip
jimmyzho Apr 15, 2026
19655aa
ci
jimmyzho Apr 16, 2026
0cf9d6f
Merge branch 'main' into skip-fmhav2
jimmyzho Apr 21, 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
1 change: 1 addition & 0 deletions flashinfer/prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4458,6 +4458,7 @@ def trtllm_fmha_v2_prefill(
query.dtype == torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else False
)
if is_e4m3:
logging.warning("The FP8 (e4m3) kernels are currently known to hang on SM90.")
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The warning about FP8 kernels hanging on SM90 is issued before the check for SM120 (Blackwell) support. If a user is on an SM120 device, they will see this warning before receiving a ValueError stating that FP8 is not yet supported on their architecture. This is confusing as the hang is specific to SM90. It would be better to move the warning after the SM120 check so it only appears for SM90 users.

if is_sm12x_supported(query.device):
raise ValueError(
"FP8 (e4m3) is not yet supported for FMHAv2 on SM120 (Blackwell). "
Expand Down
26 changes: 5 additions & 21 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 @@ -3,10 +3,6 @@
import math
from typing import Optional, Tuple, Union

pytestmark = pytest.mark.skip(
reason="todo(jimmyzho): temporarily skip this test due to hangs"
)

import flashinfer
from flashinfer.prefill import fmha_v2_prefill_deepseek
from tests.utils_fp8 import to_float8
Expand Down Expand Up @@ -831,20 +827,8 @@ def test_trtllm_fmha_v2_prefill(
pos_encoding_mode: str,
save_softmax_stats: bool,
) -> None:
# skip bs=16, q_heads=4, kv_heads=4, head_dim=256, dtype=float8_e4m3fn if packed/contiguous and sliding window due to bug
if (
batch_size == 16
and num_kv_heads == 4
and head_dim == 256
and dtype == torch.float8_e4m3fn
and input_layout in ["PACKED_QKV", "CONTIGUOUS_Q_KV"]
and mask_mode == "SLIDING_WINDOW"
):
pytest.skip("Skip due to bug in fp8 sliding window")
if mask_mode == "SLIDING_WINDOW":
pytest.skip("todo(jimmyzho): temporarily skip sliding window test due to hang")
if dtype == torch.float8_e4m3fn and o_dtype == torch.float8_e4m3fn:
pytest.skip("todo(jimmyzho): temporarily skip fp8 tests due to hang")
if dtype == torch.float8_e4m3fn:
pytest.skip("FP8 (e4m3) FMHA v2 kernels are known to hang on SM90")
run_trtllm_fmha_v2_prefill_case(
input_layout=input_layout,
batch_size=batch_size,
Expand Down Expand Up @@ -875,7 +859,7 @@ def test_trtllm_fmha_v2_prefill(
[
(torch.float16, torch.float16),
(torch.bfloat16, torch.bfloat16),
(torch.float8_e4m3fn, torch.bfloat16),
(torch.float8_e4m3fn, torch.float16),
],
)
@pytest.mark.parametrize(
Expand Down Expand Up @@ -905,6 +889,8 @@ def test_trtllm_fmha_v2_prefill_skip_softmax(
rtol: float,
atol: float,
) -> None:
if dtype == torch.float8_e4m3fn:
pytest.skip("FP8 (e4m3) FMHA v2 kernels are known to hang on SM90")
run_trtllm_fmha_v2_prefill_case(
input_layout=input_layout,
batch_size=batch_size,
Expand Down Expand Up @@ -963,8 +949,6 @@ def test_trtllm_fmha_v2_prefill_attention_sinks(

if not is_sm90a_supported(torch.device("cuda")):
pytest.skip("FMHA v2 requires SM90+ (Hopper) GPUs.")
if mask_mode == "SLIDING_WINDOW":
pytest.skip("todo(jimmyzho): temporarily skip sliding window test due to hang")
torch.manual_seed(42)
device = torch.device("cuda")

Expand Down
Loading