mstar.engine.resources.rms_norm

Backend-neutral RMSNorm operator.

mstar.engine.resources.rms_norm.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.

Parameters:

• x (Tensor)

• weight (Tensor)

• eps (float)

Return type:

Tensor

mstar.engine.resources.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

Modules

rms_norm

Backend-neutral RMSNorm custom operator.