| |
| """Higgs delayed-codebook benchmark against eager/compile and raw op.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| def elapsed_us(fn, warmup, iterations): |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| start = torch.cuda.Event(enable_timing=True) |
| end = torch.cuda.Event(enable_timing=True) |
| start.record() |
| for _ in range(iterations): |
| fn() |
| end.record() |
| end.synchronize() |
| return start.elapsed_time(end) * 1000.0 / iterations |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=("source", "installed"), default="source") |
| parser.add_argument("--artifact") |
| parser.add_argument("--warmup", type=int, default=50) |
| parser.add_argument("--iterations", type=int, default=500) |
| args = parser.parse_args() |
| if args.backend == "source": |
| tests = Path(__file__).resolve().parents[1] / "tests" |
| sys.path.insert(0, str(tests)) |
| from test_audio_codebook_primitives import load_source_ops |
| ops = load_source_ops() |
| else: |
| if args.artifact: |
| sys.path.insert(0, args.artifact) |
| ops = importlib.import_module("audio_codebook_primitives") |
|
|
| c, v, h, delay, boc = 8, 1026, 1024, 7, 1024 |
| logits = torch.randn((c, v), device="cuda", dtype=torch.bfloat16) |
| codebook = torch.randn((c, v, h), device="cuda", dtype=torch.bfloat16) |
| index = torch.arange(c, device="cuda") |
| active = index <= delay |
| boc_tensor = torch.full((c,), boc, device="cuda", dtype=torch.int64) |
|
|
| def eager(): |
| codes = torch.where(active, logits.argmax(dim=1), boc_tensor) |
| embedding = codebook[index, codes].float().sum(dim=0).bfloat16() |
| return codes, embedding |
|
|
| compiled = torch.compile(eager, fullgraph=True) |
| codes = torch.empty(c, device="cuda", dtype=torch.int64) |
| embedding = torch.empty(h, device="cuda", dtype=torch.bfloat16) |
|
|
| def wrapper(): |
| return ops.delayed_codebook_argmax_embed_bf16( |
| logits, codebook, delay=delay, boc=boc, |
| codes=codes, embedding=embedding |
| ) |
|
|
| namespace = ops.ops |
|
|
| def raw(): |
| namespace.delayed_codebook_argmax_embed_bf16( |
| logits, codebook, delay, boc, codes, embedding |
| ) |
|
|
| expected = eager() |
| actual = wrapper() |
| torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) |
| torch.testing.assert_close(actual[1], expected[1], rtol=0, atol=0) |
|
|
| rows = { |
| "torch_eager_us": elapsed_us(eager, args.warmup, args.iterations), |
| "torch_compile_us": elapsed_us(compiled, args.warmup, args.iterations), |
| "hub_wrapper_us": elapsed_us(wrapper, args.warmup, args.iterations), |
| "raw_native_op_us": elapsed_us(raw, args.warmup, args.iterations), |
| } |
| for name, value in rows.items(): |
| print(f"{name}={value:.3f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
