File size: 4,980 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Prepare pinned image overlays offline. Writes only to an explicit empty output
directory; never deploys or edits the image snapshot. Run before any coordinator
window. The generated standard verifier accepts actual FP32 q, not log(q)/T.
"""
from __future__ import annotations
import argparse
import ast
import hashlib
import json
from pathlib import Path
UTILS = 'vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py'
RUNNER = 'vllm/v1/worker/gpu/model_runner.py'
PINS = {UTILS: 'ef1d40bf31140eb644a8f3b7a003373209d6a8e297f059e70ea66cf4972d8d90',
RUNNER: 'af9be025ff8a17a5aabd452fc3e4b987abe64fb28d852c569e1372d6f1e4d3a6'}
ABI = 'pp3-dyadic24-standard-v1'
def replace_once(text: str, old: str, new: str) -> str:
if text.count(old) != 1:
raise ValueError(f'probability overlay source mismatch: {old[:80]!r}')
return text.replace(old, new)
def utilities(text: str) -> str:
text = 'import os\nREMOTE_DRAFT_PROBS_ABI = ' + repr(ABI) + '\n' + text
old = ' if HAS_DRAFT_LOGITS:\n # draft_logits is stored pre-temperature, so apply scale first.'
new = ''' if HAS_DRAFT_LOGITS == 2:
q = tl.load(draft_logits_ptr + req_state_idx * draft_logits_stride_0
+ draft_step * draft_logits_stride_1 + token,
mask=mask, other=0.).to(tl.float32)
draft_log_prob = tl.log(q)
draft_lse = 0.
elif HAS_DRAFT_LOGITS:
# draft_logits is stored pre-temperature, so apply scale first.'''
text = replace_once(text, old, new)
text = replace_once(text, ' if HAS_DRAFT_LOGITS:\n # Get local draft max',
' if HAS_DRAFT_LOGITS == 1:\n # Get local draft max')
text = replace_once(text, ' draft_log_probs = draft_logits - draft_lse', ''' draft_log_probs = draft_logits - draft_lse
if HAS_DRAFT_LOGITS == 2:
q = tl.load(draft_logits_ptr + req_state_idx * draft_logits_stride_0
+ resample_idx * draft_logits_stride_1 + block,
mask=mask, other=0.).to(tl.float32)
draft_log_probs = tl.log(q)''')
text = replace_once(text, ' # Compute the per-vocab-block logits stats,', ''' if os.environ.get("VLLM_DSPARK_SAMPLED_TOP32", "0") == "1":
if use_block_verification or synthetic_conditional_rates is not None or watermark:
raise ValueError("PP3 q ABI supports standard verification only, without watermarking")
if draft_logits is None or draft_logits.dtype != torch.float32 or draft_logits.shape[-1] != target_logits.shape[-1]:
raise ValueError("PP3 standard verification needs full-vocab FP32 probabilities")
has_draft_logits = 2
# Compute the per-vocab-block logits stats,''')
return text
def runner(text: str) -> str:
text = 'REMOTE_DRAFT_PROBS_ABI = ' + repr(ABI) + '\n' + text
return replace_once(text, ' sampler_output: SamplerOutput | None\n', ''' # PP3 failed/missing-q rounds are target-only. Mask verification, never
# reinterpret sampled IDs as point-mass proposals. Tensor values stay on GPU.
pp3_valid = getattr(self.speculator, "sampled_valid", None)
if pp3_valid is not None and input_batch.num_draft_tokens > 0:
bad = ((input_batch.expanded_local_pos > 0)
& ~pp3_valid[input_batch.expanded_idx_mapping.long()])
invalid_drafts = bad if invalid_drafts is None else invalid_drafts | bad
sampler_output: SamplerOutput | None
''')
def prepare(source: Path, output: Path) -> dict[str, object]:
if output.exists():
raise ValueError('output must not exist; refuse overwriting any prepared overlay')
texts = {}
for name, pin in PINS.items():
data = (source / name).read_bytes()
if hashlib.sha256(data).hexdigest() != pin:
raise ValueError(f'image snapshot differs from pinned PP3 source: {name}')
texts[name] = utilities(data.decode()) if name == UTILS else runner(data.decode())
ast.parse(texts[name])
compile(texts[name], name, 'exec')
result: dict[str, object] = {'abi': ABI, 'source_sha256': PINS,
'output_sha256': {name: hashlib.sha256(text.encode()).hexdigest()
for name, text in texts.items()}}
output.mkdir(parents=True)
for name, text in texts.items():
path = output / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(text)
(output / 'receipt.json').write_text(json.dumps(result, indent=2)+'\n')
return result
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--source', type=Path, required=True)
p.add_argument('--output', type=Path, required=True)
a = p.parse_args()
print(json.dumps(prepare(a.source, a.output), indent=2))
if __name__ == '__main__':
main()
|