Skip to content

vllm_gaudi.extension.utils

FP8Matmul

Bases: Module

Source code in vllm_gaudi/extension/utils.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
class FP8Matmul(torch.nn.Module):

    def __init__(
        self,
        scale_input=1.0,
        scale_other=1.0,
    ):
        super().__init__()
        self.scale_input = scale_input
        self.scale_other = scale_other

    def quant_input(self, x, scale):
        return torch.ops.hpu.cast_to_fp8_v2(x, scale, False, False, torch.float8_e4m3fn)[0]

    def matmul_fp8(self, x, other, out_dtype, scale_input_inv=None, scale_other_inv=None):
        return torch.ops.hpu.fp8_gemm_v2(
            A=x,
            trans_A=False,
            B=other,
            trans_B=False,
            D=None,
            out_dtype=out_dtype,
            A_scale_inv=scale_input_inv,
            B_scale_inv=scale_other_inv,
            bias=None,
            accumulate=False,
        )

    def forward(self, input, other):
        qinput = self.quant_input(input, self.scale_input)
        qother = self.quant_input(other, self.scale_other)
        output = self.matmul_fp8(
            qinput,
            qother,
            out_dtype=torch.bfloat16,
            scale_input_inv=1.0 / self.scale_input,
            scale_other_inv=1.0 / self.scale_other,
        )
        return output

scale_input instance-attribute 

scale_input = scale_input

scale_other instance-attribute 

scale_other = scale_other

__init__ 

__init__(scale_input=1.0, scale_other=1.0)
Source code in vllm_gaudi/extension/utils.py 
100
101
102
103
104
105
106
107
def __init__(
    self,
    scale_input=1.0,
    scale_other=1.0,
):
    super().__init__()
    self.scale_input = scale_input
    self.scale_other = scale_other

forward 

forward(input, other)
Source code in vllm_gaudi/extension/utils.py 
126
127
128
129
130
131
132
133
134
135
136
def forward(self, input, other):
    qinput = self.quant_input(input, self.scale_input)
    qother = self.quant_input(other, self.scale_other)
    output = self.matmul_fp8(
        qinput,
        qother,
        out_dtype=torch.bfloat16,
        scale_input_inv=1.0 / self.scale_input,
        scale_other_inv=1.0 / self.scale_other,
    )
    return output

matmul_fp8 

matmul_fp8(
    x,
    other,
    out_dtype,
    scale_input_inv=None,
    scale_other_inv=None,
)
Source code in vllm_gaudi/extension/utils.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
def matmul_fp8(self, x, other, out_dtype, scale_input_inv=None, scale_other_inv=None):
    return torch.ops.hpu.fp8_gemm_v2(
        A=x,
        trans_A=False,
        B=other,
        trans_B=False,
        D=None,
        out_dtype=out_dtype,
        A_scale_inv=scale_input_inv,
        B_scale_inv=scale_other_inv,
        bias=None,
        accumulate=False,
    )

quant_input 

quant_input(x, scale)
Source code in vllm_gaudi/extension/utils.py 
109
110
def quant_input(self, x, scale):
    return torch.ops.hpu.cast_to_fp8_v2(x, scale, False, False, torch.float8_e4m3fn)[0]

Matmul

Bases: Module

Source code in vllm_gaudi/extension/utils.py 
24
25
26
27
28
29
30
class Matmul(torch.nn.Module):

    def __init__(self):
        super(Matmul, self).__init__()

    def forward(self, x, y, **kwargs):
        return torch.matmul(x, y, **kwargs)

__init__ 

__init__()
Source code in vllm_gaudi/extension/utils.py 
26
27
def __init__(self):
    super(Matmul, self).__init__()

forward 

forward(x, y, **kwargs)
Source code in vllm_gaudi/extension/utils.py 
29
30
def forward(self, x, y, **kwargs):
    return torch.matmul(x, y, **kwargs)

ModuleFusedSDPA

Bases: Module

Source code in vllm_gaudi/extension/utils.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
class ModuleFusedSDPA(torch.nn.Module):

    def __init__(self, fusedSDPA):
        super().__init__()
        assert fusedSDPA is not None, f'fusedSDPA kernel is None'
        self._hpu_kernel_fsdpa = fusedSDPA

    def forward(
        self,
        query,
        key,
        value,
        attn_mask,
        dropout_p,
        is_causal,
        scale,
        softmax_mode,
        recompute_mode,
        valid_sequence_lengths,
        padding_side="left",
        window_size=None,
    ):
        if window_size is not None:
            return self._hpu_kernel_fsdpa.apply(query, key, value, attn_mask, dropout_p, is_causal, scale, softmax_mode,
                                                recompute_mode, valid_sequence_lengths, padding_side, False, False,
                                                window_size)
        else:
            return self._hpu_kernel_fsdpa.apply(query, key, value, attn_mask, dropout_p, is_causal, scale, softmax_mode,
                                                recompute_mode, valid_sequence_lengths, padding_side)

