Skip to content

vllm_gaudi.ops.hpu_layernorm

HPUGemmaRMSNorm

Bases: GemmaRMSNorm

Source code in vllm_gaudi/ops/hpu_layernorm.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
@GemmaRMSNorm.register_oot
class HPUGemmaRMSNorm(GemmaRMSNorm):

    def forward_oot(
        self,
        x: torch.Tensor,
        residual: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        from vllm_gaudi.extension.kernels import rms_norm
        HPUFusedRMSNorm = rms_norm()
        # GemmaRMSNorm uses (1 + w) instead of w
        gemma_weight = self.weight + 1.0
        if residual is not None:
            orig_shape = x.shape
            residual = residual + x.reshape(residual.shape)
            # Note: HPUFusedRMSNorm requires 3D tensors as inputs
            x = HPUFusedRMSNorm.apply(residual, gemma_weight, self.variance_epsilon)
            return x.reshape(orig_shape), residual

        x = HPUFusedRMSNorm.apply(x, gemma_weight, self.variance_epsilon)
        return x

forward_oot 

forward_oot(
    x: Tensor, residual: Optional[Tensor] = None
) -> Union[Tensor, tuple[Tensor, Tensor]]
Source code in vllm_gaudi/ops/hpu_layernorm.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def forward_oot(
    self,
    x: torch.Tensor,
    residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    from vllm_gaudi.extension.kernels import rms_norm
    HPUFusedRMSNorm = rms_norm()
    # GemmaRMSNorm uses (1 + w) instead of w
    gemma_weight = self.weight + 1.0
    if residual is not None:
        orig_shape = x.shape
        residual = residual + x.reshape(residual.shape)
        # Note: HPUFusedRMSNorm requires 3D tensors as inputs
        x = HPUFusedRMSNorm.apply(residual, gemma_weight, self.variance_epsilon)
        return x.reshape(orig_shape), residual

    x = HPUFusedRMSNorm.apply(x, gemma_weight, self.variance_epsilon)
    return x

HPURMSNorm

Bases: RMSNorm

Source code in vllm_gaudi/ops/hpu_layernorm.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
@RMSNorm.register_oot
class HPURMSNorm(RMSNorm):

    def forward_oot(
        self,
        x: torch.Tensor,
        residual: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        from vllm_gaudi.extension.kernels import rms_norm
        HPUFusedRMSNorm = rms_norm()
        if residual is not None:
            orig_shape = x.shape
            residual = residual + x.reshape(residual.shape)
            # Note: HPUFusedRMSNorm requires 3D tensors as inputs
            x = HPUFusedRMSNorm.apply(residual, self.weight, self.variance_epsilon)
            return x.reshape(orig_shape), residual

        x = HPUFusedRMSNorm.apply(x, self.weight, self.variance_epsilon)
        return x

forward_oot 

forward_oot(
    x: Tensor, residual: Optional[Tensor] = None
) -> Union[Tensor, tuple[Tensor, Tensor]]
Source code in vllm_gaudi/ops/hpu_layernorm.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def forward_oot(
    self,
    x: torch.Tensor,
    residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    from vllm_gaudi.extension.kernels import rms_norm
    HPUFusedRMSNorm = rms_norm()
    if residual is not None:
        orig_shape = x.shape
        residual = residual + x.reshape(residual.shape)
        # Note: HPUFusedRMSNorm requires 3D tensors as inputs
        x = HPUFusedRMSNorm.apply(residual, self.weight, self.variance_epsilon)
        return x.reshape(orig_shape), residual

    x = HPUFusedRMSNorm.apply(x, self.weight, self.variance_epsilon)
    return x