xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 20:55:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add a 1-line docstring to explain why calling context_attention_fwd twice in test_prefix_prefill.py (#2553)

Browse Source
This commit is contained in:
Jason Zhu 2024-01-22 14:47:25 -08:00 committed by GitHub
parent 63e835cbcc
commit 7a0b011dd5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/kernels/test_prefix_prefill.py
View File

@ -125,6 +125,7 @@ def test_contexted_kv_attention(
v_cache = v_cache.view(-1, block_size, num_heads, v_cache = v_cache.view(-1, block_size, num_heads,
head_size).permute(0, 2, 3, 1).contiguous() head_size).permute(0, 2, 3, 1).contiguous()
# Warm up the Triton kernel by calling it once before actually measuring generation time
context_attention_fwd(query, k, v, output, k_cache, v_cache, block_table, context_attention_fwd(query, k, v, output, k_cache, v_cache, block_table,
b_start_loc, b_seq_len, b_ctx_len, max_input_len) b_start_loc, b_seq_len, b_ctx_len, max_input_len)
torch.cuda.synchronize() torch.cuda.synchronize()