From 41e7612dee44e262ba2f474546a5144c449c6596 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 25 Aug 2026 16:43:58 -0700 Subject: [PATCH] [Model] Support Nemotron 3.5 Lightning speculative decoding (#36186) Co-authored-by: Ryan Stewart --- .../NVIDIA/Nemotron3.5-Lightning.mdx | 2 +- python/sglang/srt/arg_groups/overrides.py | 28 ++- .../sglang/srt/arg_groups/speculative_hook.py | 30 +-- .../srt/layers/quantization/modelopt_quant.py | 17 +- .../sglang/srt/model_executor/model_runner.py | 9 + python/sglang/srt/models/dflash.py | 188 ++++++++++++--- python/sglang/srt/models/dspark.py | 77 ++++++- python/sglang/srt/models/nemotron_h.py | 43 +++- python/sglang/srt/server_args.py | 13 +- python/sglang/srt/speculative/dflash_utils.py | 42 ++++ .../srt/speculative/dflash_worker_v2.py | 15 +- .../dspark_components/dspark_config.py | 16 +- .../dspark_components/dspark_draft.py | 68 +++++- .../dspark_components/dspark_draft_sampler.py | 24 +- .../dspark_components/dspark_worker_v2.py | 10 +- python/sglang/test/kits/eval_accuracy_kit.py | 2 + .../models_e2e/test_nvidia_nemotron_3_nano.py | 115 ++++++++-- test/registered/unit/test_model_overrides.py | 215 +++++++++++++++++- 18 files changed, 792 insertions(+), 122 deletions(-) diff --git a/docs/cookbook/autoregressive/NVIDIA/Nemotron3.5-Lightning.mdx b/docs/cookbook/autoregressive/NVIDIA/Nemotron3.5-Lightning.mdx index 92fd60a97..4ab95c627 100644 --- a/docs/cookbook/autoregressive/NVIDIA/Nemotron3.5-Lightning.mdx +++ b/docs/cookbook/autoregressive/NVIDIA/Nemotron3.5-Lightning.mdx @@ -19,7 +19,7 @@ For all methods and hardware platforms, see the [official SGLang installation gu ```bash Command pip install --upgrade pip pip install uv -SGLANG_BUILD_RUST_EXTS=none uv pip install --prerelease=allow 'git+https://github.com/sgl-project/sglang.git@refs/pull/33554/head#subdirectory=python' +SGLANG_BUILD_RUST_EXTS=none uv pip install --prerelease=allow 'git+https://github.com/sgl-project/sglang.git#subdirectory=python' ``` Then run the **Python** output of the command panel below in that environment. diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 6f29998b3..2f1654528 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1482,8 +1482,32 @@ def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict: else: overrides["moe_runner_backend"] = "flashinfer_cutlass" - if is_sm100_supported() and server_args.attention_backend is None: - overrides["attention_backend"] = "flashinfer" + if is_blackwell_supported() and server_args.is_attention_backend_not_set(): + if server_args.speculative_algorithm is not None: + speculative_algorithm = server_args.speculative_algorithm.upper() + if is_sm100_supported() and server_args.speculative_eagle_topk in ( + None, + 1, + ): + overrides["attention_backend"] = "trtllm_mha" + if server_args.page_size is None: + overrides["page_size"] = 64 + if server_args.mamba_radix_cache_strategy == "auto": + overrides["mamba_radix_cache_strategy"] = "extra_buffer" + if ( + server_args.speculative_draft_attention_backend is None + and speculative_algorithm in ("EAGLE", "NEXTN", "DSPARK") + ): + overrides["speculative_draft_attention_backend"] = "trtllm_mha" + else: + overrides["attention_backend"] = "triton" + if ( + server_args.speculative_draft_attention_backend is None + and speculative_algorithm in ("EAGLE", "NEXTN", "DFLASH", "DSPARK") + ): + overrides["speculative_draft_attention_backend"] = "flashinfer" + elif is_sm100_supported(): + overrides["attention_backend"] = "trtllm_mha" return overrides diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 7b831be96..905dbcb54 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -447,6 +447,21 @@ def _handle_dspark(server_args: ServerArgs) -> None: speculative_eagle_topk=1, ) + from sglang.srt.speculative.dspark_components.dspark_config import ( + DEFAULT_DSPARK_GAMMA, + read_draft_checkpoint_config, + ) + + draft_config = None + try: + draft_config = read_draft_checkpoint_config(server_args=server_args) + except Exception as e: + logger.warning( + "Failed to read DSpark draft config; preserving explicit/default " + "gamma resolution. Error: %s", + e, + ) + gamma: Optional[int] = None if server_args.speculative_dspark_block_size is not None: if int(server_args.speculative_dspark_block_size) <= 0: @@ -456,19 +471,8 @@ def _handle_dspark(server_args: ServerArgs) -> None: ) gamma = int(server_args.speculative_dspark_block_size) else: - from sglang.srt.speculative.dspark_components.dspark_config import ( - DEFAULT_DSPARK_GAMMA, - read_draft_checkpoint_gamma, - ) - - try: - gamma = read_draft_checkpoint_gamma(server_args=server_args) - except Exception as e: - logger.warning( - "Failed to read DSpark gamma from draft model config; " - "cannot cross-check --speculative-num-draft-tokens. Error: %s", - e, - ) + if draft_config is not None: + gamma = draft_config.resolve_gamma(default=None) if gamma is None and server_args.speculative_num_draft_tokens is None: gamma = DEFAULT_DSPARK_GAMMA logger.warning( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 96fbd2819..38508638d 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1407,6 +1407,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): "format is experimental and subject to change." ) self.is_awq = is_awq + self.is_w4a16 = False self.group_size = group_size if not is_checkpoint_nvfp4_serialized: if use_per_token_activation: @@ -1561,9 +1562,10 @@ class ModelOptFp4Config(ModelOptQuantConfig): "Expected either flat format (config.json) or nested format (hf_quant_config.json)." ) - if quant_method not in ["FP8", "NVFP4", "NVFP4_AWQ"]: + if quant_method not in ["FP8", "NVFP4", "NVFP4_AWQ", "W4A16_NVFP4"]: raise ValueError( - "ModelOpt currently only supports: FP8, NVFP4, NVFP4_AWQ " + "ModelOpt currently only supports: FP8, NVFP4, NVFP4_AWQ, " + "W4A16_NVFP4 " "quantizations in sglang. Please check the " "quantization config for your model's configuration." ) @@ -1579,14 +1581,17 @@ class ModelOptFp4Config(ModelOptQuantConfig): "NVFP4 quantization requires group_size and exclude_modules " "specified in the quantization config" ) - return cls( + quant_config = cls( is_checkpoint_nvfp4_serialized, kv_cache_quant_algo, group_size, exclude_modules, config.get("packed_modules_mapping"), is_awq="AWQ" in quant_method, + use_per_token_activation=(False if quant_method == "W4A16_NVFP4" else None), ) + quant_config.is_w4a16 = quant_method == "W4A16_NVFP4" + return quant_config def get_quant_method(self, layer: torch.nn.Module, prefix: str): from sglang.srt.layers.linear import LinearBase @@ -1606,7 +1611,11 @@ class ModelOptFp4Config(ModelOptQuantConfig): return self._get_quant_method( layer, prefix, - Linear=ModelOptFp4LinearMethod, + Linear=( + ModelOptNvFp4A16LinearMethod + if self.is_w4a16 + else ModelOptFp4LinearMethod + ), Moe=ModelOptNvFp4FusedMoEMethod, ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9ad07a41d..8c4d18f61 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -807,6 +807,15 @@ class ModelRunner: ) -> int: """Logits rows per decode batch slot.""" if self.spec_algorithm.is_speculative(): + if self.spec_algorithm.is_dspark() and self.is_draft_worker: + from sglang.srt.speculative.dspark_components.dspark_config import ( + get_dspark_sample_from_anchor, + ) + + if not get_dspark_sample_from_anchor(self.model_config.hf_config): + if num_draft_tokens is None: + num_draft_tokens = get_spec().speculative_num_draft_tokens + return int(num_draft_tokens) return resolve_num_tokens_per_req( phase="target_verify", spec_algorithm=self.spec_algorithm, diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index b7de868d7..3a88ad392 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -1,7 +1,6 @@ # Adapted from the DFlash reference implementation (HF) but implemented with -# SGLang primitives (RadixAttention + SGLang KV cache). This model intentionally -# does not include token embeddings or an LM head; DFlash uses the target model's -# embedding/lm_head. +# SGLang primitives (RadixAttention + SGLang KV cache). Most drafts borrow the +# target embedding and LM head; Nemotron 3.5 drafts carry their own embedding. from __future__ import annotations @@ -21,6 +20,7 @@ from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import ( @@ -29,18 +29,23 @@ from sglang.srt.layers.logits_processor import ( ) from sglang.srt.layers.radix_attention import AttentionType, RadixAttention from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) from sglang.srt.models.utils import apply_qk_norm -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.speculative.dflash_utils import ( can_dflash_slice_qkv_weight, get_dflash_attention_sliding_window_size, get_dflash_layer_types, is_dense_head_weight, + is_nemotron_35_draft_config, parse_dflash_draft_config, ) -from sglang.srt.utils import is_npu +from sglang.srt.utils import is_npu, set_weight_attrs from sglang.srt.utils.common import get_compiler_backend from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -62,6 +67,24 @@ def _radix_topk(scores: torch.Tensor, k: int) -> Tuple[torch.Tensor, torch.Tenso return torch.topk(scores, k, dim=-1) +def _logical_linear_weight_shape( + param: torch.Tensor, + loaded_weight: torch.Tensor, + *, + output_features: int, +) -> Tuple[int, ...]: + """Return a checkpoint linear weight shape in logical elements.""" + loaded_shape = tuple(loaded_weight.shape) + pack_factor = getattr(param, "pack_factor", None) + if pack_factor is None or loaded_shape != tuple(param.shape): + return loaded_shape + + logical_numel = int(loaded_weight.numel() * pack_factor) + if logical_numel % output_features == 0: + return (output_features, logical_numel // output_features) + return (logical_numel,) + + def _project_candidate_logits( hidden: torch.Tensor, lm_head: nn.Module, *, num_org: int, use_quant_head: bool ) -> torch.Tensor: @@ -117,7 +140,9 @@ def _get_dflash_layer_attention_params( class DFlashAttention(nn.Module): - def __init__(self, config, layer_id: int, quant_config=None) -> None: + def __init__( + self, config, layer_id: int, quant_config=None, prefix: str = "" + ) -> None: super().__init__() hidden_size = int(config.hidden_size) tp_size = int(get_parallel().tp_size) @@ -160,14 +185,14 @@ class DFlashAttention(nn.Module): total_num_kv_heads=self.total_num_kv_heads, bias=attention_bias, quant_config=quant_config, - prefix="qkv_proj", + prefix=f"{prefix}.qkv_proj" if prefix else "qkv_proj", ) self.o_proj = RowParallelLinear( self.total_num_heads * head_dim, hidden_size, bias=attention_bias, quant_config=quant_config, - prefix="o_proj", + prefix=f"{prefix}.o_proj" if prefix else "o_proj", ) # Per-head Q/K RMSNorm, matching HF Qwen3. @@ -201,6 +226,24 @@ class DFlashAttention(nn.Module): self.sliding_window_size, self.attn_type = _get_dflash_layer_attention_params( config, layer_id ) + self.attention_sink_bias = None + if is_nemotron_35_draft_config(config) and bool( + getattr(config, "attention_sink_bias", False) + ): + draft_attention_backend = get_spec().speculative_draft_attention_backend + if draft_attention_backend != "trtllm_mha": + raise ValueError( + "Nemotron 3.5 DSpark attention sinks require " + "--speculative-draft-attention-backend trtllm_mha, " + f"got {draft_attention_backend!r}." + ) + self.attention_sink_bias = nn.Parameter( + torch.empty(self.num_heads, dtype=torch.float32), requires_grad=False + ) + set_weight_attrs( + self.attention_sink_bias, + {"weight_loader": sharded_weight_loader(0)}, + ) self.attn = RadixAttention( num_heads=self.num_heads, head_dim=head_dim, @@ -259,7 +302,12 @@ class DFlashAttention(nn.Module): q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = apply_qk_norm(q, k, self.q_norm, self.k_norm, self.head_dim) q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v, forward_batch) + if self.attention_sink_bias is None: + attn_output = self.attn(q, k, v, forward_batch) + else: + attn_output = self.attn( + q, k, v, forward_batch, sinks=self.attention_sink_bias + ) attn_output = self.apply_attention_output(attn_output, hidden_states) output, _ = self.o_proj(attn_output) return output @@ -428,17 +476,25 @@ class DFlashDecoderLayer(nn.Module): attention_conv: Optional[DFlashGroupedConv] = None, mlp_conv: Optional[DFlashGroupedConv] = None, quant_config=None, + prefix: str = "", ) -> None: super().__init__() hidden_size = int(config.hidden_size) rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-6)) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + attention_prefix = f"{prefix}.self_attn" if prefix else "" self.self_attn = self.attention_cls( - config=config, layer_id=layer_id, quant_config=quant_config + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=attention_prefix, ) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = DFlashMLP(config=config, quant_config=quant_config) + mlp_prefix = f"{prefix}.mlp" if prefix else "" + self.mlp = DFlashMLP( + config=config, quant_config=quant_config, prefix=mlp_prefix + ) self.attention_conv = attention_conv self.mlp_conv = mlp_conv @@ -487,7 +543,7 @@ class DFlashDecoderLayer(nn.Module): class DFlashDraftModel(nn.Module): - """SGLang DFlash draft model (no embedding / lm_head weights). + """SGLang DFlash draft model with an optional Nemotron embedding. The checkpoint provides: - transformer weights for `layers.*` @@ -510,6 +566,16 @@ class DFlashDraftModel(nn.Module): ) self.block_size = draft_config.resolve_block_size(default=16) self.candidate_selector: Optional[nn.Module] = None + self.is_nemotron_35_draft = is_nemotron_35_draft_config(config) + self.embed_tokens: Optional[VocabParallelEmbedding] = None + if self.is_nemotron_35_draft: + embed_prefix = f"{prefix}.embed_tokens" if prefix else "embed_tokens" + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + hidden_size, + quant_config=quant_config, + prefix=embed_prefix, + ) def grouped_conv(): if not draft_config.conv_kernel_size: @@ -529,6 +595,11 @@ class DFlashDraftModel(nn.Module): attention_conv=grouped_conv(), mlp_conv=grouped_conv(), quant_config=quant_config, + prefix=( + (f"{prefix}.layers.{i}" if prefix else f"layers.{i}") + if self.is_nemotron_35_draft + else "" + ), ) for i in range(num_layers) ] @@ -550,9 +621,19 @@ class DFlashDraftModel(nn.Module): num_context_features = len(target_layer_ids) self.num_context_features = int(num_context_features) - self.fc = nn.Linear( - self.num_context_features * hidden_size, hidden_size, bias=False - ) + if self.is_nemotron_35_draft: + fc_prefix = f"{prefix}.fc" if prefix else "fc" + self.fc = ReplicatedLinear( + self.num_context_features * hidden_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=fc_prefix, + ) + else: + self.fc = nn.Linear( + self.num_context_features * hidden_size, hidden_size, bias=False + ) self.hidden_norm = RMSNorm(hidden_size, eps=rms_norm_eps) def set_block_size(self, block_size: int) -> None: @@ -571,6 +652,9 @@ class DFlashDraftModel(nn.Module): def get_attention_sliding_window_size(self) -> Optional[int]: return get_dflash_attention_sliding_window_size(self.config) + def get_input_embeddings(self) -> Optional[VocabParallelEmbedding]: + return self.embed_tokens + def prepare_context_hidden_for_kv( self, layer: DFlashDecoderLayer, ctx_hidden: torch.Tensor ) -> torch.Tensor: @@ -578,7 +662,9 @@ class DFlashDraftModel(nn.Module): def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: """Project concatenated target-layer hidden states into draft hidden_size.""" - expected = int(self.fc.in_features) + expected = int( + self.fc.input_size if self.is_nemotron_35_draft else self.fc.in_features + ) if target_hidden.ndim != 2 or int(target_hidden.shape[-1]) != expected: raise ValueError( "DFLASH target_hidden feature dim mismatch. " @@ -588,7 +674,10 @@ class DFlashDraftModel(nn.Module): "This usually means the target model is capturing a different number of layer features than " "the draft checkpoint/config expects." ) - return self.hidden_norm(self.fc(target_hidden)) + projected = self.fc(target_hidden) + if self.is_nemotron_35_draft: + projected = projected[0] + return self.hidden_norm(projected) @torch.no_grad() def forward( @@ -601,7 +690,9 @@ class DFlashDraftModel(nn.Module): pp_proxy_tensors=None, ) -> LogitsProcessorOutput: if input_embeds is None: - if hasattr(self, "forward_embed"): + if self.embed_tokens is not None: + input_embeds = self.embed_tokens(input_ids) + elif hasattr(self, "forward_embed"): input_embeds = self.forward_embed(input_ids) else: raise ValueError( @@ -679,16 +770,33 @@ class DFlashDraftModel(nn.Module): # Ignore unexpected weights (e.g., HF rotary caches). continue param = params_dict[resolved_name] - if resolved_name.endswith("fc.weight") and tuple( - loaded_weight.shape - ) != tuple(param.shape): - raise ValueError( - "DFLASH fc.weight shape mismatch. This usually means the draft checkpoint's " - "number of context features (K) does not match this config. " - f"Expected fc.weight.shape={tuple(param.shape)} " - f"(num_context_features={self.num_context_features}, hidden_size={int(self.config.hidden_size)}), " - f"but got {tuple(loaded_weight.shape)} for weight '{name}'." - ) + if resolved_name.endswith("fc.weight"): + if self.is_nemotron_35_draft: + expected_shape = ( + int(self.config.hidden_size), + int(self.num_context_features * self.config.hidden_size), + ) + loaded_shape = _logical_linear_weight_shape( + param, + loaded_weight, + output_features=expected_shape[0], + ) + shape_matches = loaded_shape == expected_shape or ( + getattr(param, "pack_factor", None) is None + and tuple(loaded_weight.shape) == tuple(param.shape) + ) + else: + expected_shape = tuple(param.shape) + loaded_shape = tuple(loaded_weight.shape) + shape_matches = loaded_shape == expected_shape + if not shape_matches: + raise ValueError( + "DFLASH fc.weight shape mismatch. This usually means the draft checkpoint's " + "number of context features (K) does not match this config. " + f"Expected fc.weight.shape={expected_shape} " + f"(num_context_features={self.num_context_features}, hidden_size={int(self.config.hidden_size)}), " + f"but got {loaded_shape} for weight '{name}'." + ) weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) @@ -696,8 +804,15 @@ class DFlashDraftModel(nn.Module): class DFlashLagunaAttention(DFlashAttention): """Laguna DFlash attention with the trained Laguna softplus gate.""" - def __init__(self, config, layer_id: int, quant_config=None) -> None: - super().__init__(config=config, layer_id=layer_id, quant_config=quant_config) + def __init__( + self, config, layer_id: int, quant_config=None, prefix: str = "" + ) -> None: + super().__init__( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=prefix, + ) hidden_size = int(config.hidden_size) total_num_heads = self.total_num_heads gating = normalize_gating(getattr(config, "gating", True)) @@ -716,7 +831,7 @@ class DFlashLagunaAttention(DFlashAttention): g_out, bias=False, quant_config=quant_config, - prefix="g_proj", + prefix=f"{prefix}.g_proj" if prefix else "g_proj", ) def apply_attention_output( @@ -764,7 +879,9 @@ class DFlashLagunaForCausalLM(DFlashDraftModel): return layer.input_layernorm(ctx_hidden) def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: - expected = int(self.fc.in_features) + expected = int( + self.fc.input_size if self.is_nemotron_35_draft else self.fc.in_features + ) if target_hidden.ndim != 2 or int(target_hidden.shape[-1]) != expected: raise ValueError( "Laguna DFLASH target_hidden feature dim mismatch. " @@ -783,7 +900,10 @@ class DFlashLagunaForCausalLM(DFlashDraftModel): for i, norm in enumerate(self.aux_hidden_norms): normed[:, i, :] = norm(slices[:, i, :]) fused = normed.reshape(target_hidden.shape[0], -1) - return self.hidden_norm(self.fc(fused)) + projected = self.fc(fused) + if self.is_nemotron_35_draft: + projected = projected[0] + return self.hidden_norm(projected) @torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index 69ecbfdd0..a0447e606 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -9,11 +9,13 @@ from torch import nn from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather from sglang.srt.environ import envs +from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.dflash import DFlashDraftModel from sglang.srt.speculative.dflash_utils import can_dflash_slice_qkv_weight from sglang.srt.speculative.dspark_components.dspark_config import ( + get_dspark_sample_from_anchor, parse_dspark_draft_config, ) from sglang.srt.speculative.ragged_verify import ( @@ -142,6 +144,39 @@ class VanillaMarkov(nn.Module): ) +class Nemotron35VanillaMarkov(VanillaMarkov): + """Checkpoint-quantized Markov head used only by Nemotron 3.5 DSpark.""" + + def __init__( + self, + *, + vocab_size: int, + markov_rank: int, + quant_config, + prefix: str, + ) -> None: + nn.Module.__init__(self) + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + if self.markov_rank <= 0: + raise ValueError( + "Nemotron35VanillaMarkov requires markov_rank > 0, " + f"got {self.markov_rank}." + ) + self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) + self.markov_w2 = ReplicatedLinear( + self.markov_rank, + self.vocab_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.markov_w2" if prefix else "markov_w2", + ) + + def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor: + bias, _ = self.markov_w2(latent_states) + return bias + + class GatedMarkovHead(VanillaMarkov): markov_head_type = "gated" @@ -298,6 +333,22 @@ def build_markov_head(config) -> Optional[nn.Module]: raise ValueError(f"Unsupported DSpark markov_head_type={markov_head_type!r}.") +def build_nemotron_35_markov_head(config, quant_config, prefix: str) -> nn.Module: + markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower() + if markov_head_type != "vanilla": + raise ValueError( + "Nemotron 3.5 DSpark requires markov_head_type='vanilla', " + f"got {markov_head_type!r}." + ) + markov_prefix = f"{prefix}.markov_head" if prefix else "markov_head" + return Nemotron35VanillaMarkov( + vocab_size=int(config.vocab_size), + markov_rank=int(config.markov_rank), + quant_config=quant_config, + prefix=markov_prefix, + ) + + class DSparkConfidenceHead(nn.Module): def __init__( @@ -365,11 +416,7 @@ def build_confidence_head(config) -> Optional[nn.Module]: ) -_DSPARK_SKIPPED_WEIGHT_PREFIXES = ( - "embed_tokens.", - "lm_head.", - "rotary_emb.", -) +_DSPARK_SKIPPED_WEIGHT_PREFIXES = ("lm_head.", "rotary_emb.") class DSparkDraftMixin: @@ -385,14 +432,21 @@ class DSparkDraftMixin: f"got markov_rank={dspark_config.markov_rank}." ) self.gamma = int(dspark_config.resolve_gamma(default=self.block_size)) - self.markov_head = build_markov_head(config) + self.sample_from_anchor = get_dspark_sample_from_anchor(config) + if self.is_nemotron_35_draft: + self.markov_head = build_nemotron_35_markov_head( + config, quant_config, prefix + ) + else: + self.markov_head = build_markov_head(config) self.confidence_head = build_confidence_head(config) self.lm_head: Optional[nn.Module] = None def attach_shared_modules( self, *, embed_tokens: nn.Module, lm_head: nn.Module ) -> None: - self.embed_tokens = embed_tokens + if not self.is_nemotron_35_draft: + self.embed_tokens = embed_tokens self.lm_head = lm_head def forward_embed(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -430,7 +484,14 @@ class DSparkDraftMixin: backbone_weights = [] params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: - if any(name.startswith(p) for p in _DSPARK_SKIPPED_WEIGHT_PREFIXES): + normalized_name = name.removeprefix("model.") + if any( + normalized_name.startswith(p) for p in _DSPARK_SKIPPED_WEIGHT_PREFIXES + ): + continue + if normalized_name.startswith("embed_tokens.") and not ( + self.is_nemotron_35_draft + ): continue if name.startswith("confidence_head."): if self.confidence_head is None: diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 7fa2f1d6b..803589f33 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -833,6 +833,7 @@ class NemotronHModel(nn.Module): self.norm_f = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) else: self.norm_f = PPMissingLayer(return_tuple=True) + self.layers_to_capture: set[int] = set() def forward( self, @@ -853,7 +854,17 @@ class NemotronHModel(nn.Module): hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] + aux_hidden_states = [] for i in range(self.start_layer, self.end_layer): + if i in self.layers_to_capture: + if residual is not None and getattr( + hidden_states, "_sglang_needs_allreduce_fusion", False + ): + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states._sglang_needs_allreduce_fusion = False + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) layer = self.layers[i] if not isinstance(layer, Layers): raise ValueError(f"Unknown layer type: {type(layer)}") @@ -867,7 +878,18 @@ class NemotronHModel(nn.Module): return PPProxyTensors( {"hidden_states": hidden_states, "residual": residual} ) + if self.end_layer in self.layers_to_capture: + if residual is not None and getattr( + hidden_states, "_sglang_needs_allreduce_fusion", False + ): + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states._sglang_needs_allreduce_fusion = False + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) hidden_states, _ = self.norm_f(hidden_states, residual) + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states @@ -962,6 +984,7 @@ class NemotronHForCausalLM(nn.Module): self.lm_head.weight.copy_(emb_token_weight) self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = False def _init_model( self, @@ -1082,9 +1105,16 @@ class NemotronHForCausalLM(nn.Module): hidden_states = self.model.forward( input_ids, positions, forward_batch, pp_proxy_tensors, input_embeds ) + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states if self.pp_group.is_last_rank: return self.logits_processor( - input_ids, hidden_states, self.lm_head, forward_batch + input_ids, + hidden_states, + self.lm_head, + forward_batch, + aux_hidden_states, ) else: return hidden_states @@ -1106,6 +1136,17 @@ class NemotronHForCausalLM(nn.Module): torch.cuda.empty_cache() torch.cuda.synchronize() + def set_dflash_layers_to_capture(self, layer_ids: list[int]): + if not self.pp_group.is_last_rank: + return + if layer_ids is None: + raise ValueError( + "DFLASH requires explicit layer_ids for aux hidden capture." + ) + + self.capture_aux_hidden_states = True + self.model.layers_to_capture = {layer_id + 1 for layer_id in layer_ids} + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False ) -> None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 769a3b404..6cc87be0c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2249,7 +2249,10 @@ class ServerArgs: ] = "prefill" speculative_draft_attention_backend: A[ Optional[str], - "Attention backend for speculative decoding drafting.", + Arg( + help="Attention backend for speculative decoding drafting.", + resolvable=True, + ), NS("spec"), ] = None speculative_draft_kv_cache_dtype: A[ @@ -6069,14 +6072,6 @@ class ServerArgs: logger.info( f"Using {attention_backend} as attention backend for {model_arch}." ) - elif model_arch in ["NemotronHForCausalLM", "NemotronHPuzzleForCausalLM"]: - # Quantization / MoE runner / attention backend defaults moved to - # the override registry (arg_groups/overrides.py: - # _nemotron_h_overrides). - assert resolved_view(self).attention_backend != "triton", ( - "NemotronHForCausalLM does not support triton attention backend," - "as the first layer might not be an attention layer" - ) elif model_arch in [ "Qwen3MoeForCausalLM", "Qwen3VLMoeForConditionalGeneration", diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index f8ce936c5..db852e486 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -413,6 +413,8 @@ def get_dflash_attention_sliding_window_size(config: Any) -> Optional[int]: sliding_window = _cfg_get( text_config, "sliding_window", _cfg_get(config, "sliding_window") ) + if sliding_window is None and is_nemotron_35_draft_config(config): + sliding_window = _get_dflash_config(config).get("swa_window_size") if sliding_window is None: raise ValueError( "DFLASH sliding_attention layers require config.sliding_window." @@ -463,6 +465,46 @@ def _get_dflash_config(config: Any) -> dict: return {} +def is_nemotron_35_draft_config(config: Any) -> bool: + """Identify the published Nemotron 3.5 DFlash/DSpark draft layout. + + Keep the non-anchor query layout and checkpoint-local vocabulary modules + scoped to this structurally distinct family instead of changing every + DFlash/DSpark checkpoint that happens to expose one of these fields. + """ + architectures = _cfg_get(config, "architectures", None) or [] + if not {"DFlashDraftModel", "Qwen3DSparkModel"}.intersection(architectures): + return False + if not bool(_cfg_get(config, "has_embed_tokens", False)): + return False + if bool(_cfg_get(config, "has_lm_head", False)): + return False + + quant_config = _cfg_get(config, "quantization_config", None) or {} + if _cfg_get(quant_config, "quant_algo", None) != "W4A16_NVFP4": + return False + + dflash_config = _get_dflash_config(config) + target_layer_ids = dflash_config.get( + "target_layer_ids", _cfg_get(config, "target_layer_ids", None) + ) + aux_layer_ids = _cfg_get(config, "eagle_aux_hidden_state_layer_ids", None) + if not target_layer_ids or not aux_layer_ids: + return False + if len(target_layer_ids) != len(aux_layer_ids): + return False + if any( + int(target) + 1 != int(aux) + for target, aux in zip(target_layer_ids, aux_layer_ids) + ): + return False + + sample_from_anchor = dflash_config.get( + "sample_from_anchor", _cfg_get(config, "sample_from_anchor", True) + ) + return sample_from_anchor is False + + def _parse_optional_int( value: Any, *, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 33bcdbda9..7da219523 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -178,6 +178,17 @@ def _commit_accept(candidates, accept_len, bonus_tokens): return out_tokens, accept_len.to(torch.int32) + 1 +def _resolve_dflash_embedding_module(draft_model, target_model): + if getattr(draft_model, "is_nemotron_35_draft", False): + embed_module = draft_model.get_input_embeddings() + if embed_module is None: + raise RuntimeError( + "Nemotron 3.5 DFLASH draft requires its checkpoint embedding." + ) + return embed_module + return target_model.get_input_embeddings() + + def _is_all_greedy(sampling_info) -> bool: return sampling_info is None or sampling_info.is_all_greedy @@ -1753,7 +1764,9 @@ class DFlashWorkerV2(BaseSpecWorker): # --- 1) Draft a fixed block with the draft model. target_model = self.target_worker.model_runner.model - embed_module = unwrap_lora_layer(target_model.get_input_embeddings()) + embed_module = unwrap_lora_layer( + _resolve_dflash_embedding_module(self.draft_model, target_model) + ) lm_head = unwrap_lora_layer(getattr(target_model, "lm_head", None)) if lm_head is None or not ( hasattr(lm_head, "weight") diff --git a/python/sglang/srt/speculative/dspark_components/dspark_config.py b/python/sglang/srt/speculative/dspark_components/dspark_config.py index 55b699f3d..068e3ff7c 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_config.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_config.py @@ -25,6 +25,11 @@ SUPPORTED_DSPARK_MARKOV_HEAD_TYPES = ("vanilla", "gated", "rnn") DSV4_DRAFT_ATTENTION_BACKEND = "dsv4" +def get_dspark_sample_from_anchor(draft_hf_config: Any) -> bool: + """Return whether a DSpark checkpoint samples the anchor query row.""" + return bool(_cfg_get(draft_hf_config, "sample_from_anchor", True)) + + def draft_is_deepseek_v4(*, server_args: ServerArgs) -> bool: from sglang.srt.configs.model_config import is_deepseek_v4 from sglang.srt.utils.hf_transformers_utils import get_config @@ -128,9 +133,8 @@ def resolve_runtime_config( ) -def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]: - """Load the draft checkpoint's hf config and read its DSpark gamma - (block_size). Raises on config-load failure; callers pick the fallback. +def read_draft_checkpoint_config(*, server_args: ServerArgs) -> DSparkDraftConfig: + """Load and normalize the DSpark draft checkpoint configuration. Reads the *resolving* configuration, not the bags: the speculative hook calls this from inside resolution, where no bag exists yet -- and the @@ -148,7 +152,11 @@ def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]: revision=resolving.speculative_draft_model_revision, model_override_args=json.loads(resolving.json_model_override_args), ) - return parse_dspark_draft_config(draft_hf_config=draft_hf_config).resolve_gamma( + return parse_dspark_draft_config(draft_hf_config=draft_hf_config) + + +def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]: + return read_draft_checkpoint_config(server_args=server_args).resolve_gamma( default=None ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 5514c7d49..4c1983a1b 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -81,6 +81,27 @@ class DraftProposal(msgspec.Struct, frozen=True): folded: bool = False +def select_draft_hidden_without_anchor( + hidden_states: torch.Tensor, + *, + bs: int, + gamma: int, +) -> tuple[torch.Tensor, torch.Tensor]: + query_token_num = gamma + 1 + expected_rows = bs * query_token_num + if hidden_states.shape[0] != expected_rows: + raise RuntimeError( + f"DSpark draft returned {hidden_states.shape[0]} hidden rows, " + f"expected {expected_rows}." + ) + hidden_by_query = hidden_states.view(bs, query_token_num, *hidden_states.shape[1:]) + selected = hidden_by_query[:, 1:].contiguous() + return ( + selected.view(bs * gamma, *hidden_states.shape[1:]), + selected.view(bs, gamma, -1), + ) + + def make_next_draft_input( *, bonus_tokens: torch.Tensor, @@ -178,6 +199,8 @@ class DraftBlockProposer: self.draft_model = draft_model self.draft_model_runner = draft_model_runner self.gamma = gamma + self.sample_from_anchor = bool(draft_model.sample_from_anchor) + self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1 self._mask_token_id = mask_token_id self._draft_block_spec_info = draft_block_spec_info self._draft_sampler = None @@ -202,7 +225,11 @@ class DraftBlockProposer: target_model, sampling_info, ) -> DraftProposal: - embed_module = unwrap_lora_layer(target_model.get_input_embeddings()) + embed_module = unwrap_lora_layer( + self.draft_model.embed_tokens + if not self.sample_from_anchor + else target_model.get_input_embeddings() + ) draft_sampler = self._draft_sampler all_greedy = sampling_info is None or sampling_info.is_all_greedy fwd = self._run_forward( @@ -276,8 +303,13 @@ class DraftBlockProposer: markov_head=self.draft_model.markov_head, device=device, ) + proposal_block_ids = ( + draft_block_ids + if self.sample_from_anchor + else draft_block_ids[:, : self.gamma].contiguous() + ) return DraftProposal( - draft_block_ids=draft_block_ids, + draft_block_ids=proposal_block_ids, draft_block=draft_block, draft_hidden=fwd.draft_hidden_3d, confidence=folded_confidence, @@ -321,16 +353,20 @@ class DraftBlockProposer: sampling_info=None, ) -> DraftForwardResult: gamma = self.gamma + query_token_num = self.query_token_num prefix_lens = batch.seq_lens positions_2d = verify_window.positions_2d verify_cache_loc_2d = verify_window.verify_cache_loc_2d draft_block_ids = torch.full( - (bs, gamma), int(self._mask_token_id), dtype=torch.long, device=device + (bs, query_token_num), + int(self._mask_token_id), + dtype=torch.long, + device=device, ) draft_block_ids[:, 0].copy_(draft_input.bonus_tokens.view(-1)) - draft_positions = positions_2d[:, :gamma].reshape(-1) - draft_cache_loc = verify_cache_loc_2d[:, :gamma].reshape(-1) + draft_positions = positions_2d[:, :query_token_num].reshape(-1) + draft_cache_loc = verify_cache_loc_2d[:, :query_token_num].reshape(-1) draft_owns_embed = envs.SGLANG_DSPARK_EMBED_IN_GRAPH.get() and hasattr( self.draft_model, "forward_embed" @@ -341,7 +377,7 @@ class DraftBlockProposer: draft_input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1]) if batch.seq_lens_cpu is not None: - draft_seq_lens_cpu = batch.seq_lens_cpu + gamma + draft_seq_lens_cpu = batch.seq_lens_cpu + query_token_num draft_seq_lens_sum = int(draft_seq_lens_cpu.sum()) elif draft_input.nxt_kv_lens_cpu is not None: draft_seq_lens_cpu = draft_input.nxt_kv_lens_cpu @@ -349,7 +385,7 @@ class DraftBlockProposer: else: raise RuntimeError("DSpark decode expected batch.seq_lens_cpu, got None") - draft_num_tokens = bs * gamma + draft_num_tokens = bs * query_token_num draft_forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, batch_size=bs, @@ -381,10 +417,24 @@ class DraftBlockProposer: raw_hidden = logits_output.hidden_states if raw_hidden is None: raise RuntimeError("DSpark draft model returned no hidden states.") - draft_hidden_3d = raw_hidden.view(bs, gamma, -1) + if self.sample_from_anchor: + expected_rows = bs * gamma + if raw_hidden.shape[0] != expected_rows: + raise RuntimeError( + f"DSpark draft returned {raw_hidden.shape[0]} hidden rows, " + f"expected {expected_rows}." + ) + model_hidden = raw_hidden + draft_hidden_3d = raw_hidden.view(bs, gamma, -1) + else: + model_hidden, draft_hidden_3d = select_draft_hidden_without_anchor( + raw_hidden, + bs=bs, + gamma=gamma, + ) return DraftForwardResult( draft_block_ids=draft_block_ids, - raw_hidden=raw_hidden, + raw_hidden=model_hidden, draft_hidden_3d=draft_hidden_3d, can_run_graph=draft_out.can_run_graph, ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index df5926cc0..1b7c3392f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py @@ -9,6 +9,9 @@ from sglang.kernels.ops.speculative.dspark.dspark_draft_model import ( SampleStepTokens, ) from sglang.srt.environ import DsparkFoldedSampling, envs +from sglang.srt.speculative.dspark_components.dspark_draft import ( + select_draft_hidden_without_anchor, +) from sglang.srt.utils import get_available_gpu_memory logger = logging.getLogger(__name__) @@ -49,6 +52,8 @@ class DsparkDraftSampler: self.model = model self.markov_head = model.markov_head self.gamma = int(gamma) + self.sample_from_anchor = bool(model.sample_from_anchor) + self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1 max_bs = int(max_bs) if out is not None: assert out.shape == (max_bs * self.gamma,) and out.dtype == torch.int64 @@ -100,10 +105,19 @@ class DsparkDraftSampler: self.greedy_mask[:bs].copy_((sampling_info.top_ks <= 1).view(-1)[:bs]) def __call__(self, hidden_states, input_ids): - bs = hidden_states.shape[0] // self.gamma - base_logits, confidence_tap = self.model.compute_base_logits(hidden_states) + bs = hidden_states.shape[0] // self.query_token_num + if self.sample_from_anchor: + model_hidden = hidden_states + sample_hidden = hidden_states.view(bs, self.gamma, -1) + else: + model_hidden, sample_hidden = select_draft_hidden_without_anchor( + hidden_states, + bs=bs, + gamma=self.gamma, + ) + base_logits, confidence_tap = self.model.compute_base_logits(model_hidden) base_logits = base_logits.view(bs, self.gamma, -1) - anchor = input_ids.view(bs, self.gamma)[:, 0] + anchor = input_ids.view(bs, self.query_token_num)[:, 0] if self.folded_sampling: @@ -125,7 +139,7 @@ class DsparkDraftSampler: draft_tokens, corrected_logits = self.markov_head.sample_block( base_logits, first_prev_tokens=anchor, - hidden_states=hidden_states.view(bs, self.gamma, -1), + hidden_states=sample_hidden, sampler=sampler, ) self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1)) @@ -135,7 +149,7 @@ class DsparkDraftSampler: ) if self.confidence_out is not None: confidence = self.confidence_fn( - draft_hidden=hidden_states.view(bs, self.gamma, -1), + draft_hidden=sample_hidden, anchor_tokens=anchor, draft_tokens=draft_tokens, confidence_tap=confidence_tap, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index adcd6c42f..f200c34dc 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -162,18 +162,22 @@ class DSparkWorkerV2(BaseSpecWorker): ) self.gamma = runtime_config.gamma self.verify_num_draft_tokens = runtime_config.verify_num_draft_tokens + self.sample_from_anchor = bool(self.draft_model.sample_from_anchor) + self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1 self.speculative_num_draft_tokens = self.verify_num_draft_tokens self._mask_token_id = runtime_config.mask_token_id if self.ps.tp_rank == 0: logger.info( "Initialized DSpark draft runner. attention_backend=%s, model=%s, " - "gamma=%s, verify_num_draft_tokens=%s, mask_token_id=%s, " - "markov_head=%s", + "gamma=%s, verify_num_draft_tokens=%s, query_token_num=%s, " + "sample_from_anchor=%s, mask_token_id=%s, markov_head=%s", bundle.resolved_attention_backend, self.draft_model.__class__.__name__, self.gamma, self.verify_num_draft_tokens, + self.query_token_num, + self.sample_from_anchor, self._mask_token_id, type(self.draft_model.markov_head).__name__, ) @@ -182,7 +186,7 @@ class DSparkWorkerV2(BaseSpecWorker): length=self.verify_num_draft_tokens, device=self.device ) self._draft_block_spec_info = make_draft_block_spec_info( - draft_token_num=int(self.gamma), device=self.device + draft_token_num=int(self.query_token_num), device=self.device ) if getattr(self.draft_model, "uses_own_vocab_modules", False): diff --git a/python/sglang/test/kits/eval_accuracy_kit.py b/python/sglang/test/kits/eval_accuracy_kit.py index ae9509c2a..857899a11 100644 --- a/python/sglang/test/kits/eval_accuracy_kit.py +++ b/python/sglang/test/kits/eval_accuracy_kit.py @@ -192,6 +192,7 @@ class GSM8KMixin: gsm8k_num_shots: int = 5 # run_eval backend only gsm8k_backend: str = "run_eval" # "run_eval" | "sgl_eval" gsm8k_thinking: bool = False # sgl_eval backend + gsm8k_max_tokens: Optional[int] = None # sgl_eval backend gsm8k_n_repeats: int = 1 # sgl_eval backend def test_gsm8k(self): @@ -213,6 +214,7 @@ class GSM8KMixin: num_examples=num_examples, num_threads=self.gsm8k_num_threads, thinking=self.gsm8k_thinking, + max_tokens=self.gsm8k_max_tokens, accept_length_thres=self.gsm8k_accept_length_thres, ) else: diff --git a/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py b/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py index 87e2f708f..7042936a6 100644 --- a/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py +++ b/test/registered/models_e2e/test_nvidia_nemotron_3_nano.py @@ -1,33 +1,110 @@ +"""B200 NVFP4 E2E coverage for NVIDIA Nemotron 3.5 Lightning. + +The three cases exercise the production NVFP4 checkpoint without speculation, +with DFlash, and with DSpark. MTP is already covered by the Nemotron model +family tests; the external-draft paths are the new coverage in this file. +""" + import unittest +from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kits.lm_eval_kit import LMEvalMixin -from sglang.test.server_fixtures.default_fixture import DefaultServerBase - -register_cuda_ci( - est_time=190, - stage="base-b", - runner_config="2-gpu-large", +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, + try_cached_model, ) -NEMOTRON_3_NANO_THINKING_ARGS = [ - "--trust-remote-code", +register_cuda_ci(est_time=500, stage="extra-b", runner_config="4-gpu-b200") + +MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" +DFLASH_DRAFT_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash" +DSPARK_DRAFT_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DSpark" + +SERVER_LAUNCH_TIMEOUT = 3600 +GSM8K_SCORE_THRESHOLD = 0.80 + +BASE_ARGS = [ + "--mamba-backend", + "flashinfer", + "--mamba-ssm-dtype", + "float16", + "--enable-mamba-cache-stochastic-rounding", + "--mamba-cache-philox-rounds", + "5", + "--mem-fraction-static", + "0.85", + "--cuda-graph-max-bs-decode", + "16", + "--reasoning-parser", + "nemotron_3", "--tool-call-parser", "qwen3_coder", - "--reasoning-parser", - "deepseek-r1", ] -class TestNvidiaNemotron3Nano30BFP8(LMEvalMixin, DefaultServerBase): - """Test Nemotron-3-Nano-30B FP8 model with lm-eval GSM8K evaluation.""" +class _Nemotron35LightningServer: + speculative_args: list[str] = [] + model = try_cached_model(MODEL) + base_url = DEFAULT_URL_FOR_TEST + gsm8k_backend = "sgl_eval" + gsm8k_thinking = True + gsm8k_num_examples = 200 + gsm8k_num_threads = 32 + gsm8k_max_tokens = 16384 + gsm8k_score_threshold = GSM8K_SCORE_THRESHOLD - model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" - model_config_name = "lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml" - other_args = [ - "--tp-size", - "2", - ] + NEMOTRON_3_NANO_THINKING_ARGS + @classmethod + def setUpClass(cls): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=BASE_ARGS + cls.speculative_args, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + +class TestNvidiaNemotron35LightningNVFP4( + _Nemotron35LightningServer, GSM8KMixin, CustomTestCase +): + """Normal autoregressive serving.""" + + +class TestNvidiaNemotron35LightningNVFP4DFlash( + _Nemotron35LightningServer, GSM8KMixin, CustomTestCase +): + """DFlash with the published W4A16 draft checkpoint.""" + + speculative_args = [ + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + DFLASH_DRAFT_MODEL, + "--speculative-dflash-block-size", + "6", + ] + + +class TestNvidiaNemotron35LightningNVFP4DSpark( + _Nemotron35LightningServer, GSM8KMixin, CustomTestCase +): + """DSpark with the published bonus-anchor W4A16 draft checkpoint.""" + + speculative_args = [ + "--speculative-algorithm", + "DSPARK", + "--speculative-draft-model-path", + DSPARK_DRAFT_MODEL, + "--speculative-dspark-block-size", + "3", + ] if __name__ == "__main__": diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index be61c8945..5c3f38bcc 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -97,11 +97,25 @@ class TestModelOverridableWhitelist(CustomTestCase): "enable_aiter_allreduce_fusion", "enable_symm_mem", "speculative_attention_mode", + "speculative_draft_attention_backend", } ), ) +class TestDSparkCheckpointConfig(CustomTestCase): + def test_sample_from_anchor_is_read_from_checkpoint_config(self): + from sglang.srt.speculative.dspark_components.dspark_config import ( + get_dspark_sample_from_anchor, + ) + + config = SimpleNamespace( + architectures=["UnrelatedDSparkModel"], sample_from_anchor=False + ) + self.assertFalse(get_dspark_sample_from_anchor(config)) + self.assertTrue(get_dspark_sample_from_anchor(SimpleNamespace())) + + class _IsolatedRegistry(CustomTestCase): """Run each test against empty registries (they are process-global).""" @@ -709,6 +723,14 @@ class TestGoldenModelOverrides(_IsolatedPublish): moe_runner_backend="auto", moe_a2a_backend="none", attention_backend=None, + prefill_attention_backend=None, + decode_attention_backend=None, + speculative_algorithm=None, + speculative_eagle_topk=None, + speculative_draft_attention_backend=None, + page_size=None, + mamba_radix_cache_strategy="auto", + is_attention_backend_not_set=lambda: True, get_model_config=lambda: model_config, ), hf_config, @@ -731,13 +753,16 @@ class TestGoldenModelOverrides(_IsolatedPublish): } ) - with patch.object(overrides_module, "is_sm100_supported", return_value=True): + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=True), + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + ): self.assertEqual( _nemotron_h_overrides(server_args, hf_config), { "quantization": "modelopt_mixed", "moe_runner_backend": "marlin", - "attention_backend": "flashinfer", + "attention_backend": "trtllm_mha", }, ) @@ -758,16 +783,169 @@ class TestGoldenModelOverrides(_IsolatedPublish): } ) - with patch.object(overrides_module, "is_sm100_supported", return_value=True): + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=True), + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + ): self.assertEqual( _nemotron_h_overrides(server_args, hf_config), { "quantization": "modelopt_mixed", "moe_runner_backend": "flashinfer_trtllm", - "attention_backend": "flashinfer", + "attention_backend": "trtllm_mha", }, ) + def test_nemotron_h_speculation_uses_arch_specific_attention_on_blackwell(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + cases = { + True: { + "attention_backend": "trtllm_mha", + "page_size": 64, + "mamba_radix_cache_strategy": "extra_buffer", + "speculative_draft_attention_backend": "trtllm_mha", + }, + False: { + "attention_backend": "triton", + "speculative_draft_attention_backend": "flashinfer", + }, + } + for is_sm100, expected in cases.items(): + with self.subTest(is_sm100=is_sm100): + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = "EAGLE" + + with ( + patch.object( + overrides_module, + "is_blackwell_supported", + return_value=True, + ), + patch.object( + overrides_module, + "is_sm100_supported", + return_value=is_sm100, + ), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + for key, value in expected.items(): + self.assertEqual(overrides[key], value) + + def test_nemotron_h_sm100_speculative_draft_backend_matrix(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + for algorithm in ("EAGLE", "NEXTN", "DSPARK"): + with self.subTest(algorithm=algorithm): + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = algorithm + with ( + patch.object( + overrides_module, "is_blackwell_supported", return_value=True + ), + patch.object( + overrides_module, "is_sm100_supported", return_value=True + ), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + self.assertEqual(overrides["attention_backend"], "trtllm_mha") + self.assertEqual( + overrides["speculative_draft_attention_backend"], + "trtllm_mha", + ) + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = "DFLASH" + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=True), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + self.assertEqual(overrides["attention_backend"], "trtllm_mha") + self.assertNotIn("speculative_draft_attention_backend", overrides) + + def test_nemotron_h_sm100_speculation_preserves_explicit_cache_and_draft(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = "DSPARK" + server_args.page_size = 128 + server_args.mamba_radix_cache_strategy = "extra_buffer_lazy" + server_args.speculative_draft_attention_backend = "flashinfer" + + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=True), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + + self.assertEqual(overrides["attention_backend"], "trtllm_mha") + self.assertNotIn("page_size", overrides) + self.assertNotIn("mamba_radix_cache_strategy", overrides) + self.assertNotIn("speculative_draft_attention_backend", overrides) + + def test_nemotron_h_sm100_topk_tree_falls_back_to_triton(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = "EAGLE" + server_args.speculative_eagle_topk = 4 + + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=True), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + + self.assertEqual(overrides["attention_backend"], "triton") + self.assertEqual(overrides["speculative_draft_attention_backend"], "flashinfer") + self.assertNotIn("page_size", overrides) + self.assertNotIn("mamba_radix_cache_strategy", overrides) + + def test_nemotron_h_target_only_sm120_defers_to_generic_attention_default(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=False), + ): + self.assertNotIn( + "attention_backend", _nemotron_h_overrides(server_args, hf_config) + ) + + def test_nemotron_h_target_only_sm100_uses_trtllm_mha(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=True), + ): + self.assertEqual( + _nemotron_h_overrides(server_args, hf_config)["attention_backend"], + "trtllm_mha", + ) + + def test_nemotron_h_explicit_split_attention_backend_wins(self): + from sglang.srt.arg_groups.overrides import _nemotron_h_overrides + + server_args, hf_config = self._nemotron_h_args(quantized_layers={}) + server_args.speculative_algorithm = "DFLASH" + server_args.prefill_attention_backend = "triton" + server_args.speculative_draft_attention_backend = "fa3" + server_args.is_attention_backend_not_set = lambda: False + + with ( + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + patch.object(overrides_module, "is_sm100_supported", return_value=True), + ): + overrides = _nemotron_h_overrides(server_args, hf_config) + self.assertNotIn("attention_backend", overrides) + self.assertNotIn("speculative_draft_attention_backend", overrides) + def test_nemotron_h_w4a16_moe_rejects_a2a_backend(self): from sglang.srt.arg_groups.overrides import _nemotron_h_overrides @@ -1334,20 +1512,36 @@ class TestGoldenModelOverrides(_IsolatedPublish): moe_runner_backend="auto", moe_a2a_backend="none", attention_backend=None, + prefill_attention_backend=None, + decode_attention_backend=None, + speculative_algorithm=None, + speculative_eagle_topk=None, + speculative_draft_attention_backend=None, + page_size=None, + mamba_radix_cache_strategy="auto", get_model_config=lambda: mc, ) defaults.update(kw) - return SimpleNamespace(**defaults) + args = SimpleNamespace(**defaults) + args.is_attention_backend_not_set = lambda: ( + args.attention_backend is None + and args.prefill_attention_backend is None + and args.decode_attention_backend is None + ) + return args hf = _hf() - with patch.object(overrides_module, "is_sm100_supported", return_value=True): + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=True), + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + ): # modelopt checkpoint: quant algo resolution + sm100 defaults self.assertEqual( _nemotron_h_overrides(_args("modelopt", hf), hf), { "quantization": "modelopt_fp4", "moe_runner_backend": "flashinfer_trtllm", - "attention_backend": "flashinfer", + "attention_backend": "trtllm_mha", }, ) hf_mixed = _hf("MIXED_PRECISION") @@ -1383,7 +1577,10 @@ class TestGoldenModelOverrides(_IsolatedPublish): ) hf_without_quant_cfg = _hf(include_quantization_config=False) - with patch.object(overrides_module, "is_sm100_supported", return_value=True): + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=True), + patch.object(overrides_module, "is_blackwell_supported", return_value=True), + ): for modelopt_quantization in ("modelopt_fp8", "modelopt_fp4"): with self.subTest(modelopt_quantization=modelopt_quantization): self.assertEqual( @@ -1394,7 +1591,7 @@ class TestGoldenModelOverrides(_IsolatedPublish): { "quantization": modelopt_quantization, "moe_runner_backend": "flashinfer_trtllm", - "attention_backend": "flashinfer", + "attention_backend": "trtllm_mha", }, )