99 lines
3.1 KiB
Bash
Executable File
99 lines
3.1 KiB
Bash
Executable File
#!/bin/bash
|
|
# probe_moe_fused_breakdown.sh — Time each step of MoE decode
|
|
|
|
python3 << 'PY'
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import time
|
|
|
|
K = 8
|
|
H = 4096
|
|
I = 2752
|
|
|
|
x = torch.randn(1, H, dtype=torch.float16, device='cuda')
|
|
w13 = torch.randn(K, 2*I, H, dtype=torch.float16, device='cuda')
|
|
w2 = torch.randn(K, H, I, dtype=torch.float16, device='cuda')
|
|
ws = torch.softmax(torch.randn(K, device='cuda'), 0).half()
|
|
|
|
def time_fn(fn, name, iters=50):
|
|
for _ in range(5): fn()
|
|
torch.cuda.synchronize()
|
|
t0 = time.time()
|
|
for _ in range(iters): fn()
|
|
torch.cuda.synchronize()
|
|
ms = (time.time() - t0) / iters * 1000
|
|
print(f" {name}: {ms:.3f} ms")
|
|
return ms
|
|
|
|
print("=== Individual operation timings ===")
|
|
|
|
# Single F.linear (one expert FC1)
|
|
time_fn(lambda: F.linear(x, w13[0]), "F.linear FC1 1 expert (1,H)x(2I,H)")
|
|
|
|
# Single F.linear (one expert FC2)
|
|
gu = F.linear(x, w13[0])
|
|
g, u = gu.chunk(2, dim=-1)
|
|
a = F.silu(g) * u
|
|
time_fn(lambda: F.linear(a, w2[0]), "F.linear FC2 1 expert (1,I)x(H,I)")
|
|
|
|
# silu * mul
|
|
time_fn(lambda: F.silu(gu[:,:I]) * gu[:,I:], "silu*mul (1, I)")
|
|
|
|
# 8x F.linear loop (baseline)
|
|
def flinear_loop():
|
|
outs = []
|
|
for i in range(K):
|
|
gu = F.linear(x, w13[i])
|
|
g, u = gu.chunk(2, dim=-1)
|
|
a = F.silu(g) * u
|
|
outs.append(F.linear(a, w2[i]))
|
|
return sum(outs[i] * ws[i] for i in range(K))
|
|
time_fn(flinear_loop, "F.linear loop 8 experts (FULL)")
|
|
|
|
# torch.mm loop (no F.linear overhead)
|
|
def mm_loop():
|
|
outs = []
|
|
for i in range(K):
|
|
gu = torch.mm(x, w13[i].t())
|
|
g, u = gu.chunk(2, dim=-1)
|
|
a = F.silu(g) * u
|
|
outs.append(torch.mm(a, w2[i].t()))
|
|
return sum(outs[i] * ws[i] for i in range(K))
|
|
time_fn(mm_loop, "torch.mm loop 8 experts (FULL)")
|
|
|
|
# Batched via cublasHgemmStridedBatched (pre-gathered weights)
|
|
# First gather weights contiguously
|
|
print("\n=== Batched approaches ===")
|
|
|
|
# Measure gather cost
|
|
time_fn(lambda: w13.reshape(K, 2*I*H), "w13 reshape (view, should be free)")
|
|
|
|
# cublasHgemmStridedBatched via torch.bmm
|
|
x_exp = x.expand(K, 1, H).contiguous()
|
|
w13_t = w13.transpose(1, 2).contiguous() # (K, H, 2I)
|
|
time_fn(lambda: w13.transpose(1, 2).contiguous(), "w13 transpose+contiguous (K,2I,H)->(K,H,2I)")
|
|
time_fn(lambda: torch.bmm(x_exp, w13_t), "torch.bmm FC1 (K,1,H)x(K,H,2I)")
|
|
|
|
# What if weights are pre-transposed?
|
|
print("\n=== Pre-transposed weights (no runtime copy) ===")
|
|
w13_pre = w13.transpose(1, 2).contiguous() # (K, H, 2I) — do this once at model load
|
|
w2_pre = w2.transpose(1, 2).contiguous() # (K, I, H)
|
|
time_fn(lambda: torch.bmm(x_exp, w13_pre), "torch.bmm FC1 pre-transposed")
|
|
|
|
gu = torch.bmm(x_exp, w13_pre).squeeze(1)
|
|
g, u = gu.chunk(2, dim=-1)
|
|
act = F.silu(g) * u
|
|
act_3d = act.unsqueeze(1)
|
|
time_fn(lambda: torch.bmm(act_3d, w2_pre), "torch.bmm FC2 pre-transposed")
|
|
|
|
def bmm_fused_pretransposed():
|
|
gu = torch.bmm(x_exp, w13_pre).squeeze(1)
|
|
g, u = gu.chunk(2, dim=-1)
|
|
a = F.silu(g) * u
|
|
eo = torch.bmm(a.unsqueeze(1), w2_pre).squeeze(1)
|
|
return (eo * ws.unsqueeze(1)).sum(0, True)
|
|
time_fn(bmm_fused_pretransposed, "bmm full MoE (pre-transposed)")
|
|
|
|
print("\n=== Summary ===")
|
|
PY
|