[16/17] DLA public integration and reshape cleanup

CHANGELOG.rst

@@ -34,6 +34,8 @@ Changelog
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
+- Add the ordered DLA compatibility pipeline through the ``modelopt.onnx.graph_surgery.make_dla_compatible``
+  API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype.
 
 **Backward Breaking Changes**
 

modelopt/onnx/graph_surgery/__init__.py

@@ -23,13 +23,15 @@
 - Converting model precision (e.g., FP16 to BF16)
 - Transposing DequantizeLinear weights for column-major storage optimization
 - Graph cleanup and optimization
+- Making models compatible with the NVIDIA DLA accelerator
 
 Example usage:
     >>> from modelopt.onnx.graph_surgery import (
     ...     replace_attention_with_gqa,
     ...     convert_fp16_to_bf16,
     ...     transpose_dequantize_linear_weights,
     ...     add_cross_kv_to_encoder,
+    ...     make_dla_compatible,
     ... )
     >>> # Replace attention with GQA for LLMs (FP16 model)
     >>> replace_attention_with_gqa(
@@ -62,18 +64,25 @@
     ...     model_path="model_quantized.onnx",
     ...     output_path="model_quantized_transposed.onnx",
     ... )
+    >>> # Apply the ordered DLA compatibility pipeline
+    >>> make_dla_compatible(
+    ...     model_path="model.onnx",
+    ...     output_path="model_dla.onnx",
+    ... )
 """
 
 from .dq_transpose import transpose_dequantize_linear_weights
 from .encoder_cross_kv import add_cross_kv_to_encoder
 from .gqa_replacement import replace_attention_with_gqa
+from .make_dla_compatible import dla_make_dla_compatible as make_dla_compatible
 from .utils.dtype_conversion import convert_fp16_to_bf16
 
 _SURGERY_REGISTRY = {
     "replace-gqa": replace_attention_with_gqa,
     "add-cross-kv": add_cross_kv_to_encoder,
     "convert-bf16": convert_fp16_to_bf16,
     "transpose-dq": transpose_dequantize_linear_weights,
+    "make-dla-compatible": make_dla_compatible,
 }
 
 
@@ -132,6 +141,7 @@ def run_graph_surgery(
     "add_cross_kv_to_encoder",
     "convert_fp16_to_bf16",
     "get_available_surgeries",
+    "make_dla_compatible",
     "replace_attention_with_gqa",
     "run_graph_surgery",
     "transpose_dequantize_linear_weights",

modelopt/onnx/graph_surgery/__main__.py

@@ -50,6 +50,12 @@
         --input model_quantized.onnx \
         --output model_quantized_transposed.onnx
 
+Make a model compatible with the NVIDIA DLA accelerator::
+
+    python -m modelopt.onnx.graph_surgery make-dla-compatible \
+        --input model.onnx \
+        --output model_dla.onnx
+
 Analyze attention pattern::
 
     python -m modelopt.onnx.graph_surgery analyze \
@@ -84,6 +90,9 @@ def main():
   Transpose DequantizeLinear weights:
     python -m modelopt.onnx.graph_surgery transpose-dq -i model_quantized.onnx -o model_transposed.onnx
 
+  Make a model DLA-compatible:
+    python -m modelopt.onnx.graph_surgery make-dla-compatible -i model.onnx -o model_dla.onnx
+
   Analyze attention pattern:
     python -m modelopt.onnx.graph_surgery analyze -i model.onnx --layer 0
         """,
@@ -235,6 +244,47 @@ def main():
         "-q", "--quiet", action="store_true", help="Suppress progress messages"
     )
 
