diff --git a/docs/src/index.md b/docs/src/index.md
index 608cfc4..cb59743 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -42,6 +42,7 @@ utils/mixed_precision
 utils/activation_comparison
 utils/casting
 utils/coreai_compression
+utils/composite_op_quantization
 ```
 
 ```{toctree}
diff --git a/docs/src/quantization/overview.md b/docs/src/quantization/overview.md
index 45286f2..753a539 100644
--- a/docs/src/quantization/overview.md
+++ b/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). This happens when the `coreai-torch` conversion needs to "externalize" certain sub-modules to map them to _composite ops_ for better runtime performance.
+- 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
 
diff --git a/docs/src/utils/composite_op_quantization.md b/docs/src/utils/composite_op_quantization.md
new file mode 100644
index 0000000..5978792
--- /dev/null
+++ b/docs/src/utils/composite_op_quantization.md
@@ -0,0 +1,251 @@
+# 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_.
+Refer to the [Externalization](https://apple.github.io/coreai-torch/main/guides/externalization.html) guide for details.
+Here, we will discuss the steps required to quantize a model in `graph` mode using `coreai-opt`.
+
+`graph`-mode quantization invokes `torch.export.export` under the hood, which decomposes a submodule's `forward` into aten ops.
+In order to preserve the composite op structure during this process for externalization, the following APIs are provided:
+
+- `_patch_model_for_externalization`: Patch the model **before** `quantizer.prepare`, so that the composite op call sites survive export and all subsequent quantization passes as opaque nodes.
+- `_subexport_and_restore`: The submodule bodies of the composite ops themselves are then exported and restored before lowering to `CoreAI`.
+
+Quantization treats each composite op as opaque, i.e., no fake-quantize op is placed inside the composite body.
+The composite's input and output boundary can still be quantized, see [Quantizing the composite op boundary](#quantizing-the-composite-op-boundary) below for details.
+
+:::{warning}
+The externalization APIs used below, `_patch_model_for_externalization` and `_subexport_and_restore` in `coreai-torch` are currently experimental.
+:::
+
+```mermaid
+---
+title: "Graph mode Quantization Workflow with Externalization"
+---
+flowchart LR
+    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
+
+`_patch_model_for_externalization` replaces the `forward` of every matching submodule in the model with a `torch.library.custom_op`, in place.
+Call it before constructing the `Quantizer`.
+The example below demonstrates this using the same `RMSNormComposite` op from the [Externalization](https://apple.github.io/coreai-torch/main/guides/externalization.html) guide, however, the same process applies for all composite ops with their respective `ExternalizeSpec`s.
+
+```python
+import torch
+import torch.nn as nn
+from coreai_torch import ExternalizeSpec, _patch_model_for_externalization
+
+
+# The composite op
+class RMSNormComposite(nn.Module):
+    def __init__(self, axes=-1, eps=1e-5, version=1):
+        super().__init__()
+        self.axes = axes
+        self.eps = eps
+        self.version = version
+
+    def forward(self, input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+        x_f32 = input.to(torch.float32)
+        inv_rms = torch.rsqrt((x_f32 * x_f32).mean(self.axes, keepdim=True) + self.eps)
+        return (input * inv_rms).to(input.dtype) * scale
+
+
+# A model that uses the composite op
+class Model(nn.Module):
+    def __init__(self, dim=32):
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+        self.norm = RMSNormComposite()
+        self.norm_weight = nn.Parameter(torch.ones(dim))
+        self.out = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        return self.out(self.norm(self.proj(x), self.norm_weight))
+
+
+model = Model().eval()
+example_inputs = (torch.randn(1, 32),)
+
+# Patch the model in-place
+# to externalize the RMSNormComposite
+_patch_model_for_externalization(
+    model,
+    targets=[
+        ExternalizeSpec(
+            target_class=RMSNormComposite,
+            composite_op_name="rms_norm",
+            composite_attrs=["axes", "eps", "version"],
+        )
+    ],
+)
+```
+
+## Step 2: Prepare, calibrate and finalize
+
+Nothing about the quantizer configuration or the calibration workflow changes.
+The composite op holds no weights of its own here, so weight quantization applies to the surrounding `Linear` layers only.
+
+```python
+import coreai_opt as opt
+from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig
+from coreai_opt.quantization.spec import (
+    default_activation_quantization_spec,
+    default_weight_quantization_spec,
+)
+
+global_config = ModuleQuantizerConfig(
+    op_state_spec={"weight": default_weight_quantization_spec()},
+    op_input_spec={"*": default_activation_quantization_spec()},
+    op_output_spec={"*": default_activation_quantization_spec()},
+)
+quant_config = QuantizerConfig(global_config=global_config)
+
+quantizer = Quantizer(model, quant_config)
+prepared_model = quantizer.prepare(example_inputs)
+
+with quantizer.calibration_mode():
+    for batch in calibration_dataloader:
+        prepared_model(batch)
+
+final_model = quantizer.finalize(backend=opt.ExportBackend.CoreAI)
+```
+
+## Step 3: Export and convert to Core AI
+
+After quantization is complete and the model is finalized, `_subexport_and_restore` API exports each patched composite op and restores the original `forward` method in the model.
+Note that the first argument to `_subexport_and_restore` is the original module that was patched in Step 1, not the finalized `GraphModule`.
+
+```python
+import coreai_torch
+from coreai_torch import TorchConverter, _subexport_and_restore
+
+exported_program = torch.export.export(final_model, example_inputs).run_decompositions(
+    coreai_torch.get_decomp_table()
+)
+externalized = _subexport_and_restore(model, exported_program)
+
+coreai_program = (
+    TorchConverter()
+    .add_exported_program(
+        exported_program, _externalized_exported_programs=externalized
+    )
+    .to_coreai()
+)
+```
+
+In the Core AI graph, the composite op is emitted as a separate private graph that `@main` reaches through `coreai.invoke`:
+
+```text
+// composite op body
+coreai.graph private noinline @norm_57e2d4a8(%arg0: tensor<1x32xf32> {coreai.name = "input"}, %arg1: tensor<32xf32> {coreai.name = "scale"}) -> (tensor<1x32xf32>) attributes {composite_decl = ...} {
+  %2 = coreai.decomposable.broadcasting_mul %0, %1 : (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
+  %4 = coreai.reduce_mean %2, %3 : (tensor<1x32xf32>, tensor<1xsi32>) -> tensor<1x1xf32>
+  %8 = coreai.decomposable.broadcasting_add %6, %7 : (tensor<1x1xf32>, tensor<f32>) -> tensor<1x1xf32>
+  %9 = coreai.rsqrt %8 : tensor<1x1xf32> -> tensor<1x1xf32>
+  %12 = coreai.decomposable.broadcasting_mul %10, %11 : (tensor<1x32xf32>, tensor<1x1xf32>) -> tensor<1x32xf32>
+  %15 = coreai.decomposable.broadcasting_mul %13, %14 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  coreai.output %15 : tensor<1x32xf32>
+}
+
+coreai.graph @main(%arg0: tensor<1x32xf32> {coreai.name = "x"}) -> (tensor<1x32xf32>) {
+  %44 = coreai.decomposable.broadcasting_add %43, %2 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  %53 = coreai.quantize %44, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %62 = coreai.dequantize %53, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+
+  // externalized composite op invocation
+  %63 = coreai.invoke @norm_57e2d4a8(%62, %0) : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  %72 = coreai.quantize %63, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %81 = coreai.dequantize %72, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+  %84 = coreai.decomposable.broadcasting_batch_matmul %81, %83 : (tensor<1x32xf32>, tensor<32x32xf32>) -> tensor<1x32xf32>
+}
+```
+
+(`coreai.cast`, `coreai.constant` and `coreai.reshape` ops omitted above for brevity.)
+
+The composite body carries no `coreai.quantize` or `coreai.dequantize` op and stays in full precision.
+
+## Quantizing the composite op boundary
+
+The `coreai.quantize` pairs surrounding the `coreai.invoke` above come from the global config. They are the output quantizer of the preceding `Linear` and the input quantizer of the following one.
+The composite op boundary itself is not targeted by the global config.
+
+To target the boundary specifically, use `module_input_spec` and `module_output_spec` on a {class}`~coreai_opt.quantization.config.ModuleQuantizerConfig` scoped by `module_type_configs` or `module_name_configs`.
+
+To see this in isolation, the following example uses a model with the composite op alone and specifies a module level spec to quantize it's boundary.
+
+```python
+from coreai_opt.quantization.spec import (
+    PerTensorGranularity,
+    QuantizationScheme,
+    QuantizationSpec,
+)
+
+
+class RMSNormOnly(nn.Module):
+    def __init__(self, dim=32):
+        super().__init__()
+        self.norm = RMSNormComposite()
+        self.norm_weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return self.norm(x, self.norm_weight)
+
+
+boundary_spec = QuantizationSpec(
+    dtype=torch.int8,
+    qscheme=QuantizationScheme.SYMMETRIC,
+    granularity=PerTensorGranularity(),
+)
+quant_config = QuantizerConfig(
+    module_type_configs={
+        RMSNormComposite: ModuleQuantizerConfig(
+            module_input_spec={"*": boundary_spec},
+            module_output_spec={"*": boundary_spec},
+        )
+    },
+)
+```
+
+Running the same patch, prepare, calibrate, finalize and convert steps as above, gives a `@main` graph containing just the boundary quantization and the composite call.
+
+```text
+coreai.graph private noinline @norm_20ea9665(%arg0: tensor<1x32xf32> {coreai.name = "input"}, %arg1: tensor<32xf32> {coreai.name = "scale"}) -> (tensor<1x32xf32>) attributes {composite_decl = ...} {
+  %2 = coreai.decomposable.broadcasting_mul %0, %1 : (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
+  %4 = coreai.reduce_mean %2, %3 : (tensor<1x32xf32>, tensor<1xsi32>) -> tensor<1x1xf32>
+  %8 = coreai.decomposable.broadcasting_add %6, %7 : (tensor<1x1xf32>, tensor<f32>) -> tensor<1x1xf32>
+  %9 = coreai.rsqrt %8 : tensor<1x1xf32> -> tensor<1x1xf32>
+  %12 = coreai.decomposable.broadcasting_mul %10, %11 : (tensor<1x32xf32>, tensor<1x1xf32>) -> tensor<1x32xf32>
+  %15 = coreai.decomposable.broadcasting_mul %13, %14 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  coreai.output %15 : tensor<1x32xf32>
+}
+
+coreai.graph @main(%arg0: tensor<1x32xf32> {coreai.name = "x"}) -> (tensor<1x32xf32>) {
+
+  // Input boundary quantizers for the composite op
+  %13 = coreai.quantize %arg0, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %22 = coreai.dequantize %13, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+
+  // externalized composite op invocation
+  %23 = coreai.invoke @norm_20ea9665(%22, %0) : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+
+  // Output boundary quantizers for the composite op
+  %32 = coreai.quantize %23, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %41 = coreai.dequantize %32, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+  coreai.output %41 : tensor<1x32xf32>
+}
+```
+
+(`coreai.cast`, `coreai.constant` and `coreai.reshape` ops omitted above for brevity.)
+
+## Notes
+
+- The same set of APIs and steps apply for Quantization Aware Training in `graph` mode as well.
diff --git a/pyproject.toml b/pyproject.toml
index 2a172f9..5f1d5d8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,7 +57,7 @@ name = "Apple Core AI Optimization Team"
 [project.optional-dependencies]
 coreai = [
     "coreai-core==1.0.0b2",
-    "coreai-torch==0.4.1",
+    "coreai-torch>=0.4.1",
     "scikit-learn>=1.7.2",
 ]
 coreml = [
@@ -205,6 +205,12 @@ conflicts = [
     ],
 ]
 [tool.uv.sources]
+# TEMPORARY: resolve coreai-torch from the module externalization API branch
+# instead of the PyPI release. That branch adds the `_patch_model_for_externalization`
+# and `_subexport_and_restore` entry points used by the externalization tests.
+# It lives on a fork; apple/coreai-torch does not carry it yet. Drop this entry
+# (and re-pin the versions above) once the work lands upstream and is released.
+coreai-torch = { git = "https://github.com/gokulkrishna98/coreai-torch.git", branch = "dev/gokul/module-externalization-api" }
 torch = [
     { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
     { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
diff --git a/tests/_test_artifacts/mnist/mnist_composite_rmsnorm_pretrained_1epoch_08132026.pt b/tests/_test_artifacts/mnist/mnist_composite_rmsnorm_pretrained_1epoch_08132026.pt
new file mode 100644
index 0000000..b707e36
Binary files /dev/null and b/tests/_test_artifacts/mnist/mnist_composite_rmsnorm_pretrained_1epoch_08132026.pt differ
diff --git a/tests/conftest.py b/tests/conftest.py
index 71e4afa..a441b67 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -32,6 +32,7 @@
     "tests.models.mnist",
     "tests.models.resnet",
     "tests.models.simple",
+    "tests.models.composite",
 ]
 
 _DEFAULT_SEED: int = 42
diff --git a/tests/export/export_utils.py b/tests/export/export_utils.py
index d622324..5203702 100644
--- a/tests/export/export_utils.py
+++ b/tests/export/export_utils.py
@@ -436,10 +436,14 @@ def convert(
         self,
         traced_model: torch.export.ExportedProgram,
         input_data: torch.Tensor,
+        externalized_model: Any = None,
         **kwargs: Any,
     ) -> AIProgram:
         _, _ = input_data, kwargs
-        coreai_program = self._lower_to_coreai(traced_model)
+        coreai_program = self._lower_to_coreai(
+            traced_model,
+            externalized_model=externalized_model,
+        )
         assert type(coreai_program) is AIProgram
 
         return coreai_program
@@ -526,10 +530,25 @@ def _verify_custom_ops_in_torch_program(
     @staticmethod
     def _lower_to_coreai(
         exported_program: torch.export.ExportedProgram,
+        externalized_model: Any = None,
     ) -> AIProgram:
-        """Lower exported program to Core AI."""
+        """Lower exported program to Core AI.
+
+        Args:
+            exported_program: The exported program to lower.
+            externalized_model: Optional ``torch.nn.Module`` that was marked in
+                place by ``coreai_torch._patch_model_for_externalization``.
+        """
         converter = coreai_torch.TorchConverter()
-        converter.add_exported_program(exported_program)
+        externalized_exported_programs = (
+            coreai_torch._subexport_and_restore(externalized_model, exported_program)
+            if externalized_model is not None
+            else None
+        )
+        converter.add_exported_program(
+            exported_program,
+            _externalized_exported_programs=externalized_exported_programs,
+        )
         return converter.to_coreai()
 
 
@@ -554,6 +573,7 @@ def convert_and_verify(
     expected_ops: Mapping[str, int],
     export_backend: ExportBackend,
     prepared_model_output: torch.Tensor | tuple[torch.Tensor, ...],
+    externalized_model: torch.nn.Module | None = None,
     snr_thresh: float = 20.0,
     psnr_thresh: float = 22.0,
     skip_finalized_model_verify: bool = False,
@@ -568,6 +588,9 @@ def convert_and_verify(
         export_backend: Target inference stack (CoreML or CoreAI)
         prepared_model_output: Pre-computed reference output from the prepared
             PyTorch model (single tensor or tuple).
+        externalized_model: Optional ``torch.nn.Module`` that was patched in place by
+            ``coreai_torch._patch_model_for_externalization``. Only supported by the
+            CoreAI backend.
         snr_thresh: Minimum acceptable SNR value
         psnr_thresh: Minimum acceptable PSNR value
         skip_finalized_model_verify: If True, skip forward pass verification on
@@ -578,9 +601,20 @@ def convert_and_verify(
     Returns:
         The converted model in the specified format
 
+    Raises:
+        ValueError: If externalized_model is given for a non-CoreAI backend.
+
     """
     converter = create_converter(export_backend)
 
+    if externalized_model is not None:
+        if export_backend is not ExportBackend.CoreAI:
+            msg = (
+                f"externalized_model is only supported by the CoreAI backend, got {export_backend}"
+            )
+            raise ValueError(msg)
+        converter_kwargs["externalized_model"] = externalized_model
+
     # Run finalized model forward pass BEFORE tracing. torch.export.export()
     # (called in trace) may mutate the model (e.g., strip parametrizations on
     # older PyTorch versions), so the forward pass must happen first.
diff --git a/tests/export/test_graph_mode_mlir_export.py b/tests/export/test_graph_mode_mlir_export.py
index 6d6ea7a..7de0585 100644
--- a/tests/export/test_graph_mode_mlir_export.py
+++ b/tests/export/test_graph_mode_mlir_export.py
@@ -10,6 +10,7 @@
 
 import pytest
 import torch
+from coreai_torch import ExternalizeSpec, _patch_model_for_externalization
 
 from coreai_opt import ExportBackend
 from coreai_opt.palettization.kmeans import KMeansPalettizer
@@ -28,7 +29,15 @@
 from tests.fixtures.compression import ParametrizedP4A8CompressionConfigs
 from tests.fixtures.fp4 import ParametrizedFP4Configs
 from tests.fixtures.fp8 import ParametrizedFP8Configs
-from tests.fixtures.quantization import ParametrizedQuantConfigs
+from tests.fixtures.quantization import (
+    ParametrizedQuantConfigs,
+    make_graph_mode_module_boundary_config,
+    make_quant_config,
+)
+from tests.models.composite import (
+    CompositeSDPAModel,
+    sdpa_externalize_spec,
+)
 
 from . import export_utils
 
@@ -39,6 +48,8 @@ def _run_graph_mode_mlir_export_test_ex(
     config: QuantizerConfig,
     expected_ops: Mapping[str, int],
     model_dtype: torch.dtype | None = None,
+    calibrate: bool = False,
+    externalized_model: torch.nn.Module | None = None,
 ) -> None:
     """Run graph-mode Core AI export test with expanded configuration parameters.
 
@@ -48,6 +59,10 @@ def _run_graph_mode_mlir_export_test_ex(
         config: graph-mode quantization configuration
         model_dtype: Model dtype (float16, float32, bfloat16, or None for no conversion)
         expected_ops: Expected operation counts in converted model
+        calibrate: If True, run one calibration pass under
+            ``quantizer.calibration_mode()`` before the reference forward.
+        externalized_model: The model patched in place by
+            ``coreai_torch._patch_model_for_externalization``.
     """
     if model_dtype is not None:
         model = model.to(dtype=model_dtype)
@@ -57,6 +72,10 @@ def _run_graph_mode_mlir_export_test_ex(
     quantizer = Quantizer(model, config)
     prepared_model = quantizer.prepare((input_data,))
 
+    if calibrate:
+        with quantizer.calibration_mode(), torch.no_grad():
+            prepared_model(input_data)
+
     with torch.no_grad():
         prepared_model_output = prepared_model(input_data)
 
@@ -68,6 +87,7 @@ def _run_graph_mode_mlir_export_test_ex(
         expected_ops=expected_ops,
         export_backend=ExportBackend.CoreAI,
         prepared_model_output=prepared_model_output,
+        externalized_model=externalized_model,
     )
 
 
@@ -438,3 +458,75 @@ def test_integer_quant_minval_export(
             "dequantize": 4 if has_activation_quant else 0,
         },
     )
+
+
+# Composite-op externalize export coverage
+
+
+@pytest.mark.parametrize("config_kind", ["w8", "w8a8", "w8a8-boundary"])
+@pytest.mark.parametrize(
+    # (model, externalize spec, composite submodule path, expected coreai.quantize
+    # count per config kind)
+    "model_cls, externalize_spec, composite_module, expected_quantize_counts",
+    [
+        pytest.param(
+            CompositeSDPAModel,
+            sdpa_externalize_spec(),
+            "composite",
+            {"w8": 0, "w8a8": 4, "w8a8-boundary": 8},
+            id="sdpa",
+        ),
+    ],
+)
+def test_composite_externalize_export(
+    model_cls: type[torch.nn.Module],
+    externalize_spec: ExternalizeSpec,
+    composite_module: str,
+    expected_quantize_counts: Mapping[str, int],
+    config_kind: str,
+) -> None:
+    """End-to-end CoreAI export of a model with an externalized composite op.
+
+    Marks the composite op for externalization, runs graph-mode PTQ, then lowers the
+    finalized graph to a .aimodel and runs it. ``convert_and_verify`` handles SNR /
+    PSNR on the runtime output and op-count verification on the exported program
+    (``constexpr_blockwise_shift_scale`` for weight quantizers and ``quantize`` /
+    ``dequantize`` for activation quantizers).
+
+    Three configs:
+
+    - ``w8`` / ``w8a8``: global config only.
+    - ``w8a8-boundary``: adds a module-scoped ``module_input_spec`` /
+      ``module_output_spec`` so the composite op's i/o edges are quantized
+    """
+    model = model_cls().eval().half()
+    input_data = torch.randn(2, 4, 32, dtype=torch.float16)
+
+    _patch_model_for_externalization(model, [externalize_spec])
+
+    if config_kind == "w8a8-boundary":
+        # uint8 boundary edges stay distinguishable from the int8 global ones.
+        config = make_graph_mode_module_boundary_config(
+            module_boundary_dtype=torch.uint8,
+            module_name=composite_module,
+        )
+    else:
+        config = make_quant_config(
+            weight_dtype=torch.int8,
+            act_dtype=torch.int8 if config_kind == "w8a8" else None,
+            execution_mode="graph",
+        )
+
+    expected_quantize_count = expected_quantize_counts[config_kind]
+    _run_graph_mode_mlir_export_test_ex(
+        model=model,
+        input_data=input_data,
+        config=config,
+        expected_ops={
+            "constexpr_blockwise_shift_scale": 2,
+            "quantize": expected_quantize_count,
+            "dequantize": expected_quantize_count,
+        },
+        calibrate=config_kind != "w8",
+        externalized_model=model,
+    )
diff --git a/tests/fixtures/quantization.py b/tests/fixtures/quantization.py
index 4ee165b..84ea8f7 100644
--- a/tests/fixtures/quantization.py
+++ b/tests/fixtures/quantization.py
@@ -81,6 +81,65 @@ def _spec(dtype: torch.dtype | str) -> QuantizationSpec:
     )
 
 
+def make_graph_mode_module_boundary_config(
+    *,
+    module_boundary_dtype: torch.dtype,
+    module_name: str | None = None,
+    module_type: type | None = None,
+    module_input_spec: dict | None = None,
+    global_dtype: torch.dtype = torch.int8,
+) -> QuantizerConfig:
+    """Build a graph-mode config that also quantizes one module's own i/o boundary.
+
+    The global part comes from ``make_quant_config``; this adds a module-scoped
+    ``module_input_spec`` / ``module_output_spec`` on top of it.
+
+    Args:
+        module_boundary_dtype: Activation dtype for the boundary spec.
+        module_name: Target the module at this path (``module_name_configs``).
+        module_type: Target modules of this type (``module_type_configs``).
+            Exactly one of module_name / module_type must be given.
+        module_input_spec: Override the boundary input spec, e.g.
+            ``{0: spec, 2: spec}`` to select individual positional args.
+            Defaults to the ``"*"`` wildcard over every boundary input.
+        global_dtype: Weight and activation dtype for the global config.
+
+    Returns:
+        QuantizerConfig: the global config plus a module-scoped boundary spec.
+
+    Raises:
+        ValueError: If not exactly one of module_name / module_type is provided.
+    """
+    if (module_name is None) == (module_type is None):
+        msg = "pass exactly one of module_name / module_type"
+        raise ValueError(msg)
+
+    def _boundary_spec() -> QuantizationSpec:
+        return QuantizationSpec(
+            dtype=module_boundary_dtype,
+            qscheme=QuantizationScheme.SYMMETRIC,
+            granularity=PerTensorGranularity(),
+        )
+
+    boundary_config = ModuleQuantizerConfig(
+        module_input_spec=module_input_spec or {"*": _boundary_spec()},
+        module_output_spec={"*": _boundary_spec()},
+    )
+    scope = (
+        {"module_name_configs": {module_name: boundary_config}}
+        if module_name is not None
+        else {"module_type_configs": {module_type: boundary_config}}
+    )
+    base = make_quant_config(
+        weight_dtype=global_dtype, act_dtype=global_dtype, execution_mode="graph"
+    )
+    return QuantizerConfig(
+        global_config=base.global_config,
+        execution_mode="graph",
+        **scope,
+    )
+
+
 @dataclass
 class ParametrizedQuantConfigs:
     """Container for parametrized Eager and PT2E quantization configs.
diff --git a/tests/models/composite.py b/tests/models/composite.py
new file mode 100644
index 0000000..2bb8985
--- /dev/null
+++ b/tests/models/composite.py
@@ -0,0 +1,132 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Test models with composite ops, for externalization test coverage."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from tests.utils import test_artifact_path
+
+# Externalize specs shared by the externalization test modules.
+
+
+def rmsnorm_externalize_spec():
+    """ExternalizeSpec targeting the RMSNormImpl composite."""
+    from coreai_torch import ExternalizeSpec  # noqa: PLC0415
+    from coreai_torch.composite_ops import RMSNormImpl  # noqa: PLC0415
+
+    return ExternalizeSpec(
+        target_class=RMSNormImpl,
+        composite_op_name="rms_norm",
+        composite_attrs=["axes", "eps"],
+    )
+
+
+def sdpa_externalize_spec():
+    """ExternalizeSpec targeting the SDPA composite."""
+    from coreai_torch import ExternalizeSpec  # noqa: PLC0415
+    from coreai_torch.composite_ops import SDPA  # noqa: PLC0415
+
+    return ExternalizeSpec(
+        target_class=SDPA,
+        composite_op_name="scaled_dot_product_attention",
+        composite_attrs=["scale", "is_causal", "window_size"],
+    )
+
+
+class MNISTCompositeRMSNormModel(nn.Module):
+    """Tiny MNIST classifier with an embedded RMSNormImpl composite op.
+
+    Architecture: Flatten -> Linear(28*28, 128) -> RMSNormImpl ->
+    ReLU -> Linear(128, 10) -> LogSoftmax.
+    """
+
+    def __init__(
+        self,
+        hidden: int = 128,
+        num_classes: int = 10,
+        eps: float = 1e-5,
+    ) -> None:
+        from coreai_torch.composite_ops import RMSNormImpl  # noqa: PLC0415
+
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.fc1 = nn.Linear(28 * 28, hidden)
+        self.norm = RMSNormImpl(eps=eps)
+        self.scale = nn.Parameter(torch.ones(hidden))
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden, num_classes)
+        self.softmax = nn.LogSoftmax(dim=-1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.flatten(x)
+        x = self.fc1(x)
+        x = self.norm(x, self.scale)
+        x = self.relu(x)
+        x = self.fc2(x)
+        return self.softmax(x)
+
+
+# Function scoped so each test gets an independently mutable model.
+@pytest.fixture(scope="function")
+def mnist_composite_rmsnorm_pretrained_model() -> MNISTCompositeRMSNormModel:
+    """Load the committed 1-epoch MNISTCompositeRMSNormModel checkpoint.
+
+    Trained with seed 42 for one epoch over the MNIST train split: Adam,
+    lr=1e-3, batch_size=128, shuffle=False, nll_loss.
+    """
+    model = MNISTCompositeRMSNormModel()
+    model.load_state_dict(
+        torch.load(
+            test_artifact_path("mnist/mnist_composite_rmsnorm_pretrained_1epoch_08132026.pt")
+        )
+    )
+    return model
+
+
+@pytest.fixture
+def mnist_composite_rmsnorm_example_input() -> torch.Tensor:
+    """A canonical example-input tensor matching MNIST shape, for prepare()."""
+    return torch.ones(1, 1, 28, 28, dtype=torch.float32)
+
+
+class CompositeRMSNormOnlyModel(nn.Module):
+    """A single RMSNormImpl composite and nothing else."""
+
+    def __init__(self, dim: int = 32, eps: float = 1e-5) -> None:
+        from coreai_torch.composite_ops import RMSNormImpl  # noqa: PLC0415
+
+        super().__init__()
+        self.norm = RMSNormImpl(eps=eps)
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.norm(x, self.scale)
+
+
+class CompositeSDPAModel(nn.Module):
+    """proj -> SDPA(composite) -> output_proj, fp16, single-head fake-dim."""
+
+    def __init__(self, dim: int = 32) -> None:
+        from coreai_torch.composite_ops import SDPA  # noqa: PLC0415
+
+        super().__init__()
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.composite = SDPA(scale=None, is_causal=True)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        # Insert a single fake head dim so SDPA sees rank-4 inputs.
+        q = q.unsqueeze(1)
+        k = k.unsqueeze(1)
+        v = v.unsqueeze(1)
+        attn = self.composite(q, k, v).squeeze(1)
+        return self.out(attn)
diff --git a/tests/quantization/test_composite_op_externalize.py b/tests/quantization/test_composite_op_externalize.py
new file mode 100644
index 0000000..87c4583
--- /dev/null
+++ b/tests/quantization/test_composite_op_externalize.py
@@ -0,0 +1,296 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Externalize-specific structural tests for _patch_model_for_externalization
+in presence of coreai-opt graph mode quantization.
+
+Test structural assertions: after ``_patch_model_for_externalization``
+patches a composite submodule's forward into a ``torch.library.custom_op``,
+the resulting opaque call_function node survives Graph-mode ``prepare`` + ``finalize``.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+from coreai_torch import ExternalizeSpec, _patch_model_for_externalization
+
+from coreai_opt import ExportBackend
+from coreai_opt.quantization import (
+    Quantizer,
+    QuantizerConfig,
+)
+from coreai_opt.quantization.spec import (
+    PerTensorGranularity,
+    QuantizationScheme,
+    QuantizationSpec,
+)
+from tests.fixtures.quantization import (
+    make_graph_mode_module_boundary_config,
+    make_quant_config,
+)
+from tests.models.composite import (
+    CompositeRMSNormOnlyModel,
+    CompositeSDPAModel,
+    rmsnorm_externalize_spec,
+    sdpa_externalize_spec,
+)
+from tests.test_utils.general import (
+    assert_single_call_function_node,
+    get_quantize_dtype,
+    is_coreai_dequantize,
+    is_coreai_quantize,
+)
+
+
+@pytest.mark.parametrize(
+    "quantize_activations",
+    [
+        pytest.param(False, id="w8-weight-only"),
+        pytest.param(True, id="w8a8"),
+    ],
+)
+def test_composite_op_survives_prepare_and_finalize(
+    quantize_activations: bool,
+) -> None:
+    """The externalized composite must remain a single opaque
+    call_function node end-to-end, under both w8 and w8a8.
+    """
+    model = CompositeSDPAModel().eval().half()
+    sample = torch.randn(2, 4, 32, dtype=torch.float16)
+
+    _patch_model_for_externalization(model, [sdpa_externalize_spec()])
+    op_name = model.composite._externalize_op_name
+    target_substr = f"coreai_torch_ext.{op_name}"
+
+    quantizer = Quantizer(
+        model,
+        make_quant_config(
+            weight_dtype=torch.int8,
+            act_dtype=torch.int8 if quantize_activations else None,
+            execution_mode="graph",
+        ),
+    )
+    prepared = quantizer.prepare((sample,))
+    assert_single_call_function_node(prepared, target_substr, stage="prepared")
+
+    finalized = quantizer.finalize(backend=ExportBackend.CoreAI)
+    assert_single_call_function_node(finalized, target_substr, stage="finalized")
+
+
+# Composite op I/O boundary quantization
+
+
+class TestCompositeOpIOQuantization:
+    """Ensure externalized composite's I/O boundary can be quantized
+    via a module-level config, by name and by type.
+
+    A global config quantizes the rest of the model with the default activation
+    dtype (int8) and the composite config provides a distinct boundary dtype
+    (uint8) on the composite's edges. Module config outranks global, so the
+    composite boundary must carry the composite dtype while every other quantized
+    edge carries the global dtype.
+    """
+
+    @staticmethod
+    def _composite_act_spec(boundary_dtype: torch.dtype) -> QuantizationSpec:
+        return QuantizationSpec(
+            dtype=boundary_dtype,
+            qscheme=QuantizationScheme.SYMMETRIC,
+            granularity=PerTensorGranularity(),
+        )
+
+    @staticmethod
+    def _config(
+        spec: ExternalizeSpec,
+        module_name: str,
+        target_by: str,
+        boundary_dtype: torch.dtype,
+        module_input_spec: dict | None = None,
+    ) -> QuantizerConfig:
+        return make_graph_mode_module_boundary_config(
+            module_boundary_dtype=boundary_dtype,
+            module_name=module_name if target_by == "name" else None,
+            module_type=spec.target_class if target_by == "type" else None,
+            module_input_spec=module_input_spec,
+        )
+
+    def _quantize_with_externalization_and_verify(
+        self,
+        model: nn.Module,
+        sample: torch.Tensor,
+        spec: ExternalizeSpec,
+        module_name: str,
+        config: QuantizerConfig,
+    ) -> tuple[torch.fx.GraphModule, str]:
+        """Run graph-mode PTQ on the model and verify the composite op stays opaque.
+
+        Patches the composite for externalization, prepares, finalizes, and asserts
+        the composite is still a single opaque call_function node after each of
+        those two stages.
+
+        Returns the finalized graph and the substring identifying that node.
+        """
+        _patch_model_for_externalization(model, [spec])
+        op_name = model.get_submodule(module_name)._externalize_op_name
+        target_substr = f"coreai_torch_ext.{op_name}"
+
+        quantizer = Quantizer(model, config)
+        prepared = quantizer.prepare((sample,))
+        assert_single_call_function_node(prepared, target_substr, stage="prepared")
+
+        finalized = quantizer.finalize(backend=ExportBackend.CoreAI)
+        assert_single_call_function_node(finalized, target_substr, stage="finalized")
+        return finalized, target_substr
+
+    def _assert_boundary_quantized(
+        self,
+        finalized: torch.fx.GraphModule,
+        target_substr: str,
+        num_tensor_inputs: int,
+        boundary_dtype: torch.dtype,
+    ) -> None:
+        composite = assert_single_call_function_node(finalized, target_substr, stage="finalized")
+
+        # A composite's non-tensor captured attributes appear either as baked-in
+        # constants (SDPA's scale / is_causal / window_size) or as a get_attr arg
+        # (RMSNorm's scale), and neither is a quantized activation edge, so
+        # filter get_attr out rather than indexing fixed arg positions.
+        tensor_inputs = [
+            a for a in composite.args if isinstance(a, torch.fx.Node) and a.op != "get_attr"
+        ]
+        assert len(tensor_inputs) == num_tensor_inputs, (
+            f"Expected {num_tensor_inputs} tensor inputs to {composite.name}, "
+            f"got {[n.name for n in tensor_inputs]}"
+        )
+        for act_input in tensor_inputs:
+            assert is_coreai_dequantize(act_input.target)
+            assert get_quantize_dtype(act_input.args[0]) == boundary_dtype
+
+        users = list(composite.users)
+        assert len(users) == 1
+        assert is_coreai_quantize(users[0].target)
+        assert get_quantize_dtype(users[0]) == boundary_dtype
+
+        composite_dtype_quant = [
+            n
+            for n in finalized.graph.nodes
+            if is_coreai_quantize(n.target) and get_quantize_dtype(n) == boundary_dtype
+        ]
+        assert len(composite_dtype_quant) == num_tensor_inputs + 1
+
+    @pytest.mark.parametrize("target_by", ["name", "type"])
+    @pytest.mark.parametrize(
+        # (model class, externalize spec, submodule attribute name, tensor input count).
+        # Both models default to dim=32 and accept the same rank-3 fp16 sample.
+        "model_cls, spec, module_name, num_tensor_inputs",
+        [
+            pytest.param(
+                CompositeRMSNormOnlyModel,
+                rmsnorm_externalize_spec(),
+                "norm",
+                1,
+                id="rmsnorm-only",
+            ),
+            pytest.param(
+                CompositeSDPAModel,
+                sdpa_externalize_spec(),
+                "composite",
+                3,
+                id="sdpa-qkv",
+            ),
+        ],
+    )
+    def test_composite_boundary_quantized(
+        self,
+        model_cls: type[nn.Module],
+        spec: ExternalizeSpec,
+        module_name: str,
+        num_tensor_inputs: int,
+        target_by: str,
+    ) -> None:
+        # uint8 boundary edges stay distinguishable from the int8 global ones.
+        boundary_dtype = torch.uint8
+        model = model_cls().eval().half()
+        sample = torch.randn(2, 4, 32, dtype=torch.float16)
+        config = self._config(spec, module_name, target_by, boundary_dtype)
+        finalized, target_substr = self._quantize_with_externalization_and_verify(
+            model, sample, spec, module_name, config
+        )
+        self._assert_boundary_quantized(finalized, target_substr, num_tensor_inputs, boundary_dtype)
+
+    @pytest.mark.parametrize("target_by", ["name", "type"])
+    def test_composite_boundary_input_index_selects_those_args(self, target_by: str) -> None:
+        """Integer keys in ``module_input_spec`` quantize exactly those positional args
+        for composite ops.
+
+        The unselected input is left unquantized rather than falling
+        back to the global spec, because the composite is opaque to the
+        op-pattern annotator and only a module-level config reaches its edges.
+        """
+        # uint8 boundary edges stay distinguishable from the int8 global ones.
+        boundary_dtype = torch.uint8
+        quantized_indices = (0, 2)
+        num_tensor_inputs = 3
+        model = CompositeSDPAModel().eval().half()
+        sample = torch.randn(2, 4, 32, dtype=torch.float16)
+        spec = sdpa_externalize_spec()
+        config = self._config(
+            spec,
+            "composite",
+            target_by,
+            boundary_dtype,
+            module_input_spec={
+                i: self._composite_act_spec(boundary_dtype) for i in quantized_indices
+            },
+        )
+
+        finalized, target_substr = self._quantize_with_externalization_and_verify(
+            model, sample, spec, "composite", config
+        )
+
+        composite = assert_single_call_function_node(finalized, target_substr, stage="finalized")
+        tensor_inputs = [
+            a for a in composite.args if isinstance(a, torch.fx.Node) and a.op != "get_attr"
+        ]
+        assert len(tensor_inputs) == num_tensor_inputs, (
+            f"Expected {num_tensor_inputs} tensor inputs to {composite.name}, "
+            f"got {[n.name for n in tensor_inputs]}"
+        )
+
+        for index, act_input in enumerate(tensor_inputs):
+            if index in quantized_indices:
+                assert is_coreai_dequantize(act_input.target), (
+                    f"input {index} was selected by module_input_spec but is not fed by "
+                    f"a dequantize: {act_input.target}"
+                )
+                input_dtype = get_quantize_dtype(act_input.args[0])
+                assert input_dtype == boundary_dtype, (
+                    f"input {index} was selected by module_input_spec but is quantized as "
+                    f"{input_dtype}, expected the composite dtype {boundary_dtype}"
+                )
+            else:
+                assert not is_coreai_dequantize(act_input.target), (
+                    f"input {index} was not selected by module_input_spec but is fed by "
+                    f"a dequantize: {act_input.target}"
+                )
+
+        # module_output_spec stays the wildcard, so the composite's consumer is
+        # quantized with the composite dtype regardless of which inputs were selected.
+        consumers = list(composite.users)
+        assert len(consumers) == 1, (
+            f"expected the composite to have exactly one consumer, got "
+            f"{[n.name for n in consumers]}"
+        )
+        consumer = consumers[0]
+        assert is_coreai_quantize(consumer.target), (
+            f"the composite's consumer is not a quantize node: {consumer.target}"
+        )
+        consumer_dtype = get_quantize_dtype(consumer)
+        assert consumer_dtype == boundary_dtype, (
+            f"the composite's output is quantized as {consumer_dtype}, expected the "
+            f"composite dtype {boundary_dtype}"
+        )
diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py
index 901cb24..ef95594 100644
--- a/tests/quantization/test_graph_mode_quantizer.py
+++ b/tests/quantization/test_graph_mode_quantizer.py
@@ -50,6 +50,7 @@
     StaticQParamsCalculator,
 )
 from tests.models.simple import SimpleModel
+from tests.test_utils.general import get_fake_quant_nodes
 
 
 @pytest.fixture
@@ -88,16 +89,6 @@ def weight_only_config():
     )
 
 
-def get_fake_quant_nodes(model: torch.fx.GraphModule) -> list[torch.fx.Node]:
-    """
-    Returns list of fake quant nodes present in the input model
-    """
-    fake_quant_nodes = [
-        node for node in model.graph.nodes if "activation_post_process" in node.name
-    ]
-    return fake_quant_nodes
-
-
 class TestGraphModeQuantizer:
     """Test cases for GraphQuantizer class."""
 
@@ -1223,125 +1214,6 @@ def test_train_eval_compatible_with_context_managers(
         assert prepared_model.training
 
 
-class TestCompositeOpQuantization:
-    """
-    Tests that models with COREAI CompositeOps can be quantized
-    """
-
-    class SDPAModule(torch.nn.Module):
-        def forward(self, query, key, value):
-            from coreai_torch.composite_ops._sdpa import (  # noqa: PLC0415
-                scaled_dot_product_attention as _scaled_dot_product_attention,
-            )
-
-            return _scaled_dot_product_attention(query, key, value, is_causal=True)
-
-    class SimpleSDPAModel(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.proj = torch.nn.Linear(16, 48)
-            self.sdpa = TestCompositeOpQuantization.SDPAModule()
-
-        def forward(self, x):
-            qkv = self.proj(x)
-            q, k, v = qkv.chunk(3, dim=-1)
-            b, s, _ = q.shape
-            q = q.reshape(b, 1, s, 16)
-            k = k.reshape(b, 1, s, 16)
-            v = v.reshape(b, 1, s, 16)
-            return self.sdpa(q, k, v)
-
-    @pytest.fixture
-    def model(self):
-        return self.SimpleSDPAModel()
-
-    @pytest.fixture
-    def example_input(self):
-        return torch.randn(1, 4, 16)
-
-    @pytest.mark.xfail(reason="tracked by coreai-torch issue #309")
-    def test_composite_op_io_quantization(self, model, example_input):
-        """
-        Verify CompositeOps boundaries can be quantized
-        """
-        qspec_dict = {
-            "dtype": "int8",
-            "qscheme": "symmetric",
-            "granularity": {"type": "per_tensor"},
-        }
-        config = QuantizerConfig.from_dict(
-            {
-                "quantization_config": {
-                    "global_config": {"op_state_spec": {"weight": qspec_dict}},
-                    "module_name_configs": {
-                        "sdpa": {
-                            "op_input_spec": None,
-                            "op_output_spec": None,
-                            "op_state_spec": None,
-                            "module_input_spec": {
-                                0: qspec_dict,
-                                1: qspec_dict,
-                                2: qspec_dict,
-                            },
-                            "module_output_spec": {
-                                "*": qspec_dict,
-                            },
-                        }
-                    },
-                }
-            }
-        )
-
-        quantizer = Quantizer(model, config)
-        prepared_model = quantizer.prepare((example_input,))
-        assert isinstance(prepared_model, torch.fx.GraphModule)
-
-        def _is_composite_op(node):
-            return (
-                node.op == "call_function"
-                and isinstance(node.target, torch._ops.OpOverload)
-                and node.target.namespace == "CompositeOps"
-            )
-
-        # Find the SDPA node
-        sdpa_nodes = [
-            n
-            for n in prepared_model.graph.nodes
-            if n.op == "call_function"
-            and isinstance(n.target, torch._ops.OpOverload)
-            and n.target._opname == "scaled_dot_product_attention"
-        ]
-        assert len(sdpa_nodes) == 1, f"Expected 1 SDPA node, got {len(sdpa_nodes)}"
-        sdpa_node = sdpa_nodes[0]
-
-        composite_op_input_nodes = [arg for arg in sdpa_node.args if _is_composite_op(arg)]
-        composite_op_output_nodes = [user for user in sdpa_node.users if _is_composite_op(user)]
-
-        assert len(composite_op_input_nodes) == 3, (
-            f"Expected 3 custom input nodes (q/k/v), got {len(composite_op_input_nodes)}"
-        )
-        assert len(composite_op_output_nodes) == 1, "Expected 1 composite op output node"
-
-        for node in composite_op_input_nodes + composite_op_output_nodes:
-            assert "name" in node.kwargs, f"kwargs not restored for {node.name}"
-            assert "op_name" in node.kwargs, f"kwargs not restored for {node.name}"
-            assert node.kwargs["op_name"] == "scaled_dot_product_attention"
-
-        for node in composite_op_input_nodes:
-            input_node = node.args[0]
-            assert "activation_post_process" in input_node.name, (
-                f"Expected fake quant before {node.name} "
-                f"(input={node.kwargs['name']}), "
-                f"got {input_node.name} instead"
-            )
-
-        for node in composite_op_output_nodes:
-            users = list(node.users.keys())
-            assert any("activation_post_process" in u.name for u in users), (
-                f"Expected fake quant after {node.name}, got users: {[u.name for u in users]}"
-            )
-
-
 class TestFP4MLIRExportValidation:
     """Test that FP4 export validation rejects unsupported configurations."""
 
diff --git a/tests/quantization/test_graph_mode_quantizer_mnist.py b/tests/quantization/test_graph_mode_quantizer_mnist.py
index adcc3ce..a23b5fc 100644
--- a/tests/quantization/test_graph_mode_quantizer_mnist.py
+++ b/tests/quantization/test_graph_mode_quantizer_mnist.py
@@ -7,6 +7,7 @@
 
 import pytest
 import torch
+from coreai_torch import _patch_model_for_externalization
 
 import tests.utils as utils
 from coreai_opt import ExportBackend
@@ -20,6 +21,9 @@
     PerChannelGranularity,
     PerTensorGranularity,
 )
+from tests.export import export_utils
+from tests.models.composite import rmsnorm_externalize_spec
+from tests.test_utils.general import assert_single_call_function_node
 
 image_size = 28
 batch_size = 128
@@ -327,3 +331,97 @@ def test_weight_and_activation_qat_mnist(mnist_pretrained_model, mnist_dataset,
 
     # Accuracy before and after finalize should match
     assert post_qat_accuracy == finalized_accuracy
+
+
+@pytest.mark.slow
+@pytest.mark.seed
+def test_weight_and_activation_qat_mnist_with_externalized_composite(
+    mnist_composite_rmsnorm_pretrained_model,
+    mnist_composite_rmsnorm_example_input,
+    mnist_dataset,
+):
+    """QAT on an MNIST classifier with an externalized RMSNormImpl composite.
+
+    Graph-mode QAT flow: baseline accuracy -> mark
+    for externalization -> prepare -> post-prepare drop -> train under
+    ``training_mode()`` -> post-QAT recovery -> finalize -> finalized
+    accuracy matches post-QAT accuracy.
+    """
+    train_loader, test_loader = utils.setup_data_loaders(mnist_dataset, batch_size)
+
+    model = mnist_composite_rmsnorm_pretrained_model
+    accuracy = utils.eval_model(model, test_loader)
+    assert accuracy > 92.0, (
+        f"expect pretrained MNIST-composite model accuracy > 92%, got {accuracy:.2f}%"
+    )
+
+    _patch_model_for_externalization(model, [rmsnorm_externalize_spec()])
+    op_name = model.norm._externalize_op_name
+    target_substr = f"coreai_torch_ext.{op_name}"
+
+    config = QuantizerConfig(global_config=ModuleQuantizerConfig())
+    quantizer = Quantizer(model, config)
+
+    prepared_model = quantizer.prepare(
+        example_inputs=(mnist_composite_rmsnorm_example_input,),
+    )
+
+    post_prepare_accuracy = utils.eval_model(prepared_model, test_loader)
+    assert post_prepare_accuracy < 90.0, (
+        f"Expect accuracy to drop below 90% after preparation with an all ones data sample; "
+        f"got {post_prepare_accuracy:.2f}% (baseline {accuracy:.2f}%)"
+    )
+
+    optimizer = torch.optim.Adam(prepared_model.parameters(), eps=1e-3, weight_decay=1e-4)
+    with quantizer.training_mode():
+        for batch_idx, (data, target) in enumerate(train_loader):
+            utils.train_step(
+                prepared_model,
+                optimizer,
+                train_loader,
+                data,
+                target,
+                batch_idx,
+                epoch=0,
+            )
+
+    post_qat_accuracy = utils.eval_model(prepared_model, test_loader)
+    assert post_qat_accuracy > 94.0, (
+        f"Expect accuracy to climb above 94% after QAT, got {post_qat_accuracy:.2f}%"
+    )
+
+    finalized_model = quantizer.finalize(backend=ExportBackend.CoreAI)
+    finalized_accuracy = utils.eval_model(finalized_model, test_loader)
+    assert post_qat_accuracy == finalized_accuracy, (
+        f"post-QAT accuracy ({post_qat_accuracy:.2f}%) must match "
+        f"post-finalize accuracy ({finalized_accuracy:.2f}%)"
+    )
+
+    assert_single_call_function_node(finalized_model, target_substr, stage="finalized")
+
+    # Lower to .aimodel and verify the runtime output matches the
+    # finalized torch output. expected_ops pins the weight constexpr and
+    # activation q-dq node counts graph mode inserts for the
+    # flatten -> Linear -> composite -> ReLU -> Linear -> LogSoftmax graph
+    # under the default w8a8 (int8) global config:
+    #   - 2 constexpr_blockwise_shift_scale: the two Linear weights.
+    #   - 5 quantize / 5 dequantize: the five annotated activation edges,
+    #     namely the model input, fc1's input and output, and fc2's input
+    #     and output.
+    # ReLU, LogSoftmax and the composite carry no quantization annotation
+    # and contribute no q-dq nodes.
+    sample_input = mnist_composite_rmsnorm_example_input
+    with torch.no_grad():
+        prepared_for_export_output = prepared_model(sample_input)
+    export_utils.convert_and_verify(
+        finalized_model=finalized_model,
+        input_data=sample_input,
+        expected_ops={
+            "constexpr_blockwise_shift_scale": 2,
+            "quantize": 5,
+            "dequantize": 5,
+        },
+        export_backend=ExportBackend.CoreAI,
+        prepared_model_output=prepared_for_export_output,
+        externalized_model=model,
+    )
diff --git a/tests/test_utils/general.py b/tests/test_utils/general.py
index b6ba93d..6a42f42 100644
--- a/tests/test_utils/general.py
+++ b/tests/test_utils/general.py
@@ -9,6 +9,7 @@
 from decimal import ROUND_FLOOR, Decimal
 
 import torch
+from torch.ops import coreai
 
 COREAI_AVAILABLE = importlib.util.find_spec("coreai") is not None
 
@@ -70,7 +71,6 @@ def compute_snr_psnr(
 
     Returns:
         Tuple of (SNR, PSNR) values
-
     """
     assert len(data) == len(reference), f"Tensor length mismatch: {len(data)} vs {len(reference)}"
 
@@ -113,3 +113,51 @@ def verify_snr_psnr(
 
     if snr <= snr_thresh or psnr <= psnr_thresh:
         raise SNRBelowThresholdError(snr, psnr, snr_thresh, psnr_thresh, prefix)
+
+
+def get_fake_quant_nodes(model: torch.fx.GraphModule) -> list[torch.fx.Node]:
+    """Return the activation_post_process observer nodes in a prepared graph."""
+    return [node for node in model.graph.nodes if "activation_post_process" in node.name]
+
+
+def assert_single_call_function_node(
+    gm: torch.fx.GraphModule, target_substr: str, *, stage: str = ""
+) -> torch.fx.Node:
+    """Assert exactly one call_function node's target contains a substring.
+
+    Current usage: externalized composite ops get a process-unique op name (sanitized module
+    path plus a uuid4 suffix), so there is no stable target object to compare
+    against and a substring match on ``str(node.target)`` is required.
+
+    Args:
+        gm: The graph module to search
+        target_substr: Substring to match against ``str(node.target)``
+        stage: Optional pipeline stage name, used only in the error message
+
+    Returns:
+        The single matching node
+    """
+    matches = [
+        n for n in gm.graph.nodes if n.op == "call_function" and target_substr in str(n.target)
+    ]
+    where = f" in the {stage.upper()} graph" if stage else ""
+    assert len(matches) == 1, (
+        f"Expected exactly one call_function matching '{target_substr}'"
+        f"{where}; got {[n.name for n in matches]}"
+    )
+    return matches[0]
+
+
+def is_coreai_quantize(target: object) -> bool:
+    """Whether an FX node target is the ``coreai::quantize`` op."""
+    return target is coreai.quantize or target is coreai.quantize.default
+
+
+def is_coreai_dequantize(target: object) -> bool:
+    """Whether an FX node target is the ``coreai::dequantize`` op."""
+    return target is coreai.dequantize or target is coreai.dequantize.default
+
+
+def get_quantize_dtype(node: torch.fx.Node) -> torch.dtype | None:
+    """Return the quantized dtype carried in an FX node's args, else None."""
+    return next((a for a in node.args if isinstance(a, torch.dtype)), None)