Skip to content

vllm_gaudi.v1.attention.backends.hpu_attn

logger module-attribute 

logger = logger()

HPUAttentionBackendV1

Bases: HPUAttentionBackend

Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
@register_backend(AttentionBackendEnum.CUSTOM, "HPU_ATTN_V1")
class HPUAttentionBackendV1(HPUAttentionBackend):

    @staticmethod
    def get_name() -> str:
        return "CUSTOM"

    @staticmethod
    def get_impl_cls() -> type["AttentionImpl"]:
        return HPUAttentionImpl

    @staticmethod
    def get_metadata_cls() -> type["AttentionMetadata"]:
        return HPUAttentionMetadataV1

    @staticmethod
    def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
        # for mamba models we don't split block size across kernels
        # kernel_block_sizes in InputBatch are the same as block_sizes
        return [128]

get_impl_cls staticmethod 

get_impl_cls() -> type[AttentionImpl]
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
28
29
30
@staticmethod
def get_impl_cls() -> type["AttentionImpl"]:
    return HPUAttentionImpl

get_metadata_cls staticmethod 

get_metadata_cls() -> type[AttentionMetadata]
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
32
33
34
@staticmethod
def get_metadata_cls() -> type["AttentionMetadata"]:
    return HPUAttentionMetadataV1

get_name staticmethod 

get_name() -> str
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
24
25
26
@staticmethod
def get_name() -> str:
    return "CUSTOM"

get_supported_kernel_block_size staticmethod 

get_supported_kernel_block_size() -> list[
    Union[int, MultipleOf]
]
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
36
37
38
39
40
@staticmethod
def get_supported_kernel_block_size() -> list[Union[int, MultipleOf]]:
    # for mamba models we don't split block size across kernels
    # kernel_block_sizes in InputBatch are the same as block_sizes
    return [128]

HPUAttentionMetadataV1 dataclass

Bases: HPUAttentionMetadata

Metadata for HPUAttentionbackend.

Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@dataclass
class HPUAttentionMetadataV1(HPUAttentionMetadata):
    # TODO(kwisniewski98): for now, in V1 input positions are not provided
    # which needs to be fixed in the future, as we need to support MLA
    """Metadata for HPUAttentionbackend."""
    is_prompt: bool
    attn_bias: Optional[torch.Tensor]
    seq_lens_tensor: Optional[torch.Tensor]
    context_lens_tensor: Optional[torch.Tensor]
    query_start_loc: Optional[torch.Tensor] = None
    query_start_loc_p: Optional[torch.Tensor] = None
    padding_mask_flat: Optional[torch.Tensor] = None

    def seq_len(self):
        return self.slot_mapping.size(-1)

    def num_blocks(self):
        if self.block_list is None:
            return 0
        return self.block_list.numel()

    @classmethod
    def make_prefill_metadata(cls,
                              attn_bias,
                              block_list,
                              context_lens_tensor,
                              seq_lens_tensor,
                              slot_mapping,
                              block_size,
                              prep_initial_states=None,
                              has_initial_states_p=None,
                              last_chunk_indices_p=None,
                              state_indices_tensor=None,
                              query_start_loc=None,
                              padding_mask_flat=None):
        return cls(is_prompt=True,
                   block_list=block_list,
                   block_mapping=None,
                   block_usage=None,
                   block_groups=None,
                   attn_bias=attn_bias,
                   alibi_blocks=None,
                   context_lens_tensor=context_lens_tensor,
                   seq_lens_tensor=seq_lens_tensor,
                   input_positions=None,
                   slot_mapping=slot_mapping,
                   block_size=block_size,
                   prep_initial_states=prep_initial_states,
                   has_initial_states_p=has_initial_states_p,
                   last_chunk_indices_p=last_chunk_indices_p,
                   state_indices_tensor=state_indices_tensor,
                   query_start_loc=query_start_loc,
                   query_start_loc_p=query_start_loc,
                   padding_mask_flat=padding_mask_flat)

    @classmethod
    def make_decode_metadata(cls,
                             block_list,
                             block_usage,
                             block_groups,
                             input_positions,
                             slot_mapping,
                             block_size,
                             window_block_list,
                             window_block_usage,
                             window_block_groups,
                             chunked_block_list,
                             chunked_block_usage,
                             chunked_block_groups,
                             state_indices_tensor=None,
                             query_start_loc=None,
                             seq_lens_tensor=None):
        return cls(is_prompt=False,
                   block_mapping=None,
                   alibi_blocks=None,
                   attn_bias=None,
                   seq_lens_tensor=seq_lens_tensor,
                   context_lens_tensor=None,
                   block_list=block_list,
                   block_usage=block_usage,
                   block_groups=block_groups,
                   window_block_list=window_block_list,
                   window_block_usage=window_block_usage,
                   window_block_groups=window_block_groups,
                   chunked_block_list=chunked_block_list,
                   chunked_block_usage=chunked_block_usage,
                   chunked_block_groups=chunked_block_groups,
                   input_positions=input_positions,
                   slot_mapping=slot_mapping,
                   block_size=block_size,
                   prep_initial_states=None,
                   state_indices_tensor=state_indices_tensor,
                   query_start_loc=query_start_loc,
                   query_start_loc_p=query_start_loc)

attn_bias instance-attribute 

attn_bias: Optional[Tensor]

context_lens_tensor instance-attribute 

context_lens_tensor: Optional[Tensor]

is_prompt instance-attribute 

is_prompt: bool

padding_mask_flat class-attribute instance-attribute 

padding_mask_flat: Optional[Tensor] = None

query_start_loc class-attribute instance-attribute 

query_start_loc: Optional[Tensor] = None

query_start_loc_p class-attribute instance-attribute 

