Enable validated prefix caching and record DGX Spark optimization benchmarks

scripts/benchmark.py

@@ -0,0 +1,128 @@
+"""Synthetic streaming benchmark; run inside serving container, writes JSONL to stdout.
+
+No external datasets or user prompts. Usage: python3 benchmark.py --label eager
+TTFT includes reasoning tokens; first_content_s separately measures visible answer.
+Decode tok/s is an approximation from total completion tokens and streaming time.
+"""
+import argparse
+import concurrent.futures
+import hashlib
+import json
+import random
+import re
+import threading
+import time
+import urllib.request
+from pathlib import Path
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--label', required=True)
+parser.add_argument('--long-max-tokens', type=int, default=1024)
+parser.add_argument('--phase', choices=['all', 'short', 'long'], default='all')
+args = parser.parse_args()
+BASE = 'http://127.0.0.1:8000'
+KEY = Path('/run/secrets/qwen_api_key').read_text().strip()
+HEADERS = {'Authorization': 'Bearer ' + KEY, 'Content-Type': 'application/json'}
+
+def emit(record):
+    print(json.dumps({'label': args.label, **record}, ensure_ascii=False), flush=True)
+
+def metrics():
+    request = urllib.request.Request(BASE + '/metrics', headers=HEADERS)
+    with urllib.request.urlopen(request, timeout=10) as response:
+        text = response.read().decode()
+    values = {}
+    for line in text.splitlines():
+        if line.startswith('#') or 'prefix_cache_' not in line:
+            continue
+        name = line.split('{')[0].split()[0]
+        if name.endswith('_total'):
+            values[name] = values.get(name, 0) + float(line.split()[-1])
+    return values
+
+def run(name, prompt, expected=None, barrier=None, max_tokens=512):
+    payload = {'model': 'qwen3.8-flash-next', 'messages': [{'role': 'user', 'content': prompt}],
+               'temperature': 0, 'seed': 42, 'max_tokens': max_tokens,
+               'reasoning_effort': 'low', 'stream': True,
+               'stream_options': {'include_usage': True}}
+    request = urllib.request.Request(BASE + '/v1/chat/completions',
+                                     data=json.dumps(payload).encode(), headers=HEADERS)
+    if barrier:
+        barrier.wait()
+    start = time.perf_counter()
+    first = first_content = last = None
+    content, reasoning = '', ''
+    usage, finish = {}, None
+    with urllib.request.urlopen(request, timeout=900) as response:
+        for raw in response:
+            line = raw.decode().strip()
+            if not line.startswith('data: ') or line == 'data: [DONE]':
+                continue
+            item = json.loads(line[6:])
+            if item.get('usage'):
+                usage = item['usage']
+            for choice in item.get('choices', []):
+                delta = choice.get('delta', {})
+                answer = delta.get('content') or ''
+                thought = delta.get('reasoning') or delta.get('reasoning_content') or ''
+                if answer or thought:
+                    now = time.perf_counter()
+                    first = first if first is not None else now
+                    last = now
+                if answer:
+                    first_content = first_content if first_content is not None else now
+                content += answer
+                reasoning += thought
+                finish = choice.get('finish_reason') or finish
+    elapsed = time.perf_counter() - start
+    count = usage.get('completion_tokens', 0)
+    record = {'test': name, 'max_tokens': max_tokens, 'elapsed_s': round(elapsed, 4),
+              'ttft_s': round(first-start, 4) if first else None,
+              'first_content_s': round(first_content-start, 4) if first_content else None,
+              'decode_tps_approx': round((count-1)/(last-first), 3) if last and first and last>first else None,
+              'usage': usage, 'finish_reason': finish, 'content': content,
+              'output_sha256': hashlib.sha256((reasoning+'\0'+content).encode()).hexdigest(),
+              'prompt_sha256': hashlib.sha256(prompt.encode()).hexdigest(),
+              'correct': bool(content.strip()) and (expected is None or expected in content)}
+    emit(record)
+    if not record['correct'] or finish != 'stop':
+        raise AssertionError('Failed answer or truncated output: ' + name)
+    return record
+
+emit({'event': 'start', 'phase': args.phase})
+if args.phase in ('all', 'short'):
+    run('warmup', '计算17乘19，只输出结果。', '323')
+    prompt = '用中文写一段约150字的说明，解释数据库索引为什么能加快查询，以及它对写入有什么影响。'
+    for i in range(3):
+        run('short_'+str(i), prompt)
+    run('code', '写一个Python函数is_even(n)，判断整数是否为偶数，只输出代码。', '%')
+    for concurrency in [2, 4]:
+        barrier = threading.Barrier(concurrency)
+        start = time.perf_counter()
+        with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
+            futures = [pool.submit(run, f'concurrent_{concurrency}_{i}',
+                        f'场景编号{i}。'+prompt, None, barrier) for i in range(concurrency)]
+            results = [f.result() for f in futures]
+        seconds = time.perf_counter()-start
+        emit({'test': 'concurrency_summary', 'concurrency': concurrency,
+              'elapsed_s': round(seconds, 3),
+              'aggregate_completion_tps': round(sum(r['usage']['completion_tokens'] for r in results)/seconds, 3)})
+if args.phase in ('all', 'long'):
+    from transformers import AutoTokenizer
+    snapshot = '/root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47'
+    tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True)
+    rng = random.Random(42)
+    filler = '\n'.join(f'Record {i}: item={rng.randrange(100000,999999)}; status=archived; category={i%17}.' for i in range(6000))
+    ids = tokenizer.encode(filler, add_special_tokens=False)
+    for size in [8192, 32768]:
+        text = tokenizer.decode(ids[:size-150])
+        pos = len(text)//2
+        text = text[:pos]+'\n唯一有效的校验口令是：青松739251。\n'+text[pos:]
+        prompt = '从下面的档案里找出唯一有效的校验口令。只输出口令，不要解释。\n<档案>\n'+text+'\n</档案>\n请输出校验口令。'
+        before = metrics()
+        a = run(f'long_{size}_cold', prompt, '青松739251', max_tokens=args.long_max_tokens)
+        b = run(f'long_{size}_repeat', prompt, '青松739251', max_tokens=args.long_max_tokens)
+        after = metrics()
+        emit({'test':f'prefix_{size}', 'same_output':a['output_sha256']==b['output_sha256'],
+              'metrics_delta':{key:after[key]-before.get(key,0) for key in after}})
+emit({'event':'BENCHMARK_PASS'})

