InplaceAddRmsNormFusionPass

Description

Scenario 1: Converts AddRmsNorm to InplaceAddRmsNorm by reusing the x1 input address as the y output address and the x2 input address as the x output address.

Scenario 2: Fuses AddRmsNorm and Cast into the AddRmsNormCast operator if the output y of AddRmsNorm is connected to only two outputs and the first output is connected to the Cast operator (to cast the output y from float16/bfloat16 to float32).

Restrictions

• General restrictions: Only inference is supported.
• Restrictions on fusion into AddRmsNormCast:
  • The input x1 can only be of the float16 or bfloat16 type.
  • Only casting from float16 or bfloat16 to float32 is supported.
  • The output must not contain rstd either before or after fusion.
• Restrictions on fusion into InplaceAddRmsNorm: In scenarios not involving Atlas 350 Accelerator Card, the second output of AddRmsNorm must not have a successor node.

Availability

Atlas A2 training product/Atlas A2 inference product

Atlas A3 training product/Atlas A3 inference product

Atlas inference product

Atlas 350 Accelerator Card