mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-07-09 11:27:14 +08:00
precommit
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
parent
7c95fd8279
commit
1429a5e9a8
@ -110,7 +110,12 @@ def vit_flash_attn_wrapper(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def apply_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, softmax_scale: float | None = None) -> torch.Tensor:
|
def apply_sdpa(
|
||||||
|
q: torch.Tensor,
|
||||||
|
k: torch.Tensor,
|
||||||
|
v: torch.Tensor,
|
||||||
|
softmax_scale: float | None = None,
|
||||||
|
) -> torch.Tensor:
|
||||||
"""
|
"""
|
||||||
Input shape:
|
Input shape:
|
||||||
(batch_size x seq_len x num_heads x head_size)
|
(batch_size x seq_len x num_heads x head_size)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user