_hpu_kernel_fsdpa instance-attribute 

_hpu_kernel_fsdpa = fusedSDPA

__init__ 

__init__(fusedSDPA)
Source code in vllm_gaudi/extension/utils.py 
141
142
143
144
def __init__(self, fusedSDPA):
    super().__init__()
    assert fusedSDPA is not None, f'fusedSDPA kernel is None'
    self._hpu_kernel_fsdpa = fusedSDPA

forward 

forward(
    query,
    key,
    value,
    attn_mask,
    dropout_p,
    is_causal,
    scale,
    softmax_mode,
    recompute_mode,
    valid_sequence_lengths,
    padding_side="left",
    window_size=None,
)
Source code in vllm_gaudi/extension/utils.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def forward(
    self,
    query,
    key,
    value,
    attn_mask,
    dropout_p,
    is_causal,
    scale,
    softmax_mode,
    recompute_mode,
    valid_sequence_lengths,
    padding_side="left",
    window_size=None,
):
    if window_size is not None:
        return self._hpu_kernel_fsdpa.apply(query, key, value, attn_mask, dropout_p, is_causal, scale, softmax_mode,
                                            recompute_mode, valid_sequence_lengths, padding_side, False, False,
                                            window_size)
    else:
        return self._hpu_kernel_fsdpa.apply(query, key, value, attn_mask, dropout_p, is_causal, scale, softmax_mode,
                                            recompute_mode, valid_sequence_lengths, padding_side)

Softmax

Bases: Module

Source code in vllm_gaudi/extension/utils.py 
33
34
35
36
37
38
39
class Softmax(torch.nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x, dim=None, inv_head=None):
        return torch.softmax(x, dim)

__init__ 

__init__()
Source code in vllm_gaudi/extension/utils.py 
35
36
def __init__(self):
    super().__init__()

forward 

forward(x, dim=None, inv_head=None)
Source code in vllm_gaudi/extension/utils.py 
38
39
def forward(self, x, dim=None, inv_head=None):
    return torch.softmax(x, dim)

VLLMFP8KVCache

Bases: VLLMKVCache

Source code in vllm_gaudi/extension/utils.py 
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
class VLLMFP8KVCache(VLLMKVCache):

    def __init__(self, input_scale=1.0):
        super().__init__()
        self.use_contiguous_pa = get_config().use_contiguous_pa
        self.input_scale = input_scale
        self.output_scale = 1.0 / self.input_scale

    def quant_input(self, input):
        return torch.ops.hpu.cast_to_fp8_v2(input, self.input_scale, False, False, torch.float8_e4m3fn)[0]

    def dequant_output(self, output):
        return torch.ops.hpu.cast_from_fp8(output, self.output_scale, torch.bfloat16)

    def forward(self, input, *args, **kwargs):
        qinput = self.quant_input(input)
        return super().forward(qinput, *args, **kwargs)

    def fetch_from_cache(self, quant_cache, blocks, permutations=None):
        if permutations:
            output_cache = super().fetch_from_cache(quant_cache, blocks, permutations)
            for i in range(len(output_cache)):
                output_cache[i] = self.dequant_output(output_cache[i])
            return output_cache
        output_cache = super().fetch_from_cache(quant_cache, blocks)
        return self.dequant_output(output_cache)

input_scale instance-attribute 

input_scale = input_scale

output_scale instance-attribute 

output_scale = 1.0 / input_scale

use_contiguous_pa instance-attribute 

use_contiguous_pa = use_contiguous_pa

__init__ 

__init__(input_scale=1.0)
Source code in vllm_gaudi/extension/utils.py 
72
73
74
75
76
def __init__(self, input_scale=1.0):
    super().__init__()
    self.use_contiguous_pa = get_config().use_contiguous_pa
    self.input_scale = input_scale
    self.output_scale = 1.0 / self.input_scale

dequant_output 

dequant_output(output)
Source code in vllm_gaudi/extension/utils.py 
81
82
def dequant_output(self, output):
    return torch.ops.hpu.cast_from_fp8(output, self.output_scale, torch.bfloat16)

fetch_from_cache 

fetch_from_cache(quant_cache, blocks, permutations=None)
Source code in vllm_gaudi/extension/utils.py 
88
89
90
91
92
93
94
95
def fetch_from_cache(self, quant_cache, blocks, permutations=None):
    if permutations:
        output_cache = super().fetch_from_cache(quant_cache, blocks, permutations)
        for i in range(len(output_cache)):
            output_cache[i] = self.dequant_output(output_cache[i])
        return output_cache
    output_cache = super().fetch_from_cache(quant_cache, blocks)
    return self.dequant_output(output_cache)

