diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py index 368852a1d..ac517945d 100644 --- a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py @@ -17,6 +17,7 @@ import torch from vllm.lora.layers import BaseLayerWithLoRA import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.models.sensenova_u1.paged_decode import ( BLOCK_SIZE, BUCKETS, @@ -169,7 +170,8 @@ def test_length_lives_in_device_tensors(): # --------------------------------------------------------------------------- -# The capture contract. These need CUDA and the bundled flash-attention kernel. +# The capture contract. Graph tests need a torch.cuda-compatible accelerator; +# paged-attention equivalence also needs the bundled kernel's paged API. # --------------------------------------------------------------------------- cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -181,8 +183,7 @@ vllm_flash_attn_only = pytest.mark.skipif( @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_paged_attention_matches_sdpa_over_the_same_buffer(): import torch.nn.functional as F @@ -223,8 +224,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_captured_graph_follows_a_later_length(): """One capture has to serve every length in the bucket. @@ -285,8 +285,7 @@ def test_a_captured_graph_follows_a_later_length(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_captured_graph_writes_each_step_to_its_own_slot(): """The write slot has to be a device tensor, not a Python index. @@ -347,8 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_whole_decode_loop_matches_the_unpaged_path(): """Per-step kernel equivalence is not loop equivalence. @@ -413,8 +411,7 @@ class _StubLM(torch.nn.Module): @cuda_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch): """``_decode_context`` builds a runner per request, so a private capture pool hands every request its own arena and never gives it back. Measured on a @@ -556,10 +553,12 @@ def test_a_kernel_missing_any_kwarg_we_pass_is_not_supported(monkeypatch): def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch): """A positional standard argument binds to whatever a future wheel puts in that slot, and the probe above only checks that the names exist.""" - seen = {} + seen_args: tuple[object, ...] = () + seen_kwargs: dict[str, object] = {} - def recorder(*args, **kwargs): - seen["args"], seen["kwargs"] = args, kwargs + def recorder(*args: object, **kwargs: object): + nonlocal seen_args, seen_kwargs + seen_args, seen_kwargs = args, kwargs return torch.zeros(1, N_HEADS, HEAD_DIM) monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder) @@ -572,8 +571,8 @@ def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch): torch.zeros(1, KV_HEADS, 1, HEAD_DIM), 1.0, ) - assert len(seen["args"]) == 3, "only q, k and v may be positional" - assert set(seen["kwargs"]) == { + assert len(seen_args) == 3, "only q, k and v may be positional" + assert set(seen_kwargs) == { "max_seqlen_q", "cu_seqlens_q", "max_seqlen_k", @@ -694,8 +693,7 @@ def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypat @cuda_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_second_request_replays_the_first_capture(monkeypatch): """Reuse is only worth anything if the capture survives it.""" dev = torch.device("cuda")