diff --git a/benchmarks/bench_offload_cache_copy.py b/benchmarks/bench_offload_cache_copy.py index a2434a66..7bc227af 100644 --- a/benchmarks/bench_offload_cache_copy.py +++ b/benchmarks/bench_offload_cache_copy.py @@ -195,9 +195,25 @@ def print_table( miss_rates: list[float], repeat: int, device: torch.device, -) -> None: +) -> bool: + """False when the geometry is refused and nothing was measured; the caller tallies those.""" per_expert = expert_bytes(profile) cache_gib = cache_slots * per_expert / 2**30 + # Build before printing: a refused geometry would otherwise leave a header with no rows + # under it, which is what the abort already looked like. `validate_rebuild` runs in + # __post_init__ ahead of the allocations, so the refused path costs nothing. + # + # Catch rather than pre-check: the rule is the pool's and covers both the num_experts + # floor and the nvfp4_marlin slot cap (992 -- moe_align_block_size needs + # round_up(experts, 32) < 1024). Re-deriving that here would duplicate the rule, not just + # the number. Note this is ValueError only: a device OOM is a RuntimeError and still + # aborts, which is correct -- it is not a statement about this geometry being illegal. + try: + cache = make_cache(profile, cache_slots, device) + except ValueError as exc: + print(f"\n{name} @ cache_slots={cache_slots} ({cache_gib:.1f} GiB): SKIPPED -- {exc}") + return False + print( f"\n{name} ({profile.quant_format}, {len(bank_specs(profile))} banks, " f"L={profile.layers} E={profile.experts} k={profile.topk}) " @@ -207,7 +223,6 @@ def print_table( print("bs active miss_rate misses time_ms copy_MiB bw_GBps tok_ms") print("-- ------ --------- ------ ------- -------- ------- ------") - cache = make_cache(profile, cache_slots, device) for batch_size in batch_sizes: for miss_rate in miss_rates: active, misses, time_ms = time_case( @@ -223,26 +238,44 @@ def print_table( ) del cache torch.cuda.empty_cache() + return True -def main() -> None: +def main() -> int: args = parse_args() assert torch.cuda.is_available(), "CUDA is required" torch.cuda.set_device(args.device) device = torch.device("cuda") print("gpu", torch.cuda.get_device_name(device), flush=True) - for name in args.models: - profile = MODELS[name] - slot_counts = args.cache_slots or [ - profile.experts, - int(0.4 * profile.layers * profile.experts), - ] - for cache_slots in slot_counts: - print_table( - name, profile, cache_slots, args.batch_sizes, args.miss_rates, args.repeat, device - ) + plan = [ + (name, MODELS[name], cache_slots) + for name in args.models + for cache_slots in ( + args.cache_slots + or [MODELS[name].experts, int(0.4 * MODELS[name].layers * MODELS[name].experts)] + ) + ] + skipped = [] + for name, profile, cache_slots in plan: + if print_table( + name, profile, cache_slots, args.batch_sizes, args.miss_rates, args.repeat, device + ): + continue + skipped.append(f"{name}@{cache_slots}") + if not skipped: + return 0 + + print(f"\nskipped {len(skipped)} of {len(plan)} combinations: {', '.join(skipped)}") + # A derived default that one backend cannot satisfy is expected output, not a failed + # run -- the no-arg sweep would otherwise be red forever. Asking for a geometry by hand + # and getting nothing is a different thing, and so is measuring nothing at all. + if len(skipped) == len(plan): + return 1 + if args.cache_slots is not None: + return 1 + return 0 if __name__ == "__main__": - main() + raise SystemExit(main())