55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
import sys, os, torch, time
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from cutedsl.moe_pipeline import run_nvfp4_moe, run_nvfp4_moe_fused, quantize_weight
|
|
|
|
torch.manual_seed(42)
|
|
device = "cuda"
|
|
num_experts = 3
|
|
hidden = 7168
|
|
intermediate = 3072
|
|
n_tokens = 128
|
|
|
|
l1_weights = [torch.randn(2*intermediate, hidden, dtype=torch.bfloat16, device=device) for _ in range(num_experts)]
|
|
l2_weights = [torch.randn(hidden, intermediate, dtype=torch.bfloat16, device=device) for _ in range(num_experts)]
|
|
|
|
l1_fp4, l1_sf, l1_gs = [], [], []
|
|
l2_fp4, l2_sf, l2_gs = [], [], []
|
|
for l1_w, l2_w in zip(l1_weights, l2_weights):
|
|
fp4, sf, gs = quantize_weight(l1_w)
|
|
l1_fp4.append(fp4); l1_sf.append(sf); l1_gs.append(gs)
|
|
fp4, sf, gs = quantize_weight(l2_w)
|
|
l2_fp4.append(fp4); l2_sf.append(sf); l2_gs.append(gs)
|
|
|
|
weights = {
|
|
"l1_fp4": l1_fp4, "l1_sf": l1_sf, "l1_gs": l1_gs,
|
|
"l2_fp4": l2_fp4, "l2_sf": l2_sf, "l2_gs": l2_gs,
|
|
}
|
|
|
|
hidden_states = torch.randn(n_tokens, hidden, dtype=torch.bfloat16, device=device) * 0.1
|
|
expert_ids = torch.zeros(n_tokens, 1, dtype=torch.int32, device=device)
|
|
expert_weights = torch.ones(n_tokens, 1, dtype=torch.float32, device=device)
|
|
expert_indices = [0, 1, 2]
|
|
|
|
# Warmup
|
|
_ = run_nvfp4_moe(hidden_states, expert_ids, expert_weights, weights, expert_indices)
|
|
_ = run_nvfp4_moe_fused(hidden_states, expert_ids, expert_weights, weights, expert_indices)
|
|
torch.cuda.synchronize()
|
|
|
|
# Benchmark
|
|
N = 50
|
|
t0 = time.perf_counter()
|
|
for _ in range(N):
|
|
out1 = run_nvfp4_moe(hidden_states, expert_ids, expert_weights, weights, expert_indices)
|
|
torch.cuda.synchronize()
|
|
t1 = time.perf_counter()
|
|
|
|
for _ in range(N):
|
|
out2 = run_nvfp4_moe_fused(hidden_states, expert_ids, expert_weights, weights, expert_indices)
|
|
torch.cuda.synchronize()
|
|
t2 = time.perf_counter()
|
|
|
|
print(f"Non-fused: {(t1-t0)/N*1000:.2f} ms/iter")
|
|
print(f"Fused: {(t2-t1)/N*1000:.2f} ms/iter")
|
|
print(f"Speedup: {(t1-t0)/(t2-t1):.2f}x")
|
|
print(f"Output match: {torch.allclose(out1.float(), out2.float(), atol=1.0)}")
|