Jiangyun ZhuandSamit
98fe15762d
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: samithuang <285365963@qq.com> Co-authored-by: Samit <285365963@qq.com>
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import (
|
|
AttentionMetadata,
|
|
)
|
|
from vllm_omni.diffusion.attention.selector import get_attn_backend
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_heads: int,
|
|
head_size: int,
|
|
causal: bool,
|
|
softmax_scale: float,
|
|
num_kv_heads: int | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.attn_backend = get_attn_backend(-1)
|
|
self.attn_impl_cls = self.attn_backend.get_impl_cls()
|
|
self.attention = self.attn_impl_cls(
|
|
num_heads=num_heads,
|
|
head_size=head_size,
|
|
softmax_scale=softmax_scale,
|
|
causal=causal,
|
|
num_kv_heads=num_kv_heads,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AttentionMetadata = None,
|
|
) -> torch.Tensor:
|
|
# shape: (batch_size, seq_len, num_heads, head_size)
|
|
attn_output = self.attention.forward(query, key, value, attn_metadata)
|
|
return attn_output
|