[Example]: Calibration-free FP8/NVFP4 PTQ for speculative-decoding drafters

examples/specdec_bench/run.py

@@ -305,7 +305,7 @@ def run_simple(args):
         type=str,
         required=False,
         default="EAGLE3",
-        choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "NONE"],
+        choices=["EAGLE3", "EAGLE", "DRAFT_TARGET", "NGRAM", "MTP", "DFLASH", "DSPARK", "NONE"],
         help="Speculative algorithm to use",
     )
     parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory")

examples/specdec_bench/specdec_bench/models/auto_deploy.py

@@ -93,6 +93,8 @@ def create_auto_deploy_model(model_path: str, max_concurrent_requests: int, kwar
         )
     elif speculative_algorithm == "NONE":
         specdec = None
+    elif speculative_algorithm == "DSPARK":
+        raise NotImplementedError("DSPARK is only supported by --engine VLLM.")
 
     max_num_tokens = kwargs.get("max_num_tokens", 8192)
 

examples/specdec_bench/specdec_bench/models/sglang.py

@@ -49,6 +49,8 @@ def __init__(
             speculative_algorithm = "LOOKAHEAD"
         elif speculative_algorithm == "NONE":
             speculative_algorithm = None
+        elif speculative_algorithm == "DSPARK":
+            raise NotImplementedError("DSPARK is only supported by --engine VLLM.")
 
         engine_kwargs = {
             "model_path": model_dir,

examples/specdec_bench/specdec_bench/models/trtllm_torch_api.py

@@ -125,6 +125,8 @@ def create_executor(model_path: str, max_concurrent_requests, kwargs):
         )
     elif kwargs.get("speculative_algorithm", None) == "NONE":
         specdec = None
+    elif kwargs.get("speculative_algorithm", None) == "DSPARK":
+        raise NotImplementedError("DSPARK is only supported by --engine VLLM.")
     else:
         specdec = None
 

examples/specdec_bench/specdec_bench/models/vllm.py

@@ -28,6 +28,16 @@
     vllm = None
 
 
+# Forwarded from ``--runtime_params`` ``engine_args.<key>``; extend as needed.
+PASSTHROUGH_ENGINE_ARGS = (
+    "mamba_backend",
+    "mamba_ssm_cache_dtype",
+    "mamba_cache_mode",
+    "mamba_cache_philox_rounds",
+    "enable_mamba_cache_stochastic_rounding",
+)
+
+
 class VLLMModel(Model):
     # Cross-engine ``--max_seq_len`` (run.py) lands in kwargs under the
     # vLLM-native name ``max_model_len`` (see run.py's ``_MAX_SEQ_LEN_KEY``)
@@ -114,6 +124,13 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs
                 "model": kwargs.get("draft_model_dir"),
                 "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 8),
             }
+        elif kwargs.get("speculative_algorithm") == "DSPARK":
+            specdec = {
+                "method": "dspark",
+                "model": kwargs.get("draft_model_dir"),
+                "num_speculative_tokens": kwargs.get("speculative_num_draft_tokens", 7),
+                "draft_sample_method": kwargs.get("draft_sample_method", "greedy"),
+            }
         elif kwargs.get("speculative_algorithm") == "NONE":
             specdec = None
 
@@ -133,8 +150,9 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs
             max_num_seqs=max_concurrent_requests * num_speculative_tokens,
             skip_tokenizer_init=False,
             async_scheduling=kwargs.get("async_scheduling", True),
-            enforce_eager=False,
+            enforce_eager=kwargs.get("enforce_eager", False),
             max_model_len=kwargs.get("max_model_len"),
+            **{key: kwargs[key] for key in PASSTHROUGH_ENGINE_ARGS if kwargs.get(key) is not None},
         )
         self.engine_args = engine_args
         self.model = AsyncLLM.from_engine_args(engine_args)