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

small fix

Browse Source
This commit is contained in:
Cyril Vallez 2025-01-17 16:05:47 +00:00
parent b03d7ae951
commit b40c889360
No known key found for this signature in database
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
server/text_generation_server/models/transformers_flash_causal_lm.py
View File

@ -37,6 +37,7 @@ def tgi_flash_attention_forward(
query_states: torch.Tensor, query_states: torch.Tensor,
key_states: torch.Tensor, key_states: torch.Tensor,
value_states: torch.Tensor, value_states: torch.Tensor,
attention_mask: Optional[torch.Tensor], # This needs to stay as it is passed as a positional arg in transformers
kv_cache: List[KVCache], kv_cache: List[KVCache],
kv_head_mapping: torch.Tensor, kv_head_mapping: torch.Tensor,
slots: torch.Tensor, slots: torch.Tensor,