update flowchart

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

Author: pkmandke

docs/src/quantization/overview.md

@@ -204,7 +204,7 @@ The two modes are expected to produce very similar models for weight-only quanti
 A few scenarios where `eager` mode may need to be used instead of `graph`:
 
 - If you run into any errors during the `prepare` call which, under the hood, invokes the `torch.export.export` and `torchao`'s `prepare_qat_pt2e`/`convert_pt2e` APIs. See [Graph Mode Troubleshooting](../debugging/graph_mode_troubleshooting.md) for common export errors and workarounds before falling back to eager mode.
-- When `torch.nn.Module` needs to be provided as an input, instead of `ExportedProgram` to the conversion API of [coreai-torch](https://github.com/apple/coreai-torch). Note that models whose submodules must be "externalized" to map them to _composite ops_ for better runtime performance can still be quantized in graph mode. See [Quantizing Models with Core AI Composite Ops](../utils/composite_op_quantization.md).
+- When `torch.nn.Module` needs to be provided as an input, instead of `ExportedProgram` to the conversion API of [coreai-torch](https://github.com/apple/coreai-torch). Note that models whose submodules must be "externalized" to map them to _composite ops_ for better runtime performance can still be quantized in graph mode. See [Quantizing Models with Core AI Composite Ops in Graph Mode](../utils/composite_op_quantization.md).
 
 #### Weights and activations quantization
 

docs/src/utils/composite_op_quantization.md

@@ -1,4 +1,4 @@
-# Quantizing Models with Core AI Composite Ops
+# Quantizing Models with Core AI Composite Ops in Graph Mode
 
 Core AI recognizes certain well-known building blocks, such as SDPA or RMSNorm, as _composite ops_ and applies optimized implementations for them.
 `coreai-torch` establishes those boundaries through _externalization_.
@@ -20,17 +20,17 @@ The externalization APIs used below, `_patch_model_for_externalization` and `_su
 
 ```mermaid
 ---
-title: "Composite Op Quantization Workflow"
+title: "Graph mode Quantization Workflow with Externalization"
 ---
 flowchart LR
-    model["Full Precision<br>Model"] --> patch["Patch model"]
-    patch --> prepare["quantizer.prepare(...)"]
-    prepare --> calibrate["Calibrate"]
-    calibrate --> qfin["quantizer.finalize(...)"]
-    qfin --> export["torch.export.export(...)"]
-    export --> sub["Sub-export<br>and restore"]
-    sub --> convert["TorchConverter().to_coreai()"]
+    model["Full Precision<br>Model"] --> patch["Patch Model for<br>Externalization"]
+    patch --> prepare["Prepare and<br>Calibrate"]
+    prepare --> qfin["Finalize and<br>Export"]
+    qfin --> sub["Sub-export<br>and Restore"]
+    sub --> convert["Convert to<br>Core AI"]
     style model fill:#f9f9f9,stroke:#999
+    style patch fill:#e8f0fe,stroke:#4285f4
+    style sub fill:#e8f0fe,stroke:#4285f4
 ```
 
 ## Step 1: Patch the model before prepare
@@ -239,4 +239,4 @@ coreai.graph @main(%arg0: tensor<1x32xf32> {coreai.name = "x"}) -> (tensor<1x32x
 
 ## Notes
 
-- The same set of APIs and workflows can be followed for Quantization Aware Training in `graph` mode as well.
+- The same set of APIs and steps apply for Quantization Aware Training in `graph` mode as well.