Skip to content

vllm_gaudi.attention.ops.hpu_paged_attn

_PARTITION_SIZE module-attribute 

_PARTITION_SIZE = 512

HPUPageAttentionInputBuilderBase dataclass

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
45
46
47
@dataclass
class HPUPageAttentionInputBuilderBase:
    pass

__init__ 

__init__() -> None

HPUPagedAttention

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
 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
class HPUPagedAttention:

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return list(range(1, 257))

    @classmethod
    def supports_attn_type(cls, attn_type: str) -> bool:
        """CPU attention supports decoder and encoder-only attention."""
        from vllm.attention.backends.abstract import AttentionType

        return attn_type in (
            AttentionType.DECODER,
            AttentionType.ENCODER,
            AttentionType.ENCODER_ONLY,
        )

    @staticmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> tuple[int, ...]:
        return (num_blocks * block_size, num_kv_heads, head_size)

    @staticmethod
    def split_kv_cache(
        kv_cache: tuple,
        num_kv_heads: int,
        head_size: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        key_cache = kv_cache[0]
        value_cache = kv_cache[1]
        k_scales = kv_cache[2]
        v_scales = kv_cache[3]
        return key_cache, value_cache, k_scales, v_scales

    @staticmethod
    def write_to_paged_cache(key: torch.Tensor, value: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor,
                             slot_mapping: torch.Tensor, kv_cache_dtype: str, is_prompt: bool) -> None:
        cache_ops.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, is_prompt)

    @staticmethod
    def forward_decode(**kwargs) -> torch.Tensor:
        if kwargs.get("kv_lora_rank"):
            return ops.flat_pa_mla(**kwargs)
        return ops.flat_pa(**kwargs)

    @staticmethod
    def swap_blocks(
        src_kv_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
        dst_kv_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
        src_to_dsts: torch.Tensor,
    ) -> None:
        src_key_cache = src_kv_cache[0]
        dst_key_cache = dst_kv_cache[0]
        cache_ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dsts)

        src_value_cache = src_kv_cache[1]
        dst_value_cache = dst_kv_cache[1]
        cache_ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dsts)

        src_key_scales = src_kv_cache[2]
        dst_key_scales = dst_kv_cache[2]
        src_value_scales = src_kv_cache[3]
        dst_value_scales = dst_kv_cache[3]
        if src_key_scales is not None:
            cache_ops.swap_blocks(src_key_scales, dst_key_scales, src_to_dsts)
        if src_value_scales is not None:
            cache_ops.swap_blocks(src_value_scales, dst_value_scales, src_to_dsts)

    @staticmethod
    def copy_blocks(
        kv_caches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
        src_to_dsts: torch.Tensor,
    ) -> None:
        key_caches = [kv_cache[0] for kv_cache in kv_caches]
        value_caches = [kv_cache[1] for kv_cache in kv_caches]
        key_scales = [kv_cache[2] for kv_cache in kv_caches]
        value_scales = [kv_cache[3] for kv_cache in kv_caches]
        cache_ops.copy_blocks(key_caches, value_caches, key_scales, value_scales, src_to_dsts)

copy_blocks staticmethod 

copy_blocks(
    kv_caches: list[tuple[Tensor, Tensor, Tensor, Tensor]],
    src_to_dsts: Tensor,
) -> None
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
122
123
124
125
126
127
128
129
130
131
@staticmethod
def copy_blocks(
    kv_caches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
    src_to_dsts: torch.Tensor,
) -> None:
    key_caches = [kv_cache[0] for kv_cache in kv_caches]
    value_caches = [kv_cache[1] for kv_cache in kv_caches]
    key_scales = [kv_cache[2] for kv_cache in kv_caches]
    value_scales = [kv_cache[3] for kv_cache in kv_caches]
    cache_ops.copy_blocks(key_caches, value_caches, key_scales, value_scales, src_to_dsts)

forward_decode staticmethod 

forward_decode(**kwargs) -> Tensor
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
93
94
95
96
97
@staticmethod
def forward_decode(**kwargs) -> torch.Tensor:
    if kwargs.get("kv_lora_rank"):
        return ops.flat_pa_mla(**kwargs)
    return ops.flat_pa(**kwargs)

get_kv_cache_shape staticmethod 

get_kv_cache_shape(
    num_blocks: int,
    block_size: int,
    num_kv_heads: int,
    head_size: int,
) -> tuple[int, ...]
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
67
68
69
70
71
72
73
74
@staticmethod
def get_kv_cache_shape(
    num_blocks: int,
    block_size: int,
    num_kv_heads: int,
    head_size: int,
) -> tuple[int, ...]:
    return (num_blocks * block_size, num_kv_heads, head_size)

get_supported_head_sizes staticmethod 

get_supported_head_sizes() -> list[int]
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
52
53
54
@staticmethod
def get_supported_head_sizes() -> list[int]:
    return list(range(1, 257))

split_kv_cache staticmethod 

