[Spec] Mamba2 support in target models (#13434)

2025-12-06 00:50:46 +08:00
parent 05284378d6
commit 889b46ea50
6 changed files with 398 additions and 112 deletions
@@ -24,5 +24,57 @@ class TestNvidiaNemotronNanoV2NVFP4(GSM8KMixin, CustomTestCase):
    other_args = ["--max-mamba-cache-size", "256"]


+class TestNvidiaNemotronNanoV2SpeculativeDecoding(GSM8KMixin, CustomTestCase):
+    accuracy = 0.87
+    model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+    other_args = [
+        "--speculative-algorithm",
+        "STANDALONE",
+        "--speculative-num-steps",
+        "2",
+        "--speculative-eagle-topk",
+        "3",
+        "--speculative-num-draft-tokens",
+        "5",
+        "--speculative-draft-model-path",
+        "meta-llama/Llama-3.2-1B",
+        "--speculative-draft-load-format",
+        "dummy",
+        "--max-running-requests",
+        "8",
+        "--max-total-tokens",
+        "2048",
+        "--json-model-override-args",
+        '{"vocab_size": 131072}',
+    ]
+
+
+class TestNvidiaNemotronNanoV2SpeculativeDecodingBF16Cache(GSM8KMixin, CustomTestCase):
+    accuracy = 0.87
+    model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+    other_args = [
+        "--speculative-algorithm",
+        "STANDALONE",
+        "--speculative-num-steps",
+        "2",
+        "--speculative-eagle-topk",
+        "3",
+        "--speculative-num-draft-tokens",
+        "5",
+        "--speculative-draft-model-path",
+        "meta-llama/Llama-3.2-1B",
+        "--speculative-draft-load-format",
+        "dummy",
+        "--max-running-requests",
+        "8",
+        "--max-total-tokens",
+        "2048",
+        "--json-model-override-args",
+        '{"vocab_size": 131072}',
+        "--mamba-ssm-dtype",
+        "bfloat16",
+    ]
+
+
 if __name__ == "__main__":
    unittest.main()