[Kernel][Attention] Separate Attention.kv_scale into k_scale and v_scale (#6081)

2024-07-16 18:31:32 -04:00
parent 160e1d8c99
commit 978aed5300
33 changed files with 317 additions and 185 deletions
--- a/benchmarks/kernels/benchmark_paged_attention.py
+++ b/benchmarks/kernels/benchmark_paged_attention.py
@@ -100,7 +100,7 @@ def main(
        start_time = time.perf_counter()

        # Using default kv_scale
-        kv_scale = 1.0
+        k_scale = v_scale = 1.0

        for _ in range(num_iters):
            if version == "v1":
@@ -117,7 +117,8 @@ def main(
                    max_seq_len,
                    alibi_slopes,
                    kv_cache_dtype,
-                    kv_scale,
+                    k_scale,
+                    v_scale,
                )
            elif version == "v2":
                ops.paged_attention_v2(
@@ -136,7 +137,8 @@ def main(
                    max_seq_len,
                    alibi_slopes,
                    kv_cache_dtype,
-                    kv_scale,
+                    k_scale,
+                    v_scale,
                )
            else:
                raise ValueError(f"Invalid version: {version}")