Fix

build/torch27-cxx11-rocm63-x86_64-linux/residual_rms_rocm/__init__.py

@@ -1,3 +0,0 @@
-from .wrapped_rms import residual_rms, reference_residual_rms
-
-__all__ = ["residual_rms", "reference_residual_rms"]

build/torch27-cxx11-rocm63-x86_64-linux/residual_rms_rocm/_ops.py

@@ -1,9 +0,0 @@
-import torch
-from . import _residual_rms_rocm_7d048af
-ops = torch.ops._residual_rms_rocm_7d048af
-
-def add_op_namespace_prefix(op_name: str):
-    """
-    Prefix op by namespace.
-    """
-    return f"_residual_rms_rocm_7d048af::{op_name}"

build/torch27-cxx11-rocm63-x86_64-linux/residual_rms_rocm/wrapped_rms.py

@@ -1,171 +0,0 @@
-from typing import Tuple, Optional
-import torch
-from torch import Tensor
-
-from ._ops import ops
-
-
-_HIGHEST_RESIDUAL_RMS_MODE = 3
-
-
-def residual_rms_checks(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    scale_tensor: Tensor,
-    epsilon: float,
-    next_buffer: Tensor,
-) -> None:
-    # Check shapes
-    assert input.dim() == 2, f"Expected input to have 2 dimensions but got {input.dim() = } instead."
-    assert residual.shape == input.shape, \
-        f"Expected input and residual to have same shape but got {input.shape = } and {residual.shape = }"
-    assert weight.shape == (input.size(1), ), \
-        f"Expected weight to have shape {(input.size(1), ) = } but got {weight.shape = }"
-    # Check devices
-    device = input.device
-    assert device.type == "cuda", f"Expected input.device to be of type cuda, but got {device.type = } instead."
-    assert residual.device == device, f"Expected {residual.device = } to be the same as {input.device = }"
-    if scale_tensor is not None:
-        assert scale_tensor.device == device, f"Expected {scale_tensor.device = } to be the same as {input.device = }"
-    assert next_buffer.device == device, f"Expected {next_buffer.device = } to be the same as {input.device = }"
-    # Check layouts
-    assert input.is_contiguous(), f"Expected input to be contiguous but got {input.stride() = }"
-    assert residual.is_contiguous(), f"Expected residual to be contiguous but got {residual.stride() = }"
-    # Check scalars
-    assert epsilon > 0, f"Expected RMS epsilon to be > 0 to avoid division by zero, but got {epsilon = }"
-
-
-def residual_rms_choose_mode(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    next_buffer: Tensor,
-    mode: int,
-) -> int:
-    cols_is_multiple_of_8 = (input.size(1) % 8 == 0) and (next_buffer.size(1) % 8 == 0)
-    tensors_are_16b_aligned = all([x.data_ptr() % 16 == 0 for x in [input, residual, weight]])
-    if mode == -1:
-        mode = _HIGHEST_RESIDUAL_RMS_MODE if (tensors_are_16b_aligned and cols_is_multiple_of_8) else 0
-    elif mode > 0:
-        assert tensors_are_16b_aligned, (
-            f"Requested a {mode = } > 0 requires tensors to be 16 bits aligned but got {input.data_ptr() % 16 = }, "
-            f"{residual.data_ptr() % 16 = }, {weight.data_ptr() % 16 = }"
-        )
-        assert cols_is_multiple_of_8, f"Requested {mode = } requires {input.size(1) = } to be a multiple of 8."
-    return mode
-
-
-def infer_num_threads(rows: int, num_threads: int) -> int:
-    # Error case
-    if num_threads < 0 or num_threads > 1024:
-        raise ValueError(f"{num_threads = } is not between 0 and 1024")
-    # Case: num_threads was specified
-    elif num_threads != 0:
-        return num_threads
-    # Otherwise, we branch upon the number of rows
-    if rows <= 16:
-        return 1024
-    if rows <= 32:
-        return 768
-    if rows <= 64:
-        return 1024
-    if rows <= 256:
-        return 960
-    return 1024
-
-## Main kernel
-def residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor] = None,
-    next_buffer: Optional[Tensor] = None,
-    num_threads: int = 0,
-    force_scalar: bool = False,
-) -> Tuple[Tensor, Tensor]:
-    """Kernel that fuses a residual connection, an RMS normalization and a conversion to fp8. The resdiual argument is
-    modified inplace (residual <- input + residual).
-    Args:
-        - input: a fp16 tensor of shape (rows, cols) in row-major format
-        - residual: a fp16 tensor of shape (rows, cols) in row-major format
-        - weight: a fp16 tensor of shape (cols, ) in row-major format which contains the weight of the RMS norm
-        - epsilon: the small epsilon used inside the RMS norm to avoid division by zero
-        - scale_tensor: a fp32 one-item tensor to divide the output of the RMS norm before their conversion to fp8. If
-            set to None, then the output dtype is fp16
-        - next_buffer: an optional tensor of shape (rows, .) to initialize to zero if the output dtype in fp8
-        - num_threads: the number of threads per block in the kernel. Default value is 0, which then defaults to 1024
-    Outputs:
-        - an fp8 tensor of shape (rows, cols) in row-major format
-        - the residual modified in place
-    """
-    if next_buffer is None:
-        next_buffer = torch.empty(size=(input.size(0), 0), device=input.device, dtype=torch.float16)
-
-    residual_rms_checks(input, residual, weight, scale_tensor, epsilon, next_buffer)
-    num_threads = infer_num_threads(input.size(0), num_threads)
-
-    if scale_tensor is not None:
-        output = torch.empty(size=input.shape, dtype=torch.float8_e4m3fnuz, device=input.device)
-    else:
-        # TODO: here, we could use input as the output tensor
-        output = torch.empty(size=input.shape, dtype=torch.float16, device=input.device)
-    ops.residual_rms(
-        input=input,
-        residual=residual,
-        weight=weight,
-        scale_tensor=scale_tensor,
-        epsilon=epsilon,
-        output=output,
-        next_buffer=next_buffer,
-        num_threads=num_threads,
-        force_scalar=force_scalar,
-    )
-    return output, residual
-
-## Reference implementation
-def fp8_quantize(
-    x_full_precision: Tensor,
-    scale: Tensor,
-) -> Tuple[Tensor, Tensor]:
-    finfo = torch.finfo(torch.float8_e4m3fn)
-    x_quantized = (x_full_precision * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max)
-    x_quantized = x_quantized.to(torch.float8_e4m3fn)
-    weight_as_int8 = x_quantized.view(torch.int8)
-    ROCM_FP8_NAN_AS_INT = -128
-    mask = weight_as_int8 == ROCM_FP8_NAN_AS_INT
-    weight_as_int8[mask] = 0
-    x_quantized = weight_as_int8.view(torch.float8_e4m3fnuz)
-    return x_quantized, scale * 2.0
-
-def reference_residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor],
-    next_buffer: Optional[Tensor] = None,
-) -> Tuple[Tensor, Tensor, float]:
-    """Reference for the residual_rms operation. Check its docstring for more details, the only difference here is that
-    the scale needs to be passed a tensor and not a float."""
-    assert input.dtype == torch.float16, f"Expected torch.float16 but got {input.dtype = }"
-    assert residual.dtype == torch.float16, f"Expected torch.float16 but got {residual.dtype = }"
-    input += residual
-    residual = input
-    input = reference_rms(input, epsilon)
-    if weight.dtype in [torch.float16, torch.bfloat16]:
-        input = input.to(weight.dtype)
-    input = weight * input
-    if scale_tensor is not None:
-        qinput, scale_tensor = fp8_quantize(input, scale_tensor)
-        if next_buffer is not None:
-            next_buffer.fill_(0)
-    else:
-        qinput = input
-    return qinput, residual, scale_tensor
-
-def reference_rms(x: Tensor, eps: float) -> Tensor:
-    x = x.to(torch.float32)
-    variance = x.pow(2).mean(-1, keepdim=True)
-    return x * torch.rsqrt(variance + eps)

