Merge branch 'chenjiel/support_mtp_qwen_next' of github.com:NVIDIA/TensorRT-Model-Optimizer into chenjiel/support_mtp_qwen_next

Author: cjluo-nv

.github/workflows/example_tests.yml

@@ -86,14 +86,14 @@ jobs:
       pip_install_extras: "[hf,dev-test]"
       runner: linux-amd64-gpu-h100-latest-2
 
-  ##### Speculative Decoding Example Tests (requires 25.08 image) #####
+  ##### Speculative Decoding Example Tests (requires 26.01 image) #####
   speculative-decoding-pr:
     needs: [check-file-changes, wait-checks]
     if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
     uses: ./.github/workflows/_example_tests_runner.yml
     secrets: inherit
     with:
-      docker_image: "nvcr.io/nvidia/pytorch:25.08-py3"
+      docker_image: "nvcr.io/nvidia/pytorch:26.01-py3"
       example: speculative_decoding
       pip_install_extras: "[hf,dev-test]"
       runner: linux-amd64-gpu-l4-latest-1
@@ -103,7 +103,7 @@ jobs:
     uses: ./.github/workflows/_example_tests_runner.yml
     secrets: inherit
     with:
-      docker_image: "nvcr.io/nvidia/pytorch:25.08-py3"
+      docker_image: "nvcr.io/nvidia/pytorch:26.01-py3"
       example: speculative_decoding
       pip_install_extras: "[hf,dev-test]"
       runner: linux-amd64-gpu-h100-latest-2

examples/speculative_decoding/eagle_utils.py

@@ -31,7 +31,6 @@
 from packaging.version import Version
 from PIL import Image
 from scripts.ar_validate import validate_ar
-from torch.distributed.tensor.experimental._attention import _SDPAMerger
 from torch.utils.data import Dataset
 from transformers import AutoProcessor, Trainer, TrainerCallback
 from transformers.trainer_pt_utils import LabelSmoother
@@ -581,7 +580,7 @@ def on_step_end(self, args, state, control, **kwargs):
 def get_patched_templated_ring_attn(orig_templated_attn: Callable):
     """
     Return patched version of
-    torch.distributed.tensor.experimental._attention._templated_ring_attention
+    torch.distributed.tensor.experimental._context_parallel._attention._templated_ring_attention
     to support TTT.
     """
 
@@ -630,7 +629,7 @@ def _get_sharded_ttt_msk(i, rank, size, q_len, ttt_step, dtype):
         return attn_bias
 
     def patched_templated_attn(*args, **kwargs):
-        """Patched version of torch.distributed.tensor.experimental._attention._templated_ring_attention."""
+        """Patched version of _templated_ring_attention."""
         # Get original attention op
         # Sensitive to impl of _templated_ring_attention
         original_op = args[2]
@@ -678,40 +677,35 @@ def patch_ring_attention_for_ttt():
     """Patch torch ring attention to support context parallelism for TTT."""
     # Torch Ring Attention only supports no mask or causal mask. We apply the following patches to enable TTT mask.
 
-    if not (
-        Version(torch.__version__) > Version("2.7.1")
-        and Version(torch.__version__) < Version("2.9.0")
-    ):
+    if Version(torch.__version__) < Version("2.10.0"):
         raise RuntimeError(
-            f"Context parallel TTT only supported for PyTorch 2.8.0 now. "
+            f"Context parallel TTT only supported for PyTorch >= 2.10.0. "
             f"Got {torch.__version__}. "
-            f"Please use nvcr.io/nvidia/pytorch:25.08-py3 or torch 2.8.0 or cp_size=1."
+            f"Please use torch 2.10.0 or cp_size=1."
         )
 
+    from torch.distributed.tensor.experimental._context_parallel import _attention
+
     # 1. Disable load balance, which is designed for causal mask.
     # This affect how buffers are sharded. So need to be done permanently before accelerate/hf trainer init.
-    torch.distributed.tensor.experimental._attention._cp_options.enable_load_balance = False
+    _attention._cp_options.enable_load_balance = False
 
     # 2. Patch templated ring attention for TTT mask.
-    original_templated_ring_attention = (
-        torch.distributed.tensor.experimental._attention._templated_ring_attention
-    )
-    original_templated_ring_attention_backward = (
-        torch.distributed.tensor.experimental._attention._templated_ring_attention_backward
-    )
-    torch.distributed.tensor.experimental._attention._templated_ring_attention = (
-        get_patched_templated_ring_attn(original_templated_ring_attention)
+    original_templated_ring_attention = _attention._templated_ring_attention
+    original_templated_ring_attention_backward = _attention._templated_ring_attention_backward
+    _attention._templated_ring_attention = get_patched_templated_ring_attn(
+        original_templated_ring_attention
     )
-    torch.distributed.tensor.experimental._attention._templated_ring_attention_backward = (
-        get_patched_templated_ring_attn(original_templated_ring_attention_backward)
+    _attention._templated_ring_attention_backward = get_patched_templated_ring_attn(
+        original_templated_ring_attention_backward
     )
 
     # 3. Patch merger to skip the blank shard to avoid difference in output.
-    original_sdpa_merger_step = _SDPAMerger.step
+    original_sdpa_merger_step = _attention._SDPAMerger.step
 
     def patched_sdpa_merger_step(self, out: torch.Tensor, lse: torch.Tensor, partial: bool):
         if lse.sum() <= 0:
             return
         return original_sdpa_merger_step(self, out, lse, partial)
 
-    _SDPAMerger.step = patched_sdpa_merger_step
+    _attention._SDPAMerger.step = patched_sdpa_merger_step

experimental/README.md

@@ -0,0 +1,135 @@
+# Experimental Optimization Techniques
+
+Experimental optimization algorithms and research prototypes under active development.
+
+## Purpose
+
+For new optimization techniques (quantization, pruning, sparsity, etc.) that are:
+
+- Novel or research-stage algorithms
+- Not yet production-ready
+- May have unstable APIs
+
+**⚠️ Warning**: Experimental features are not guaranteed to work across releases. APIs may change or features may be removed without notice. Use at your own risk.
+
+## Requirements
+
+Each experimental technique must include:
+
+- **README.md** - Explains what the technique does, how to use it, current status, model support, and references
+- **Working code** - Clear, readable implementation
+- **Comprehensive tests** - Good test coverage demonstrating correctness
+- **Detailed documentation** - Clear docs on usage, APIs, and behavior
+- **Example** - Demonstrating usage
+- **Model support list** - Which models/frameworks are supported
+- **Deployment info** - Supported deployment frameworks (TensorRT-LLM, vLLM, SGLang, etc.) and whether custom kernels are required
+- **requirements.txt** - Additional dependencies beyond base modelopt
+- **License headers** - Apache 2.0 headers on all Python files
+
+## Example Structures
+
+Organize your code however makes sense. Here are some examples:
+
+**Simple flat structure:**
+
+```text
+experimental/my_technique/
+├── README.md
+├── requirements.txt
+├── my_technique.py
+├── test_my_technique.py
+└── example.py
+```
+
+**Package structure:**
+
+```text
+experimental/my_technique/
+├── README.md
+├── requirements.txt
+├── my_technique/
+│   ├── __init__.py
+│   ├── core.py
+│   └── config.py
+├── tests/
+│   └── test_core.py
+└── examples/
+    └── example_usage.py
+```
+
+## Quality Standards
+
+Experimental code must meet quality standards:
+
+- Comprehensive test coverage required
+- Clear documentation required
+- Pass all pre-commit checks
+
+## PR Guidelines
+
+Keep PRs focused and reviewable:
+
+- **Split large features**: Break complex techniques into multiple PRs if needed
+- **Reasonable scope**: PRs with tens of thousands of lines are difficult to review
+- **Incremental development**: Consider submitting core functionality first, then enhancements
+- If your technique is large, discuss the implementation plan in an issue first
+
+## Example Documentation Template
+
+Your technique's README.md should include:
+
+```markdown
+# Your Technique Name
+
+Brief description of the optimization technique.
+
+## Model Support
+
+| Model/Framework | Supported | Notes |
+|-----------------|-----------|-------|
+| LLMs (Llama, GPT, etc.) | ✅ | Tested on Llama 3.1 |
+| Diffusion Models | ❌ | Not yet supported |
+| Vision Models | ✅ | Experimental |
+
+## Deployment
+
+| Framework | Supported | Notes |
+|-----------|-----------|-------|
+| TensorRT-LLM | ✅ | Requires custom kernel |
+| vLLM | ❌ | Not yet supported |
+| SGLang | ✅ | Uses standard ops |
+
+## Usage
+
+\`\`\`python
+from experimental.my_technique import my_optimize
+...
+\`\`\`
+
+## Status
+
+Current state: Prototype
+
+Known issues:
+- Issue 1
+- Issue 2
+
+## References
+
+- [Paper](link)
+- [Code repository](link)
+- [Project page](link)
+- [Related work](link)
+```
+
+## Path to Production
+
+When a technique is ready for production (proven effective, stable API, full tests, comprehensive docs), it can be promoted to the main `modelopt` package.
+
+**Contributors**: Open an issue proposing graduation with evidence of effectiveness and stability.
+
+**Users**: If you find an experimental feature valuable, open a GitHub issue requesting promotion to production. User demand is a key signal for production readiness.
+
+## Questions?
+
+Open a GitHub issue with `[experimental]` prefix.

experimental/__init__.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Experimental optimization techniques for Model Optimizer.
+
+This package contains experimental and research-stage optimization algorithms
+that are under active development. APIs may change without notice.
+
+Warning:
+    Code in this package is experimental and not covered by semantic versioning.
+    Use at your own risk in production environments.
+"""
+
+import warnings
+
+warnings.warn(
+    "The 'experimental' package contains unstable APIs that may change. "
+    "Use at your own risk in production environments.",
+    FutureWarning,
+    stacklevel=2,
+)
+
+__all__ = []

modelopt/onnx/autocast/referencerunner.py

@@ -89,7 +89,16 @@ def _validate_inputs(self, data_loader):
             if sorted(self.input_names) != sorted(data_loader[0].keys()):
                 raise ValueError("Input names from ONNX model do not match provided input names.")
             for inp_name, inp_shape in data_loader[0].items():
-                if self.input_shapes[inp_name] != list(inp_shape.shape):
+                # Get model and data shapes as numpy arrays
+                inp_shape_model = np.array(self.input_shapes[inp_name])
+                inp_shape_data = np.array(inp_shape.shape)
+                # Compare input rank
+                raise_value_error = len(inp_shape_model) != len(inp_shape_data)
+                if not raise_value_error:
+                    # Compare input shape, skipping check for unknown dimensions
+                    mask = inp_shape_model > 0
+                    raise_value_error = np.any(inp_shape_model[mask] != inp_shape_data[mask])
+                if raise_value_error:
                     raise ValueError(
                         f"Input shape from '{inp_name}' does not match provided input shape: "
                         f"{self.input_shapes[inp_name]} vs {list(inp_shape.shape)}. "

modelopt/onnx/op_types.py

@@ -367,3 +367,22 @@ def get_set_ops():
         "Unique",
         "NonZero",
     }
+
+
+def get_symmetric_ops():
+    """Returns set of commutative/symmetric operations where operand order doesn't matter."""
+    return {
+        "Add",
+        "Mul",
+        "And",
+        "Or",
+        "Xor",
+        "Equal",
+        "Max",
+        "Min",
+        "Sum",
+        "Mean",
+        "BitwiseAnd",
+        "BitwiseOr",
+        "BitwiseXor",
+    }