Skip to content
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

Faster Custom Paged Attention kernels #372

Open
wants to merge 14 commits into
base: main
Choose a base branch
from
Open

Faster Custom Paged Attention kernels #372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
42d4c5c
integrate new cpa kernel, update tests and benchmark
sanyalington Jan 17, 2025
ba5df8b
added comments to mfma4 kernel
sanyalington Jan 20, 2025
4fdcd75
further comments for mfma16 kernel
sanyalington Jan 20, 2025
bf47251
Merge branch 'main' into shsanyal_cpa_main_integration
sanyalington Jan 21, 2025
45945a7
Merge branch 'main' into shsanyal_cpa_main_integration
gshtras Jan 21, 2025
72cf38e
Merge remote-tracking branch 'origin/main' into shsanyal_cpa_main_int…
gshtras Jan 22, 2025
6f8e708
clang-format
gshtras Jan 22, 2025
7aa547f
Lint
gshtras Jan 22, 2025
62755f3
Merge branch 'main' into shsanyal_cpa_main_integration
gshtras Jan 23, 2025
4d0166c
add flag for logits rtz conversion and disable by default
sanyalington Jan 23, 2025
52a0b95
lint
sanyalington Jan 23, 2025
5510e8c
Merge branch 'main' into shsanyal_cpa_main_integration
shajrawi Jan 24, 2025
49dfc1d
[Bugfix]: Fix paged attention unit tests of https://github.com/ROCm/v…
tjtanaa Jan 28, 2025
e2fe1db
Merge remote-tracking branch 'origin/main' into shsanyal_cpa_main_int…
gshtras Jan 29, 2025
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
16 changes: 11 additions & 5 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,8 +9,9 @@
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser,
create_kv_caches_with_random)

NUM_BLOCKS = 1024 * 1024
NUM_BLOCKS = 128 * 1024
PARTITION_SIZE = 512
PARTITION_SIZE_ROCM = 256


@torch.inference_mode()
Expand Down Expand Up @@ -78,9 +79,12 @@ def main(
# Prepare for the paged attention kernel.
output = torch.empty_like(query)
if version == "v2":
if current_platform.is_rocm() and not args.custom_paged_attn:
if current_platform.is_rocm():
global PARTITION_SIZE
PARTITION_SIZE = 1024
if not args.custom_paged_attn:
PARTITION_SIZE = 1024
else:
PARTITION_SIZE = PARTITION_SIZE_ROCM
num_partitions = ((max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE)
tmp_output = torch.empty(
size=(num_seqs, num_query_heads, num_partitions, head_size),
Expand Down Expand Up @@ -163,6 +167,8 @@ def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
kv_cache_dtype,
k_scale,
v_scale,
None,
PARTITION_SIZE,
)
else:
raise ValueError(f"Invalid version: {version}")
Expand All @@ -176,13 +182,13 @@ def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
# Warmup.
print("Warming up...")
run_benchmark = run_cuda_benchmark
run_benchmark(num_iters=3, profile=False)
run_benchmark(num_iters=500, profile=False)

# Benchmark.
if do_profile:
latency = run_benchmark(num_iters=1, profile=True)
else:
latency = run_benchmark(num_iters=1000, profile=False)
latency = run_benchmark(num_iters=10000, profile=False)
print(f"Kernel running time: {latency * 1000000:.3f} us")


Expand Down
Loading