[OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [1/4]

modelopt/torch/models/__init__.py

@@ -70,13 +70,15 @@
     deepseek,
     deepseek_v3,
     deepseek_v4,
+    falcon,
     gemma,
     gemma2,
     gemma3,
     gemma4,
     gemma4_text,
     gpt_oss,
     llama,
+    llama4,
     minimax,
     mixtral,
     phimoe,
@@ -87,4 +89,7 @@
     qwen3_5_moe,
     qwen3_moe,
     qwen3_next,
+    qwen3_vl_moe,
+    step3p5,
+    step3p7,
 )

modelopt/torch/models/falcon/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+"""Per-model specs for HF model type ``falcon``."""
+
+from . import specs

modelopt/torch/models/falcon/specs.py

@@ -0,0 +1,24 @@
+# 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.
+
+"""Falcon specs (HF model type ``falcon``)."""
+
+from ..specs import ModelSpec, register
+
+__all__: list[str] = []
+
+# Dense, and nothing model-specific for export. Older remote-code Falcon checkpoints
+# (RWForCausalLM) report other model types and are handled by the HF quantization plugin alone.
+register(ModelSpec(model_type="falcon", min_transformers_version="4.57"))

modelopt/torch/models/llama4/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+"""Per-model specs for HF model type ``llama4``."""
+
+from . import specs

modelopt/torch/models/llama4/specs.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""Llama4 specs (HF model type ``llama4``)."""
+
+from ..specs import ModelSpec, MoESpec, register
+
+__all__: list[str] = []
+
+# Llama4TextExperts is fused: one module holding 3-D gate_up_proj and down_proj
+# parameters, run through torch.bmm. There is no (gate, up) pair left to fuse, and the
+# grouped-export path (get_experts_list) does not apply. The block also carries a
+# shared_expert MLP, which is an ordinary dense MLP, not part of this layout.
+register(
+    ModelSpec(
+        model_type="llama4",
+        min_transformers_version="4.57",
+        moe_spec=MoESpec(
+            block_names=("Llama4TextMoe",),
+            expert_linear_names=("gate_up_proj", "down_proj"),
+            fused_expert_names=True,
+        ),
+    )
+)

modelopt/torch/models/qwen3_vl_moe/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+"""Per-model specs for HF model type ``qwen3_vl_moe``."""
+
+from . import specs

modelopt/torch/models/qwen3_vl_moe/specs.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""Qwen3-VL-MoE specs (HF model type ``qwen3_vl_moe``)."""
+
+from ..specs import ModelSpec, MoESpec, register
+
+__all__: list[str] = []
+
+# Qwen3VLMoeTextExperts is fused on every supported transformers: 3-D gate_up_proj and
+# down_proj parameters (transposed before 5.12, standard layout from 5.12 on). Before 5.12
+# the PTQ wrapper unrolls them into gate_proj/up_proj/down_proj ModuleLists (see
+# quantization/plugins/huggingface.py); that rewrite is not iterable per expert, so no export path groups it.
+register(
+    ModelSpec(
+        model_type="qwen3_vl_moe",
+        min_transformers_version="4.57",
+        moe_spec=MoESpec(
+            block_names=("Qwen3VLMoeTextSparseMoeBlock",),
+            expert_linear_names=("gate_up_proj", "down_proj"),
+            fused_expert_names=True,
+        ),
+    )
+)

modelopt/torch/models/step3p5/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+"""Per-model specs for HF model type ``step3p5``."""
+
+from . import specs

modelopt/torch/models/step3p5/specs.py

@@ -0,0 +1,30 @@
+# 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.
+
+"""Step-3.5 specs (HF model type ``step3p5``)."""
+
+from ..specs import ModelSpec, register
+
+__all__: list[str] = []
+
+# Native in transformers since 5.16, where the step3p7 package also registers this text model
+# type; older Step-3.5 checkpoints ship their own modeling code (trust_remote_code).
+#
+# No MoESpec. The native MoE block (Step3p7SparseMoeBlock, with fused gate_up_proj/down_proj
+# experts) is already recognized by is_moe by name and handled by the generic fused-experts
+# path. The remote-code layout keeps its routed experts as expert-indexed MoELinear
+# projections on the MoE MLP itself, with no `experts` container; declaring that block would
+# make is_moe claim it and send AWQ export into get_experts_list, which does not support it.
+register(ModelSpec(model_type="step3p5", min_transformers_version="5.16"))

modelopt/torch/models/step3p7/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+"""Per-model specs for HF model type ``step3p7``."""
+
+from . import specs

modelopt/torch/models/step3p7/specs.py

@@ -0,0 +1,28 @@
+# 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.
+
+"""Step-3.7 specs (HF model type ``step3p7``).
+
+Shares Step-3.5's architecture; its PTQ modeling (in the HF quantization plugin)
+matches every Step revision.
+"""
+
+from ..specs import ModelSpec, register
+
+__all__: list[str] = []
+
+# Native in transformers since 5.16. No MoESpec, for the same reasons as step3p5 (see
+# step3p5/specs.py).
+register(ModelSpec(model_type="step3p7", min_transformers_version="5.16"))

tests/unit/torch/models/test_model_specs.py

@@ -552,6 +552,7 @@ def test_fused_experts_shortcut_ignores_unrelated_attributes():
         ("gate_proj", "up_proj"),
     ),
     ("gpt_oss", ("GptOssMLP", "GptOssMoE"), ("gate_up_proj", "down_proj"), True, None),
+    ("llama4", ("Llama4TextMoe",), ("gate_up_proj", "down_proj"), True, None),
     ("minimax", ("MiniMaxSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")),
     ("mixtral", ("MixtralSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")),
     ("phimoe", ("PhimoeSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")),
@@ -584,6 +585,13 @@ def test_fused_experts_shortcut_ignores_unrelated_attributes():
         False,
         ("gate_proj", "up_proj"),
     ),
+    (
+        "qwen3_vl_moe",
+        ("Qwen3VLMoeTextSparseMoeBlock",),
+        ("gate_up_proj", "down_proj"),
+        True,
+        None,
+    ),
 ]
 
 
@@ -717,6 +725,7 @@ def test_iterable_experts_matches_pre_refactor_support():
         "gemma4": "Gemma4ForConditionalGeneration",
         "gemma4_text": "Gemma4TextForCausalLM",
         "gpt_oss": "GptOssForCausalLM",
+        "llama4": "Llama4ForConditionalGeneration",
         "minimax": "MiniMaxForCausalLM",
         "mixtral": "MixtralForCausalLM",
         "phimoe": "PhimoeForCausalLM",
@@ -725,6 +734,7 @@ def test_iterable_experts_matches_pre_refactor_support():
         "qwen3_5_moe": "Qwen3_5MoeForCausalLM",
         "qwen3_moe": "Qwen3MoeForCausalLM",
         "qwen3_next": "Qwen3NextForCausalLM",
+        "qwen3_vl_moe": "Qwen3VLMoeForConditionalGeneration",
     }
     moe_specs = [s for s in get_specs() if _layouts(s)]
     assert {s.model_type for s in moe_specs} == set(root_class_names), (