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

Co-authored-by: Ryan Stewart <rystewart@nvidia.com>

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
 
 

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(

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,
         )
 

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,

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)

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:

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:

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",

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,
     *,

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")

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
     )
 

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,
         )

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,

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):

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: