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

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
This commit is contained in:
Thomas Parnell 2024-07-13 03:30:54 +02:00 committed by GitHub
parent a27f87da34
commit e1684a766a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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),