@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3ccc1ff495
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 10:00 UTC

[Eval][CI] Add basic mrcr eval to tests/evals/ (#40164)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed May 1, 2026 at 12:00 UTC 3ccc1ff4958dd07dbffeaa1c48463325c892b518
8 files changed +531
.buildkite/test_areas/lm_eval.yaml
+7
@@ -137,3 +137,10 @@ steps:
137 commands:
138 - uv pip install --system 'gpt-oss[eval]==0.0.5'
139 - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
140 +
141 +- label: MRCR Eval Small Models
142 + timeout_in_minutes: 30
143 + source_file_dependencies:
144 + - tests/evals/mrcr/
145 + commands:
146 + - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
tests/evals/mrcr/README.md new
+44
@@ -0,0 +1,44 @@
1 +# MRCR Long-Context Accuracy Evaluation
2 +
3 +Smoke test for long-context behavior using OpenAI's public [`openai/mrcr`](https://huggingface.co/datasets/openai/mrcr) dataset. The model sees a long chat with several near-duplicate "needles" and must reproduce a specific earlier assistant turn verbatim, prepended with a random anti-guessing string.
4 +
5 +**Scoring:** if the response doesn't start with `random_string_to_prepend`, score is 0; otherwise the prefix is stripped and the mean `SequenceMatcher.ratio()` against the reference answer is reported.
6 +
7 +## Usage
8 +
9 +```bash
10 +# Pytest (spawns the server)
11 +pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \
12 + --config-list-file=configs/models-small.txt
13 +
14 +# Standalone (server already running; model and context auto-discovered)
15 +vllm serve Qwen/Qwen3-0.6B --reasoning-parser qwen3 --port 8000
16 +python tests/evals/mrcr/mrcr_eval.py --port 8000
17 +```
18 +
19 +## Configuration
20 +
21 +```yaml
22 +model_name: "Qwen/Qwen3-0.6B"
23 +# Per-needle thresholds catch bucket-specific regressions (sliding window,
24 +# chunked prefill, prefix cache) that an aggregate can hide. A scalar
25 +# (e.g. `match_ratio_threshold: 0.20`) is also accepted and checked against
26 +# the mean match ratio.
27 +match_ratio_threshold:
28 + 2: 0.30
29 + 4: 0.15
30 + 8: 0.10
31 +num_samples: 30
32 +needles: [2, 4, 8]
33 +# max_prompt_tokens: 32768 # Optional; defaults to server max_model_len - max_tokens - 256
34 +max_tokens: 2048
35 +concurrency: 8
36 +server_args: "--max-model-len 32768 --reasoning-parser qwen3"
37 +```
38 +
39 +## Notes
40 +
41 +- Samples stream from three parquet shards (`{N}needle/{N}needle_0.parquet`); only the first few row groups are fetched, not the full 1.4 GB repo.
42 +- `max_prompt_tokens` defaults to `max_model_len - max_tokens - 256`, i.e. fills whatever context the server advertises. Set `--max-model-len` on the server to control the smoke-test context length; override `--max-prompt-tokens` on the client to cap below that.
43 +- Sample length is pre-filtered by `n_chars × 4 ≤ max_prompt_tokens`, then verified via the server's `/tokenize` endpoint under the actual chat template.
44 +- Reasoning models: start the server with `--reasoning-parser <name>` (e.g. `qwen3`, `deepseek_r1`) so `<think>` goes to `message.reasoning_content` and doesn't contaminate the scored answer.
tests/evals/mrcr/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/evals/mrcr/configs/Qwen3.5-4B.yaml new
+7
@@ -0,0 +1,7 @@
1 +model_name: "Qwen/Qwen3.5-4B"
2 +needles: [2, 4, 8]
3 +match_ratio_threshold:
4 + 2: 0.99
5 + 4: 0.84
6 + 8: 0.76
7 +server_args: "--max-model-len 128K --reasoning-parser qwen3"
tests/evals/mrcr/configs/models-small.txt new
+1
@@ -0,0 +1 @@
1 +Qwen3.5-4B.yaml
tests/evals/mrcr/conftest.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from pathlib import Path
5 +
6 +
7 +def pytest_addoption(parser):
8 + """Add custom command line options."""
9 + parser.addoption(
10 + "--config-list-file",
11 + default="configs/models-small.txt",
12 + help="File containing list of config files to test",
13 + )
14 +
15 +
16 +def pytest_generate_tests(metafunc):
17 + """Generate test parameters from config files."""
18 + if "config_filename" in metafunc.fixturenames:
19 + config_list_file = metafunc.config.getoption("--config-list-file")
20 +
21 + config_list_path = Path(config_list_file)
22 + if not config_list_path.is_absolute():
23 + test_dir_path = Path(__file__).parent / config_list_file
24 + if test_dir_path.exists():
25 + config_list_path = test_dir_path
26 + else:
27 + config_list_path = Path.cwd() / config_list_file
28 +
29 + print(f"Looking for config list at: {config_list_path}")
30 +
31 + config_files = []
32 + if config_list_path.exists():
33 + config_dir = config_list_path.parent
34 + with open(config_list_path) as f:
35 + for line in f:
36 + line = line.strip()
37 + if line and not line.startswith("#"):
38 + config_path = config_dir / line
39 + if config_path.exists():
40 + config_files.append(config_path)
41 + print(f" ✓ Found: {config_path}")
42 + else:
43 + print(f" ✗ Missing: {config_path}")
44 + else:
45 + print(f"Config list file not found: {config_list_path}")
46 +
47 + if config_files:
48 + metafunc.parametrize(
49 + "config_filename",
50 + config_files,
51 + ids=[config_file.stem for config_file in config_files],
52 + )
53 + else:
54 + print("No config files found, test will be skipped")
tests/evals/mrcr/mrcr_eval.py new
+333
@@ -0,0 +1,333 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""MRCR long-context evaluation for vLLM's OpenAI-compatible server.
5 +
6 +Streams samples from `openai/mrcr` on HuggingFace, sends chat completions to
7 +the server, and scores each response with a prefix-gated SequenceMatcher ratio
8 +against the reference answer.
9 +"""
10 +
11 +import argparse
12 +import asyncio
13 +import json
14 +import time
15 +from difflib import SequenceMatcher
16 +
17 +import aiohttp
18 +import numpy as np
19 +import requests
20 +from tqdm.asyncio import tqdm
21 +
22 +DATASET_REPO = "openai/mrcr"
23 +NEEDLE_SHARDS = {
24 + 2: "2needle/2needle_0.parquet",
25 + 4: "4needle/4needle_0.parquet",
26 + 8: "8needle/8needle_0.parquet",
27 +}
28 +# Reserve headroom for chat-template tokens on top of the messages.
29 +PROMPT_SAFETY_BUFFER = 256
30 +# Pre-filter heuristic before the authoritative /tokenize check.
31 +CHARS_PER_TOKEN = 4
32 +# Skip chain-of-thought on reasoning models; ignored by non-reasoning templates.
33 +DEFAULT_EXTRA_BODY: dict = {"chat_template_kwargs": {"enable_thinking": False}}
34 +
35 +
36 +def discover_server_model(base_url: str) -> tuple[str, int | None]:
37 + """Return (model_id, max_model_len) from /v1/models."""
38 + resp = requests.get(f"{base_url}/v1/models", timeout=30)
39 + resp.raise_for_status()
40 + data = resp.json().get("data", [])
41 + if not data:
42 + raise RuntimeError(f"No models advertised at {base_url}/v1/models")
43 + entry = data[0]
44 + return entry["id"], entry.get("max_model_len")
45 +
46 +
47 +def count_chat_tokens(base_url: str, model: str, messages: list[dict]) -> int:
48 + """Return the chat-template-rendered token count via /tokenize."""
49 + resp = requests.post(
50 + f"{base_url}/tokenize",
51 + json={"model": model, "messages": messages, "add_generation_prompt": True},
52 + timeout=120,
53 + )
54 + resp.raise_for_status()
55 + return int(resp.json()["count"])
56 +
57 +
58 +def _load_mrcr_samples(
59 + needles: list[int],
60 + max_prompt_tokens: int,
61 + num_samples: int,
62 + seed: int,
63 + base_url: str,
64 + model_name: str,
65 +) -> list[dict]:
66 + """Stream MRCR samples balanced across needle buckets, token-verified."""
67 + try:
68 + from datasets import load_dataset
69 + except ImportError as e:
70 + raise ImportError(
71 + "MRCR eval requires `datasets`. Install with: uv pip install datasets"
72 + ) from e
73 +
74 + max_chars = max_prompt_tokens * CHARS_PER_TOKEN
75 + per_bucket = num_samples // len(needles)
76 + leftover = num_samples - per_bucket * len(needles)
77 +
78 + samples: list[dict] = []
79 + for idx, n in enumerate(needles):
80 + if n not in NEEDLE_SHARDS:
81 + raise ValueError(f"Unsupported needle count {n}")
82 + target = per_bucket + (1 if idx < leftover else 0)
83 + if target == 0:
84 + continue
85 +
86 + ds = load_dataset(
87 + DATASET_REPO,
88 + data_files=NEEDLE_SHARDS[n],
89 + split="train",
90 + streaming=True,
91 + ).shuffle(seed=seed + n, buffer_size=16)
92 +
93 + taken = 0
94 + for row in ds:
95 + if int(row.get("n_chars", 0)) > max_chars:
96 + continue
97 + prompt = row["prompt"]
98 + messages = json.loads(prompt) if isinstance(prompt, str) else list(prompt)
99 + n_tokens = count_chat_tokens(base_url, model_name, messages)
100 + if n_tokens > max_prompt_tokens:
101 + continue
102 + samples.append(
103 + {
104 + "messages": messages,
105 + "answer": row["answer"],
106 + "random_string_to_prepend": row["random_string_to_prepend"],
107 + "n_needles": int(row["n_needles"]),
108 + "n_tokens": n_tokens,
109 + }
110 + )
111 + taken += 1
112 + if taken >= target:
113 + break
114 +
115 + if taken < target:
116 + print(f"Warning: only {taken}/{target} samples for n_needles={n}")
117 +
118 + if not samples:
119 + raise RuntimeError("No MRCR samples fit; loosen max_prompt_tokens.")
120 + return samples
121 +
122 +
123 +def score_mrcr(response: str, answer: str, random_prefix: str) -> float:
124 + """Prefix-gated SequenceMatcher ratio; 0 if the prefix is missing."""
125 + if not response.startswith(random_prefix):
126 + return 0.0
127 + stripped = response[len(random_prefix) :]
128 + return SequenceMatcher(a=answer, b=stripped, autojunk=False).ratio()
129 +
130 +
131 +async def _call_chat(
132 + session: aiohttp.ClientSession,
133 + url: str,
134 + model: str,
135 + messages: list[dict],
136 + max_tokens: int,
137 + temperature: float,
138 + seed: int | None,
139 + extra_body: dict,
140 +) -> tuple[str, int]:
141 + data = {
142 + "model": model,
143 + "messages": messages,
144 + "temperature": temperature,
145 + "max_tokens": max_tokens,
146 + **extra_body,
147 + }
148 + if seed is not None:
149 + data["seed"] = seed
150 + try:
151 + async with session.post(f"{url}/v1/chat/completions", json=data) as resp:
152 + resp.raise_for_status()
153 + result = await resp.json()
154 + text = result["choices"][0]["message"]["content"] or ""
155 + return text, result.get("usage", {}).get("completion_tokens", 0)
156 + except Exception as e:
157 + print(f"chat request failed: {e}")
158 + return "", 0
159 +
160 +
161 +def evaluate_mrcr(
162 + model_name: str | None = None,
163 + num_samples: int = 40,
164 + needles: list[int] | None = None,
165 + max_prompt_tokens: int | None = None,
166 + max_tokens: int = 2048,
167 + host: str = "http://127.0.0.1",
168 + port: int = 8000,
169 + temperature: float = 0.0,
170 + seed: int | None = 42,
171 + concurrency: int = 8,
172 + extra_body: dict | None = None,
173 +) -> dict:
174 + """Run MRCR against a vLLM server; auto-discovers model and context."""
175 + needles = needles or [2, 4, 8]
176 + extra_body = DEFAULT_EXTRA_BODY if extra_body is None else extra_body
177 + base_url = f"{host}:{port}"
178 +
179 + discovered_model, server_max_len = discover_server_model(base_url)
180 + if model_name is None:
181 + model_name = discovered_model
182 + if max_prompt_tokens is None:
183 + if server_max_len is None:
184 + raise RuntimeError(
185 + "Server did not advertise max_model_len; pass --max-prompt-tokens."
186 + )
187 + max_prompt_tokens = max(512, server_max_len - max_tokens - PROMPT_SAFETY_BUFFER)
188 + print(
189 + f"Model: {model_name} | max_prompt_tokens={max_prompt_tokens} "
190 + f"(server max_model_len={server_max_len}, max_tokens={max_tokens})"
191 + )
192 +
193 + samples = _load_mrcr_samples(
194 + needles=needles,
195 + max_prompt_tokens=max_prompt_tokens,
196 + num_samples=num_samples,
197 + seed=seed or 0,
198 + base_url=base_url,
199 + model_name=model_name,
200 + )
201 + tok_counts = [s["n_tokens"] for s in samples]
202 + print(
203 + f"Loaded {len(samples)} samples (needles={needles}, "
204 + f"tokens={min(tok_counts)}-{max(tok_counts)})"
205 + )
206 +
207 + async def run():
208 + sem = asyncio.Semaphore(concurrency)
209 + responses = [""] * len(samples)
210 + out_tokens = [0] * len(samples)
211 +
212 + async def one(session, i):
213 + async with sem:
214 + text, toks = await _call_chat(
215 + session=session,
216 + url=base_url,
217 + model=model_name,
218 + messages=samples[i]["messages"],
219 + max_tokens=max_tokens,
220 + temperature=temperature,
221 + seed=seed,
222 + extra_body=extra_body,
223 + )
224 + responses[i] = text
225 + out_tokens[i] = toks
226 +
227 + timeout = aiohttp.ClientTimeout(total=1800)
228 + async with aiohttp.ClientSession(timeout=timeout) as session:
229 + await tqdm.gather(
230 + *[one(session, i) for i in range(len(samples))], desc="MRCR"
231 + )
232 + return responses, out_tokens
233 +
234 + tic = time.perf_counter()
235 + responses, out_tokens = asyncio.run(run())
236 + latency = time.perf_counter() - tic
237 +
238 + scores = np.array(
239 + [
240 + score_mrcr(r, s["answer"], s["random_string_to_prepend"])
241 + for r, s in zip(responses, samples)
242 + ]
243 + )
244 + prefix_hits = np.array(
245 + [
246 + r.startswith(s["random_string_to_prepend"])
247 + for r, s in zip(responses, samples)
248 + ]
249 + )
250 + per_needle = {
251 + f"match_ratio_n{n}": float(
252 + scores[np.array([s["n_needles"] == n for s in samples])].mean()
253 + )
254 + for n in needles
255 + if any(s["n_needles"] == n for s in samples)
256 + }
257 +
258 + total_out = int(sum(out_tokens))
259 + return {
260 + "model": model_name,
261 + "match_ratio": float(scores.mean()),
262 + "prefix_hit_rate": float(prefix_hits.mean()),
263 + "per_needle": per_needle,
264 + "num_samples": len(samples),
265 + "latency": latency,
266 + "total_output_tokens": total_out,
267 + "tokens_per_second": total_out / latency if latency > 0 else 0.0,
268 + "max_tokens": max_tokens,
269 + "needles": needles,
270 + "max_prompt_tokens": max_prompt_tokens,
271 + }
272 +
273 +
274 +def main() -> None:
275 + p = argparse.ArgumentParser(description="MRCR evaluation for vLLM serve")
276 + p.add_argument("--model", default=None, help="Default: discovered from /v1/models")
277 + p.add_argument("--num-samples", type=int, default=40)
278 + p.add_argument(
279 + "--needles", type=int, nargs="+", default=[2, 4, 8], choices=[2, 4, 8]
280 + )
281 + p.add_argument(
282 + "--max-prompt-tokens",
283 + type=int,
284 + default=None,
285 + help="Default: server max_model_len - max_tokens - buffer",
286 + )
287 + p.add_argument("--max-tokens", type=int, default=2048)
288 + p.add_argument("--host", default="http://127.0.0.1")
289 + p.add_argument("--port", type=int, default=8000)
290 + p.add_argument("--temperature", type=float, default=0.0)
291 + p.add_argument("--seed", type=int, default=42)
292 + p.add_argument("--concurrency", type=int, default=8)
293 + p.add_argument(
294 + "--extra-body",
295 + default=None,
296 + help="JSON merged into each request. "
297 + "Pass '{}' to disable the default enable_thinking=false.",
298 + )
299 + p.add_argument("--save-results", default=None)
300 + args = p.parse_args()
301 +
302 + extra_body = json.loads(args.extra_body) if args.extra_body else None
303 +
304 + result = evaluate_mrcr(
305 + model_name=args.model,
306 + num_samples=args.num_samples,
307 + needles=args.needles,
308 + max_prompt_tokens=args.max_prompt_tokens,
309 + max_tokens=args.max_tokens,
310 + host=args.host,
311 + port=args.port,
312 + temperature=args.temperature,
313 + seed=args.seed,
314 + concurrency=args.concurrency,
315 + extra_body=extra_body,
316 + )
317 +
318 + print("\nResults:")
319 + print(f" match_ratio: {result['match_ratio']:.4f}")
320 + print(f" prefix_hit_rate: {result['prefix_hit_rate']:.4f}")
321 + for k, v in result["per_needle"].items():
322 + print(f" {k}: {v:.4f}")
323 + print(f" samples: {result['num_samples']}")
324 + print(f" latency: {result['latency']:.1f}s")
325 + print(f" output tok/s: {result['tokens_per_second']:.1f}")
326 +
327 + if args.save_results:
328 + with open(args.save_results, "w") as f:
329 + json.dump(result, f, indent=2)
330 +
331 +
332 +if __name__ == "__main__":
333 + main()
tests/evals/mrcr/test_mrcr_correctness.py new
+83
@@ -0,0 +1,83 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +MRCR long-context accuracy test.
5 +
6 +Usage:
7 + pytest -s -v tests/evals/mrcr/test_mrcr_correctness.py \
8 + --config-list-file=configs/models-small.txt
9 +"""
10 +
11 +import shlex
12 +
13 +import yaml
14 +
15 +from tests.utils import RemoteOpenAIServer
16 +
17 +from .mrcr_eval import evaluate_mrcr
18 +
19 +
20 +def _split_host_port(url: str, default_port: int = 8000) -> tuple[str, int]:
21 + if "://" in url:
22 + url = url.split("://", 1)[1]
23 + host_port = url.split("/", 1)[0]
24 + if ":" in host_port:
25 + host, p = host_port.split(":", 1)
26 + return f"http://{host}", int(p)
27 + return f"http://{host_port}", default_port
28 +
29 +
30 +def test_mrcr_correctness(config_filename):
31 + cfg = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
32 +
33 + server_args = shlex.split(cfg.get("server_args", ""))
34 + server_args += ["--trust-remote-code", "--disable-uvicorn-access-log"]
35 +
36 + print(
37 + f"MRCR eval for {cfg['model_name']} (threshold {cfg['match_ratio_threshold']})"
38 + )
39 +
40 + with RemoteOpenAIServer(
41 + cfg["model_name"],
42 + server_args,
43 + env_dict=cfg.get("env"),
44 + max_wait_seconds=cfg.get("startup_max_wait_seconds", 600),
45 + ) as server:
46 + host, port = _split_host_port(server.url_for("v1"))
47 + results = evaluate_mrcr(
48 + model_name=cfg.get("model_name"),
49 + num_samples=cfg.get("num_samples", 40),
50 + needles=cfg.get("needles", [2, 4, 8]),
51 + max_prompt_tokens=cfg.get("max_prompt_tokens"),
52 + max_tokens=cfg.get("max_tokens", 2048),
53 + host=host,
54 + port=port,
55 + concurrency=cfg.get("concurrency", 8),
56 + extra_body=cfg.get("extra_body"),
57 + )
58 +
59 + threshold = cfg["match_ratio_threshold"]
60 + tol = cfg.get("tolerance", 0.05)
61 +
62 + print(f" match_ratio: {results['match_ratio']:.4f}")
63 + print(f" prefix_hit_rate: {results['prefix_hit_rate']:.4f}")
64 + for k, v in results["per_needle"].items():
65 + print(f" {k}: {v:.4f}")
66 +
67 + failures: list[str] = []
68 + if isinstance(threshold, dict):
69 + for n, expected in threshold.items():
70 + key = f"match_ratio_n{int(n)}"
71 + measured = results["per_needle"].get(key)
72 + if measured is None:
73 + failures.append(f"{key}: no samples collected")
74 + elif measured < expected - tol:
75 + failures.append(f"{key}: {measured:.4f} < {expected:.4f} - {tol:.4f}")
76 + else:
77 + measured = results["match_ratio"]
78 + if measured < threshold - tol:
79 + failures.append(
80 + f"match_ratio: {measured:.4f} < {threshold:.4f} - {tol:.4f}"
81 + )
82 +
83 + assert not failures, "MRCR thresholds failed: " + "; ".join(failures)