[Bugfix] Fix hard-coded value of x in context_attention_fwd (#6373)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>

[Bugfix] Fix hard-coded value of x in context_attention_fwd (#6373)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
e1684a76 · Thomas Parnell · GitHub · a27f87da · e1684a76
Unverified Commit e1684a76 authored Jul 13, 2024 by Thomas Parnell Committed by GitHub Jul 12, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

vllm/attention/ops/prefix_prefill.py vllm/attention/ops/prefix_prefill.py +2 -2

No files found.
--- a/vllm/attention/ops/prefix_prefill.py
+++ b/vllm/attention/ops/prefix_prefill.py
@@ -718,7 +718,7 @@ if triton.__version__ >= "2.1.0":
                b_ctx_len,
                alibi_slopes,
                v_cache.shape[3],
-                8,
+                k_cache.shape[4],
                o,
                b_loc.stride(0),
                b_loc.stride(1),
@@ -768,7 +768,7 @@ if triton.__version__ >= "2.1.0":
            b_seq_len,
            b_ctx_len,
            v_cache.shape[3],
-            8,
+            k_cache.shape[4],
            o,
            b_loc.stride(0),
            b_loc.stride(1),