minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[Model] Support Nemotron 3.5 Lightning speculative decoding (#36186)

Browse Source 
Co-authored-by: Ryan Stewart <rystewart@nvidia.com>
This commit is contained in:
Baizhou Zhang
2026-08-25 16:43:58 -07:00
committed by GitHub
co-authored by Ryan Stewart
parent 2d88c79b3e
commit 41e7612dee
18 changed files with 792 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cookbook/autoregressive/NVIDIA/Nemotron3.5-Lightning.mdx
+1 -1
View File
@@ -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.
python/sglang/srt/arg_groups/overrides.py
+26 -2
View File
@@ -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
python/sglang/srt/arg_groups/speculative_hook.py
+17 -13
View File
@@ -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(
python/sglang/srt/layers/quantization/modelopt_quant.py
+13 -4
View File
@@ -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,
)
python/sglang/srt/model_executor/model_runner.py
+9
View File
@@ -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,
python/sglang/srt/models/dflash.py
+154 -34
View File
@@ -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)
python/sglang/srt/models/dspark.py
+69 -8
View File
@@ -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:
python/sglang/srt/models/nemotron_h.py
+42 -1
View File
@@ -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:
python/sglang/srt/server_args.py
+4 -9
View File
@@ -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",
python/sglang/srt/speculative/dflash_utils.py
+42
View File
@@ -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,
*,
python/sglang/srt/speculative/dflash_worker_v2.py
+14 -1
View File
@@ -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")
python/sglang/srt/speculative/dspark_components/dspark_config.py
+12 -4
View File
@@ -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
)
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+59 -9
View File
@@ -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,
)
python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py
+19 -5
View File
@@ -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,
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+7 -3
View File
@@ -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):
python/sglang/test/kits/eval_accuracy_kit.py
+2
View File
@@ -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:
test/registered/models_e2e/test_nvidia_nemotron_3_nano.py
+96 -19
View File
@@ -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__":
test/registered/unit/test_model_overrides.py
+206 -9
View File
@@ -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",
},
)