+    # Make DLA compatible subcommand
+    dla_parser = subparsers.add_parser(
+        "make-dla-compatible",
+        help="Make a model compatible with the NVIDIA DLA accelerator",
+        description=(
+            "Apply the ordered DLA graph transforms and validate the resulting ONNX model. "
+            "Where lowering assumes finite FP32 branches and may change the sign of zero."
+        ),
+    )
+    dla_parser.add_argument("-i", "--input", required=True, help="Input ONNX model path")
+    dla_parser.add_argument("-o", "--output", required=True, help="Output ONNX model path")
+    dla_parser.add_argument(
+        "--no-external-data",
+        action="store_true",
+        help="Don't save weights as external data",
+    )
+    dla_parser.add_argument(
+        "--external-data-name",
+        type=str,
+        default=None,
+        help="Name for external data file",
+    )
+    qdq_group = dla_parser.add_mutually_exclusive_group()
+    qdq_group.add_argument(
+        "--qdq-quantized-dtypes",
+        nargs="+",
+        default=None,
+        metavar="DTYPE",
+        help=(
+            "Zero-point storage dtypes whose Q/DQ pairs are stripped in the early "
+            "QDQ-removal step (e.g. uint16 int16). Accepts TensorProto names or ints. "
+            "Default: UINT16 and INT16."
+        ),
+    )
+    dla_parser.add_argument("-q", "--quiet", action="store_true", help="Suppress progress messages")
+    qdq_group.add_argument(
+        "--remove-all-qdq",
+        action="store_true",
+        help="Remove QDQ pairs for every dtype, matching the C++ remove-all mode",
+    )
+
     # Analyze subcommand
     analyze_parser = subparsers.add_parser(
         "analyze",
@@ -305,6 +355,20 @@ def main():
             verbose=not args.quiet,
         )
 
+    elif args.command == "make-dla-compatible":
+        from . import run_graph_surgery
+
+        run_graph_surgery(
+            "make-dla-compatible",
+            model_path=args.input,
+            output_path=args.output,
+            use_external_data=not args.no_external_data,
+            external_data_name=args.external_data_name,
+            verbose=not args.quiet,
+            qdq_quantized_dtypes=args.qdq_quantized_dtypes,
+            remove_all_qdq=args.remove_all_qdq,
+        )
+
     elif args.command == "analyze":
         from .gqa_replacement import analyze_attention_pattern
 

modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py

@@ -851,10 +851,18 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto:
     for o in graph.output:
         used.add(o.name)
 
+    if model.ir_version >= 4:
+        used.update(value.name for value in graph.input)
+
     kept = [init for init in graph.initializer if init.name in used]
+    discarded = {init.name for init in graph.initializer if init.name not in used}
     removed = len(graph.initializer) - len(kept)
     graph.ClearField("initializer")
     graph.initializer.extend(kept)
+    if model.ir_version <= 3:
+        inputs = [value for value in graph.input if value.name not in discarded]
+        graph.ClearField("input")
+        graph.input.extend(inputs)
     logger.debug("remove_unused_initializers: removed %d unused initializer(s).", removed)
     return model
 

modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Collapse reshape chains at the final, inferred stage of the DLA pipeline."""
+
+from itertools import pairwise
+
+import numpy as np
+from onnx import TensorProto, helper
+
+from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes
+from .dla_constants_to_initializers import _add_legacy_initializer_inputs
+from .dla_graph_cleanup import _subgraph_captures
+from .dla_indexing import _operator_nodes
+from .dla_rank_conversion import _runtime_input_names
+
+__all__ = []
+
+
+def _apply_remove_reshapes(model):
+    """Collapse fixed reshape chains while retaining rank-sensitive consumers and runtime controls."""
+    graph = model.graph
+    cache = GraphCache(graph)
+    runtime_inputs = _runtime_input_names(model)
+    reshapes = _operator_nodes(model, {"Reshape"})
+    captured = _subgraph_captures(graph)
+    positions = {id(node): index for index, node in enumerate(reshapes)}
+    chains, tails = [], {}
+    for node in reshapes:
+        producer = cache.get_producer(node.input[0])
+        chain = tails.pop(id(producer), None)
+        if chain is None:
+            chain = []
+            chains.append(chain)
+        chain.append(node)
+        consumers = cache.get_consumers(node.output[0])
+        if (
+            node.output[0] not in cache.graph_output_names | captured
+            and len(consumers) == 1
+            and consumers[0][1] == 0
+        ):
+            tails[id(node)] = chain
+    replacements = {}
+    for chain in chains:
+        first, last = chain[0], chain[-1]
+        if any(
+            node.input[1] not in cache.init_names or node.input[1] in runtime_inputs
+            for node in chain
+        ):
+            continue
+        input_shape = cache.get_shape(first.input[0])
+        output_shape = cache.get_shape(last.output[0])
+        if (
+            not input_shape
+            or not output_shape
+            or any(d <= 0 for d in [*input_shape, *output_shape])
+        ):
+            continue
+        removable = input_shape == output_shape
+        if removable and last.output[0] not in cache.graph_output_names | captured:
+            for consumer, port in cache.get_consumers(last.output[0]):
+                consumer.input[port] = first.input[0]
+        elif input_shape == output_shape:
+            replacements[id(last)] = [
+                helper.make_node(
+                    "Identity",
+                    [first.input[0]],
+                    list(last.output),
+                    name=cache.unique_name((last.name or last.output[0]) + "_identity"),
+                    domain=last.domain,
+                )
+            ]
+        elif len(chain) > 1:
+            if all(cache.get_shape(node.output[0]) == input_shape for node in chain[:-1]) and any(
+                positions[id(right)] != positions[id(left)] + 1 for left, right in pairwise(chain)
+            ):
+                # Interleaved identity prefixes need only bypassing; retain the existing target and attributes.
+                last.input[0] = first.input[0]
+                for node in chain[:-1]:
+                    replacements[id(node)] = []
+                continue
+            target = add_initializer(
+                graph,
+                cache.unique_name(last.name + "_reshape_shape"),
+                np.asarray(output_shape, dtype=np.int64),
+                dtype=TensorProto.INT64,
+                with_value_info=False,
+            )
+            replacements[id(last)] = [
+                helper.make_node(
+                    "Reshape",
+                    [first.input[0], target],
+                    list(last.output),
+                    name=cache.unique_name((last.name or last.output[0]) + "_reshape"),
+                    domain=last.domain,
+                )
+            ]
+        else:
+            continue
+        for node in chain:
+            replacements.setdefault(id(node), [])
+    batch_replace_nodes(
+        graph,
+        {
+            index: ([node], replacements[id(node)])
+            for index, node in enumerate(graph.node)
+            if id(node) in replacements
+        },
+    )
+    _add_legacy_initializer_inputs(model)
+    return model

modelopt/onnx/graph_surgery/make_dla_compatible.py

@@ -68,6 +68,7 @@
 )
 from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin
 from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model
+from .dla_transforms.dla_remove_reshapes import _apply_remove_reshapes
 from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes
 from .dla_transforms.dla_reshape_patterns import (
     _apply_reshape_qdq_reduce_sum,
@@ -82,11 +83,13 @@
     GraphCache,
     infer_shapes,
     parse_qdq_quantized_dtype_list,
+    save_onnx_model,
 )
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
-__all__ = []
+
+__all__ = ["DLA_TRANSFORMS", "dla_make_dla_compatible"]
 
 
 def _remove_deqlin_and_unused_initializers(model):
@@ -147,6 +150,7 @@ def _apply_constants_to_initializers(model, **options):
     ("Greater", _apply_greater),
     ("ShapeInference", infer_shapes),
     ("TopK", partial(_apply_topk, rank_converted=True)),
+    ("RemoveReshapes", _apply_remove_reshapes),
     ("RemoveUnusedInitializers", _remove_unused_initializers),
 )
 
@@ -295,3 +299,50 @@ def _transform_make_dla_compatible(
     if verbose:
         logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started)
     return model
+
+
+def dla_make_dla_compatible(
+    model_path: str,
+    output_path: str,
+    *,
+    use_external_data: bool = True,
+    external_data_name: str | None = None,
+    verbose: bool = True,
+    qdq_quantized_dtypes: Iterable[int | str] | None = None,
+    remove_all_qdq: bool = False,
+) -> onnx.ModelProto:
+    """Load an ONNX model, apply the ordered DLA transforms, and save the result.
+
+    Arithmetic lowering of zero-branch ``Where`` assumes finite FP32 branch tensors.
+    NaN and infinity values at those nodes are outside the supported input contract.
+    Masked multiplication can change the sign of zero even for finite inputs; validate
+    sign-sensitive downstream operations before using this conversion.
+
+    Args:
+        model_path: Input file; external weights are loaded from its directory.
+        output_path: Destination ONNX file.
+        use_external_data: Store large tensor payloads in a sidecar file.
+        external_data_name: Sidecar filename; defaults to output filename plus "_data".
+        verbose: Log each transform stage.
+        qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed.
+            Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names.
+        remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag.
+            Cannot be combined with qdq_quantized_dtypes.
+
+    Returns:
+        The transformed model. Invalid models and failed transformations raise.
+    """
+    model = _transform_make_dla_compatible(
+        onnx.load(model_path, load_external_data=True),
+        verbose=verbose,
+        qdq_quantized_dtypes=qdq_quantized_dtypes,
+        remove_all_qdq=remove_all_qdq,
+    )
+    save_onnx_model(
+        model,
+        output_path,
+        use_external_data=use_external_data,
+        external_data_name=external_data_name,
+        verbose=verbose,
+    )
+    return model