xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-08-14 03:34:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

yapf

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-26 05:27:38 +00:00
parent d16a348477
commit 4ea41d01a9
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
vllm/model_executor/models/jax/ops/flash_attn.py
View File

@ -8,6 +8,7 @@ _DEFAULT_BLOCK_SIZES = {
"block_b": 2, "block_b": 2,
} }
def flash_attn( def flash_attn(
q: jax.Array, # [batch, seq_len, num_heads, head_size] q: jax.Array, # [batch, seq_len, num_heads, head_size]
k: jax.Array, # [batch, seq_len, num_heads, head_size] k: jax.Array, # [batch, seq_len, num_heads, head_size]