[TRANSFORMATIONS] Preserve bias output precision during activation scaling

[andrew-k-park]

Details:

• Fix ScaleDownSingleLayer incorrectly restoring the MatMul output precision when a supported bias Add extends the scaled region.
• Original graph

  flowchart LR
      A["MatMul<br/>output: FP16"] --> B["TypeRelaxed Add<br/>output contract: FP32"]
      Bias["Bias<br/>FP32"] --> B
      B --> C["GELU<br/>FP32"]
      C --> D["Downstream MatMul<br/>activation: FP32<br/>weights: FP32"]
      W["Decompressed weights<br/>FP32"] --> D
  

  Loading
  • The first MatMul computes in FP16, while its type-relaxed bias Add intentionally exposes an FP32 output.
  • The downstream GELU and MatMul therefore receive FP32 activations.
• Graph before the fix

  flowchart LR
      A["MatMul<br/>output: FP16"] --> B["TypeRelaxed Add<br/>output: FP32"]
      Bias["Scaled bias<br/>FP32"] --> B
      B --> X["Convert<br/>FP32 to FP16"]
      X --> S["Scale-up Multiply<br/>FP16"]
      S --> C["GELU<br/>FP16"]
      C --> D["Downstream MatMul<br/>INVALID"]
      W["Decompressed weights<br/>FP32"] --> D
  
      style C fill:#fff3cd,stroke:#b58105,color:#24292f
      style W fill:#fff3cd,stroke:#b58105,color:#24292f
      style D fill:#ffebe9,stroke:#cf222e,stroke-width:2px,color:#24292f
  

  Loading
  • ScaleDownSingleLayer captured output_prec from the first MatMul before detecting that the bias Add extended the scaled region.
  • It consequently inserted an FP16 conversion after the FP32 Add and changed the downstream activation to FP16.
  • The downstream MatMul was left with incompatible input types:

    Arguments do not have the same element type
    (arg0 element type: f16, arg1 element type: f32)
    

• Graph after the fix

  flowchart LR
      A["MatMul<br/>output: FP16"] --> B["TypeRelaxed Add<br/>output: FP32"]
      Bias["Scaled bias<br/>FP32"] --> B
      B --> S["Scale-up Multiply<br/>FP32"]
      S --> C["GELU<br/>FP32"]
      C --> D["Downstream MatMul<br/>VALID"]
      W["Decompressed weights<br/>FP32"] --> D
  
      style S fill:#ddf4ff,stroke:#0969da,color:#24292f
      style D fill:#dafbe1,stroke:#1a7f37,stroke-width:2px,color:#24292f
  

  Loading
  • When a supported bias Add extends the scaled region, use the Add's original output element type as the scale-up restoration precision:

    if (has_bias) {
        auto add = child_node->shared_from_this();
        output_prec = add->get_output_element_type(0);
        // ...
    }

  • This preserves the FP32 output contract of the type-relaxed Add while retaining FP16 computation and activation scaling in the first MatMul.

Tickets:

• CVS-194558

AI Assistance:

• AI assistance used: yes
• AI was used to help analyze transformation-pass traces, isolate the minimal failing graph, draft the source change and regression test. Human validation included reviewing the source diff, building the targets, running common transformation and GPU functional tests, and completing the end-to-end WWB evaluation

[Copilot]

🟡 Changes recommended

The regression test must verify that the transformation actually executes.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Fixes activation scaling to preserve a type-relaxed bias Add’s output precision.

Changes:

• Uses the bias Add’s output type for scale restoration.
• Adds an FP32 regression scenario.

File summaries

File	Description
activations_scaling.cpp	Preserves bias Add output precision.
activations_scaling_test.cpp	Adds TypeRelaxed Add coverage.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.