[None][feat] Add AD custom model for Qwen2 family (#219)

* [None][feat] Add AD custom model for Qwen2 family

Add a lean prefill-only custom model implementation for the Qwen2 architecture
family, covering Qwen2.5, QwQ, DeepSeek-R1-Distill-Qwen, and related models.

The Qwen2 architecture is similar to Llama3 (GQA + SwiGLU MLP + RMSNorm + RoPE)
with the key difference that Q/K/V projections have bias=True while O and MLP
projections have bias=False.

Uses AD canonical ops:
- torch_rmsnorm for normalization
- torch_rope_with_explicit_cos_sin for rotary embeddings
- torch_attention for grouped-query attention (no repeat_kv needed)

Files:
- tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen2.py
- tests/unittest/auto_deploy/singlegpu/models/test_qwen2_modeling.py

Registered for Qwen2Config covering all qwen2-based models in the registry:
Qwen2.5-{0.5B,1.5B,3B,7B,14B,72B}, Qwen2.5-Coder, Qwen2.5-Math, QwQ-32B,
DeepSeek-R1-Distill-Qwen-{1.5B,7B,14B,32B}, Skywork-SWE-32B,
OpenReasoning-Nemotron-32B, r1-1776-distill-qwen-32b.

Signed-off-by: Lucas Liebenwein <lliebenwein@nvidia.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>

* [None][fix] Fix causal mask in Qwen2 standalone attention/decoder-layer tests

HF eager_attention_forward skips masking when attention_mask=None, making
the HF reference run non-causal (full) attention. The custom AD attention
uses is_causal=True. This mismatch caused RMSE ratio ~0.93 in the
standalone attention and decoder-layer equivalence tests.

Fix: build an additive causal mask and pass it to the HF attention call
in test_qwen2_attention_equivalence and test_qwen2_decoder_layer_equivalence.

The full-model tests are unaffected since Qwen2Model.forward builds
the causal mask internally for both HF and custom paths.

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

* [None][fix] Use _get_hf_rotary_class() helper instead of redundant module-level import

Remove the module-level try/except import of HFQwen2RotaryEmbedding and use
the HFRotary variable from _get_hf_rotary_class() already called in the test.

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

---------

Signed-off-by: Lucas Liebenwein <lliebenwein@nvidia.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Co-authored-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Author: 2 people

tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py

@@ -15,6 +15,7 @@
 from .modeling_nemotron_flash import NemotronFlashForCausalLM, NemotronFlashPreTrainedTokenizerFast
 from .modeling_nemotron_h import NemotronHForCausalLM
 from .modeling_olmo3 import Olmo3ForCausalLM
+from .modeling_qwen2 import Qwen2ForCausalLM
 from .modeling_qwen3 import Qwen3ForCausalLM
 from .modeling_qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5ForConditionalGeneration
 from .modeling_qwen3_5_moe import Qwen3_5MoeForCausalLM, Qwen3_5MoeForConditionalGeneration
@@ -44,6 +45,7 @@
     "NemotronFlashPreTrainedTokenizerFast",
     "NemotronHForCausalLM",
     "Olmo3ForCausalLM",
+    "Qwen2ForCausalLM",
     "Qwen3ForCausalLM",
     "Qwen3_5ForCausalLM",
     "Qwen3_5ForConditionalGeneration",

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen2.py

@@ -0,0 +1,353 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-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.
+
+"""Slimmed down PyTorch Qwen2 model implementation for auto_deploy export.
+
+Source:
+https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct
+
+This implementation differs from the original HuggingFace version in the following ways:
+* Simplified for prefill-only inference (no KV caching)
+* Uses auto_deploy custom ops for export compatibility
+* Removed flash attention variants (uses torch_attention custom op)
+* Removed gradient checkpointing and training code paths
+* Removed attention dropout (inference only)
+* Removed sliding window attention (AD manages masking)
+* No repeat_kv — AD attention ops handle GQA natively
+
+The Qwen2 family (Qwen2.5, QwQ, DeepSeek-R1-Distill-Qwen, etc.) shares a single
+architecture with GQA, SwiGLU MLP, RMSNorm, and RoPE. Key difference from Llama:
+Q/K/V projections have bias=True while O and MLP projections have bias=False.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.utils import ModelOutput
+
+from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory
+
+
+class Qwen2RMSNorm(nn.Module):
+    """RMS Normalization for Qwen2 using AutoDeploy torch_rmsnorm reference op."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return torch.ops.auto_deploy.torch_rmsnorm(
+            hidden_states, self.weight, self.variance_epsilon
+        )
+
+
+class Qwen2RotaryEmbedding(nn.Module):
+    """Rotary Position Embedding for Qwen2 family.
+
+    Supports all rope types (default, linear, dynamic, etc.) via
+    transformers ROPE_INIT_FUNCTIONS. Precomputes and caches cos/sin values.
+    Slices by position_ids once and returns pre-sliced cos/sin to all layers.
+
+    Uses _ad_ prefix for buffer names to work with AutoDeploy's lift_to_meta.
+    """
+
+    def __init__(self, config: Qwen2Config):
+        super().__init__()
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            rope_type = config.rope_scaling.get(
+                "rope_type", config.rope_scaling.get("type", "default")
+            )
+        else:
+            rope_type = "default"
+
+        inv_freq, self.attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, device=None)
+
+        max_pos = config.max_position_embeddings
+        t = torch.arange(max_pos, dtype=inv_freq.dtype)
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("_ad_cos_cached", emb.cos() * self.attention_scaling, persistent=False)
+        self.register_buffer("_ad_sin_cached", emb.sin() * self.attention_scaling, persistent=False)
+
+    def forward(
+        self, x: torch.Tensor, position_ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        cos = self._ad_cos_cached.to(dtype=x.dtype, device=x.device)
+        sin = self._ad_sin_cached.to(dtype=x.dtype, device=x.device)
+        return cos[position_ids], sin[position_ids]
+
+
+class Qwen2MLP(nn.Module):
+    """MLP layer for Qwen2 (SwiGLU activation). All projections have bias=False."""
+
+    def __init__(self, config: Qwen2Config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Qwen2Attention(nn.Module):
+    """Grouped Query Attention for Qwen2.
+
+    Uses AD canonical ops for attention and RoPE. GQA is handled natively
+    by torch_attention — no repeat_kv needed. Q/K/V projections have bias=True,
+    O projection has bias=False.
+    """
+
+    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.scaling = self.head_dim ** (-0.5)
+
+        # Qwen2: Q/K/V have bias=True, O has bias=False
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=True)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=True)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, q_len, _ = hidden_states.size()
+
+        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+
+        q, k = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin(
+            q,
+            k,
+            cos,
+            sin,
+            2,  # unsqueeze_dim=2 for BSND layout
+        )
+
+        attn_output = torch.ops.auto_deploy.torch_attention(
+            q,
+            k,
+            v,
+            None,  # attn_mask
+            0.0,  # dropout_p
+            True,  # is_causal
+            self.scaling,
+            None,  # sinks
+            None,  # sliding_window
+            None,  # logit_cap
+            "bsnd",
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim)
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output
+
+
+class Qwen2DecoderLayer(nn.Module):
+    """Transformer decoder layer for Qwen2."""
+
+    def __init__(self, config: Qwen2Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = Qwen2Attention(config, layer_idx=layer_idx)
+        self.mlp = Qwen2MLP(config)
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, position_embeddings)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+@dataclass
+class Qwen2ModelOutput(ModelOutput):
+    """Output for Qwen2Model."""
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class Qwen2CausalLMOutput(ModelOutput):
+    """Output for Qwen2ForCausalLM."""
+
+    logits: Optional[torch.FloatTensor] = None
+
+
+class Qwen2PreTrainedModel(PreTrainedModel):
+    """Base class for Qwen2 models."""
+
+    config_class = Qwen2Config
+    base_model_prefix = "model"
+    _no_split_modules = ["Qwen2DecoderLayer"]
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class Qwen2Model(Qwen2PreTrainedModel):
+    """Qwen2 transformer decoder model."""
+
+    def __init__(self, config: Qwen2Config):
+        super().__init__(config)
+        self.config = config
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen2DecoderLayer(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen2RotaryEmbedding(config)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Qwen2ModelOutput:
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("Cannot specify both input_ids and inputs_embeds")
+        elif input_ids is None and inputs_embeds is None:
+            raise ValueError("Must specify either input_ids or inputs_embeds")
+
+        assert position_ids is not None, "position_ids must be provided for AD export"
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        inputs_embeds = inputs_embeds.to(self.norm.weight.dtype)
+
+        position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
+
+        hidden_states = inputs_embeds
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(hidden_states, position_embeddings)
+
+        hidden_states = self.norm(hidden_states)
+
+        return Qwen2ModelOutput(last_hidden_state=hidden_states)
+
+
+class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin):
+    """Qwen2 model with language modeling head."""
+
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.model = Qwen2Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> Qwen2CausalLMOutput:
+        assert position_ids is not None, "position_ids must be provided for AD export"
+        outputs = self.model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states).float()
+
+        return Qwen2CausalLMOutput(logits=logits)
+
+
+# Register with AutoModelForCausalLMFactory
+AutoModelForCausalLMFactory.register_custom_model_cls("Qwen2Config", Qwen2ForCausalLM)