70 lines
2.1 KiB
Bash
Executable File
70 lines
2.1 KiB
Bash
Executable File
#!/bin/bash
|
|
# probe_moe_fused_breakdown.sh — Time each step of MoE decode
|
|
|
|
python3 << 'PY'
|
|
import torch
|
|
import time
|
|
import importlib.util
|
|
|
|
spec = importlib.util.spec_from_file_location('m',
|
|
'qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/corex_batched_gemm.so')
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
|
|
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("=== Step-by-step breakdown ===")
|
|
|
|
# Step 0: expand+contiguous
|
|
time_fn(lambda: x.expand(K, 1, H).contiguous(), "expand+contiguous (1,H)->(K,1,H)")
|
|
|
|
# Step 1: batched GEMM FC1 only
|
|
x_exp = x.expand(K, 1, H).contiguous()
|
|
time_fn(lambda: mod.batched_gemm_fp16(x_exp, w13), "batched_gemm FC1 (K,1,H)x(K,2I,H)")
|
|
|
|
# Step 2: silu * mul
|
|
gate_up = mod.batched_gemm_fp16(x_exp, w13).squeeze(1)
|
|
chunks = gate_up.chunk(2, dim=1)
|
|
time_fn(lambda: torch.sigmoid(chunks[0]) * chunks[0] * chunks[1], "silu*mul (K,I)")
|
|
|
|
# Step 3: batched GEMM FC2 only
|
|
act = (torch.sigmoid(chunks[0]) * chunks[0] * chunks[1]).unsqueeze(1)
|
|
time_fn(lambda: mod.batched_gemm_fp16(act, w2), "batched_gemm FC2 (K,1,I)x(K,H,I)")
|
|
|
|
# Step 4: weighted reduction
|
|
eo = mod.batched_gemm_fp16(act, w2).squeeze(1)
|
|
time_fn(lambda: (eo * ws.unsqueeze(1)).sum(0, True), "weighted_sum")
|
|
|
|
# Full fused
|
|
time_fn(lambda: mod.moe_decode_fused(x, w13, w2, ws), "moe_decode_fused (full)")
|
|
|
|
# Comparison: 8x F.linear loop
|
|
import torch.nn.functional as F
|
|
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 (baseline)")
|
|
PY
|