File size: 3,133 Bytes
6e668dc | 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 | """Stable long-context needle retrieval on the vLLM SERVING path.
Same construction as the throwaway /tmp probe, but with enough trials per length
to get a low-noise bf16-vs-int4_kivi comparison (the n=5 probe could not resolve
a ~20% K-error change from sampling noise). Exact-integer graded.
Run once per dtype:
KVD=auto -> bf16 KV cache (ceiling)
KVD=int4_kivi -> our custom INT4-KIVI backend
TRIALS and LENGTHS are env-overridable.
"""
import os, json, random, time
from vllm import LLM, SamplingParams
KVD = os.environ.get("KVD", "int4_kivi")
TRIALS = int(os.environ.get("TRIALS", "20"))
LENGTHS = [int(x) for x in os.environ.get("LENGTHS", "8000,16000,32000").split(",")]
DEPTH = 0.06 # needle near the start (in a frozen, fully-quantized region)
llm = LLM(model="poolside/Laguna-XS.2", dtype="bfloat16", kv_cache_dtype=KVD,
gpu_memory_utilization=0.55, max_model_len=34000, enforce_eager=True)
tok = llm.get_tokenizer()
def filler(rng, n):
out = []
for i in range(n):
k = rng.randint(2, 99)
out.append(f"def _f{i}(x):\n # assorted helper {i}\n return x * {k} + {i % 7}\n")
return "\n".join(out)
def build(length_tok, depth, magic_id, magic_val, rng):
needle = f"MAGIC_SEED_{magic_id} = {magic_val} # unique constant, remember it\n\n"
blocks, ntok, placed = [], 0, False
while ntok < length_tok:
if not placed and ntok >= depth * length_tok:
blocks.append(needle); placed = True
b = filler(rng, 20)
blocks.append(b)
ntok += len(tok(b)["input_ids"])
if not placed:
blocks.insert(1, needle)
body = "\n".join(blocks)
q = (f"\n\n# End of module. Recall the unique constant defined far above.\n"
f"# Write only the integer value assigned to MAGIC_SEED_{magic_id}.\n"
f"# MAGIC_SEED_{magic_id} =")
return body + q
prompts, meta = [], []
rng = random.Random(0)
for L in LENGTHS:
for t in range(TRIALS):
mid = f"{L}_{t}"
val = rng.randint(1000003, 9999991)
prompts.append(build(L, DEPTH, mid, val, rng))
meta.append((L, val))
ctx_tok = [len(tok(p)["input_ids"]) for p in prompts]
print(f"[{KVD}] {len(prompts)} prompts ({TRIALS}/len), ctx {min(ctx_tok)}..{max(ctx_tok)}")
sp = SamplingParams(temperature=0.0, max_tokens=12)
t0 = time.time()
outs = llm.generate(prompts, sp)
gen_s = time.time() - t0
res = {L: {"pass": 0, "n": 0} for L in LENGTHS}
for (L, val), o in zip(meta, outs):
txt = o.outputs[0].text.replace(",", "").replace("_", "")
res[L]["pass"] += int(str(val) in txt)
res[L]["n"] += 1
print(f"=== [{KVD}] needle retrieval ({TRIALS} trials/len) ===")
for L in LENGTHS:
r = res[L]
print(f" len~{L:>6}: {r['pass']:>2}/{r['n']:>2} ({100*r['pass']/r['n']:.0f}%)")
tot_p = sum(r['pass'] for r in res.values()); tot_n = sum(r['n'] for r in res.values())
print(f"[{KVD}] TOTAL {tot_p}/{tot_n} ({100*tot_p/tot_n:.0f}%) gen {gen_s:.0f}s")
json.dump({"kvd": KVD, "res": res, "gen_s": gen_s, "trials": TRIALS},
open(f"/tmp/needle_serving_{KVD}.json", "w"))
print(f"NEEDLE_SERVING DONE [{KVD}]")
|