From 3633f8b0cfef2de3d31a451f7ae5187a0fc9a2af Mon Sep 17 00:00:00 2001 From: Zijian Zhang <35801754+futrime@users.noreply.github.com> Date: Sun, 9 Nov 2025 17:32:47 +0800 Subject: [PATCH] Add Jet-Nemotron (#12448) --- docs/supported_models/generative_models.md | 1 + python/sglang/srt/configs/__init__.py | 2 + python/sglang/srt/configs/jet_nemotron.py | 74 +++ .../sglang/srt/model_executor/model_runner.py | 3 +- python/sglang/srt/models/jet_nemotron.py | 596 ++++++++++++++++++ .../sglang/srt/utils/hf_transformers_utils.py | 2 + python/sglang/test/test_utils.py | 2 +- 7 files changed, 678 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/configs/jet_nemotron.py create mode 100644 python/sglang/srt/models/jet_nemotron.py diff --git a/docs/supported_models/generative_models.md b/docs/supported_models/generative_models.md index f08cb847b..2e4651b35 100644 --- a/docs/supported_models/generative_models.md +++ b/docs/supported_models/generative_models.md @@ -59,3 +59,4 @@ in the GitHub search bar. | **Llama Nemotron Ultra** (v1, NVIDIA) | `nvidia/Llama-3_1-Nemotron-Ultra-253B-v1` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. | | **NVIDIA Nemotron Nano 2.0** | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | The [NVIDIA Nemotron](https://www.nvidia.com/en-us/ai-data-science/foundation-models/nemotron/) family of multimodal models provides state-of-the-art reasoning models specifically designed for enterprise-ready AI agents. `Nemotron-Nano-9B-v2` is a hybrid Mamba-Transformer language model designed to increase throughput for reasoning workloads while achieving state-of-the-art accuracy compared to similarly-sized models. | | **StarCoder2** (3B-15B) | `bigcode/starcoder2-7b` | StarCoder2 is a family of open large language models (LLMs) specialized for code generation and understanding. It is the successor to StarCoder, jointly developed by the BigCode project (a collaboration between Hugging Face, ServiceNow Research, and other contributors). | +| **Jet-Nemotron** | `jet-ai/Jet-Nemotron-2B` | Jet-Nemotron is a new family of hybrid-architecture language models that surpass state-of-the-art open-source full-attention language models, while achieving significant efficiency gains. | diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 2e5a3ba9e..690a1e3eb 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/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", ] diff --git a/python/sglang/srt/configs/jet_nemotron.py b/python/sglang/srt/configs/jet_nemotron.py new file mode 100644 index 000000000..c05a2ec10 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 963e2cd81..6e3a2ef6c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/models/jet_nemotron.py b/python/sglang/srt/models/jet_nemotron.py new file mode 100644 index 000000000..918cdc65b --- /dev/null +++ b/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 diff --git a/python/sglang/srt/utils/hf_transformers_utils.py b/python/sglang/srt/utils/hf_transformers_utils.py index 30d0dc16d..fbc34e906 100644 --- a/python/sglang/srt/utils/hf_transformers_utils.py +++ b/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 = { diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 42af14537..822c7e12c 100644 --- a/python/sglang/test/test_utils.py +++ b/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"