Skip to content

vllm_gaudi.extension.kernels

_kernel 

_kernel(name)
Source code in vllm_gaudi/extension/kernels.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def _kernel(name):

    def loader(fn):
        result = []

        def loader_impl():
            if not result:
                try:
                    result.append(fn())
                except (ImportError, AttributeError):
                    from .utils import logger
                    logger().warning(f"Could not import HPU {name} kernel. "
                                     "vLLM will use native implementation")
                    result.append(None)
            return result[0]

        return loader_impl

    return loader

block_softmax_adjustment 

block_softmax_adjustment()
Source code in vllm_gaudi/extension/kernels.py 
42
43
44
45
@_kernel("block_softmax_adjustment")
def block_softmax_adjustment():
    import torch
    return torch.ops.hpu.block_softmax_adjustment

fsdpa 

fsdpa()
Source code in vllm_gaudi/extension/kernels.py 
30
31
32
33
@_kernel("FusedSDPA")
def fsdpa():
    from habana_frameworks.torch.hpex.kernels import FusedSDPA
    return FusedSDPA

rms_norm 

rms_norm()
Source code in vllm_gaudi/extension/kernels.py 
36
37
38
39
@_kernel("FusedRMSNorm")
def rms_norm():
    from habana_frameworks.torch.hpex.normalization import FusedRMSNorm
    return FusedRMSNorm

softmax_fa2 

softmax_fa2()
Source code in vllm_gaudi/extension/kernels.py 
48
49
50
51
@_kernel("softmax_fa2")
def softmax_fa2():
    import torch
    return torch.ops.hpu.softmax_fa2