query_start_loc_p: Optional[Tensor] = None

seq_lens_tensor instance-attribute 

seq_lens_tensor: Optional[Tensor]

__init__ 

__init__(
    block_list: Optional[Tensor],
    block_mapping: Optional[Tensor],
    block_usage: Optional[Tensor],
    block_groups: Optional[Tensor],
    alibi_blocks: Optional[Tensor],
    is_prompt: bool,
    block_size: int,
    prep_initial_states: bool,
    slot_mapping: Tensor,
    attn_bias: Optional[Tensor],
    seq_lens_tensor: Optional[Tensor],
    context_lens_tensor: Optional[Tensor],
    input_positions: Tensor,
    seq_lens: Optional[list[int]] = None,
    encoder_seq_lens: Optional[list[int]] = None,
    encoder_seq_lens_tensor: Optional[Tensor] = None,
    max_encoder_seq_len: Optional[int] = None,
    cross_block_list: Optional[Tensor] = None,
    cross_slot_mapping: Optional[Tensor] = None,
    cross_block_mapping: Optional[Tensor] = None,
    cross_block_groups: Optional[Tensor] = None,
    cross_block_usage: Optional[Tensor] = None,
    cross_attn_bias: Optional[Tensor] = None,
    window_block_list: Optional[Tensor] = None,
    window_slot_mapping: Optional[Tensor] = None,
    window_block_mapping: Optional[Tensor] = None,
    window_block_groups: Optional[Tensor] = None,
    window_block_usage: Optional[Tensor] = None,
    window_attn_bias: Optional[Tensor] = None,
    chunked_slot_mapping: Optional[Tensor] = None,
    chunked_attn_bias: Optional[Tensor] = None,
    chunked_block_mapping: Optional[Tensor] = None,
    chunked_block_list: Optional[Tensor] = None,
    chunked_block_groups: Optional[Tensor] = None,
    chunked_block_usage: Optional[Tensor] = None,
    has_initial_states_p: Optional[Tensor] = None,
    last_chunk_indices_p: Optional[Tensor] = None,
    state_indices_tensor: Optional[Tensor] = None,
    query_start_loc: Optional[Tensor] = None,
    query_start_loc_p: Optional[Tensor] = None,
    padding_mask_flat: Optional[Tensor] = None,
) -> None

make_decode_metadata classmethod 

make_decode_metadata(
    block_list,
    block_usage,
    block_groups,
    input_positions,
    slot_mapping,
    block_size,
    window_block_list,
    window_block_usage,
    window_block_groups,
    chunked_block_list,
    chunked_block_usage,
    chunked_block_groups,
    state_indices_tensor=None,
    query_start_loc=None,
    seq_lens_tensor=None,
)
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@classmethod
def make_decode_metadata(cls,
                         block_list,
                         block_usage,
                         block_groups,
                         input_positions,
                         slot_mapping,
                         block_size,
                         window_block_list,
                         window_block_usage,
                         window_block_groups,
                         chunked_block_list,
                         chunked_block_usage,
                         chunked_block_groups,
                         state_indices_tensor=None,
                         query_start_loc=None,
                         seq_lens_tensor=None):
    return cls(is_prompt=False,
               block_mapping=None,
               alibi_blocks=None,
               attn_bias=None,
               seq_lens_tensor=seq_lens_tensor,
               context_lens_tensor=None,
               block_list=block_list,
               block_usage=block_usage,
               block_groups=block_groups,
               window_block_list=window_block_list,
               window_block_usage=window_block_usage,
               window_block_groups=window_block_groups,
               chunked_block_list=chunked_block_list,
               chunked_block_usage=chunked_block_usage,
               chunked_block_groups=chunked_block_groups,
               input_positions=input_positions,
               slot_mapping=slot_mapping,
               block_size=block_size,
               prep_initial_states=None,
               state_indices_tensor=state_indices_tensor,
               query_start_loc=query_start_loc,
               query_start_loc_p=query_start_loc)

make_prefill_metadata classmethod 

make_prefill_metadata(
    attn_bias,
    block_list,
    context_lens_tensor,
    seq_lens_tensor,
    slot_mapping,
    block_size,
    prep_initial_states=None,
    has_initial_states_p=None,
    last_chunk_indices_p=None,
    state_indices_tensor=None,
    query_start_loc=None,
    padding_mask_flat=None,
)
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@classmethod
def make_prefill_metadata(cls,
                          attn_bias,
                          block_list,
                          context_lens_tensor,
                          seq_lens_tensor,
                          slot_mapping,
                          block_size,
                          prep_initial_states=None,
                          has_initial_states_p=None,
                          last_chunk_indices_p=None,
                          state_indices_tensor=None,
                          query_start_loc=None,
                          padding_mask_flat=None):
    return cls(is_prompt=True,
               block_list=block_list,
               block_mapping=None,
               block_usage=None,
               block_groups=None,
               attn_bias=attn_bias,
               alibi_blocks=None,
               context_lens_tensor=context_lens_tensor,
               seq_lens_tensor=seq_lens_tensor,
               input_positions=None,
               slot_mapping=slot_mapping,
               block_size=block_size,
               prep_initial_states=prep_initial_states,
               has_initial_states_p=has_initial_states_p,
               last_chunk_indices_p=last_chunk_indices_p,
               state_indices_tensor=state_indices_tensor,
               query_start_loc=query_start_loc,
               query_start_loc_p=query_start_loc,
               padding_mask_flat=padding_mask_flat)

num_blocks 

num_blocks()
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
59
60
61
62
def num_blocks(self):
    if self.block_list is None:
        return 0
    return self.block_list.numel()

seq_len 

seq_len()
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
56
57
def seq_len(self):
    return self.slot_mapping.size(-1)