| """Long-context CODING benchmark on the vLLM SERVING path: executed HumanEval |
| pass@1 for bf16 vs our custom INT4-KIVI KV cache. |
| |
| Unlike the exact-integer needle (which a single quantized-logit digit-flip can |
| fail), this measures *executed* pass@1 -- the model must emit a function that |
| actually passes the hidden unit tests. That is the metric KV-cache quant has to |
| preserve, and it is far less brittle than digit-exact recall. |
| |
| Two regimes, identical prompts for both dtypes (so it is apples-to-apples): |
| * short : plain HumanEval (KV cache ~hundreds of tokens at decode start). |
| * long : the SAME problems with a long, in-distribution Python-source prefix |
| prepended (real ``transformers`` modeling_*.py), so the model must |
| decode while attending back over a fully-quantized long context. |
| |
| Run once per dtype (mirrors needle_serving.py): |
| KVD=auto -> bf16 KV cache (ceiling) |
| KVD=int4_kivi -> our custom INT4-KIVI backend |
| N, PREFIX_TOKENS, MAXNEW are env-overridable. Greedy (temperature 0). |
| """ |
| import glob |
| import json |
| import os |
| import re |
| import subprocess |
| import sys |
| import tempfile |
| import time |
|
|
| from datasets import load_dataset |
| from vllm import LLM, SamplingParams |
|
|
| ROOT = "/home/alex/poolside-hackathon-kv-quant" |
| MODEL = "poolside/Laguna-XS.2" |
|
|
| KVD = os.environ.get("KVD", "int4_kivi") |
| N = int(os.environ.get("N", "20")) |
| PREFIX_TOKENS = int(os.environ.get("PREFIX_TOKENS", "12000")) |
| MAXNEW = int(os.environ.get("MAXNEW", "256")) |
|
|
|
|
| |
| |
| |
| def extract_code(response: str, prompt: str) -> str: |
| m = re.search(r"```(?:python)?\n(.*?)```", response, re.DOTALL) |
| if m: |
| block = m.group(1) |
| if prompt.split("def ", 1)[-1].split("(")[0].strip() in block: |
| return block |
| return prompt + block |
| lines = response.splitlines() |
| body, in_body = [], False |
| for line in lines: |
| if not in_body and (line.startswith(" ") or line.startswith("\t")): |
| in_body = True |
| if in_body: |
| if line.startswith("def ") and body: |
| break |
| body.append(line) |
| if body: |
| return prompt + "\n".join(body) + "\n" |
| return prompt + response |
|
|
|
|
| def run_tests(solution_code: str, test_code: str, entry_point: str): |
| full = solution_code + "\n\n" + test_code + f"\ncheck({entry_point})\n" |
| with tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False) as f: |
| f.write(full) |
| fname = f.name |
| try: |
| r = subprocess.run( |
| [sys.executable, fname], capture_output=True, text=True, timeout=10 |
| ) |
| if r.returncode == 0: |
| return True, "" |
| return False, (r.stderr or r.stdout).strip()[-300:] |
| except subprocess.TimeoutExpired: |
| return False, "timeout" |
| except Exception as e: |
| return False, str(e) |
| finally: |
| try: |
| os.unlink(fname) |
| except OSError: |
| pass |
|
|
|
|
| |
| |
| |
| def build_prefix_text(tok, target_tokens: int): |
| files = [] |
| for venv in (".venv-vllm", ".venv"): |
| files = sorted(glob.glob(f"{ROOT}/{venv}/**/transformers/**/modeling_*.py", |
| recursive=True)) |
| if files: |
| break |
| if not files: |
| files = sorted(glob.glob(f"{ROOT}/**/*.py", recursive=True)) |
| texts = [] |
| for f in files: |
| try: |
| texts.append(open(f).read()) |
| except OSError: |
| continue |
| ids = tok("\n\n".join(texts))["input_ids"] |
| if len(ids) >= target_tokens: |
| return tok.decode(ids[:target_tokens]), target_tokens |
| ids = tok("\n\n".join(texts))["input_ids"] |
| return tok.decode(ids), len(ids) |
|
|
|
|
| |
| SYS_MSG = ( |
| "You are a Python coding assistant. Complete the function below. Return a " |
| "fenced ```python``` code block containing the complete function (including " |
| "signature and docstring)." |
| ) |
|
|
|
|
| def build_msgs(prompt, prefix_text): |
| if prefix_text: |
| user = ("Here is some reference Python source code for context. You do " |
| "not need to use it; it is provided only as background.\n\n" |
| f"```python\n{prefix_text}\n```\n\n" |
| "Now, ignoring the reference above, complete this Python " |
| f"function:\n\n```python\n{prompt}```") |
| else: |
| user = f"Complete this Python function:\n\n```python\n{prompt}```" |
| return [{"role": "system", "content": SYS_MSG}, |
| {"role": "user", "content": user}] |
|
|
|
|
| |
| llm = LLM(model=MODEL, dtype="bfloat16", kv_cache_dtype=KVD, |
| gpu_memory_utilization=0.55, max_model_len=PREFIX_TOKENS + 2048, |
| enforce_eager=True) |
| tok = llm.get_tokenizer() |
|
|
| prefix_text, prefix_tok = build_prefix_text(tok, PREFIX_TOKENS) |
| ds = load_dataset("openai/openai_humaneval", split="test").select(range(N)) |
| sp = SamplingParams(temperature=0.0, max_tokens=MAXNEW) |
|
|
| summary = {} |
| for regime in ("short", "long"): |
| pfx = prefix_text if regime == "long" else None |
| convs = [build_msgs(p["prompt"], pfx) for p in ds] |
| t0 = time.time() |
| outs = llm.chat(convs, sp, add_generation_prompt=True) |
| gen_s = time.time() - t0 |
| ctx = [len(o.prompt_token_ids) for o in outs] |
|
|
| npass = 0 |
| for prob, o in zip(ds, outs): |
| sol = extract_code(o.outputs[0].text, prob["prompt"]) |
| ok, _ = run_tests(sol, prob["test"], prob["entry_point"]) |
| npass += int(ok) |
| summary[regime] = {"pass": npass, "n": len(ds), |
| "ctx_min": min(ctx), "ctx_max": max(ctx), "gen_s": gen_s} |
| print(f"=== [{KVD}] {regime.upper()} HumanEval pass@1 ===") |
| print(f" {npass}/{len(ds)} ({100*npass/len(ds):.0f}%) " |
| f"ctx {min(ctx)}..{max(ctx)} gen {gen_s:.0f}s") |
|
|
| json.dump({"kvd": KVD, "prefix_tok": prefix_tok, "summary": summary}, |
| open(f"/tmp/longctx_code_serving_{KVD}.json", "w")) |
| print(f"LONGCTX_CODE_SERVING DONE [{KVD}]") |
|
|
