cast_to_fp4_benchmarks.py

gistfile1.txt

import gc
import time

import torch
import triton
import triton.language as tl

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


# ─── 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)



# ─── Inplace Triton kernel (no result tensor, reuse x) ─────────────────────

@triton.jit
def _cast_to_fp4_inplace_kernel(
    x_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(x_ptr + offsets, mask=mask, other=0.0)

    sign = tl.where(x < 0.0, -1.0, 1.0)
    x = tl.abs(x)

    x = tl.where(x > 6.0, 12.0, x) # 6   + 7 
    x = tl.where(x <= 0.25, 7.0, x) # 0   + 7
    x = tl.where(x < 0.75, 7.5, x)  # 0.5 + 7
    x = tl.where(x <= 1.25, 8.0, x) # 1   + 7
    x = tl.where(x < 1.75, 8.5, x)  # 1.5 + 7
    x = tl.where(x <= 2.5, 9, x)    # 2   + 7
    x = tl.where(x < 3.5, 10.0, x)  # 3   + 7
    x = tl.where(x <= 5.0, 11.0, x) # 4   + 7
    x = tl.where(x <= 6.0, 12.0, x) # 6   + 7 

    x *= sign
    x -= 7.0

    tl.store(x_ptr + offsets, x, mask=mask)


def inplace_triton_cast_to_fp4(x):
    shape = x.shape
    x = x.flatten()
    n = x.numel()
    block_size = 1024
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)  # noqa: E731
    _cast_to_fp4_inplace_kernel[grid](x, n, BLOCK_SIZE=block_size)
    return x.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)


@triton.jit
def _cast_to_fp4_inplace_v2_kernel(
    x_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(x_ptr + offsets, mask=mask, other=0.0)

    sign = tl.where(x < 0.0, -32.0, 32.0)
    x = tl.abs(x)

    # move all values from 0 to .25 to 0. Do this first to
    # clear up space to store the other rounded values temporarily.
    x = tl.where(x <= 0.25, 0.0, x) 

    # starting with largest bucket, round values to fp4 values divided by 32.
    # this moves each value temporarily into the 0 to .25 range so it won't be 
    # picked up by subsequent threshold checks. 
    x = tl.where(x > 5.0, 6/32, x)
    x = tl.where(x >= 3.5, 4/32, x)
    x = tl.where(x > 2.5, 3/32, x)
    x = tl.where(x >= 1.75, 2/32, x)
    x = tl.where(x > 1.25, 1.5/32, x)
    x = tl.where(x >= .75, 1/32, x)
    x = tl.where(x > .25, .5/32, x)
    # Note: fp4_val/32 is perfectly representable by any float with at least 2 mantissa bits
    #       and 3 exponent bits, so no rounding errors.

    x *= sign #  sign is actually sign(x)*32 so will rescale everything to exact fp4 

    tl.store(x_ptr + offsets, x, mask=mask)


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


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 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")

    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 inplace Triton kernel...")
    t, p = benchmark(
        inplace_triton_cast_to_fp4, test_data, "inplace_triton", warmup=True
    )
    results["Inplace Triton"] = (t, p)
    print(f"  Time: {t * 1000:.2f}ms, Peak: {p:.1f} GB")

    print("\nRunning inplace Triton kernel v2...")
    t, p = benchmark(
        inplace_triton_cast_to_fp4_v2, test_data, "inplace_triton_v2", warmup=True
    )
    results["Inplace Triton v2"] = (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()

================================================================================
SUMMARY
================================================================================
Implementation            Time (ms)       Peak (GB)       Speedup        
--------------------------------------------------------------------------------
Original (compiled)            37.29 ms        33.8 GB   baseline
Direct Triton                  30.98 ms        33.8 GB   1.20x (bf16 cast)
Kyle Triton                    30.98 ms        33.8 GB   1.20x
Inplace Triton                 30.09 ms        16.9 GB   1.24x
Inplace Triton v2              28.78 ms        16.9 GB   1.30x