Buckets:
| """vLLM helpers: letter readout, thinking-mode sampling, readout after a chain. | |
| The readout restricts the next token to the K option letters (allowed_token_ids) and asks for processed | |
| logprobs, so the K returned values are exactly log softmax over the K letter logits: one forward pass, nothing | |
| generated beyond that token. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import os | |
| import time | |
| import uuid | |
| import numpy as np | |
| from vllm import LLM, SamplingParams | |
| from vllm.sampling_params import RequestOutputKind | |
| from common import CONTEXT, LETTERS, letter_ids, render, split_chain, parse_letter | |
| # Thinking-mode sampling: Qwen's recommended settings by default; a model card's own settings via the environment | |
| # (MiniCPM5: THINK_TEMPERATURE=1.0 THINK_TOP_K=-1 THINK_MIN_P=0) | |
| THINK_SAMPLING = dict(temperature=float(os.environ.get("THINK_TEMPERATURE", "0.6")), | |
| top_p=float(os.environ.get("THINK_TOP_P", "0.95")), | |
| top_k=int(os.environ.get("THINK_TOP_K", "20")), | |
| min_p=float(os.environ.get("THINK_MIN_P", "0.0"))) | |
| def engine_kwargs(model: str, enable_lora: bool = False, gpu_mem: float = 0.90, max_lora_rank: int = 32) -> dict: | |
| kw = dict(model=model, dtype="bfloat16", max_model_len=CONTEXT, gpu_memory_utilization=gpu_mem, | |
| enable_prefix_caching=True, max_logprobs=len(LETTERS), logprobs_mode="processed_logprobs", | |
| seed=17, | |
| # the in-process async engine otherwise falls back to 128 running sequences. 256 was faster for n=4 | |
| # generation (KV ~80-99%) but thrashed the n=1 eval (KV 99.5%, preemption, 1.9k tok/s): set per stage | |
| max_num_seqs=int(os.environ.get("MAX_NUM_SEQS", "256")), | |
| disable_log_stats=False) # engine stats (printed by the progress reporter each minute) | |
| if enable_lora: | |
| kw.update(enable_lora=True, max_lora_rank=max_lora_rank, max_loras=3) | |
| return kw | |
| def _env_overrides(kw: dict) -> dict: | |
| """Serving knobs for benchmarking and per-stage tuning, all optional: | |
| GPU_MEM (fraction), MAX_BATCHED_TOKENS (per step). Scheduling only: they never change what the model computes. | |
| (Speculative decoding was considered and left out on purpose: an untested path for this hybrid model.)""" | |
| if os.environ.get("GPU_MEM"): | |
| kw["gpu_memory_utilization"] = float(os.environ["GPU_MEM"]) | |
| if os.environ.get("MAX_BATCHED_TOKENS"): | |
| kw["max_num_batched_tokens"] = int(os.environ["MAX_BATCHED_TOKENS"]) | |
| return kw | |
| def make_llm(model: str, **kw) -> LLM: | |
| return LLM(**_env_overrides(engine_kwargs(model, **kw))) | |
| def _letter_logprobs(output, ids: list[int]) -> np.ndarray: | |
| lp = output.outputs[0].logprobs[0] # dict token_id -> Logprob, the K allowed letters | |
| return np.array([lp[i].logprob if i in lp else -1e4 for i in ids], dtype=np.float64) | |
| def readout_prompts(tok, items, orders, think_chains=None) -> list[str]: | |
| """No-think prompts, or (with chains) thinking prompts continued by the chain and '</think>\\n\\n'.""" | |
| out = [] | |
| for i, (it, order) in enumerate(zip(items, orders)): | |
| if think_chains is None: | |
| out.append(render(tok, it, think=False, order=order)) | |
| else: | |
| base = render(tok, it, think=True, order=order) | |
| if not base.endswith("<think>\n"): | |
| base += "<think>\n" | |
| out.append(base + think_chains[i] + "\n</think>\n\n") | |
| return out | |
| def _strip_bos(tok, text: str) -> str: | |
| """Chat templates that write the BOS token themselves (MiniCPM, Llama-style) would get a second one when vLLM | |
| tokenizes the text; Qwen's templates have none, so this changes nothing for them.""" | |
| b = tok.bos_token | |
| return text[len(b):] if b and text.startswith(b) else text | |
| def readout(llm: LLM, tok, items, orders, lora=None, think_chains=None) -> list[np.ndarray]: | |
| """K letter log-probabilities per item (in prompt order), from one forward pass each.""" | |
| ids = letter_ids(tok) | |
| prompts = [_strip_bos(tok, p) for p in readout_prompts(tok, items, orders, think_chains)] | |
| params = [SamplingParams(max_tokens=1, temperature=1.0, allowed_token_ids=ids[:len(it["options"])], | |
| logprobs=len(it["options"])) for it in items] | |
| outs = llm.generate(prompts, params, lora_request=lora, use_tqdm=False) | |
| return [_letter_logprobs(o, ids[:len(it["options"])]) for o, it in zip(outs, items)] | |
| def chain_record(text: str, n_tokens: int, finish: str, k: int) -> dict: | |
| chain, answer = split_chain(text) | |
| letter = parse_letter(answer, k) if chain is not None else None | |
| return {"text": text, "tokens": n_tokens, "finish": finish, "closed": chain is not None, | |
| "answer": letter, "looped": chain is None} | |
| def think_sync(llm: LLM, tok, items, orders, n: int = 1, seed: int = 0, lora=None) -> list[list[dict]]: | |
| """n thinking-mode samples per item, no length cap below the context.""" | |
| prompts = [_strip_bos(tok, render(tok, it, think=True, order=o)) for it, o in zip(items, orders)] | |
| lens = [len(x) for x in tok(prompts, add_special_tokens=False)["input_ids"]] | |
| params = [SamplingParams(n=n, max_tokens=CONTEXT - L, seed=seed + i, **THINK_SAMPLING) for i, L in enumerate(lens)] | |
| outs = llm.generate(prompts, params, lora_request=lora, use_tqdm=False) | |
| return [[chain_record(c.text, len(c.token_ids), c.finish_reason, len(it["options"])) for c in o.outputs] | |
| for o, it in zip(outs, items)] | |
| def make_async_engine(model: str, **kw): | |
| from vllm import AsyncEngineArgs, AsyncLLMEngine | |
| return AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**_env_overrides(engine_kwargs(model, **kw)))) | |
| async def _final(engine, prompt, params, lora=None): | |
| """The finished RequestOutput, with all n completions (FINAL_ONLY: one output, after every sample ends).""" | |
| params.output_kind = RequestOutputKind.FINAL_ONLY | |
| final = None | |
| async for res in engine.generate(prompt, params, request_id=str(uuid.uuid4()), lora_request=lora): | |
| final = res | |
| return final | |
| async def readout_async(engine, tok, items, orders, lora=None, think_chains=None, concurrency: int = 512) -> list[np.ndarray]: | |
| """Async twin of readout(): K letter log-probabilities per item, in prompt order.""" | |
| ids = letter_ids(tok) | |
| prompts = [_strip_bos(tok, p) for p in readout_prompts(tok, items, orders, think_chains)] | |
| sem = asyncio.Semaphore(concurrency) | |
| async def one(prompt, k): | |
| async with sem: | |
| params = SamplingParams(max_tokens=1, temperature=1.0, allowed_token_ids=ids[:k], logprobs=k) | |
| return _letter_logprobs(await _final(engine, prompt, params, lora), ids[:k]) | |
| return await asyncio.gather(*[one(p, len(it["options"])) for p, it in zip(prompts, items)]) | |
| async def latency_async(engine, tok, items, orders, think: bool, lora=None) -> list[float]: | |
| """Seconds per request, one request at a time (batch size 1).""" | |
| ids = letter_ids(tok) | |
| out = [] | |
| for it, order in zip(items, orders): | |
| k = len(it["options"]) | |
| if think: | |
| prompt = _strip_bos(tok, render(tok, it, think=True, order=order)) | |
| plen = len(tok(prompt, add_special_tokens=False)["input_ids"]) | |
| params = SamplingParams(max_tokens=CONTEXT - plen, seed=17, **THINK_SAMPLING) | |
| else: | |
| prompt = _strip_bos(tok, render(tok, it, think=False, order=order)) | |
| params = SamplingParams(max_tokens=1, temperature=1.0, allowed_token_ids=ids[:k], logprobs=k) | |
| t0 = time.perf_counter() | |
| await _final(engine, prompt, params, lora) | |
| out.append(time.perf_counter() - t0) | |
| return out | |
| async def post_readout_ids(engine, tok, prompt_ids, gen_ids, k: int, lora=None) -> list[float] | None: | |
| """Letter log-probs right after a chain, fed as the exact token ids just generated (prompt + chain up to | |
| and including the line break after </think>), so the prefix cache still holds them and the pass is cheap.""" | |
| end_id = tok.convert_tokens_to_ids("</think>") | |
| if end_id not in gen_ids: | |
| return None | |
| i = gen_ids.index(end_id) | |
| sep = tok.encode("\n\n", add_special_tokens=False) | |
| tail = list(gen_ids[:i + 1 + len(sep)]) if list(gen_ids[i + 1:i + 1 + len(sep)]) == sep else list(gen_ids[:i + 1]) + sep | |
| ids = letter_ids(tok)[:k] | |
| params = SamplingParams(max_tokens=1, temperature=1.0, allowed_token_ids=ids, logprobs=k) | |
| out = await _final(engine, {"prompt_token_ids": list(prompt_ids) + tail}, params, lora) | |
| return [float(v) for v in _letter_logprobs(out, ids)] | |
| async def think_async(model_or_engine, tok, jobs, out_path: str, n: int, deadline: float, concurrency: int = 256, | |
| lora=None, gpu_mem: float = 0.90, post_readout: bool = False) -> dict: | |
| """Stream thinking samples for jobs [(item, order)], appending one JSON line per finished item. | |
| Stops taking new items at `deadline` (unix time) and lets in-flight ones finish; never cuts a chain. | |
| With post_readout, each closed chain also gets its letter readout right away ("post_lp"). | |
| """ | |
| engine = model_or_engine | |
| if isinstance(model_or_engine, str): | |
| engine = make_async_engine(model_or_engine, gpu_mem=gpu_mem) | |
| sem = asyncio.Semaphore(concurrency) | |
| stats = {"submitted": 0, "finished": 0, "gen_tokens": 0, "t0": time.time()} | |
| fout = open(out_path, "a", encoding="utf-8") | |
| progress_path = out_path.replace(".jsonl", "") + ".progress.jsonl" | |
| seen_tokens, seen_parsed, seen_looped = [], [], [] | |
| async def reporter(): | |
| """One progress line a minute: throughput, chain lengths and parse rate so far.""" | |
| with open(progress_path, "a") as pf: | |
| while True: | |
| await asyncio.sleep(60) | |
| el = time.time() - stats["t0"] | |
| toks = np.array(seen_tokens) if seen_tokens else np.zeros(1) | |
| rec = {"time": time.strftime("%H:%M:%S"), "elapsed_min": round(el / 60, 1), | |
| "submitted": stats["submitted"], "finished": stats["finished"], | |
| "in_flight": stats["submitted"] - stats["finished"], "gen_tokens": stats["gen_tokens"], | |
| "tok_per_s": round(stats["gen_tokens"] / max(el, 1e-9)), | |
| "chain_tokens_median": float(np.median(toks)), "chain_tokens_p90": float(np.percentile(toks, 90)), | |
| "chain_tokens_max": int(toks.max()), "parse_rate": float(np.mean(seen_parsed)) if seen_parsed else None, | |
| "looped": int(sum(seen_looped)), "minutes_to_deadline": round((deadline - time.time()) / 60, 1)} | |
| pf.write(json.dumps(rec) + "\n") | |
| pf.flush() | |
| print("progress " + json.dumps(rec), flush=True) | |
| try: # vLLM's own stats line (throughput, running/waiting, KV use) | |
| await engine.do_log_stats() | |
| except Exception: | |
| pass | |
| rep = asyncio.create_task(reporter()) | |
| async def one(idx, item, order): | |
| prompt = _strip_bos(tok, render(tok, item, think=True, order=order)) | |
| plen = len(tok(prompt, add_special_tokens=False)["input_ids"]) | |
| params = SamplingParams(n=n, max_tokens=CONTEXT - plen, seed=17 + idx, **THINK_SAMPLING) | |
| final = await _final(engine, prompt, params, lora) | |
| samples = [chain_record(c.text, len(c.token_ids), c.finish_reason, len(item["options"])) for c in final.outputs] | |
| if post_readout: | |
| for smp, c in zip(samples, final.outputs): | |
| smp["post_lp"] = (await post_readout_ids(engine, tok, final.prompt_token_ids, list(c.token_ids), | |
| len(item["options"]), lora)) if smp["closed"] else None | |
| rec = {"id": item["id"], "order": order, "prompt_tokens": plen, "samples": samples} | |
| fout.write(json.dumps(rec, ensure_ascii=False) + "\n") | |
| fout.flush() | |
| stats["finished"] += 1 | |
| stats["gen_tokens"] += sum(s["tokens"] for s in samples) | |
| for s in samples: | |
| seen_tokens.append(s["tokens"]) | |
| seen_parsed.append(s["answer"] is not None) | |
| seen_looped.append(s["looped"]) | |
| async def guarded(idx, item, order): | |
| try: | |
| await one(idx, item, order) | |
| finally: | |
| sem.release() | |
| tasks = [] | |
| for idx, (item, order) in enumerate(jobs): | |
| await sem.acquire() | |
| if time.time() >= deadline: | |
| sem.release() | |
| break | |
| tasks.append(asyncio.create_task(guarded(idx, item, order))) | |
| stats["submitted"] += 1 | |
| await asyncio.gather(*tasks) | |
| rep.cancel() | |
| fout.close() | |
| stats["seconds"] = time.time() - stats["t0"] | |
| stats["tok_per_s"] = stats["gen_tokens"] / max(stats["seconds"], 1e-9) | |
| return stats | |
Xet Storage Details
- Size:
- 12.9 kB
- Xet hash:
- cb8fa59385f8ae544f1fb110f40f1b33e0d0e0f5499e03e898e70033e68fa04a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.