benchmarks/vortex_monolith.py

import torch
import time
import triton
import triton.language as tl

@triton.jit
def vortex_monolith_kernel(
    X, Out, N, BLOCK_SIZE: tl.constexpr
):
    # PERSISTENT RECURENCE (The 10x Path)
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    x = tl.load(X + offsets, mask=mask)
    
    # Simulate a heavy 10-pass SSM Recurrence in Registers
    state = x
    for _ in range(10):
        state = tl.exp(state * 0.5) - 1.0
        state = tl.log(tl.abs(state) + 1.0)
        state = state * 1.1
    
    tl.store(Out + offsets, state, mask=mask)

def run_monolith():
    N = 1024 * 1024 * 64
    print("--- BLITZ VORTEX: THE 10X MONOLITH (H200) ---")
    X = torch.randn(N, device="cuda")
    Out = torch.empty_like(X)
    
    # 1. Eager (Standard "Crap" Path)
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(10):
        s = X
        for j in range(10):
            s = torch.exp(s * 0.5) - 1.0
            s = torch.log(torch.abs(s) + 1.0)
            s = s * 1.1
    torch.cuda.synchronize()
    eager_ms = (time.time() - start) / 10 * 1000
    
    # 2. Blitz Monolith (Artisan Path)
    grid = (triton.cdiv(N, 16384),)
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(10): vortex_monolith_kernel[grid](X, Out, N, BLOCK_SIZE=16384)
    torch.cuda.synchronize()
    vortex_ms = (time.time() - start) / 10 * 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_monolith()