File size: 1,681 Bytes
2811c56 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import torch
import time
import triton
import triton.language as tl
@triton.jit
def vortex_unroll_kernel(X, Out, N, BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < N
x = tl.load(X + offsets, mask=mask)
# UNROLLED ARTISAN OPS (Sm_90 Register Persistent)
y = x * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
y = y * 1.1 + 0.1
tl.store(Out + offsets, y, mask=mask)
def run_unroll():
N = 1024 * 1024 * 128
print("--- BLITZ VORTEX: THE 10X UNROLL (H200) ---")
X = torch.randn(N, device="cuda", dtype=torch.bfloat16)
Out = torch.empty_like(X)
# 1. Eager Baseline
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
y = X * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1
y = y * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1; y = y * 1.1 + 0.1
torch.cuda.synchronize()
eager_ms = (time.time() - start) / 100 * 1000
# 2. Artisan Unroll
grid = (triton.cdiv(N, 16384),)
torch.cuda.synchronize()
start = time.time()
for _ in range(100): vortex_unroll_kernel[grid](X, Out, N, BLOCK_SIZE=16384)
torch.cuda.synchronize()
vortex_ms = (time.time() - start) / 100 * 1000
print(f"Eager Latency: {eager_ms:.4f}ms")
print(f"Vortex Latency: {vortex_ms:.4f}ms")
print(f"ARTISAN SPEEDUP: {eager_ms/vortex_ms:.2f}x")
if __name__ == "__main__":
run_unroll()
|