squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix attention

Browse Source
This commit is contained in:
Woosuk Kwon 2023-02-23 22:29:46 +00:00
parent 87e0bcd426
commit ba84b8728a
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
cacheflow/models/attention.py
View File

@ -1,4 +1,4 @@
from typing import Optional, Tuple from typing import Optional
import torch import torch
import torch.nn as nn import torch.nn as nn
@ -24,8 +24,12 @@ class OPTCacheFlowAttention(nn.Module):
key: torch.Tensor, key: torch.Tensor,
value: torch.Tensor, value: torch.Tensor,
) -> None: ) -> None:
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
out = xops.memory_efficient_attention( out = xops.memory_efficient_attention(
query, key, value, attn_bias=self.attention_mask, scale=self.scale) query, key, value, attn_bias=self.attention_mask, scale=self.scale)
out = out.squeeze(0)
# FIXME(woosuk): Directly write the attention output. # FIXME(woosuk): Directly write the attention output.
output.copy_(out, non_blocking=True) output.copy_(out, non_blocking=True)