()
| 44 | |
| 45 | |
| 46 | def main(): |
| 47 | parser = argparse.ArgumentParser(description="RL-Kernel Production Benchmark Suite") |
| 48 | parser.add_argument("--vocab-size", type=int, default=128256, help="Model vocab size") |
| 49 | parser.add_argument("--seq-len", type=int, default=512, help="Sequence length") |
| 50 | parser.add_argument("--g-sizes", type=str, default="64,128,256", help="Batch sizes to test") |
| 51 | parser.add_argument("--top-k", type=int, default=50, help="Top-K sampling") |
| 52 | parser.add_argument("--top-p", type=float, default=0.9, help="Top-P sampling") |
| 53 | args = parser.parse_args() |
| 54 | |
| 55 | device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "AMD MI300X/ROCm" |
| 56 | dtype = device_ctx.get_preferred_dtype() |
| 57 | |
| 58 | logger.info("Phase 1: Benchmarking Fused Logprobs...") |
| 59 | logp_results = run_logp_perf(args, return_data=True) |
| 60 | |
| 61 | logger.info("Phase 2: Benchmarking FlashInfer Sampling...") |
| 62 | sample_results = run_sample_perf(args, return_data=True) |
| 63 | |
| 64 | logger.info("Phase 3: Testing Weight Sync Latency...") |
| 65 | bridge = IPCWeightBridge() |
| 66 | dummy_model = torch.nn.Linear(1024, 1024).to(device_ctx.device) |
| 67 | t0 = time.perf_counter() |
| 68 | handles = bridge.export_model_handles(dummy_model) |
| 69 | _ = bridge.import_model_weights(handles) |
| 70 | sync_latency = (time.perf_counter() - t0) * 1000 |
| 71 | |
| 72 | metrics = { |
| 73 | "tip": "install flashinfer and aiter for maximum throughput.", |
| 74 | "device": device_name, |
| 75 | "vocab": args.vocab_size, |
| 76 | "seq": args.seq_len, |
| 77 | "dtype": dtype, |
| 78 | "vram_saved": max([r["vram_saved_val"] for r in logp_results]), |
| 79 | "logp_speedup": logp_results[-1]["speedup"], |
| 80 | "sample_speedup": sample_results[-1]["speedup"], |
| 81 | "avg_logp_ms": sum([r["engine_ms"] for r in logp_results]) / len(logp_results), |
| 82 | "avg_sample_ms": sum([r["engine_ms"] for r in sample_results]) / len(sample_results), |
| 83 | "sync_ms": sync_latency, |
| 84 | } |
| 85 | |
| 86 | PerfReport.print_panel(metrics) |
| 87 | |
| 88 | |
| 89 | if __name__ == "__main__": |
no test coverage detected