@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 156 lines 5.44 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5
6 import pytest
7
8 from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
9 from vllm.config import CompilationConfig
10 from vllm.platforms import current_platform
11
12 from ...utils import compute_acceptance_len, compute_acceptance_rate
13
14 REGRESSION_TOLERANCE = 0.95
15
16
17 @dataclass(frozen=True)
18 class DSparkCorrectnessConfig:
19 model: str
20 draft_model: str
21 reference_accuracy: float
22 reference_acceptance_rate: float
23 reference_acceptance_len: float
24 num_speculative_tokens: int = 7
25 max_model_len: int = 4096
26 max_num_seqs: int | None = None
27 num_questions: int = 1319
28 max_tokens: int = 256
29 use_chat_completions: bool = False
30 chat_template_kwargs: dict[str, object] | None = None
31 attention_backend: str | None = None
32 gpu_memory_utilization: float = 0.92
33 enforce_eager: bool = False
34 disable_flashinfer_sampler: bool = False
35 language_model_only: bool = False
36
37
38 # References from 12 full GSM8K runs at temperature 1.0: accuracy 0.782-0.814,
39 # acceptance rate 0.418-0.434, acceptance length 3.928-4.037.
40 QWEN3_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
41 model="Qwen/Qwen3-4B-FP8",
42 draft_model="deepseek-ai/dspark_qwen3_4b_block7",
43 reference_accuracy=0.801,
44 reference_acceptance_rate=0.428,
45 reference_acceptance_len=3.994,
46 attention_backend="FLASH_ATTN",
47 )
48
49 # References from five 200-question runs at temperature 1.0: accuracy
50 # 0.900-0.955, acceptance rate 0.578-0.595, acceptance length 5.044-5.167.
51 GEMMA4_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
52 model="RedHatAI/gemma-4-12B-it-NVFP4",
53 draft_model="deepseek-ai/dspark_gemma4_12b_block7",
54 reference_accuracy=0.92,
55 reference_acceptance_rate=0.58,
56 reference_acceptance_len=5.0,
57 max_num_seqs=32,
58 num_questions=200,
59 use_chat_completions=True,
60 gpu_memory_utilization=0.85,
61 enforce_eager=True,
62 disable_flashinfer_sampler=True,
63 language_model_only=True,
64 )
65
66 # Reference from 200 GSM8K questions at temperature 1.0.
67 QWEN3_6_DSPARK_SPECULATORS = DSparkCorrectnessConfig(
68 model="RedHatAI/Qwen3.6-35B-A3B-NVFP4",
69 draft_model="RedHatAI/Qwen3.6-35B-A3B-speculator.dspark",
70 reference_accuracy=0.85,
71 reference_acceptance_rate=0.41,
72 reference_acceptance_len=4.25,
73 num_speculative_tokens=8,
74 max_num_seqs=32,
75 num_questions=200,
76 use_chat_completions=True,
77 chat_template_kwargs={"enable_thinking": False},
78 gpu_memory_utilization=0.85,
79 language_model_only=True,
80 )
81
82
83 @pytest.mark.parametrize(
84 "config",
85 [
86 pytest.param(QWEN3_DSPARK_DEEPSPEC, id="qwen3-deepspec"),
87 pytest.param(GEMMA4_DSPARK_DEEPSPEC, id="gemma4-deepspec"),
88 pytest.param(QWEN3_6_DSPARK_SPECULATORS, id="qwen3.6-speculators"),
89 ],
90 )
91 def test_dspark_correctness_and_acceptance_rate(
92 monkeypatch: pytest.MonkeyPatch,
93 config: DSparkCorrectnessConfig,
94 vllm_runner,
95 ):
96 """Guard GSM8K accuracy and acceptance metrics for DSpark models."""
97 if config.disable_flashinfer_sampler:
98 monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
99
100 speculative_config = {
101 "method": "dspark",
102 "model": config.draft_model,
103 "num_speculative_tokens": config.num_speculative_tokens,
104 "draft_sample_method": "probabilistic",
105 }
106 if config.attention_backend is not None and current_platform.is_cuda():
107 speculative_config["attention_backend"] = config.attention_backend
108
109 runner_config = {
110 "block_size": None,
111 "trust_remote_code": True,
112 "speculative_config": speculative_config,
113 "max_model_len": config.max_model_len,
114 "gpu_memory_utilization": config.gpu_memory_utilization,
115 "enforce_eager": config.enforce_eager,
116 "enable_chunked_prefill": None,
117 "enable_prefix_caching": False,
118 "disable_log_stats": False,
119 "compilation_config": CompilationConfig(),
120 }
121 if config.max_num_seqs is not None:
122 runner_config["max_num_seqs"] = config.max_num_seqs
123 if config.language_model_only:
124 runner_config["language_model_only"] = True
125
126 with vllm_runner(config.model, **runner_config) as spec_runner:
127 spec_llm = spec_runner.llm
128 results = evaluate_gsm8k_offline(
129 spec_llm,
130 num_questions=config.num_questions,
131 max_tokens=config.max_tokens,
132 temperature=1.0,
133 use_chat_completions=config.use_chat_completions,
134 chat_template_kwargs=config.chat_template_kwargs,
135 )
136 accuracy = results["accuracy"]
137 metrics = spec_llm.get_metrics()
138 acceptance_rate = compute_acceptance_rate(metrics)
139 acceptance_len = compute_acceptance_len(metrics)
140 context = f"DSpark target={config.model}, draft={config.draft_model}"
141 metrics_summary = (
142 f"gsm8k_accuracy={accuracy:.3f}, "
143 f"acceptance_rate={acceptance_rate:.3f}, "
144 f"acceptance_len={acceptance_len:.3f}"
145 )
146 print(f"{context}: {metrics_summary}")
147
148 assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE, (
149 f"{context}: {metrics_summary}"
150 )
151 assert (
152 acceptance_rate >= config.reference_acceptance_rate * REGRESSION_TOLERANCE
153 ), f"{context}: {metrics_summary}"
154 assert (
155 acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE
156 ), f"{context}: {metrics_summary}"