[Model] Support Nemotron 3.5 Lightning speculative decoding (#36186)
Co-authored-by: Ryan Stewart <rystewart@nvidia.com>
This commit is contained in:
Baizhou Zhang
co-authored by
Ryan Stewart
Ryan Stewart
parent
2d88c79b3e
commit
41e7612dee
@@ -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__":
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user