如何不使用 xformers 推理 about cogvlm2 HOT 3 CLOSED

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_heads
        head_dim = config.hidden_size // config.num_heads
        self.scale = head_dim ** -0.5
        self.query_key_value = nn.Linear(config.hidden_size, config.hidden_size * 3)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.output_dropout = torch.nn.Dropout(config.dropout_prob)
        self.dropout = torch.nn.Dropout(config.dropout_prob)

    def forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
        B, L, _ = x.shape
        qkv = self.query_key_value(x)
        qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(2, 0, 1, 3, 4)  # 3, B, L, H, D
        query, key, value = qkv[0], qkv[1], qkv[2]

        # out = xops.memory_efficient_attention(
        #     q, k, v, scale=self.scale,
        # )
        # out = self.attention(query, key, value)

        scale = 1.0 / query.shape[-1] ** 0.5
        query = query * scale
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        attn = query @ key.transpose(-2, -1)
        attn = attn.softmax(-1)
        attn = self.dropout(attn)
        attn = attn @ value
        out = attn.transpose(1, 2)

        output = self.dense(out.contiguous().view(B, L, -1))
        output = self.output_dropout(output)
        return output