95a9311696
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>
97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.attention.backends.ring import ring_kernels
|
|
from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
|
|
from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch):
|
|
expected_out = object()
|
|
expected_lse = object()
|
|
captured = {}
|
|
|
|
def fake_fa4(q, k, v, **kwargs):
|
|
captured.update(q=q, k=k, v=v, **kwargs)
|
|
return expected_out, expected_lse
|
|
|
|
monkeypatch.setattr(ring_kernels, "HAS_FA4", True)
|
|
monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4)
|
|
|
|
out, lse = ring_kernels.flash_attn4_func_forward(
|
|
"q",
|
|
"k",
|
|
"v",
|
|
softmax_scale=0.125,
|
|
causal=True,
|
|
window_size=(-1, -1),
|
|
)
|
|
|
|
assert (out, lse) == (expected_out, expected_lse)
|
|
assert captured == {
|
|
"q": "q",
|
|
"k": "k",
|
|
"v": "v",
|
|
"softmax_scale": 0.125,
|
|
"causal": True,
|
|
"window_size": (None, None),
|
|
"softcap": 0.0,
|
|
"return_lse": True,
|
|
}
|
|
|
|
|
|
def test_ring_prefers_fa4_on_blackwell(monkeypatch):
|
|
captured = {}
|
|
|
|
def fake_ring_flash_attn_func(*args, **kwargs):
|
|
captured.update(kwargs)
|
|
return "output"
|
|
|
|
monkeypatch.setattr(ring_parallel, "HAS_FA4", True)
|
|
monkeypatch.setattr(ring_parallel, "HAS_FA3", True)
|
|
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
|
|
|
|
from vllm_omni.diffusion.attention.backends import ring_flash_attn
|
|
|
|
monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func)
|
|
|
|
query = SimpleNamespace(
|
|
dtype=torch.bfloat16,
|
|
shape=(1, 16, 4, 128),
|
|
is_cuda=True,
|
|
device=torch.device("cuda:0"),
|
|
)
|
|
strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object()))
|
|
|
|
assert strategy.run_attention(query, "key", "value", None) == "output"
|
|
assert captured["attn_type"] is AttnType.FA4
|
|
|
|
|
|
def test_flashinfer_wrapper_canonicalizes_unbatched_lse_to_bhs(monkeypatch):
|
|
seq_len = num_heads = 4
|
|
query = torch.randn(1, seq_len, num_heads, 8)
|
|
unbatched_out = torch.randn(seq_len, num_heads, 8)
|
|
unbatched_lse = torch.arange(num_heads * seq_len, dtype=torch.float32).reshape(num_heads, seq_len)
|
|
assert not torch.equal(unbatched_lse, unbatched_lse.transpose(0, 1))
|
|
|
|
def fake_prefill(_q, _k, _v, **_kwargs):
|
|
return unbatched_out.clone(), unbatched_lse.clone()
|
|
|
|
monkeypatch.setattr(ring_kernels, "HAS_FLASHINFER", True)
|
|
# FlashInfer symbols are only bound when HAS_FLASHINFER was true at import.
|
|
monkeypatch.setattr(ring_kernels, "single_prefill_with_kv_cache", fake_prefill, raising=False)
|
|
monkeypatch.setattr(ring_kernels, "_LOG2_E", 1.0, raising=False)
|
|
|
|
out, lse = ring_kernels.flashinfer_attn_forward(query, query, query, softmax_scale=0.1)
|
|
|
|
assert out.shape == query.shape
|
|
assert lse.shape == (1, num_heads, seq_len)
|
|
assert torch.equal(lse[0], unbatched_lse)
|