Skip to content

Add BF16 support to custom PA #133

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
merged 12 commits into from
Aug 14, 2024
Merged

Add BF16 support to custom PA #133

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
4d16a0d
tightened atol for custom PA; enable supported head size, block sizes…
sanyalington Aug 12, 2024
f705795
update num_blocks and num_iters in benchmark PA to realistic settings
sanyalington Aug 12, 2024
b092f58
move to generic b16 type
sanyalington Aug 12, 2024
5e1ff77
bf16 first port
sanyalington Aug 12, 2024
2256ced
enabled all bf16 tests, set atol for bf16
sanyalington Aug 12, 2024
901fef9
enable custom PA for bf16 as well as block size 32 and head size 64
sanyalington Aug 12, 2024
1e85c9b
fix cast to zero in custom PA reduce
sanyalington Aug 13, 2024
afd6689
Merge branch 'main' into shsanyal_dev_PA
charlifu Aug 13, 2024
d539011
py linter fixes
sanyalington Aug 14, 2024
69489ea
clang format fixes
sanyalington Aug 14, 2024
69a6d82
div round up clang-format
sanyalington Aug 14, 2024
0a3f007
Merge branch 'main' into shsanyal_dev_PA
gshtras Aug 14, 2024
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
4 changes: 2 additions & 2 deletions benchmarks/kernels/benchmark_paged_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from vllm._custom_C import paged_attention_custom
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, create_kv_caches_with_random

NUM_BLOCKS = 1024
NUM_BLOCKS = 1024*1024
PARTITION_SIZE = 256


Expand Down Expand Up @@ -176,7 +176,7 @@ def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
if do_profile:
latency = run_benchmark(num_iters=1, profile=True)
else:
latency = run_benchmark(num_iters=100, profile=False)
latency = run_benchmark(num_iters=1000, profile=False)
print(f"Kernel running time: {latency * 1000000:.3f} us")


Expand Down
Loading