mstar.engine.resources.rms_norm.rms_norm

Backend-neutral RMSNorm custom operator.

The operator keeps fused accelerator implementations behind one stable mstar::rms_norm graph boundary. CUDA currently dispatches to FlashInfer, XPU dispatches to vllm-xpu-kernels, and other devices use the portable PyTorch implementation.

Functions

run_rms_norm(input, weight[, eps, ...])

Apply the backend-neutral mstar::rms_norm custom operator.

mstar.engine.resources.rms_norm.rms_norm.run_rms_norm(input, weight, eps=1e-6, rms_norm_dtype=None)

Apply the backend-neutral mstar::rms_norm custom operator.

Parameters:

• input (Tensor)

• weight (Tensor)

• eps (float)

• rms_norm_dtype (dtype | None)

Return type:

Tensor