split_kv_cache(
    kv_cache: tuple, num_kv_heads: int, head_size: int
) -> tuple[Tensor, Tensor, Tensor, Tensor]
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
76
77
78
79
80
81
82
83
84
85
86
@staticmethod
def split_kv_cache(
    kv_cache: tuple,
    num_kv_heads: int,
    head_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    key_cache = kv_cache[0]
    value_cache = kv_cache[1]
    k_scales = kv_cache[2]
    v_scales = kv_cache[3]
    return key_cache, value_cache, k_scales, v_scales

supports_attn_type classmethod 

supports_attn_type(attn_type: str) -> bool

CPU attention supports decoder and encoder-only attention.

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
56
57
58
59
60
61
62
63
64
65
@classmethod
def supports_attn_type(cls, attn_type: str) -> bool:
    """CPU attention supports decoder and encoder-only attention."""
    from vllm.attention.backends.abstract import AttentionType

    return attn_type in (
        AttentionType.DECODER,
        AttentionType.ENCODER,
        AttentionType.ENCODER_ONLY,
    )

swap_blocks staticmethod 

swap_blocks(
    src_kv_cache: tuple[Tensor, Tensor, Tensor, Tensor],
    dst_kv_cache: tuple[Tensor, Tensor, Tensor, Tensor],
    src_to_dsts: Tensor,
) -> None
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
@staticmethod
def swap_blocks(
    src_kv_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
    dst_kv_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
    src_to_dsts: torch.Tensor,
) -> None:
    src_key_cache = src_kv_cache[0]
    dst_key_cache = dst_kv_cache[0]
    cache_ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dsts)

    src_value_cache = src_kv_cache[1]
    dst_value_cache = dst_kv_cache[1]
    cache_ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dsts)

    src_key_scales = src_kv_cache[2]
    dst_key_scales = dst_kv_cache[2]
    src_value_scales = src_kv_cache[3]
    dst_value_scales = dst_kv_cache[3]
    if src_key_scales is not None:
        cache_ops.swap_blocks(src_key_scales, dst_key_scales, src_to_dsts)
    if src_value_scales is not None:
        cache_ops.swap_blocks(src_value_scales, dst_value_scales, src_to_dsts)

write_to_paged_cache staticmethod 

write_to_paged_cache(
    key: Tensor,
    value: Tensor,
    key_cache: Tensor,
    value_cache: Tensor,
    slot_mapping: Tensor,
    kv_cache_dtype: str,
    is_prompt: bool,
) -> None
Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
88
89
90
91
@staticmethod
def write_to_paged_cache(key: torch.Tensor, value: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor,
                         slot_mapping: torch.Tensor, kv_cache_dtype: str, is_prompt: bool) -> None:
    cache_ops.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, is_prompt)

HPUPagedAttentionMetadata dataclass

Metadata for PagedAttention.

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
18
19
20
21
22
23
24
25
@dataclass
class HPUPagedAttentionMetadata:
    """Metadata for PagedAttention."""
    block_list: Optional[torch.Tensor]
    block_mapping: Optional[torch.Tensor]
    block_usage: Optional[torch.Tensor]
    block_groups: Optional[torch.Tensor]
    alibi_blocks: Optional[torch.Tensor]

alibi_blocks instance-attribute 

alibi_blocks: Optional[Tensor]

block_groups instance-attribute 

block_groups: Optional[Tensor]

block_list instance-attribute 

block_list: Optional[Tensor]

block_mapping instance-attribute 

block_mapping: Optional[Tensor]

block_usage instance-attribute 

block_usage: Optional[Tensor]

__init__ 

__init__(
    block_list: Optional[Tensor],
    block_mapping: Optional[Tensor],
    block_usage: Optional[Tensor],
    block_groups: Optional[Tensor],
    alibi_blocks: Optional[Tensor],
) -> None

HPUPagedAttentionMetadataBuilder dataclass

Bases: AttentionMetadataBuilder

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
@dataclass
class HPUPagedAttentionMetadataBuilder(AttentionMetadataBuilder):

    def __init__(self, input_builder: "HPUPageAttentionInputBuilderBase") -> None:
        """Create the builder, remember some configuration and parameters."""
        self.input_builder = input_builder

    def prepare(self) -> None:
        """Prepare for one batch."""
        pass

    def build(self, seq_lens: list[int], query_lens: list[int], cuda_graph_pad_size: int,
              batch_size: int) -> type[HPUPagedAttentionMetadata]:
        """Build attention metadata with on-device tensors."""
        return HPUPagedAttentionMetadata

input_builder instance-attribute 

input_builder = input_builder

__init__ 

__init__(
    input_builder: HPUPageAttentionInputBuilderBase,
) -> None

Create the builder, remember some configuration and parameters.

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
31
32
33
def __init__(self, input_builder: "HPUPageAttentionInputBuilderBase") -> None:
    """Create the builder, remember some configuration and parameters."""
    self.input_builder = input_builder

build 

build(
    seq_lens: list[int],
    query_lens: list[int],
    cuda_graph_pad_size: int,
    batch_size: int,
) -> type[HPUPagedAttentionMetadata]

Build attention metadata with on-device tensors.

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
39
40
41
42
def build(self, seq_lens: list[int], query_lens: list[int], cuda_graph_pad_size: int,
          batch_size: int) -> type[HPUPagedAttentionMetadata]:
    """Build attention metadata with on-device tensors."""
    return HPUPagedAttentionMetadata

prepare 

prepare() -> None

Prepare for one batch.

Source code in vllm_gaudi/attention/ops/hpu_paged_attn.py 
35
36
37
def prepare(self) -> None:
    """Prepare for one batch."""
    pass