跳到内容

vllm_gaudi.ops.hpu_multihead_attn

HpuMultiHeadAttention

基类: MultiHeadAttention

源代码在 vllm_gaudi/ops/hpu_multihead_attn.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class HpuMultiHeadAttention(MultiHeadAttention):

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> torch.Tensor:
        """Input shape: 
        (batch_size x seq_len x hidden_size) or
        (batch_size x seq_len x num_heads x head_size)
        """
        bsz, q_len = query.size()[:2]
        kv_len = key.size(1)

        query = query.view(bsz, q_len, self.num_heads, self.head_size)
        key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
        value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)

        if (num_repeat := self.num_queries_per_kv) > 1:
            # Handle MQA and GQA
            key = torch.repeat_interleave(key, num_repeat, dim=2)
            value = torch.repeat_interleave(value, num_repeat, dim=2)

        query, key, value = (x.transpose(1, 2) for x in (query, key, value))

        from vllm_gaudi.extension.runtime import get_config

        if get_config().prompt_attn_impl == 'fsdpa_impl':

            from vllm_gaudi.extension.utils import ModuleFusedSDPA
            import vllm_gaudi.extension.kernels as kernels

            HPUFusedSDPA = kernels.fsdpa()
            fsdpa_op = ModuleFusedSDPA(HPUFusedSDPA)

            out = fsdpa_op(query,
                           key,
                           value,
                           None,
                           dropout_p=0.0,
                           is_causal=False,
                           scale=self.scale,
                           softmax_mode="fast",
                           recompute_mode=True,
                           valid_sequence_lengths=None)
        else:
            out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)

        out = out.transpose(1, 2)
        return out.reshape(bsz, q_len, -1)

forward 

forward(
    query: Tensor, key: Tensor, value: Tensor
) -> Tensor

输入形状: (batch_size x seq_len x hidden_size) 或 (batch_size x seq_len x num_heads x head_size)

源代码在 vllm_gaudi/ops/hpu_multihead_attn.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
) -> torch.Tensor:
    """Input shape: 
    (batch_size x seq_len x hidden_size) or
    (batch_size x seq_len x num_heads x head_size)
    """
    bsz, q_len = query.size()[:2]
    kv_len = key.size(1)

    query = query.view(bsz, q_len, self.num_heads, self.head_size)
    key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
    value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)

    if (num_repeat := self.num_queries_per_kv) > 1:
        # Handle MQA and GQA
        key = torch.repeat_interleave(key, num_repeat, dim=2)
        value = torch.repeat_interleave(value, num_repeat, dim=2)

    query, key, value = (x.transpose(1, 2) for x in (query, key, value))

    from vllm_gaudi.extension.runtime import get_config

    if get_config().prompt_attn_impl == 'fsdpa_impl':

        from vllm_gaudi.extension.utils import ModuleFusedSDPA
        import vllm_gaudi.extension.kernels as kernels

        HPUFusedSDPA = kernels.fsdpa()
        fsdpa_op = ModuleFusedSDPA(HPUFusedSDPA)

        out = fsdpa_op(query,
                       key,
                       value,
                       None,
                       dropout_p=0.0,
                       is_causal=False,
                       scale=self.scale,
                       softmax_mode="fast",
                       recompute_mode=True,
                       valid_sequence_lengths=None)
    else:
        out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)

    out = out.transpose(1, 2)
    return out.reshape(bsz, q_len, -1)