| |
| """RTX benchmark for grouped W4A4 expert projection launch aggregation.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| PACKAGE = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PACKAGE / "tests")) |
| from test_grouped_moe_gemv import load_source_ops, sfb_bytes |
|
|
|
|
| def elapsed_us(fn, warmup: int = 20, repeats: int = 100) -> float: |
| 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(repeats): |
| fn() |
| end.record() |
| end.synchronize() |
| return start.elapsed_time(end) * 1000.0 / repeats |
|
|
|
|
| def load_ops(backend: str, artifact: str | None): |
| if backend == "source": |
| return load_source_ops() |
| if artifact: |
| sys.path.insert(0, artifact) |
| return importlib.import_module("grouped_moe_gemv") |
|
|
|
|
| def run_case(ops, label: str, m: int, top_k: int, n: int, k: int) -> dict: |
| experts = 8 |
| x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) * 0.2 |
| packed = torch.empty((m, k // 2), device="cuda", dtype=torch.uint8) |
| sfa = torch.empty((sfb_bytes(m, k),), device="cuda", dtype=torch.uint8) |
| weights = torch.full( |
| (experts, n, k // 2), 0x11, device="cuda", dtype=torch.uint8 |
| ) |
| sfb = torch.full( |
| (experts, sfb_bytes(n, k)), 0x38, device="cuda", dtype=torch.uint8 |
| ) |
| alpha = torch.ones((experts,), device="cuda", dtype=torch.float32) |
| idx = ( |
| torch.arange(m * top_k, device="cuda", dtype=torch.int32) |
| .reshape(m, top_k) |
| .remainder(experts) |
| .contiguous() |
| ) |
| out = torch.empty((m, top_k, n), device="cuda", dtype=torch.bfloat16) |
| routed_x = x[:, None, :].expand(m, top_k, k).reshape(m * top_k, k).contiguous() |
| w4a16_out = torch.empty((m * top_k, n), device="cuda", dtype=torch.bfloat16) |
|
|
| def grouped_region(): |
| ops.quantize_activations_nvfp4_bf16(x, packed=packed, sfa=sfa) |
| ops.grouped_w4a4_gemv_bf16( |
| packed, weights, sfa, sfb, alpha, idx, out=out |
| ) |
|
|
| def grouped_kernel(): |
| ops.grouped_w4a4_gemv_bf16( |
| packed, weights, sfa, sfb, alpha, idx, out=out |
| ) |
|
|
| route_out = torch.empty((m, 1, n), device="cuda", dtype=torch.bfloat16) |
| route_idx = [idx[:, route : route + 1].contiguous() for route in range(top_k)] |
|
|
| def route_loop(): |
| ops.quantize_activations_nvfp4_bf16(x, packed=packed, sfa=sfa) |
| for route in range(top_k): |
| ops.grouped_w4a4_gemv_bf16( |
| packed, |
| weights, |
| sfa, |
| sfb, |
| alpha, |
| route_idx[route], |
| out=route_out, |
| ) |
|
|
| pair_packed = torch.empty((1, k // 2), device="cuda", dtype=torch.uint8) |
| pair_sfa = torch.empty((sfb_bytes(1, k),), device="cuda", dtype=torch.uint8) |
| pair_out = torch.empty((1, 1, n), device="cuda", dtype=torch.bfloat16) |
| pair_idx = [ |
| idx[token : token + 1, route : route + 1].contiguous() |
| for token in range(m) |
| for route in range(top_k) |
| ] |
|
|
| def pair_loop(): |
| pair = 0 |
| for token in range(m): |
| for _route in range(top_k): |
| ops.quantize_activations_nvfp4_bf16( |
| x[token : token + 1], packed=pair_packed, sfa=pair_sfa |
| ) |
| ops.grouped_w4a4_gemv_bf16( |
| pair_packed, |
| weights, |
| pair_sfa, |
| sfb, |
| alpha, |
| pair_idx[pair], |
| out=pair_out, |
| ) |
| pair += 1 |
|
|
| grouped_us = elapsed_us(grouped_region) |
| grouped_kernel_us = elapsed_us(grouped_kernel) |
| def w4a16_region(): |
| ops.grouped_w4a16_gemv_bf16( |
| routed_x, |
| weights, |
| sfb, |
| alpha, |
| idx.reshape(-1), |
| w_stride=n * k // 2, |
| sfb_stride=sfb.shape[1], |
| n=n, |
| out=w4a16_out, |
| ) |
|
|
| w4a16_us = elapsed_us(w4a16_region) |
| route_us = elapsed_us(route_loop) if top_k > 1 else grouped_us |
| pair_repeats = 30 if m * top_k >= 32 else 100 |
| pair_us = elapsed_us(pair_loop, warmup=5, repeats=pair_repeats) |
| return { |
| "label": label, |
| "M": m, |
| "top_k": top_k, |
| "pairs": m * top_k, |
| "N": n, |
| "K": k, |
| "grouped_us": grouped_us, |
| "grouped_kernel_us": grouped_kernel_us, |
| "route_loop_us": route_us, |
| "pair_loop_us": pair_us, |
| "w4a16_us": w4a16_us, |
| "vs_route_loop": route_us / grouped_us, |
| "vs_pair_loop": pair_us / grouped_us, |
| "vs_w4a16": w4a16_us / grouped_us, |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=["source", "installed"], default="source") |
| parser.add_argument("--artifact") |
| args = parser.parse_args() |
| ops = load_ops(args.backend, args.artifact) |
| cases = [ |
| ("gate_up_decode", 1, 8, 1024, 2048), |
| ("gate_up_verify", 7, 8, 1024, 2048), |
| ("down_decode", 8, 1, 2048, 512), |
| ("down_verify", 56, 1, 2048, 512), |
| ] |
| print("label,M,top_k,pairs,N,K,w4a4_region_us,w4a4_kernel_us,w4a16_us,route_loop_us,pair_loop_us,vs_w4a16,vs_route,vs_pair") |
| for case in cases: |
| result = run_case(ops, *case) |
| print( |
| f"{result['label']},{result['M']},{result['top_k']},{result['pairs']}," |
| f"{result['N']},{result['K']},{result['grouped_us']:.3f}," |
| f"{result['grouped_kernel_us']:.3f},{result['w4a16_us']:.3f}," |
| f"{result['route_loop_us']:.3f}," |
| f"{result['pair_loop_us']:.3f},{result['vs_w4a16']:.2f}x," |
| f"{result['vs_route_loop']:.2f}x,{result['vs_pair_loop']:.2f}x" |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
