Skip to content
Merged
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
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,
Expand Down Expand Up @@ -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")
Expand All @@ -181,8 +183,7 @@ def test_length_lives_in_device_tensors():

@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

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -413,8 +411,7 @@ def forward(self, input_ids, indexes, past_key_values=None, use_cache=False, pag


@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
Expand Down Expand Up @@ -556,10 +553,12 @@ def current_kernel(
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)
Expand All @@ -572,8 +571,8 @@ def recorder(*args, **kwargs):
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",
Expand Down Expand Up @@ -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")
Expand Down
Loading