[None][feat] Add AD custom model for Seed-OSS family (#238)

lucaslie · bmarimuthu-nv · commit 3e5ce2f43cf8 · 2026-03-31T16:33:17.000-07:00
* [None][feat] Add AD custom model for Seed-OSS family

Add prefill-only custom model implementation for ByteDance-Seed/Seed-OSS-36B-Instruct
using AutoDeploy canonical ops (torch_attention, torch_rope_with_explicit_cos_sin,
torch_rmsnorm). Seed-OSS is a dense Llama-style model with GQA (80 Q / 8 KV heads),
SwiGLU MLP, and attention_bias=True on Q/K/V projections.

Includes hierarchical equivalence tests (MLP, Attention, Decoder Layer, Full Model,
Export) comparing against HF reference implementation.

Signed-off-by: Lucas Liebenwein &lt;11156568+lucaslie@users.noreply.github.com&gt;

* [None][fix] Add causal mask to HF reference in attention/decoder layer tests

The HF eager attention does NOT apply causal masking when attention_mask=None,
while our custom model always uses is_causal=True. Provide explicit causal mask
to HF reference to ensure equivalent comparison.

Signed-off-by: Lucas Liebenwein &lt;11156568+lucaslie@users.noreply.github.com&gt;

---------

Signed-off-by: Lucas Liebenwein &lt;11156568+lucaslie@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py
@@ -18,6 +18,7 @@
 from .modeling_qwen2 import Qwen2ForCausalLM
 from .modeling_qwen3_5_moe import Qwen3_5MoeForCausalLM, Qwen3_5MoeForConditionalGeneration
 from .modeling_qwen3_moe import Qwen3MoeForCausalLM
+from .modeling_seed_oss import SeedOssForCausalLM
 from .modeling_skywork_r1v2 import SkyworkR1V2ForConditionalGeneration
 from .modeling_smollm3 import SmolLM3ForCausalLM
 from .modeling_starcoder2 import Starcoder2ForCausalLM
@@ -48,6 +49,7 @@
     "Qwen3_5MoeForCausalLM",
     "Qwen3_5MoeForConditionalGeneration",
     "Qwen3MoeForCausalLM",
+    "SeedOssForCausalLM",
     "SkyworkR1V2ForConditionalGeneration",
     "SmolLM3ForCausalLM",
     "Starcoder2ForCausalLM",
diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_seed_oss.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_seed_oss.py
@@ -0,0 +1,379 @@
+# 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 Seed-OSS model implementation for auto_deploy export.
+
+Source:
+https://huggingface.co/ByteDance-Seed/Seed-OSS-36B-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/residual dropout (inference only)
+
+The Seed-OSS model uses Grouped Query Attention (GQA) with standard RoPE.
+It is a dense Llama-style model with attention_bias=True and attention_out_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_utils import PreTrainedModel
+from transformers.models.seed_oss.configuration_seed_oss import SeedOssConfig
+from transformers.utils import ModelOutput
+
+from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory
+
+
+class SeedOssRMSNorm(nn.Module):
+    """RMS Normalization for Seed-OSS 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 SeedOssRotaryEmbedding(nn.Module):
+    """Rotary Position Embedding for Seed-OSS.
+
+    Simplified version that precomputes and caches cos/sin values.
+    Returns pre-sliced values indexed by position_ids.
+
+    Uses _ad_ prefix for buffer names to work with AutoDeploy's lift_to_meta.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: int = 524288,
+        base: float = 10000.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build cos/sin cache with AD-specific naming
+        self._set_cos_sin_cache(max_position_embeddings)
+
+    def _set_cos_sin_cache(self, seq_len: int):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(seq_len, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        # Use _ad_ prefix for AutoDeploy compatibility with lift_to_meta
+        self.register_buffer("_ad_cos_cached", emb.cos(), persistent=False)
+        self.register_buffer("_ad_sin_cached", emb.sin(), persistent=False)
+
+    def forward(
+        self, x: torch.Tensor, position_ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Slice cos/sin by position_ids here (once) instead of in every attention layer
+        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 SeedOssMLP(nn.Module):
+    """MLP layer for Seed-OSS (SwiGLU activation)."""
+
+    def __init__(self, config: SeedOssConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        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 SeedOssAttention(nn.Module):
+    """Grouped Query Attention for Seed-OSS with standard RoPE.
+
+    Uses attention_bias on Q/K/V projections and attention_out_bias on O projection.
+    AD canonical attention ops handle GQA natively (no repeat_kv needed).
+    """
+
+    def __init__(self, config: SeedOssConfig, 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 = config.head_dim
+        self.scaling = self.head_dim ** (-0.5)
+
+        # Q/K/V/O projections
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_out_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, q_len, _ = hidden_states.size()
+
+        # Project Q/K/V and reshape to [B, S, N, head_dim] (BSND layout)
+        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)
+
+        # Get pre-sliced cos/sin from position_embeddings (already indexed by position_ids)
+        cos, sin = position_embeddings  # [B, S, head_dim]
+
+        # Apply RoPE using custom op (BSND layout, unsqueeze_dim=2)
+        q, k = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin(
+            q,
+            k,
+            cos,
+            sin,
+            2,  # unsqueeze_dim=2 for BSND layout
+        )
+
+        # Attention using custom op with GQA support (BSND layout)
+        attn_output = torch.ops.auto_deploy.torch_attention(
+            q,  # [B, S, N, head_dim]
+            k,  # [B, S, N_kv, head_dim]
+            v,  # [B, S, N_kv, head_dim]
+            None,  # attn_mask
+            0.0,  # dropout_p
+            True,  # is_causal
+            self.scaling,  # scale
+            None,  # sinks
+            None,  # sliding_window
+            None,  # logit_cap
+            "bsnd",  # layout
+        )
+
+        # Reshape [B, S, N, head_dim] -> [B, S, N * head_dim] and project
+        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 SeedOssDecoderLayer(nn.Module):
+    """Transformer decoder layer for Seed-OSS."""
+
+    def __init__(self, config: SeedOssConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = SeedOssAttention(config, layer_idx=layer_idx)
+        self.mlp = SeedOssMLP(config)
+        self.input_layernorm = SeedOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SeedOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        # Self attention
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, position_embeddings)
+        hidden_states = residual + hidden_states
+
+        # MLP
+        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 SeedOssOutput(ModelOutput):
+    """Output for SeedOssModel."""
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class SeedOssCausalLMOutput(ModelOutput):
+    """Output for SeedOssForCausalLM."""
+
+    logits: Optional[torch.FloatTensor] = None
+
+
+class SeedOssPreTrainedModel(PreTrainedModel):
+    """Base class for Seed-OSS models."""
+
+    config_class = SeedOssConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["SeedOssDecoderLayer"]
+    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 SeedOssModel(SeedOssPreTrainedModel):
+    """Seed-OSS transformer decoder model."""
+
+    def __init__(self, config: SeedOssConfig):
+        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(
+            [SeedOssDecoderLayer(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]
+        )
+        self.norm = SeedOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Shared rotary embedding at model level
+        self.rotary_emb = SeedOssRotaryEmbedding(
+            config.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
+        )
+
+        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,
+    ) -> SeedOssOutput:
+        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)
+
+        # Cast to compute dtype (e.g., bfloat16) for FP8 models where embedding
+        # output may be FP8 but downstream ops (RMSNorm, attention) require FP16/BF16
+        inputs_embeds = inputs_embeds.to(self.norm.weight.dtype)
+
+        # Compute position embeddings once (sliced by position_ids in RoPE)
+        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 SeedOssOutput(last_hidden_state=hidden_states)
+
+
+class SeedOssForCausalLM(SeedOssPreTrainedModel, GenerationMixin):
+    """Seed-OSS model with language modeling head."""
+
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.model = SeedOssModel(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,
+    ) -> SeedOssCausalLMOutput:
+        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 SeedOssCausalLMOutput(logits=logits)
+
+
+# Register with AutoModelForCausalLMFactory
+AutoModelForCausalLMFactory.register_custom_model_cls("SeedOssConfig", SeedOssForCausalLM)
diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_seed_oss_modeling.py
