Skip to content

[Infer] add check triton and cuda version for tests #4627

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Xu-Kai merged 10 commits into from
Sep 6, 2023
Merged

[Infer] add check triton and cuda version for tests #4627

Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
c7d6d8d
fix rotary embedding
Xu-Kai Sep 4, 2023
0eadde2
merge colossal inference
Xu-Kai Sep 5, 2023
64c0782
delete print
Xu-Kai Sep 5, 2023
64d8d55
fix init seq len bug
Xu-Kai Sep 5, 2023
a255925
rename pytest
Xu-Kai Sep 5, 2023
c525e78
add benchmark for llama
Xu-Kai Sep 5, 2023
117cdf1
refactor codes
Xu-Kai Sep 5, 2023
2af8002
delete useless code
Xu-Kai Sep 5, 2023
ae4b162
Merge remote-tracking branch 'upstream/feature/colossal-inference' in…
Xu-Kai Sep 5, 2023
f3e3046
add check triton and cuda
Xu-Kai Sep 6, 2023
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
6 changes: 6 additions & 0 deletions tests/test_infer_ops/triton/test_rotary_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ def torch_rotary_emb(x, cos, sin):
return torch.cat((o0, o1), dim=-1)


@pytest.mark.skipif(not TRITON_CUDA_SUPPORT or not HAS_TRITON,
reason="triton requires cuda version to be higher than 11.4")
def test_rotary_emb():
SEQ_LEN = 1
HEAD_NUM = 32
Expand All @@ -50,3 +52,7 @@ def test_rotary_emb():
# print("max delta:", torch.max(torch.abs(y_torch - y_triton)))
# compare
assert torch.allclose(y_torch, y_triton, atol=1e-2, rtol=1e-2)


if __name__ == "__main__":
test_rotary_emb()