Add Jet-Nemotron (#12448)

python/sglang/srt/configs/__init__.py

@@ -6,6 +6,7 @@ from sglang.srt.configs.dots_vlm import DotsVLMConfig
 from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.falcon_h1 import FalconH1Config
 from sglang.srt.configs.janus_pro import MultiModalityConfig
+from sglang.srt.configs.jet_nemotron import JetNemotronConfig
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.configs.kimi_vl import KimiVLConfig
 from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
@@ -38,4 +39,5 @@ __all__ = [
     "DotsOCRConfig",
     "FalconH1Config",
     "NemotronHConfig",
+    "JetNemotronConfig",
 ]

python/sglang/srt/configs/jet_nemotron.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import Any
+
+from transformers.configuration_utils import PretrainedConfig
+
+from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
+
+
+@dataclass
+class JetBlockConfig:
+    mode: str
+    expand_v: float
+    num_heads: int
+    head_dim: int
+    norm_eps: str
+    conv_size: int
+    dconv_generator_reduction: int
+    dconv_implementation: str
+
+
+class JetNemotronConfig(PretrainedConfig):
+    model_type: str = "jet_nemotron"
+
+    efficient_attention_config: dict[str, dict[str, Any]]
+    hidden_act: str
+    hidden_size: int
+    initializer_range: float
+    intermediate_size: int
+    layer_types: list[str]
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_scaling: None
+    rope_theta: float
+
+    @property
+    def full_attention_layer_ids(self) -> list[int]:
+        return [
+            idx
+            for idx, layer_type in enumerate(self.layer_types)
+            if layer_type in ("attn", "swa")
+        ]
+
+    @property
+    def linear_layer_ids(self) -> list[int]:
+        return [
+            idx
+            for idx, layer_type in enumerate(self.layer_types)
+            if layer_type == "jet"
+        ]
+
+    @property
+    def mamba2_cache_params(self) -> Mamba2CacheParams:
+        from sglang.srt.layers.dp_attention import get_attention_tp_size
+
+        jet_block_config = JetBlockConfig(**self.efficient_attention_config["jet"])
+
+        num_heads = jet_block_config.num_heads
+        head_k_dim = jet_block_config.head_dim
+        head_v_dim = int(head_k_dim * jet_block_config.expand_v)
+        total_v_dim = num_heads * head_v_dim
+
+        shape = Mamba2StateShape.create(
+            tp_world_size=get_attention_tp_size(),
+            intermediate_size=total_v_dim,
+            n_groups=num_heads,
+            num_heads=num_heads,
+            head_dim=head_v_dim,
+            state_size=head_k_dim,
+            conv_kernel=jet_block_config.conv_size,
+        )
+
+        return Mamba2CacheParams(shape=shape, layers=self.linear_layer_ids)

python/sglang/srt/model_executor/model_runner.py

@@ -31,6 +31,7 @@ import torch.distributed as dist
 
 from sglang.srt.configs import (
     FalconH1Config,
+    JetNemotronConfig,
     KimiLinearConfig,
     NemotronHConfig,
     Qwen3NextConfig,
@@ -1374,7 +1375,7 @@ class ModelRunner:
     @property
     def hybrid_gdn_config(self):
         config = self.model_config.hf_config
-        if isinstance(config, Qwen3NextConfig):
+        if isinstance(config, Qwen3NextConfig | JetNemotronConfig):
             return config
         return None
 

python/sglang/srt/models/jet_nemotron.py

@@ -0,0 +1,596 @@
+from collections.abc import Iterable
+from typing import cast
+
+import einops
+import torch
+import torch.nn as nn
+
+from sglang.srt.configs.jet_nemotron import JetBlockConfig, JetNemotronConfig
+from sglang.srt.layers.attention.fla.fused_recurrent import (
+    fused_recurrent_gated_delta_rule_update,
+)
+from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+    MambaAttnBackendBase,
+)
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
+from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.rotary_embedding import get_rope
+from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.qwen2 import Qwen2MLP, Qwen2Model
+from sglang.srt.utils import add_prefix
+
+
+class DynamicShortConvolutionKernelGenerator(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        output_size: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.w1 = ColumnParallelLinear(
+            input_size,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("w1", prefix),
+        )
+
+        self.act = nn.SiLU()
+
+        self.w2 = ColumnParallelLinear(
+            hidden_size,
+            output_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=add_prefix("w2", prefix),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.w1(x)
+        x = self.act(x)
+        x, _ = self.w2(x)
+        return x
+
+
+class DynamicShortConvolution(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        kernel_size: int,
+        generator_input_size: int,
+        generator_reduction: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        generator_hidden_size = hidden_size // generator_reduction
+
+        self.kernel_generator = DynamicShortConvolutionKernelGenerator(
+            input_size=generator_input_size,
+            hidden_size=generator_hidden_size,
+            output_size=hidden_size * kernel_size,
+            quant_config=quant_config,
+            prefix=add_prefix("kernel_generator", prefix),
+        )
+
+        self.hidden_size = hidden_size
+        self.kernel_size = kernel_size
+
+    def forward(
+        self,
+        x: torch.Tensor,  # (cu_seq_len, hidden_size)
+        *,
+        conv_state: torch.Tensor,  # (batch_size, hidden_size, kernel_size - 1)
+        generator_input: torch.Tensor,  # (cu_seq_len, generator_input_size)
+        seq_lens: torch.Tensor,  # (batch_size,)
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: (cu_seq_len, hidden_size)
+            conv_state: (batch_size, hidden_size, kernel_size - 1)
+            generator_input: (cu_seq_len, generator_input_size)
+            seq_lens: (batch_size,)
+
+        Returns:
+            out: (cu_seq_len, hidden_size)
+            conv_state: (batch_size, hidden_size, kernel_size - 1)
+        """
+
+        x_seqs = self._continuous_to_seqs(x, seq_lens=seq_lens)
+        conv_state = einops.rearrange(conv_state, "b d k -> b k d")
+        x_seqs = [torch.cat([conv_state[i], x_seqs[i]]) for i in range(len(x_seqs))]
+        x = self._seqs_to_batch(
+            x_seqs
+        )  # (batch_size, max_seq_len + kernel_size - 1, hidden_size)
+
+        x = einops.rearrange(x, "b l d -> b d l")
+
+        new_conv_state = x[
+            :, :, -(self.kernel_size - 1) :
+        ]  # (batch_size, hidden_size, kernel_size - 1)
+
+        x = x.unfold(
+            dimension=-1, size=self.kernel_size, step=1
+        )  # (batch_size, hidden_size, max_seq_len, kernel_size)
+        x = einops.rearrange(x, "b d l k -> b l d k")
+
+        kernels = self.kernel_generator(
+            generator_input
+        )  # (cu_seq_len, hidden_size * kernel_size)
+        kernels = einops.rearrange(
+            kernels,
+            "l (d k) -> l d k",
+            d=self.hidden_size,
+            k=self.kernel_size,
+        )
+        kernels = self._seqs_to_batch(
+            self._continuous_to_seqs(kernels, seq_lens=seq_lens)
+        )  # (batch_size, max_seq_len, hidden_size, kernel_size)
+
+        out = (x * kernels).sum(dim=-1)  # (batch_size, max_seq_len, hidden_size)
+
+        out = self._batch_to_continuous(
+            out, seq_lens=seq_lens
+        )  # (cu_seq_len, hidden_size)
+
+        out = nn.functional.silu(out)
+
+        return out, new_conv_state
+
+    def _batch_to_continuous(
+        self,
+        x: torch.Tensor,
+        *,
+        seq_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.cat([x[i, -seq_lens[i] :] for i in range(seq_lens.size(0))])
+
+    def _continuous_to_seqs(
+        self,
+        x: torch.Tensor,
+        *,
+        seq_lens: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        return [
+            x[(seq_lens[:i].sum()) : (seq_lens[: i + 1].sum())]
+            for i in range(seq_lens.size(0))
+        ]
+
+    def _seqs_to_batch(
+        self,
+        seqs: list[torch.Tensor],
+    ) -> torch.Tensor:
+        return nn.utils.rnn.pad_sequence(
+            seqs,
+            batch_first=True,
+            padding_side="left",
+        )
+
+
+class JetBlock(nn.Module):
+    def __init__(
+        self,
+        config: JetNemotronConfig,
+        layer_id: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        jet_block_config = JetBlockConfig(
+            **self.config.efficient_attention_config[self.config.layer_types[layer_id]]
+        )
+
+        hidden_size = self.config.hidden_size
+        num_heads = jet_block_config.num_heads
+        head_k_dim = jet_block_config.head_dim
+        total_k_dim = num_heads * head_k_dim
+        head_v_dim = int(head_k_dim * jet_block_config.expand_v)
+        total_v_dim = num_heads * head_v_dim
+        conv_size = jet_block_config.conv_size
+
+        self.qkvabz_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [
+                total_k_dim,
+                total_k_dim,
+                total_v_dim,
+                num_heads,
+                num_heads,
+                total_v_dim,
+            ],
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("qkvabz_proj", prefix),
+        )
+
+        self.o_proj = RowParallelLinear(total_v_dim, hidden_size, bias=False)
+
+        self.A_log = nn.Parameter(torch.empty(num_heads, dtype=torch.float32))
+        self.dt_bias = nn.Parameter(torch.empty(num_heads))
+
+        self.dynamic_conv1d = DynamicShortConvolution(
+            quant_config=quant_config,
+            prefix=add_prefix("dynamic_conv1d", prefix),
+            hidden_size=total_v_dim,
+            kernel_size=conv_size,
+            generator_input_size=hidden_size,
+            generator_reduction=jet_block_config.dconv_generator_reduction,
+        )
+
+        self.o_norm = RMSNormGated(
+            head_v_dim,
+            eps=float(jet_block_config.norm_eps),
+        )
+
+        # Attributes.
+        self.conv_size = conv_size
+        self.head_k_dim = head_k_dim
+        self.head_v_dim = head_v_dim
+        self.layer_id = layer_id
+        self.num_heads = num_heads
+        self.total_k_dim = total_k_dim
+        self.total_v_dim = total_v_dim
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        assert isinstance(forward_batch.attn_backend, HybridLinearAttnBackend)
+        assert isinstance(
+            forward_batch.attn_backend.linear_attn_backend, MambaAttnBackendBase
+        )
+        linear_attn_backend = forward_batch.attn_backend.linear_attn_backend
+        forward_metadata = linear_attn_backend.forward_metadata
+        layer_cache = linear_attn_backend.req_to_token_pool.mamba2_layer_cache(
+            self.layer_id
+        )
+
+        qkvabz, _ = self.qkvabz_proj(hidden_states)
+        q, k, v, a, beta, z = qkvabz.split(
+            [
+                self.total_k_dim,
+                self.total_k_dim,
+                self.total_v_dim,
+                self.num_heads,
+                self.num_heads,
+                self.total_v_dim,
+            ],
+            dim=-1,
+        )
+
+        q = nn.functional.silu(q)
+        q = einops.rearrange(q, "l (h d) -> l h d", h=self.num_heads, d=self.head_k_dim)
+
+        k = nn.functional.silu(k)
+        k = einops.rearrange(k, "l (h d) -> l h d", h=self.num_heads, d=self.head_k_dim)
+
+        conv_cache = layer_cache.conv
+        assert isinstance(conv_cache, torch.Tensor)
+        v, new_conv_state = self.dynamic_conv1d(
+            v,
+            conv_state=conv_cache[
+                forward_metadata.mamba_cache_indices, -self.total_v_dim :, :
+            ],
+            generator_input=hidden_states,
+            seq_lens=(
+                forward_batch.extend_seq_lens
+                if forward_batch.extend_seq_lens is not None
+                else torch.ones(
+                    (forward_batch.batch_size,),
+                    dtype=torch.long,
+                )
+            ),
+        )
+        conv_cache[forward_metadata.mamba_cache_indices, -self.total_v_dim :, :] = (
+            new_conv_state
+        )
+        v = einops.rearrange(v, "l (h d) -> l h d", h=self.num_heads, d=self.head_v_dim)
+
+        g = -self.A_log.float().exp() * nn.functional.softplus(a.float() + self.dt_bias)
+
+        beta = nn.functional.sigmoid(beta)
+
+        o = fused_recurrent_gated_delta_rule_update(
+            q=q.unsqueeze(0),
+            k=k.unsqueeze(0),
+            v=v.unsqueeze(0),
+            g=g.unsqueeze(0),
+            beta=beta.unsqueeze(0),
+            initial_state_source=layer_cache.temporal,
+            initial_state_indices=forward_metadata.mamba_cache_indices,
+            cu_seqlens=cast(torch.LongTensor, forward_metadata.query_start_loc),
+            use_qk_l2norm_in_kernel=True,
+        ).squeeze(0)
+
+        z = einops.rearrange(z, "l (h d) -> l h d", h=self.num_heads)
+
+        o = self.o_norm(o, z)
+
+        o = einops.rearrange(o, "l h d -> l (h d)")
+
+        o, _ = self.o_proj(o)
+
+        return o
+
+
+class JetNemotronAttention(nn.Module):
+    def __init__(
+        self,
+        config: JetNemotronConfig,
+        layer_id: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        self.head_dim = self.config.hidden_size // self.config.num_attention_heads
+
+        self.q_size = self.config.num_attention_heads * self.head_dim
+        self.kv_size = self.config.num_key_value_heads * self.head_dim
+
+        self.qkv_proj = QKVParallelLinear(
+            self.config.hidden_size,
+            self.head_dim,
+            self.config.num_attention_heads,
+            self.config.num_key_value_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
+        )
+        self.o_proj = RowParallelLinear(
+            self.config.num_attention_heads * self.head_dim,
+            self.config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=self.config.max_position_embeddings,
+            base=int(self.config.rope_theta),
+            rope_scaling=self.config.rope_scaling,
+        )
+
+        match self.config.layer_types[layer_id]:
+            case "attn":
+                sliding_window_size = -1
+
+            case "swa":
+                sliding_window_size = self.config.efficient_attention_config["swa"][
+                    "window_size"
+                ]
+
+            case _:
+                raise NotImplementedError
+
+        self.attn = RadixAttention(
+            self.config.num_attention_heads,
+            self.head_dim,
+            self.head_dim**-0.5,
+            num_kv_heads=self.config.num_key_value_heads,
+            layer_id=layer_id,
+            sliding_window_size=sliding_window_size,
+            quant_config=quant_config,
+            prefix=add_prefix("attn", prefix),
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v, forward_batch)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class JetNemotronDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: JetNemotronConfig,
+        alt_stream: torch.cuda.Stream | None = None,
+        layer_id: int = 0,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        match config.layer_types[layer_id]:
+            case "attn" | "swa":
+                self.self_attn = JetNemotronAttention(
+                    config,
+                    quant_config=quant_config,
+                    prefix=add_prefix("self_attn", prefix),
+                    layer_id=layer_id,
+                )
+
+            case "jet":
+                self.self_attn = JetBlock(
+                    config,
+                    quant_config=quant_config,
+                    prefix=add_prefix("self_attn", prefix),
+                    layer_id=layer_id,
+                )
+
+            case _:
+                raise NotImplementedError
+
+        self.mlp = Qwen2MLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Self Attention
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            forward_batch=forward_batch,
+        )
+
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        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, None
+
+
+class JetNemotronForCausalLM(nn.Module):
+    def __init__(
+        self,
+        config: JetNemotronConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = Qwen2Model(
+            config,
+            quant_config=quant_config,
+            prefix=add_prefix("model", prefix),
+            decoder_layer_type=JetNemotronDecoderLayer,
+        )
+
+        if config.tie_word_embeddings:
+            self.lm_head = self.model.embed_tokens
+        else:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=add_prefix("lm_head", prefix),
+            )
+
+        self.logits_processor = LogitsProcessor(config)
+        self.pooler = Pooler(PoolingType.LAST, normalize=True)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: torch.Tensor | None = None,
+        get_embedding: bool = False,
+    ) -> EmbeddingPoolerOutput | LogitsProcessorOutput:
+        hidden_states = self.model(
+            input_ids,
+            positions,
+            forward_batch,
+            input_embeds,
+        )
+
+        if not get_embedding:
+            return self.logits_processor(
+                input_ids, hidden_states, self.lm_head, forward_batch
+            )
+        else:
+            return self.pooler(hidden_states, forward_batch)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        stacked_params_mapping: list[tuple[str, str, str | int]] = [
+            # (param_name, shard_weight_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+            ("qkvabz_proj", "q_proj", 0),
+            ("qkvabz_proj", "k_proj", 1),
+            ("qkvabz_proj", "v_proj", 2),
+            ("qkvabz_proj", "a_proj", 3),
+            ("qkvabz_proj", "b_proj", 4),
+            ("qkvabz_proj", "g_proj", 5),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        for weight_name, loaded_weight in weights:
+            # Handle stacked parameters first.
+            for (
+                param_name_part,
+                shard_weight_name_part,
+                shard_id,
+            ) in stacked_params_mapping:
+                if shard_weight_name_part not in weight_name.split("."):
+                    continue
+
+                param_name = weight_name.replace(
+                    shard_weight_name_part, param_name_part
+                )
+
+                if param_name not in params_dict:
+                    # Fall back to direct match if no such stacked parameter.
+                    continue
+
+                param = params_dict[param_name]
+                weight_loader = getattr(param, "weight_loader")
+                weight_loader(param, loaded_weight, shard_id)
+                break
+
+            else:
+                param_name = weight_name
+
+                param = params_dict[param_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+
+EntryClass = JetNemotronForCausalLM

python/sglang/srt/utils/hf_transformers_utils.py

@@ -44,6 +44,7 @@ from sglang.srt.configs import (
     DotsVLMConfig,
     ExaoneConfig,
     FalconH1Config,
+    JetNemotronConfig,
     KimiLinearConfig,
     KimiVLConfig,
     LongcatFlashConfig,
@@ -77,6 +78,7 @@ _CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
     DotsOCRConfig,
     NemotronHConfig,
     DeepseekVLV2Config,
+    JetNemotronConfig,
 ]
 
 _CONFIG_REGISTRY = {

python/sglang/test/test_utils.py

@@ -116,7 +116,7 @@ DEFAULT_ENABLE_THINKING_MODEL_NAME_FOR_TEST = "Qwen/Qwen3-30B-A3B"
 DEFAULT_DEEPSEEK_W4AFP8_MODEL_FOR_TEST = "Barrrrry/DeepSeek-R1-W4AFP8"
 
 # Nightly tests
-DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1 = "meta-llama/Llama-3.1-8B-Instruct,mistralai/Mistral-7B-Instruct-v0.3,deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct,google/gemma-2-27b-it"
+DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1 = "meta-llama/Llama-3.1-8B-Instruct,mistralai/Mistral-7B-Instruct-v0.3,deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct,google/gemma-2-27b-it,jet-ai/Jet-Nemotron-2B"
 DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2 = "meta-llama/Llama-3.1-70B-Instruct,mistralai/Mixtral-8x7B-Instruct-v0.1,Qwen/Qwen2-57B-A14B-Instruct"
 DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1 = "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8,neuralmagic/Mistral-7B-Instruct-v0.3-FP8,neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8,neuralmagic/gemma-2-2b-it-FP8"
 DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2 = "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8,neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8,neuralmagic/Qwen2-72B-Instruct-FP8,neuralmagic/Qwen2-57B-A14B-Instruct-FP8,neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8,zai-org/GLM-4.5-Air-FP8"