| 363 | "prompt_len": config.prompt_len, |
| 364 | "completion_len": config.completion_len, |
| 365 | "vocab_size": config.vocab_size, |
| 366 | "mask_density": config.mask_density, |
| 367 | "valid_tokens": metadata["valid_tokens"], |
| 368 | "reference_ms": f"{reference_ms:.4f}", |
| 369 | "candidate_ms": f"{candidate_ms:.4f}" if isinstance(candidate_ms, float) else candidate_ms, |
| 370 | "peak_memory_gb": f"{peak_memory_gb:.6f}", |
| 371 | "max_error": max_error, |
| 372 | "ratio_drift": ratio_drift, |
| 373 | "kl_drift": kl_drift, |
| 374 | "status": status, |
| 375 | "notes": notes, |
| 376 | } |
| 377 | |
| 378 | |
| 379 | def _write_rows(rows: list[dict[str, Any]], output: Path | None) -> None: |
| 380 | if output is None: |
| 381 | writer = csv.DictWriter(__import__("sys").stdout, fieldnames=CSV_COLUMNS) |
| 382 | writer.writeheader() |
| 383 | writer.writerows(rows) |
| 384 | return |
| 385 | |
| 386 | output.parent.mkdir(parents=True, exist_ok=True) |
| 387 | exists = output.exists() and output.stat().st_size > 0 |
| 388 | with output.open("a", newline="", encoding="utf-8") as handle: |
| 389 | writer = csv.DictWriter(handle, fieldnames=CSV_COLUMNS) |
| 390 | if not exists: |
| 391 | writer.writeheader() |
| 392 | writer.writerows(rows) |
| 393 | |
| 394 | |
| 395 | def build_arg_parser() -> argparse.ArgumentParser: |
| 396 | parser = argparse.ArgumentParser(description="RL-shaped RL-Kernel benchmark runner") |
| 397 | parser.add_argument("--case", default="selected_logprob", choices=["selected_logprob"]) |
| 398 | parser.add_argument( |
| 399 | "--candidate", |