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

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

Author: tdoublep

vllm/attention/ops/prefix_prefill.py

@@ -718,7 +718,7 @@ def context_attention_fwd(q,
                 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 @@ def context_attention_fwd(q,
             b_seq_len,
             b_ctx_len,
             v_cache.shape[3],
-            8,
+            k_cache.shape[4],
             o,
             b_loc.stride(0),
             b_loc.stride(1),