[Hardware][Intel-Gaudi] Update hpu-extension and update bucketing system for HPU device (#17186)

Signed-off-by: Agata Dobrzyniewicz <adobrzyniewicz@habana.ai>
This commit is contained in:
Agata Dobrzyniewicz
2025-04-26 14:55:14 +02:00
committed by GitHub
parent 909fdaf152
commit c48334d405
6 changed files with 128 additions and 335 deletions

View File

@@ -168,7 +168,8 @@ class RMSNorm(CustomOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
from vllm_hpu_extension.ops import HPUFusedRMSNorm
from vllm_hpu_extension.kernels import rms_norm
HPUFusedRMSNorm = rms_norm()
if HPUFusedRMSNorm is None:
return self.forward_native(x, residual)
if residual is not None: