596c16a550
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Signed-off-by: Mengjie Zhao <zmj0129@gmail.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from vllm.triton_utils import HAS_TRITON
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
|
|
|
|
_HEAD_DIM = 128
|
|
_ROTARY_DIM = 96
|
|
_EPS = 1e-5
|
|
|
|
|
|
def _reference(q, k, q_weight, k_weight, rope_table):
|
|
q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
|
|
k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
|
|
half = _ROTARY_DIM // 2
|
|
cos = rope_table[..., :half].unsqueeze(1)
|
|
sin = rope_table[..., half:].unsqueeze(1)
|
|
|
|
def apply(x):
|
|
first = x[..., :half]
|
|
second = x[..., half:_ROTARY_DIM]
|
|
return torch.cat(
|
|
(
|
|
first * cos - second * sin,
|
|
second * cos + first * sin,
|
|
x[..., _ROTARY_DIM:],
|
|
),
|
|
dim=-1,
|
|
)
|
|
|
|
return apply(q), apply(k)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
|
|
@pytest.mark.parametrize("seq_len", [1, 257, 1024])
|
|
def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
|
|
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
|
|
fused_qk_norm_rope,
|
|
)
|
|
|
|
torch.manual_seed(17)
|
|
heads = 14
|
|
qkv = torch.randn(
|
|
seq_len,
|
|
heads * _HEAD_DIM * 3,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
)
|
|
q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
|
|
k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
|
|
seq_len,
|
|
heads,
|
|
_HEAD_DIM,
|
|
)
|
|
q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
|
k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
|
freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
|
|
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
|
|
|
|
expected_q, expected_k = _reference(
|
|
q,
|
|
k,
|
|
q_weight,
|
|
k_weight,
|
|
rope_table,
|
|
)
|
|
actual_q, actual_k = fused_qk_norm_rope(
|
|
q,
|
|
k,
|
|
q_weight,
|
|
k_weight,
|
|
rope_table,
|
|
_EPS,
|
|
)
|
|
|
|
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
|
|
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
|