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

Use FlashAttention kernel

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-17 20:24:45 +00:00
parent 7e3a230c38
commit 62b870fa07
2 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/model_executor/models/jax/gemma.py
View File

@ -165,8 +165,8 @@ class Attention(nn.Module):
value_proj = jnp.repeat(value_proj, self.num_heads, axis=-2) value_proj = jnp.repeat(value_proj, self.num_heads, axis=-2)
key_proj = jnp.repeat(key_proj, self.num_heads, axis=-2) key_proj = jnp.repeat(key_proj, self.num_heads, axis=-2)
if False: if True:
# FIXME(woosuk) # FlashAttention.
output = flash_attn( output = flash_attn(
query_proj, query_proj,
key_proj, key_proj,
@ -174,6 +174,7 @@ class Attention(nn.Module):
self.sm_scale, self.sm_scale,
) )
else: else:
# Naive attention with masking.
seq_len = query_proj.shape[1] seq_len = query_proj.shape[1]
attn_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) attn_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))

4
vllm/worker/tpu_model_runner.py
View File

@ -81,6 +81,10 @@ class TPUModelRunner:
max_prompt_len = max(prompt_lens) max_prompt_len = max(prompt_lens)
assert max_prompt_len > 0 assert max_prompt_len > 0
# NOTE(woosuk): The pallas FlashAttention kernel requires the sequence
# length to be a multiple of 16. We pad the prompt length to the nearest
# multiple of 16. This is also good for performance.
max_prompt_len = (max_prompt_len + 15) // 16 * 16
input_tokens = _make_array_with_pad(input_tokens, input_tokens = _make_array_with_pad(input_tokens,
max_prompt_len, max_prompt_len,