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
41
42
43
44
45
46
47
48
49
@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_sizes() -> list[Union[int, MultipleOf]]:
        # 16 is supported for testing/smaller models; 128 is the standard HPU
        # kernel block size; 528 is required for Granite 4.0-H
        # (granitemoehybrid) without prefix caching (16-token FA alignment),
        # 768 with prefix caching (chunk-aligned).
        return [16, 128, 528, 768]

    @classmethod
    def get_preferred_block_size(cls, default_block_size: int) -> int:
        # Always prefer 128-token HPU kernel blocks. For Granite 4.0-H,
        # check_and_update_config computes the correct block size (528 or
        # 768) before update_block_size_for_backend runs.
        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_preferred_block_size classmethod 

get_preferred_block_size(default_block_size: int) -> int
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
44
45
46
47
48
49
@classmethod
def get_preferred_block_size(cls, default_block_size: int) -> int:
    # Always prefer 128-token HPU kernel blocks. For Granite 4.0-H,
    # check_and_update_config computes the correct block size (528 or
    # 768) before update_block_size_for_backend runs.
    return 128

get_supported_kernel_block_sizes staticmethod 

get_supported_kernel_block_sizes() -> list[
    Union[int, MultipleOf]
]
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
36
37
38
39
40
41
42
@staticmethod
def get_supported_kernel_block_sizes() -> list[Union[int, MultipleOf]]:
    # 16 is supported for testing/smaller models; 128 is the standard HPU
    # kernel block size; 528 is required for Granite 4.0-H
    # (granitemoehybrid) without prefix caching (16-token FA alignment),
    # 768 with prefix caching (chunk-aligned).
    return [16, 128, 528, 768]

HPUAttentionMetadataV1 dataclass

Bases: HPUAttentionMetadata

Metadata for HPUAttentionbackend.

Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
 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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@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
    blocks_caching_range: Optional[torch.Tensor] = None
    mamba_chunks_to_block_mapping: Optional[torch.Tensor] = None
    seqlens_offsets_for_blocks: 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,
                              load_indices_tensor=None,
                              store_indices_tensor=None,
                              query_start_loc=None,
                              padding_mask_flat=None,
                              blocks_caching_range=None,
                              mamba_chunks_to_block_mapping=None,
                              seqlens_offsets_for_blocks=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,
                   load_indices_tensor=load_indices_tensor,
                   store_indices_tensor=store_indices_tensor,
                   query_start_loc=query_start_loc,
                   query_start_loc_p=query_start_loc,
                   padding_mask_flat=padding_mask_flat,
                   blocks_caching_range=blocks_caching_range,
                   mamba_chunks_to_block_mapping=mamba_chunks_to_block_mapping,
                   seqlens_offsets_for_blocks=seqlens_offsets_for_blocks)

    @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,
                             load_indices_tensor=None,
                             store_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,
                   load_indices_tensor=load_indices_tensor,
                   store_indices_tensor=store_indices_tensor,
                   query_start_loc=query_start_loc,
                   query_start_loc_p=query_start_loc)

attn_bias instance-attribute 

attn_bias: Optional[Tensor]

blocks_caching_range class-attribute instance-attribute 

blocks_caching_range: Optional[Tensor] = None

context_lens_tensor instance-attribute 

context_lens_tensor: Optional[Tensor]

is_prompt instance-attribute 

is_prompt: bool

mamba_chunks_to_block_mapping class-attribute instance-attribute 

mamba_chunks_to_block_mapping: Optional[Tensor] = None

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]

seqlens_offsets_for_blocks class-attribute instance-attribute 

seqlens_offsets_for_blocks: Optional[Tensor] = None

__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,
    load_indices_tensor: Optional[Tensor] = None,
    store_indices_tensor: Optional[Tensor] = None,
    query_start_loc: Optional[Tensor] = None,
    query_start_loc_p: Optional[Tensor] = None,
    padding_mask_flat: Optional[Tensor] = None,
    blocks_caching_range: Optional[Tensor] = None,
    mamba_chunks_to_block_mapping: Optional[Tensor] = None,
    seqlens_offsets_for_blocks: 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,
    load_indices_tensor=None,
    store_indices_tensor=None,
    query_start_loc=None,
    seq_lens_tensor=None,
)
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@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,
                         load_indices_tensor=None,
                         store_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,
               load_indices_tensor=load_indices_tensor,
               store_indices_tensor=store_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,
    load_indices_tensor=None,
    store_indices_tensor=None,
    query_start_loc=None,
    padding_mask_flat=None,
    blocks_caching_range=None,
    mamba_chunks_to_block_mapping=None,
    seqlens_offsets_for_blocks=None,
)
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
 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
@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,
                          load_indices_tensor=None,
                          store_indices_tensor=None,
                          query_start_loc=None,
                          padding_mask_flat=None,
                          blocks_caching_range=None,
                          mamba_chunks_to_block_mapping=None,
                          seqlens_offsets_for_blocks=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,
               load_indices_tensor=load_indices_tensor,
               store_indices_tensor=store_indices_tensor,
               query_start_loc=query_start_loc,
               query_start_loc_p=query_start_loc,
               padding_mask_flat=padding_mask_flat,
               blocks_caching_range=blocks_caching_range,
               mamba_chunks_to_block_mapping=mamba_chunks_to_block_mapping,
               seqlens_offsets_for_blocks=seqlens_offsets_for_blocks)

num_blocks 

num_blocks()
Source code in vllm_gaudi/v1/attention/backends/hpu_attn.py 
71
72
73
74
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 
68
69
def seq_len(self):
    return self.slot_mapping.size(-1)