A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Kernel] update comment for KV shape in unified triton attn (#18099)
Signed-off-by: haochengxia <xhc_1007@163.com>
Percy committed
May 20, 2025 at 13:19 UTC
980a172474fa0f32433dda87ae1fa4aadba24c51
1 file changed
+2
-2
vllm/attention/ops/triton_unified_attention.py
+2
-2
@@ -31,8 +31,8 @@ def apply_softcap(S, x):
31
def kernel_unified_attention_2d(
32
output_ptr, # [num_tokens, num_query_heads, head_size]
33
query_ptr, # [num_tokens, num_query_heads, head_size]
34
- key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
35
- value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
34
+ key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
35
+ value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
36
block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
37
seq_lens_ptr, # [num_seqs]
38
alibi_slopes_ptr, # [num_query_heads]