Spaces:
Sleeping
Sleeping
| """Fair batch-one decode comparison: HF eager, torch.compile, and CISM. | |
| All baselines run the SAME model checkpoint in FP32 with a KV cache, greedy | |
| decoding, one torch thread, warmup discarded. torch.compile uses a manual | |
| decode loop with StaticCache (its recommended fast path), not slow | |
| use_cache=False generate(). CISM numbers come from the same machine via | |
| cism bench. Output: JSON to stdout. Exit 0 on success. | |
| """ | |
| import argparse | |
| import json | |
| import time | |
| import torch | |
| def timed(fn, warmup_tokens: int, measured_tokens: int): | |
| fn(warmup_tokens) | |
| start = time.perf_counter() | |
| generated = fn(measured_tokens) | |
| elapsed = time.perf_counter() - start | |
| return { | |
| "tokens": generated, | |
| "seconds": elapsed, | |
| "tokens_per_second": generated / elapsed if elapsed > 0 else 0, | |
| } | |
| def run_eager(model, input_ids, eos_pad, count, warmup): | |
| def loop(tokens): | |
| out = model.generate( | |
| input_ids, attention_mask=None, max_new_tokens=tokens, | |
| min_new_tokens=tokens, do_sample=False, num_beams=1, | |
| use_cache=True, pad_token_id=eos_pad, | |
| ) | |
| return tokens | |
| result = timed(loop, warmup, count) | |
| result["method"] = "transformers generate(use_cache=True) eager" | |
| return result | |
| def run_manual_eager(model, input_ids, count, warmup): | |
| """Raw decode loop without generate()'s framework overhead.""" | |
| from transformers import StaticCache | |
| cache = StaticCache(config=model.config, max_batch_size=1, | |
| max_cache_len=warmup + count + len(input_ids[0]), | |
| device="cpu", dtype=torch.float32) | |
| def loop(tokens): | |
| with torch.inference_mode(): | |
| cache.reset() | |
| generated = 0 | |
| out = model(input_ids, past_key_values=cache, use_cache=True) | |
| token = int(out.logits[0, -1].argmax()) | |
| generated += 1 | |
| one = torch.tensor([[token]]) | |
| while generated < tokens: | |
| out = model(one, past_key_values=cache, use_cache=True) | |
| token = int(out.logits[0, -1].argmax()) | |
| generated += 1 | |
| return generated | |
| result = timed(loop, warmup, count) | |
| result["method"] = "manual eager decode loop + StaticCache (no generate overhead)" | |
| return result | |
| def _dynamo_snapshot(): | |
| import torch._dynamo | |
| return {name: sum(counts.values()) for name, counts in torch._dynamo.utils.counters.items()} | |
| def _dynamo_delta(before, after): | |
| keys = set(before) | set(after) | |
| return {key: after.get(key, 0) - before.get(key, 0) for key in keys if after.get(key, 0) != before.get(key, 0)} | |
| def run_compiled(model, input_ids, count, warmup, dynamic): | |
| from transformers import StaticCache | |
| cache = StaticCache(config=model.config, max_batch_size=1, | |
| max_cache_len=warmup + count + len(input_ids[0]), | |
| device="cpu", dtype=torch.float32) | |
| compiled = torch.compile(model, dynamic=dynamic) | |
| def loop(tokens): | |
| with torch.inference_mode(): | |
| cache.reset() | |
| generated = 0 | |
| out = compiled(input_ids, past_key_values=cache, use_cache=True) | |
| token = int(out.logits[0, -1].argmax()) | |
| generated += 1 | |
| one = torch.tensor([[token]]) | |
| while generated < tokens: | |
| out = compiled(one, past_key_values=cache, use_cache=True) | |
| token = int(out.logits[0, -1].argmax()) | |
| generated += 1 | |
| return generated | |
| # Warmup must absorb ALL compilation and recompilation; counters are reset | |
| # afterwards so the measured window only reports steady-state guard hits. | |
| loop(warmup) | |
| loop(warmup) | |
| before = _dynamo_snapshot() | |
| generated = timed(loop, warmup, count) | |
| after = _dynamo_snapshot() | |
| generated["dynamo_counter_delta_during_measurement"] = _dynamo_delta(before, after) | |
| generated["method"] = f"torch.compile(model, dynamic={dynamic}) + StaticCache decode loop" | |
| return generated | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("model") | |
| parser.add_argument("--revision", default=None) | |
| parser.add_argument("--threads", type=int, default=1) | |
| parser.add_argument("--tokens", type=int, default=64) | |
| parser.add_argument("--warmup", type=int, default=32) | |
| parser.add_argument("--skip-compile", action="store_true") | |
| args = parser.parse_args() | |
| torch.set_num_threads(args.threads) | |
| torch.set_num_interop_threads(1) | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| args.model, revision=args.revision, dtype=torch.float32, | |
| attn_implementation="eager", | |
| ).to("cpu").eval() | |
| config = {"max_position_embeddings": model.config.max_position_embeddings} | |
| ids = tokenizer("The future of artificial intelligence is", return_tensors="pt").input_ids | |
| eos_pad = tokenizer.eos_token_id or 0 | |
| report = { | |
| "model": args.model, "torch": torch.__version__, "threads": args.threads, | |
| "precision": "fp32", "greedy": True, "warmup_tokens": args.warmup, | |
| "measured_tokens": args.tokens, | |
| } | |
| report["eager_generate"] = run_eager(model, ids, eos_pad, args.tokens, args.warmup) | |
| report["eager_manual_loop"] = run_manual_eager(model, ids, args.tokens, args.warmup) | |
| if args.skip_compile: | |
| report["compiled_dynamic_false"] = {"skipped": True} | |
| report["compiled_dynamic_true"] = {"skipped": True} | |
| else: | |
| for dynamic in (False, True): | |
| key = f"compiled_dynamic_{str(dynamic).lower()}" | |
| try: | |
| report[key] = run_compiled(model, ids, args.tokens, args.warmup, dynamic) | |
| except Exception as error: | |
| report[key] = {"failed": f"{type(error).__name__}: {error}"} | |
| print(json.dumps(report, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |