mstar.engine.resources.rms_norm.rms_norm

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. mstar::fused_add_rms_norm does the same for the in-place residual add + RMSNorm: FlashInfer’s fused kernel on CUDA in half precision, the unfused add and norm elsewhere.

Functions

gemma_rms_norm(x, weight, eps)

Gemma's RMSNorm, normed * (1 + weight), computed in fp32 as Hugging Face Gemma computes it and returned in x's dtype.

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

Apply the backend-neutral mstar::rms_norm custom operator.
mstar.engine.resources.rms_norm.rms_norm.gemma_rms_norm(x, weight, eps)[source]#

Gemma’s RMSNorm, normed * (1 + weight), computed in fp32 as Hugging Face Gemma computes it and returned in x’s dtype.

Parameters:
Return type:

Tensor

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

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

Parameters:
Return type:

Tensor