huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Making it work on non flash decoding.

Browse Source
This commit is contained in:
Nicolas Patry 2024-05-31 21:41:19 +00:00
parent 4293a12863
commit 66081e6ae7
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
View File

@ -194,10 +194,10 @@ class FlashLlamaAttention(torch.nn.Module):
reshape_and_cache(kv[:, 0], kv[:, 1], kv_cache[0], kv_cache[1], slots) reshape_and_cache(kv[:, 0], kv[:, 1], kv_cache[0], kv_cache[1], slots)
# output tensor # output tensor
attn_output = torch.empty_like(query)
# Prefill # Prefill
if cu_seqlen_prefill is not None: if cu_seqlen_prefill is not None:
attn_output = torch.empty_like(query)
# flash attention # flash attention
attention( attention(
query, query,
@ -211,7 +211,7 @@ class FlashLlamaAttention(torch.nn.Module):
# Decode # Decode
else: else:
attn_output = paged_attention( attn_output = paged_attention(
None, attn_output,
query, query,
kv_cache[0], kv_cache[0],
kv_cache[1], kv_cache[1],