build/torch28-cxx11-rocm63-x86_64-linux/residual_rms_rocm/__init__.py

@@ -1,3 +0,0 @@
-from .wrapped_rms import residual_rms, reference_residual_rms
-
-__all__ = ["residual_rms", "reference_residual_rms"]

build/torch28-cxx11-rocm63-x86_64-linux/residual_rms_rocm/_ops.py

@@ -1,9 +0,0 @@
-import torch
-from . import _residual_rms_rocm_7d048af
-ops = torch.ops._residual_rms_rocm_7d048af
-
-def add_op_namespace_prefix(op_name: str):
-    """
-    Prefix op by namespace.
-    """
-    return f"_residual_rms_rocm_7d048af::{op_name}"

build/torch28-cxx11-rocm63-x86_64-linux/residual_rms_rocm/wrapped_rms.py

@@ -1,171 +0,0 @@
-from typing import Tuple, Optional
-import torch
-from torch import Tensor
-
-from ._ops import ops
-
-
-_HIGHEST_RESIDUAL_RMS_MODE = 3
-
-
-def residual_rms_checks(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    scale_tensor: Tensor,
-    epsilon: float,
-    next_buffer: Tensor,
-) -> None:
-    # Check shapes
-    assert input.dim() == 2, f"Expected input to have 2 dimensions but got {input.dim() = } instead."
-    assert residual.shape == input.shape, \
-        f"Expected input and residual to have same shape but got {input.shape = } and {residual.shape = }"
-    assert weight.shape == (input.size(1), ), \
-        f"Expected weight to have shape {(input.size(1), ) = } but got {weight.shape = }"
-    # Check devices
-    device = input.device
-    assert device.type == "cuda", f"Expected input.device to be of type cuda, but got {device.type = } instead."
-    assert residual.device == device, f"Expected {residual.device = } to be the same as {input.device = }"
-    if scale_tensor is not None:
-        assert scale_tensor.device == device, f"Expected {scale_tensor.device = } to be the same as {input.device = }"
-    assert next_buffer.device == device, f"Expected {next_buffer.device = } to be the same as {input.device = }"
-    # Check layouts
-    assert input.is_contiguous(), f"Expected input to be contiguous but got {input.stride() = }"
-    assert residual.is_contiguous(), f"Expected residual to be contiguous but got {residual.stride() = }"
-    # Check scalars
-    assert epsilon > 0, f"Expected RMS epsilon to be > 0 to avoid division by zero, but got {epsilon = }"
-
-
-def residual_rms_choose_mode(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    next_buffer: Tensor,
-    mode: int,
-) -> int:
-    cols_is_multiple_of_8 = (input.size(1) % 8 == 0) and (next_buffer.size(1) % 8 == 0)
-    tensors_are_16b_aligned = all([x.data_ptr() % 16 == 0 for x in [input, residual, weight]])
-    if mode == -1:
-        mode = _HIGHEST_RESIDUAL_RMS_MODE if (tensors_are_16b_aligned and cols_is_multiple_of_8) else 0
-    elif mode > 0:
-        assert tensors_are_16b_aligned, (
-            f"Requested a {mode = } > 0 requires tensors to be 16 bits aligned but got {input.data_ptr() % 16 = }, "
-            f"{residual.data_ptr() % 16 = }, {weight.data_ptr() % 16 = }"
-        )
-        assert cols_is_multiple_of_8, f"Requested {mode = } requires {input.size(1) = } to be a multiple of 8."
-    return mode
-
-
-def infer_num_threads(rows: int, num_threads: int) -> int:
-    # Error case
-    if num_threads < 0 or num_threads > 1024:
-        raise ValueError(f"{num_threads = } is not between 0 and 1024")
-    # Case: num_threads was specified
-    elif num_threads != 0:
-        return num_threads
-    # Otherwise, we branch upon the number of rows
-    if rows <= 16:
-        return 1024
-    if rows <= 32:
-        return 768
-    if rows <= 64:
-        return 1024
-    if rows <= 256:
-        return 960
-    return 1024
-
-## Main kernel
-def residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor] = None,
-    next_buffer: Optional[Tensor] = None,
-    num_threads: int = 0,
-    force_scalar: bool = False,
-) -> Tuple[Tensor, Tensor]:
-    """Kernel that fuses a residual connection, an RMS normalization and a conversion to fp8. The resdiual argument is
-    modified inplace (residual <- input + residual).
-    Args:
-        - input: a fp16 tensor of shape (rows, cols) in row-major format
-        - residual: a fp16 tensor of shape (rows, cols) in row-major format
-        - weight: a fp16 tensor of shape (cols, ) in row-major format which contains the weight of the RMS norm
-        - epsilon: the small epsilon used inside the RMS norm to avoid division by zero
-        - scale_tensor: a fp32 one-item tensor to divide the output of the RMS norm before their conversion to fp8. If
-            set to None, then the output dtype is fp16
-        - next_buffer: an optional tensor of shape (rows, .) to initialize to zero if the output dtype in fp8
-        - num_threads: the number of threads per block in the kernel. Default value is 0, which then defaults to 1024
-    Outputs:
-        - an fp8 tensor of shape (rows, cols) in row-major format
-        - the residual modified in place
-    """
-    if next_buffer is None:
-        next_buffer = torch.empty(size=(input.size(0), 0), device=input.device, dtype=torch.float16)
-
-    residual_rms_checks(input, residual, weight, scale_tensor, epsilon, next_buffer)
-    num_threads = infer_num_threads(input.size(0), num_threads)
-
-    if scale_tensor is not None:
-        output = torch.empty(size=input.shape, dtype=torch.float8_e4m3fnuz, device=input.device)
-    else:
-        # TODO: here, we could use input as the output tensor
-        output = torch.empty(size=input.shape, dtype=torch.float16, device=input.device)
-    ops.residual_rms(
-        input=input,
-        residual=residual,
-        weight=weight,
-        scale_tensor=scale_tensor,
-        epsilon=epsilon,
-        output=output,
-        next_buffer=next_buffer,
-        num_threads=num_threads,
-        force_scalar=force_scalar,
-    )
-    return output, residual
-
-## Reference implementation
-def fp8_quantize(
-    x_full_precision: Tensor,
-    scale: Tensor,
-) -> Tuple[Tensor, Tensor]:
-    finfo = torch.finfo(torch.float8_e4m3fn)
-    x_quantized = (x_full_precision * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max)
-    x_quantized = x_quantized.to(torch.float8_e4m3fn)
-    weight_as_int8 = x_quantized.view(torch.int8)
-    ROCM_FP8_NAN_AS_INT = -128
-    mask = weight_as_int8 == ROCM_FP8_NAN_AS_INT
-    weight_as_int8[mask] = 0
-    x_quantized = weight_as_int8.view(torch.float8_e4m3fnuz)
-    return x_quantized, scale * 2.0
-
-def reference_residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor],
-    next_buffer: Optional[Tensor] = None,
-) -> Tuple[Tensor, Tensor, float]:
-    """Reference for the residual_rms operation. Check its docstring for more details, the only difference here is that
-    the scale needs to be passed a tensor and not a float."""
-    assert input.dtype == torch.float16, f"Expected torch.float16 but got {input.dtype = }"
-    assert residual.dtype == torch.float16, f"Expected torch.float16 but got {residual.dtype = }"
-    input += residual
-    residual = input
-    input = reference_rms(input, epsilon)
-    if weight.dtype in [torch.float16, torch.bfloat16]:
-        input = input.to(weight.dtype)
-    input = weight * input
-    if scale_tensor is not None:
-        qinput, scale_tensor = fp8_quantize(input, scale_tensor)
-        if next_buffer is not None:
-            next_buffer.fill_(0)
-    else:
-        qinput = input
-    return qinput, residual, scale_tensor
-
-def reference_rms(x: Tensor, eps: float) -> Tensor:
-    x = x.to(torch.float32)
-    variance = x.pow(2).mean(-1, keepdim=True)
-    return x * torch.rsqrt(variance + eps)

