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

[mypy] Enable type checking for test directory (#5017)

Browse Source
This commit is contained in:
Cyrus Leung
2024-06-15 12:45:31 +08:00
committed by GitHub
parent 1b8a0d71cf
commit 0e9164b40a
92 changed files with 509 additions and 378 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/kernels/test_cutlass.py
View File

@@ -17,13 +17,13 @@ capability = torch.cuda.get_device_capability()
capability = capability[0] * 10 + capability[1]
def to_fp8(tensor: torch.tensor):
def to_fp8(tensor: torch.Tensor):
finfo = torch.finfo(torch.float8_e4m3fn)
return torch.round(tensor.clamp(
min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)
def to_int8(tensor: torch.tensor):
def to_int8(tensor: torch.Tensor):
return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)