rms_norm_ref() — pytorch Function Reference

Architecture documentation for the rms_norm_ref() function in kernels.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 1 called by 3

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ rms_norm_ref() — pytorch Function Reference

Dependency Diagram

graph TD
  ca055034_d94c_dc37_1965_ddd3537c220d["rms_norm_ref()"]
  64a4f8ff_baf4_3c9a_cbdf_a34032e1e46b["eager()"]
  64a4f8ff_baf4_3c9a_cbdf_a34032e1e46b -->|calls| ca055034_d94c_dc37_1965_ddd3537c220d
  2b63d71e_33f7_1a73_5259_ed1e703de26b["rms_norm_ref()"]
  2b63d71e_33f7_1a73_5259_ed1e703de26b -->|calls| ca055034_d94c_dc37_1965_ddd3537c220d
  3a2f2893_5914_e379_2f43_005a0538e035["eager()"]
  3a2f2893_5914_e379_2f43_005a0538e035 -->|calls| ca055034_d94c_dc37_1965_ddd3537c220d
  2b63d71e_33f7_1a73_5259_ed1e703de26b["rms_norm_ref()"]
  ca055034_d94c_dc37_1965_ddd3537c220d -->|calls| 2b63d71e_33f7_1a73_5259_ed1e703de26b
  style ca055034_d94c_dc37_1965_ddd3537c220d fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/genai_layers/kernels.py lines 381–387

    def rms_norm_ref(self, x, w):
        x_f32 = x.float()
        return (
            x_f32
            * torch.rsqrt(torch.mean(x_f32.square(), dim=-1, keepdim=True) + 1e-6)
            * w
        ).to(x.dtype)