scripts/benchmark.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
+"${compose[@]}" exec -T vllm python3 -u - --label "${1:?Usage: benchmark.sh LABEL [all|short|long]}" --phase "${2:-all}" < scripts/benchmark.py

scripts/prefix-check.py

@@ -0,0 +1,39 @@
+"""Check cached-prefix retrieval at several document positions, including changed suffixes."""
+import json
+import re
+import urllib.request
+from pathlib import Path
+
+base = 'http://127.0.0.1:8000'
+key = Path('/run/secrets/qwen_api_key').read_text().strip()
+headers = {'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'}
+def request(path, payload=None):
+    req = urllib.request.Request(base+path, headers=headers,
+        data=json.dumps(payload).encode() if payload is not None else None)
+    with urllib.request.urlopen(req, timeout=600) as response:
+        return response.read().decode()
+def hits():
+    return sum(float(line.split()[-1]) for line in request('/metrics').splitlines()
+               if line.startswith('vllm:prefix_cache_hits_total'))
+
+rows = [f'归档记录{i}：此行仅供背景阅读，无有效预算。' for i in range(700)]
+for position, name, amount in [(20,'青松','17391'),(350,'白鹭','28647'),(670,'海棠','39583')]:
+    rows[position] = f'已核定：{name}项目预算为{amount}元。'
+document = '\n'.join(rows)
+before = hits()
+for name, expected in [('青松','17391'),('白鹭','28647'),('海棠','39583'),('青松','17391')]:
+    # Only the question suffix changes. The document prefix is identical.
+    payload = {'model':'qwen3.8-flash-next', 'temperature':0, 'seed':42,
+        'reasoning_effort':'low','max_tokens':512,
+        'messages':[{'role':'user','content':'请根据以下档案回答，忽略无效归档行。\n<档案>\n'+document+'\n</档案>\n'+name+'项目的已核定预算是多少元？只输出数字。'}]}
+    result = json.loads(request('/v1/chat/completions',payload))
+    choice = result['choices'][0]
+    content = choice['message'].get('content') or ''
+    correct = content.strip() == expected and choice['finish_reason']=='stop'
+    print(json.dumps({'project':name,'expected':expected,'content':content,
+        'correct':correct,'usage':result.get('usage')},ensure_ascii=False),flush=True)
+    assert correct, 'Cached prefix answer mismatch'
+delta = hits()-before
+print(json.dumps({'prefix_cache_hits_delta':delta}),flush=True)
+assert delta>0, 'No actual cache hits measured'
+print('PREFIX_CHECK_PASS',flush=True)

scripts/start.sh

@@ -2,8 +2,10 @@
 source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
 if [[ "${1:-}" == baseline ]]; then
   compose+=(-f configs/baseline-32k.yaml)
+elif [[ "${1:-}" == no-prefix ]]; then
+  compose+=(-f configs/no-prefix.yaml)
 elif [[ $# -gt 0 ]]; then
-  echo "Usage: $0 [baseline]" >&2; exit 2
+  echo "Usage: $0 [baseline|no-prefix]" >&2; exit 2
 fi
 "${compose[@]}" config --quiet
 "${compose[@]}" up -d --no-build vllm

scripts/summarize-benchmark.py

@@ -0,0 +1,26 @@
+"""Summarize a JSONL benchmark without interpreting stdout logs as results."""
+import argparse
+import json
+import statistics
+from pathlib import Path
+
+parser = argparse.ArgumentParser()
+parser.add_argument('files', nargs='+', type=Path)
+args = parser.parse_args()
+for path in args.files:
+    rows = []
+    for line in path.read_text().splitlines():
+        try:
+            rows.append(json.loads(line))
+        except ValueError:
+            continue
+    short = [r for r in rows if r.get('test', '').startswith('short_')]
+    summary = {'file': path.name,
+               'completed': any(r.get('event') == 'BENCHMARK_PASS' for r in rows),
+               'short_decode_tps_median': statistics.median(r['decode_tps_approx'] for r in short) if short else None,
+               'short_ttft_s_median': statistics.median(r['ttft_s'] for r in short) if short else None,
+               'all_checked_answers_correct': all(r['correct'] for r in rows if 'correct' in r),
+               'long': [{k:r[k] for k in ['test','ttft_s','elapsed_s','correct']} for r in rows if r.get('test','').startswith('long_')],
+               'concurrency': [r for r in rows if r.get('test') == 'concurrency_summary'],
+               'prefix': [r for r in rows if r.get('test','').startswith('prefix_')]}
+    print(json.dumps(summary, ensure_ascii=False, indent=2))

scripts/test-prefix.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+source "$(dirname -- "${BASH_SOURCE[0]}")/common.sh"
+"${compose[@]}" exec -T vllm python3 -u - < scripts/prefix-check.py