gistfile1.txt

import gc
import time

import torch
import triton
import triton.language as tl

SIZE = 844_000_000
device = "cuda:0" if torch.cuda.is_available() else "cpu"
N_RUNS = 20


# ─── Original implementation (from main) ────────────────────────────────────


@torch.compile
def _original_cast_to_fp4_compiled(x):
    sign = torch.sign(x)
    x = torch.abs(x)
    x[(x >= 0.0) & (x <= 0.25)] = 0.0
    x[(x > 0.25) & (x < 0.75)] = 0.5
    x[(x >= 0.75) & (x <= 1.25)] = 1.0
    x[(x > 1.25) & (x < 1.75)] = 1.5
    x[(x >= 1.75) & (x <= 2.5)] = 2.0
    x[(x > 2.5) & (x < 3.5)] = 3.0
    x[(x >= 3.5) & (x <= 5.0)] = 4.0
    x[x > 5.0] = 6.0
    return x * sign


def original_cast_to_fp4(x):
    return _original_cast_to_fp4_compiled(x.flatten()).reshape(x.shape)



# ─── Current Triton kernel (from branch) ────────────────────────────────────


@triton.jit
def _cast_to_fp4_kernel(
    input_ptr,
    output_ptr,
    n,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n

    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)

    # Extract sign and absolute value
    sign = tl.where(x < 0.0, -1.0, 1.0)
    abs_x = tl.abs(x)

    # Map absolute values to FP4 representable values
    # Using sequential tl.where for the quantization mapping
    result = tl.zeros_like(abs_x)
    result = tl.where(abs_x > 0.25, 0.5, result)
    result = tl.where(abs_x >= 0.75, 1.0, result)
    result = tl.where(abs_x > 1.25, 1.5, result)
    result = tl.where(abs_x >= 1.75, 2.0, result)
    result = tl.where(abs_x > 2.5, 3.0, result)
    result = tl.where(abs_x >= 3.5, 4.0, result)
    result = tl.where(abs_x > 5.0, 6.0, result)

    # Restore sign
    result *= sign
    tl.store(output_ptr + offsets, result, mask=mask)


def current_triton_cast_to_fp4(x):
    shape = x.shape
    x = x.flatten()
    output = torch.empty_like(x)
    n = x.numel()
    block_size = 1024
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)  # noqa: E731
    _cast_to_fp4_kernel[grid](x, output, n, BLOCK_SIZE=block_size)
    return output.reshape(shape)



# ─── Direct threshold counting Triton kernel ────────────────────────────────

@triton.jit
def _cast_to_fp4_direct_kernel(
    input_ptr,
    output_ptr,
    n,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n

    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)

    # Extract sign bit (bit 15 of bfloat16)
    sign_bit = x.to(tl.int16, bitcast=True) & (-32768)
    abs_x = tl.abs(x)

    result = tl.zeros_like(abs_x)
    result = tl.where(abs_x > 0.25, 0.5, result)
    result = tl.where(abs_x >= 0.75, 1.0, result)
    result = tl.where(abs_x > 1.25, 1.5, result)
    result = tl.where(abs_x >= 1.75, 2.0, result)
    result = tl.where(abs_x > 2.5, 3.0, result)
    result = tl.where(abs_x >= 3.5, 4.0, result)
    result = tl.where(abs_x > 5.0, 6.0, result)

    # Apply sign bit via OR (result is always positive)
    result = (result.to(tl.int16, bitcast=True) | sign_bit).to(tl.bfloat16, bitcast=True)
    tl.store(output_ptr + offsets, result, mask=mask)


def direct_triton_cast_to_fp4(x):
    shape = x.shape
    x = x.flatten()
    output = torch.empty_like(x)
    n = x.numel()
    block_size = 1024
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)  # noqa: E731
    _cast_to_fp4_direct_kernel[grid](x, output, n, BLOCK_SIZE=block_size)
    return output.reshape(shape)


# ─── Benchmarking utilities ─────────────────────────────────────────────────


def create_test_data(size, device):
    """Create test data with random values in a reasonable range."""
    x = torch.randn(size, dtype=torch.bfloat16, device=device) * 6.0
    return x.reshape(-1, 2)


def benchmark(func, test_data, name, warmup=False):
    if warmup:
        print(f"  Warming up {name}...")
        warmup_data = create_test_data(1000, device)
        for _ in range(10):
            _ = func(warmup_data.clone())
        del warmup_data
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.synchronize()
        print(f"  Warmup complete, starting benchmark...")

    times = []
    peaks = []

    for _ in range(N_RUNS):
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.reset_peak_memory_stats()

        baseline_mem = torch.cuda.memory_allocated(0)

        torch.cuda.synchronize()
        start = time.time()
        result = func(test_data.clone())
        torch.cuda.synchronize()
        elapsed = time.time() - start

        peak = (torch.cuda.max_memory_allocated(0) - baseline_mem) / 1e9

        times.append(elapsed)
        peaks.append(peak)

        del result
        torch.cuda.empty_cache()
        gc.collect()

    avg_time = sum(times) / N_RUNS
    avg_peak = sum(peaks) / N_RUNS

    return avg_time, avg_peak



def main():
    if not torch.cuda.is_available():
        print("CUDA not available")
        return

    print(f"Benchmarking cast_to_fp4 implementations ({SIZE / 1e6:.1f}M elements)\n")
    print("=" * 80)

    print("Creating test data...")
    test_data = create_test_data(SIZE, device)
    print(f"Test data: {test_data.shape}, {test_data.dtype}\n")

    results = {}

    print("Running original (compiled)...")
    t, p = benchmark(original_cast_to_fp4, test_data, "original_compiled", warmup=True)
    results["Original (compiled)"] = (t, p)
    print(f"  Time: {t * 1000:.2f}ms, Peak: {p:.1f} GB")

    print("\nRunning current Triton kernel...")
    t, p = benchmark(
        current_triton_cast_to_fp4, test_data, "current_triton", warmup=True
    )
    results["Current Triton"] = (t, p)
    print(f"  Time: {t * 1000:.2f}ms, Peak: {p:.1f} GB")

    print("\nRunning direct Triton kernel...")
    t, p = benchmark(
        direct_triton_cast_to_fp4, test_data, "direct_triton", warmup=True
    )
    results["Direct Triton"] = (t, p)
    print(f"  Time: {t * 1000:.2f}ms, Peak: {p:.1f} GB")

    del test_data

    # Summary
    baseline_time = results["Original (compiled)"][0]
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    print(
        f"{'Implementation':<25} {'Time (ms)':<15} {'Peak (GB)':<15} {'Speedup':<15}"
    )
    print("-" * 80)
    for name, (t, p) in results.items():
        speedup = (
            f"{baseline_time / t:.2f}x" if name != "Original (compiled)" else "baseline"
        )
        print(f"{name:<25} {t * 1000:>10.2f} ms  {p:>10.1f} GB   {speedup}")


if __name__ == "__main__":
    main()