forward 

forward(input, *args, **kwargs)
Source code in vllm_gaudi/extension/utils.py 
84
85
86
def forward(self, input, *args, **kwargs):
    qinput = self.quant_input(input)
    return super().forward(qinput, *args, **kwargs)

quant_input 

quant_input(input)
Source code in vllm_gaudi/extension/utils.py 
78
79
def quant_input(self, input):
    return torch.ops.hpu.cast_to_fp8_v2(input, self.input_scale, False, False, torch.float8_e4m3fn)[0]

VLLMKVCache

Bases: Module

Source code in vllm_gaudi/extension/utils.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class VLLMKVCache(torch.nn.Module):

    def __init__(self, is_v_cache: bool = False):
        super().__init__()
        self.use_contiguous_pa = get_config().use_contiguous_pa
        # is_v_cache is used in INC FP8 dynamic quantization to identify V cache
        self.is_v_cache = is_v_cache

    def forward(self, input, cache, slot_mapping, scales=None, **kwargs):
        # In cross-attention kv cache forward inputs are None in decode
        # We don't want to store them in the cache in such case
        if input is not None:
            cache.index_copy_(0, slot_mapping, input)
        return cache

    def fetch_from_cache(self, cache, blocks, scales=None, **kwargs):
        if self.use_contiguous_pa:
            return cache[:blocks.size(0)]
        else:
            return cache.index_select(0, blocks)

is_v_cache instance-attribute 

is_v_cache = is_v_cache

use_contiguous_pa instance-attribute 

use_contiguous_pa = use_contiguous_pa

__init__ 

__init__(is_v_cache: bool = False)
Source code in vllm_gaudi/extension/utils.py 
50
51
52
53
54
def __init__(self, is_v_cache: bool = False):
    super().__init__()
    self.use_contiguous_pa = get_config().use_contiguous_pa
    # is_v_cache is used in INC FP8 dynamic quantization to identify V cache
    self.is_v_cache = is_v_cache

fetch_from_cache 

fetch_from_cache(cache, blocks, scales=None, **kwargs)
Source code in vllm_gaudi/extension/utils.py 
63
64
65
66
67
def fetch_from_cache(self, cache, blocks, scales=None, **kwargs):
    if self.use_contiguous_pa:
        return cache[:blocks.size(0)]
    else:
        return cache.index_select(0, blocks)

forward 

forward(input, cache, slot_mapping, scales=None, **kwargs)
Source code in vllm_gaudi/extension/utils.py 
56
57
58
59
60
61
def forward(self, input, cache, slot_mapping, scales=None, **kwargs):
    # In cross-attention kv cache forward inputs are None in decode
    # We don't want to store them in the cache in such case
    if input is not None:
        cache.index_copy_(0, slot_mapping, input)
    return cache

align_and_pad 

align_and_pad(data, bucketing, padding_gen)
Source code in vllm_gaudi/extension/utils.py 
177
178
179
180
181
182
183
184
185
def align_and_pad(data, bucketing, padding_gen):
    bs = len(data)
    target_bs, target_len = bucketing
    if target_bs == 1 and bs > 1:
        data = [list(itertools.chain(*data))]
    data = [pad_list(x, target_len, padding_gen) for x in data]
    padding = itertools.islice(padding_gen, target_len)
    data = pad_list(data, target_bs, itertools.tee(padding, target_bs - len(data)))
    return data

get_kv_fetch_extra_args 

get_kv_fetch_extra_args(**kwargs)
Source code in vllm_gaudi/extension/utils.py 
42
43
44
45
def get_kv_fetch_extra_args(**kwargs):
    if not get_config().per_token_kv_scaling_support:
        kwargs.pop('scales', None)
    return kwargs

is_fake_hpu cached 

is_fake_hpu() -> bool
Source code in vllm_gaudi/extension/utils.py 
19
20
21
@lru_cache(maxsize=None)
def is_fake_hpu() -> bool:
    return os.environ.get('VLLM_USE_FAKE_HPU', '0') != '0'

pad_list 

pad_list(input, target_len, val_generator)
Source code in vllm_gaudi/extension/utils.py 
170
171
172
173
174
def pad_list(input, target_len, val_generator):
    padding = target_len - len(input)
    if padding > 0:
        input.extend(itertools.islice(val_generator, padding))
    return input

with_default 

with_default(value: Optional[Any], default: Any) -> Any
Source code in vllm_gaudi/extension/utils.py 
188
189
190
191
def with_default(value: Optional[Any], default: Any) -> Any:
    if value is not None:
        return value
    return default