| |
| |
| """FFN block micro-benchmark (CUDA): dense SwiGLU vs eugen_sparse (selection + kernel), forward and |
| forward+backward, bf16, optionally under torch.compile. |
| |
| python speed_check.py --tokens 65536 --d-model 512 --d-ff 4096 --group-size 32 \ |
| --active-groups 8 --samples-per-group 8 [--compile] |
| |
| This times one FFN block on random inputs. It is not a prediction of end-to-end training speed: |
| measure your own training loop. Random inputs spread tokens evenly over the groups; a trained model |
| is more skewed (popular groups), which the load-balancing bias counteracts. |
| """ |
| import argparse |
| import statistics |
| import time |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| import eugen_sparse as es |
|
|
|
|
| def timer(fn, warmup=5, iters=20): |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| ts = [] |
| for _ in range(iters): |
| t0 = time.perf_counter() |
| fn() |
| torch.cuda.synchronize() |
| ts.append((time.perf_counter() - t0) * 1000) |
| return statistics.median(ts) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--tokens', type=int, default=65536) |
| ap.add_argument('--d-model', type=int, default=512) |
| ap.add_argument('--d-ff', type=int, default=4096) |
| ap.add_argument('--group-size', type=int, default=32) |
| ap.add_argument('--active-groups', type=int, default=8) |
| ap.add_argument('--samples-per-group', type=int, default=8) |
| ap.add_argument('--compile', action='store_true') |
| a = ap.parse_args() |
| assert torch.cuda.is_available(), 'CUDA required' |
| G, kg = a.group_size, a.active_groups |
| if not es.kernel_supported(True, a.d_model, a.d_ff, G, kg): |
| raise SystemExit('shape not supported by the kernel (see kernel_supported)') |
| dt = torch.bfloat16 |
| torch.manual_seed(1) |
| x = torch.randn(a.tokens, a.d_model, device='cuda', dtype=dt) |
| gw, uw = (torch.randn(a.d_ff, a.d_model, device='cuda', dtype=dt) * a.d_model ** -0.5 for _ in range(2)) |
| dw = torch.randn(a.d_model, a.d_ff, device='cuda', dtype=dt) * a.d_ff ** -0.5 |
| dy = torch.randn_like(x) |
| sub, spg = es.make_subsample_index(a.d_ff, G, a.samples_per_group) |
| sub = sub.cuda() |
|
|
| def dense(x_, g_, u_, d_): |
| return F.linear(F.silu(F.linear(x_, g_)) * F.linear(x_, u_), d_) |
|
|
| def sparse(x_, g_, u_, d_): |
| idx = es.grouped_topk_select(x_, g_, G, kg, sub, spg) |
| return es.sparse_swiglu_ffn(x_, g_, u_, d_, idx, G) |
|
|
| print(f'tokens={a.tokens} d_model={a.d_model} d_ff={a.d_ff} group_size={G} active_groups={kg} ' |
| f'(density {kg * G / a.d_ff:.2%}) compile={a.compile} | {torch.cuda.get_device_name(0)} | ' |
| f'torch {torch.__version__}') |
| print('plan:', es.plan(a.tokens, a.d_model, a.d_ff, G, kg)) |
| res = {} |
| for name, fn in (('dense', dense), ('sparse', sparse)): |
| if a.compile: |
| fn = torch.compile(fn, fullgraph=True, dynamic=False) |
| t_f = timer(lambda: fn(x, gw, uw, dw)) |
| leaves = tuple(t.clone().requires_grad_() for t in (x, gw, uw, dw)) |
|
|
| def fb(fn=fn, l=leaves): |
| torch.autograd.grad(fn(*l), l, dy) |
| t_fb = timer(fb) |
| res[name] = (t_f, t_fb) |
| print(f' {name:6s} fwd {t_f:8.2f} ms fwd+bwd {t_fb:8.2f} ms') |
| print(f' sparse vs dense: fwd x{res["dense"][0] / res["sparse"][0]:.2f} ' |
| f'fwd+bwd x{res["dense"][1] / res["sparse"][1]:.2f} (>1 = sparse faster)') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
