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

Update window size rocm flash decoding

Browse Source
This commit is contained in:
Mohit Sharma 2025-03-14 07:50:11 +00:00
parent b30cdabf68
commit 170a12f331
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
server/text_generation_server/layers/attention/rocm.py
View File

@ -83,6 +83,8 @@ def paged_attention(
max_k = max_s max_k = max_s
import flash_attn_2_cuda import flash_attn_2_cuda
window_size_right = -1 if window_size_left == -1 else 0
if softcap is None: if softcap is None:
softcap = 0.0 softcap = 0.0
out = flash_attn_2_cuda.varlen_fwd( out = flash_attn_2_cuda.varlen_fwd(
@ -102,8 +104,8 @@ def paged_attention(
softmax_scale, softmax_scale,
False, # zero_tensors False, # zero_tensors
True, # causal True, # causal
-1, # Window_left window_size_left, # Window_left
-1, # Window right window_size_right, # Window right
softcap, softcap,
False, # return softmax False, # return softmax
None, # generator None, # generator