build/torch28-cxx11-rocm64-x86_64-linux/residual_rms_rocm/__init__.py

@@ -1,3 +0,0 @@
-from .wrapped_rms import residual_rms, reference_residual_rms
-
-__all__ = ["residual_rms", "reference_residual_rms"]

build/torch28-cxx11-rocm64-x86_64-linux/residual_rms_rocm/_ops.py

@@ -1,9 +0,0 @@
-import torch
-from . import _residual_rms_rocm_7d048af
-ops = torch.ops._residual_rms_rocm_7d048af
-
-def add_op_namespace_prefix(op_name: str):
-    """
-    Prefix op by namespace.
-    """
-    return f"_residual_rms_rocm_7d048af::{op_name}"

build/torch28-cxx11-rocm64-x86_64-linux/residual_rms_rocm/wrapped_rms.py

@@ -1,171 +0,0 @@
-from typing import Tuple, Optional
-import torch
-from torch import Tensor
-
-from ._ops import ops
-
-
-_HIGHEST_RESIDUAL_RMS_MODE = 3
-
-
-def residual_rms_checks(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    scale_tensor: Tensor,
-    epsilon: float,
-    next_buffer: Tensor,
-) -> None:
-    # Check shapes
-    assert input.dim() == 2, f"Expected input to have 2 dimensions but got {input.dim() = } instead."
-    assert residual.shape == input.shape, \
-        f"Expected input and residual to have same shape but got {input.shape = } and {residual.shape = }"
-    assert weight.shape == (input.size(1), ), \
-        f"Expected weight to have shape {(input.size(1), ) = } but got {weight.shape = }"
-    # Check devices
-    device = input.device
-    assert device.type == "cuda", f"Expected input.device to be of type cuda, but got {device.type = } instead."
-    assert residual.device == device, f"Expected {residual.device = } to be the same as {input.device = }"
-    if scale_tensor is not None:
-        assert scale_tensor.device == device, f"Expected {scale_tensor.device = } to be the same as {input.device = }"
-    assert next_buffer.device == device, f"Expected {next_buffer.device = } to be the same as {input.device = }"
-    # Check layouts
-    assert input.is_contiguous(), f"Expected input to be contiguous but got {input.stride() = }"
-    assert residual.is_contiguous(), f"Expected residual to be contiguous but got {residual.stride() = }"
-    # Check scalars
-    assert epsilon > 0, f"Expected RMS epsilon to be > 0 to avoid division by zero, but got {epsilon = }"
-
-
-def residual_rms_choose_mode(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    next_buffer: Tensor,
-    mode: int,
-) -> int:
-    cols_is_multiple_of_8 = (input.size(1) % 8 == 0) and (next_buffer.size(1) % 8 == 0)
-    tensors_are_16b_aligned = all([x.data_ptr() % 16 == 0 for x in [input, residual, weight]])
-    if mode == -1:
-        mode = _HIGHEST_RESIDUAL_RMS_MODE if (tensors_are_16b_aligned and cols_is_multiple_of_8) else 0
-    elif mode > 0:
-        assert tensors_are_16b_aligned, (
-            f"Requested a {mode = } > 0 requires tensors to be 16 bits aligned but got {input.data_ptr() % 16 = }, "
-            f"{residual.data_ptr() % 16 = }, {weight.data_ptr() % 16 = }"
-        )
-        assert cols_is_multiple_of_8, f"Requested {mode = } requires {input.size(1) = } to be a multiple of 8."
-    return mode
-
-
-def infer_num_threads(rows: int, num_threads: int) -> int:
-    # Error case
-    if num_threads < 0 or num_threads > 1024:
-        raise ValueError(f"{num_threads = } is not between 0 and 1024")
-    # Case: num_threads was specified
-    elif num_threads != 0:
-        return num_threads
-    # Otherwise, we branch upon the number of rows
-    if rows <= 16:
-        return 1024
-    if rows <= 32:
-        return 768
-    if rows <= 64:
-        return 1024
-    if rows <= 256:
-        return 960
-    return 1024
-
-## Main kernel
-def residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor] = None,
-    next_buffer: Optional[Tensor] = None,
-    num_threads: int = 0,
-    force_scalar: bool = False,
-) -> Tuple[Tensor, Tensor]:
-    """Kernel that fuses a residual connection, an RMS normalization and a conversion to fp8. The resdiual argument is
-    modified inplace (residual <- input + residual).
-    Args:
-        - input: a fp16 tensor of shape (rows, cols) in row-major format
-        - residual: a fp16 tensor of shape (rows, cols) in row-major format
-        - weight: a fp16 tensor of shape (cols, ) in row-major format which contains the weight of the RMS norm
-        - epsilon: the small epsilon used inside the RMS norm to avoid division by zero
-        - scale_tensor: a fp32 one-item tensor to divide the output of the RMS norm before their conversion to fp8. If
-            set to None, then the output dtype is fp16
-        - next_buffer: an optional tensor of shape (rows, .) to initialize to zero if the output dtype in fp8
-        - num_threads: the number of threads per block in the kernel. Default value is 0, which then defaults to 1024
-    Outputs:
-        - an fp8 tensor of shape (rows, cols) in row-major format
-        - the residual modified in place
-    """
-    if next_buffer is None:
-        next_buffer = torch.empty(size=(input.size(0), 0), device=input.device, dtype=torch.float16)
-
-    residual_rms_checks(input, residual, weight, scale_tensor, epsilon, next_buffer)
-    num_threads = infer_num_threads(input.size(0), num_threads)
-
-    if scale_tensor is not None:
-        output = torch.empty(size=input.shape, dtype=torch.float8_e4m3fnuz, device=input.device)
-    else:
-        # TODO: here, we could use input as the output tensor
-        output = torch.empty(size=input.shape, dtype=torch.float16, device=input.device)
-    ops.residual_rms(
-        input=input,
-        residual=residual,
-        weight=weight,
-        scale_tensor=scale_tensor,
-        epsilon=epsilon,
-        output=output,
-        next_buffer=next_buffer,
-        num_threads=num_threads,
-        force_scalar=force_scalar,
-    )
-    return output, residual
-
-## Reference implementation
-def fp8_quantize(
-    x_full_precision: Tensor,
-    scale: Tensor,
-) -> Tuple[Tensor, Tensor]:
-    finfo = torch.finfo(torch.float8_e4m3fn)
-    x_quantized = (x_full_precision * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max)
-    x_quantized = x_quantized.to(torch.float8_e4m3fn)
-    weight_as_int8 = x_quantized.view(torch.int8)
-    ROCM_FP8_NAN_AS_INT = -128
-    mask = weight_as_int8 == ROCM_FP8_NAN_AS_INT
-    weight_as_int8[mask] = 0
-    x_quantized = weight_as_int8.view(torch.float8_e4m3fnuz)
-    return x_quantized, scale * 2.0
-
-def reference_residual_rms(
-    input: Tensor,
-    residual: Tensor,
-    weight: Tensor,
-    epsilon: float,
-    scale_tensor: Optional[Tensor],
-    next_buffer: Optional[Tensor] = None,
-) -> Tuple[Tensor, Tensor, float]:
-    """Reference for the residual_rms operation. Check its docstring for more details, the only difference here is that
-    the scale needs to be passed a tensor and not a float."""
-    assert input.dtype == torch.float16, f"Expected torch.float16 but got {input.dtype = }"
-    assert residual.dtype == torch.float16, f"Expected torch.float16 but got {residual.dtype = }"
-    input += residual
-    residual = input
-    input = reference_rms(input, epsilon)
-    if weight.dtype in [torch.float16, torch.bfloat16]:
-        input = input.to(weight.dtype)
-    input = weight * input
-    if scale_tensor is not None:
-        qinput, scale_tensor = fp8_quantize(input, scale_tensor)
-        if next_buffer is not None:
-            next_buffer.fill_(0)
-    else:
-        qinput = input
-    return qinput, residual, scale_tensor
-
-def reference_rms(x: Tensor, eps: float) -> Tensor:
-    x = x.to(torch.float32)
-    variance = x.pow(2).mean(-1, keepdim=True)
-    return x * torch.rsqrt(variance + eps)