biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[ROCm] Fix some kernels failed unit tests (#2498)

Browse Source
This commit is contained in:
Hongxia Yang
2024-02-05 17:25:36 -05:00
committed by GitHub
parent 72d3a30c63
commit 56f738ae9b
5 changed files with 62 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/kernels/test_pos_encoding.py
View File

@@ -2,7 +2,7 @@ from typing import Optional
import pytest
import torch
from allclose_default import get_default_atol, get_default_rtol
from vllm.model_executor.layers.rotary_embedding import get_rope
IS_NEOX_STYLE = [True, False]
@@ -64,5 +64,11 @@ def test_rotary_embedding(
ref_query, ref_key = rope._forward(positions, query, key)
out_query, out_key = rope.forward(positions, query, key)
# Compare the results.
assert torch.allclose(out_query, ref_query, atol=1e-5, rtol=1e-5)
assert torch.allclose(out_key, ref_key, atol=1e-5, rtol=1e-5)
assert torch.allclose(out_query,
ref_query,
atol=get_default_atol(out_query),
rtol=get_default_rtol(out_query))
assert torch.allclose(out_key,
ref_key,
atol=get_default_atol(out_key),
rtol=get_default_rtol(out_key))