reference.py

import torch
from task import input_t, output_t
from utils import make_match_reference
from laguna_dual_gemm import custom_kernel
import numpy as np
from quack.gemm_interface import gemm_gated
import torch.nn.functional as F

# Scaling factor vector size
sf_vec_size = 16

# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b

# Helper function to convert scale factor tensor to blocked format
def to_blocked(input_matrix):
    rows, cols = input_matrix.shape

    # Please ensure rows and cols are multiples of 128 and 4 respectively
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)

    padded = input_matrix
    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)

    return rearranged.flatten()

@torch.compile
def ref_kernel(
    data: input_t,
) -> output_t:
    """
    PyTorch reference implementation of NVFP4 block-scaled dual GEMM with silu activation,
    C = silu(A @ B1) * (A @ B2).
    """
    a_ref, b1_ref, b2_ref, sfa_ref_cpu, sfb1_ref_cpu, sfb2_ref_cpu, _, _, _, c_ref = data
    
    # Get dimensions from MxNxL layout
    m, n, l = c_ref.shape

    # Call torch._scaled_mm to compute the GEMV result
    ref1 = torch.empty(
        (l, m, n),
        dtype=torch.float32,
        device="cuda",
    ).permute(1, 2, 0)
    ref2 = torch.empty(
        (l, m, n),
        dtype=torch.float32,
        device="cuda",
    ).permute(1, 2, 0)
    for l_idx in range(l):
        # Convert the scale factor tensor to blocked format
        scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
        scale_b1 = to_blocked(sfb1_ref_cpu[:, :, l_idx])
        scale_b2 = to_blocked(sfb2_ref_cpu[:, :, l_idx])
        # (m, k) @ (n, k).T -> (m, n)
        res1 = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b1_ref[:, :, l_idx].transpose(0, 1),
            scale_a.cuda(),
            scale_b1.cuda(),
            bias=None,
            out_dtype=torch.float32,
        )
        ref1[:, :, l_idx] = res1

        res2 = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b2_ref[:, :, l_idx].transpose(0, 1),
            scale_a.cuda(),
            scale_b2.cuda(),
            bias=None,
            out_dtype=torch.float32,
        )
        ref2[:, :, l_idx] = res2
    # Do silu on the first GEMM result and multiply with the second GEMM result
    c_ref = (torch.nn.functional.silu(ref1) * ref2).to(torch.float16)
    return c_ref


def generate_input(
    m: int,
    n: int,
    k: int,
    l: int,
    seed: int,
):
    """
    Generate input tensors for NVFP4 block-scaled dual GEMM with silu activation,
    C = silu(A @ B1) * (A @ B2).
    
    Args:
        m: Number of rows in matrix A
        n: Number of columns in matrix B1 and B2
        k: Number of columns in A and rows of B1 and B2
        l: Batch size
        seed: Random seed for reproducibility
    
    Returns:
        Tuple of (a, b, scale_a, scale_b, c) where:
            a: [m, k, l] - Input matrix in torch.float4e2m1fn_x2 data type
            b1: [n, k, l] - Input matrix in torch.float4e2m1fn_x2 data type
            b2: [n, k, l] - Input matrix in torch.float4e2m1fn_x2 data type
            scale_a: [m, k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b1: [n, k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b2: [n, k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_a_permuted: [32, 4, rest_m, 4, rest_k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b1_permuted: [32, 4, rest_n, 4, rest_k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b2_permuted: [32, 4, rest_n, 4, rest_k, l] - Input scale factors in torch.float8e4m3fn data type
            c: [m, n, l] - Output matrix in torch.float16 data type
    """
    torch.manual_seed(seed)
    
    def create_fp4_tensors(l, mn, k):
        # generate uint8 tensor, then convert to float4e2m1fn_x2 data type
        # generate all bit patterns
        ref_i8 = torch.randint(255, size=(l, mn, k // 2), dtype=torch.uint8, device="cuda")

        # for each nibble, only keep the sign bit and 2 LSBs
        # the possible values are [-1.5, -1, -0.5, 0, +0.5, +1, +1.5]
        ref_i8 = ref_i8 & 0b1011_1011

        return ref_i8.permute(1, 2, 0).view(torch.float4_e2m1fn_x2)

    # Generate uint8 tensor, then convert to float4e2m1fn_x2 data type
    a_ref = create_fp4_tensors(l, m, k)
    b1_ref = create_fp4_tensors(l, n, k)
    b2_ref = create_fp4_tensors(l, n, k)
    a_ref = a_ref.view(torch.float4_e2m1fn_x2)
    b1_ref = b1_ref.view(torch.float4_e2m1fn_x2)
    b2_ref = b2_ref.view(torch.float4_e2m1fn_x2)

    # Create float16 output tensor
    c_ref = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute(
        1, 2, 0
    )
    
    # Helper function to prepare the scale factor tensors for both reference
    # kernel and customize kernel. The customized data layout can be found in:
    # https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout
    def create_scale_factor_tensors(l, mn, sf_k):
        # Create the reference scale factor tensor (mn, sf_k, l) on CPU.
        ref_shape = (l, mn, sf_k)
        ref_permute_order = (1, 2, 0)
        # Init with fp32 tensor in [0,1), then convert to float8_e4m3fn
        ref_f8_random_fp32 = torch.rand(ref_shape, dtype=torch.float32, device='cuda')
        ref_f8_torch_tensor = ref_f8_random_fp32.to(dtype=torch.float8_e4m3fn)
        # permute to match ref_permute_order
        ref_f8_torch_tensor_permuted = ref_f8_torch_tensor.permute(*ref_permute_order)

        atom_m = (32, 4)
        atom_k = 4
        mma_shape = (
            l,  # batch size
            ceil_div(mn, atom_m[0] * atom_m[1]),
            ceil_div(sf_k, atom_k),
            atom_m[0],
            atom_m[1],
            atom_k,
        )

        # Reorder scale factor tensor to (32, 4, rest_m, 4, rest_k, l) layout
        # Which is needed by the CuTe customized kernel
        mma_permute_order = (3, 4, 1, 5, 2, 0)
        # Generate a random int8 tensor, then convert to float8_e4m3fn
        rand_int_tensor = torch.empty(mma_shape, dtype=torch.int8, device='cuda')
        reordered_f8_torch_tensor = rand_int_tensor.to(dtype=torch.float8_e4m3fn)
        # Permute according to mma_permute_order
        reordered_f8_torch_tensor = reordered_f8_torch_tensor.permute(*mma_permute_order)

        # GPU-side vectorized reordering (replaces slow CPU nested loops)
        # Create index grids for all dimensions
        i_idx = torch.arange(mn, device='cuda')
        j_idx = torch.arange(sf_k, device='cuda')
        b_idx = torch.arange(l, device='cuda')
        
        # Create meshgrid for all combinations of (i, j, b)
        i_grid, j_grid, b_grid = torch.meshgrid(i_idx, j_idx, b_idx, indexing='ij')
        
        # Calculate target indices in vectorized manner
        mm = i_grid // (atom_m[0] * atom_m[1])
        mm32 = i_grid % atom_m[0]
        mm4 = (i_grid % 128) // atom_m[0]
        kk = j_grid // atom_k
        kk4 = j_grid % atom_k
        
        # Perform the reordering with advanced indexing (all on GPU)
        reordered_f8_torch_tensor[mm32, mm4, mm, kk4, kk, b_grid] = ref_f8_torch_tensor_permuted[i_grid, j_grid, b_grid]
        
        return ref_f8_torch_tensor_permuted.cpu(), reordered_f8_torch_tensor

    sf_k = ceil_div(k, sf_vec_size)
    sfa_ref_cpu, sfa_ref_permuted = create_scale_factor_tensors(l, m, sf_k)
    sfb1_ref_cpu, sfb1_ref_permuted = create_scale_factor_tensors(l, n, sf_k)
    sfb2_ref_cpu, sfb2_ref_permuted = create_scale_factor_tensors(l, n, sf_k)

    return (a_ref, b1_ref, b2_ref, sfa_ref_cpu.to("cuda"), sfb1_ref_cpu.to("cuda"), sfb2_ref_cpu.to("cuda"), sfa_ref_permuted, sfb1_ref_permuted, sfb2_ref_permuted, c_ref)

def run_comprehensive_benchmark(custom_kernel, reference_fn, data, sonic_moe_layer, warmups=20, reps=100):
    # Profile Custom Kernel
    for _ in range(warmups):
        _ = custom_kernel(data)
    torch.cuda.synchronize()
    
    c_start = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    c_end = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        c_start[i].record()
        _ = custom_kernel(data)
        c_end[i].record()
    torch.cuda.synchronize()
    custom_times = [s.elapsed_time(e) for s, e in zip(c_start, c_end)]
    custom_ms = np.mean(custom_times)
    
    # Profile PyTorch Reference Path
    for _ in range(warmups):
        _ = reference_fn(data)
    torch.cuda.synchronize()
    
    ref_start = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    ref_end = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
    for i in range(reps):
        ref_start[i].record()
        _ = reference_fn(data)
        ref_end[i].record()
    torch.cuda.synchronize()
    ref_times = [s.elapsed_time(e) for s, e in zip(ref_start, ref_end)]
    ref_ms = np.mean(ref_times)
    # Profile SonicMoE Path
    # SonicMoE expects a flat token tensor [M, Hidden_Size] where M=4096
    M, N, K = 4096, 512, 2048
    sonic_x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)

    with torch.no_grad():
        for _ in range(warmups):
            _, _ = sonic_moe_layer(sonic_x, kernel_backend_moe=KernelBackendMoE.sonicmoe)
        torch.cuda.synchronize()
        
        sonic_start = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
        sonic_end = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
        for i in range(reps):
            sonic_start[i].record()
            _, _ = sonic_moe_layer(sonic_x, kernel_backend_moe=KernelBackendMoE.sonicmoe)
            sonic_end[i].record()
        torch.cuda.synchronize()
    sonic_times = [s.elapsed_time(e) for s, e in zip(sonic_start, sonic_end)]
    sonic_ms = np.mean(sonic_times)
    
    # Compute TFLOPs (4 * M * N * K ops total)
    M, N, K = 4096, 512, 2048
    flops = 4 * M * N * K
    sonic_flops = 6 * M * N * K   # Full layer (Up-projection + Down-projection)
    custom_tflops = (flops / (custom_ms / 1000.0)) / 1e12
    ref_tflops = (flops / (ref_ms / 1000.0)) / 1e12
    sonic_tflops = (sonic_flops / (sonic_ms / 1000.0)) / 1e12
    
    print(f"\n📊 B200 Performance Comparison (M={M}, N={N}, K={K}):")
    print(f"  {'Implementation':<20} | {'Latency':<12} | {'Throughput':<15}")
    print(f"  {'-'*20}-+-{'-'*12}-+-{'-'*15}")
    print(f"  {'PyTorch Reference':<20} | {ref_ms:8.4f} ms | {ref_tflops:10.2f} TFLOPs")
    print(f"  {'SonicMoE Kernel':<20} | {sonic_ms:8.4f} ms | {sonic_tflops:10.2f} TFLOPs")
    print(f"  {'Custom CuTe Kernel':<20} | {custom_ms:8.4f} ms | {custom_tflops:10.2f} TFLOPs")
    print(f"  ⚡ Speedup vs SonicMoE: {sonic_ms / custom_ms:.2f}x")

def benchmark_fair_fused_swiglu(custom_kernel, ref_kernel, data, warmups=20, reps=100):
    # Extract dimensions dynamically from the provided data tuple.
    M = data[-1].shape[0]
    N = data[-1].shape[1]
    # The float4e2m1fn_x2 datatype packs 2 elements per byte
    K = data[0].shape[1] * 2

    print(f"Setting up Fair Fused SwiGLU Benchmark (M={M}, N={N}, K={K})...")
    
    # --- 1. Setup Data for BF16 Baselines ---
    x_bf16 = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
    w_concat_bf16 = torch.randn(K, N * 2, device="cuda", dtype=torch.bfloat16)
    out_bf16 = torch.empty(M, N, device="cuda", dtype=torch.bfloat16)

    # --- 2. Setup PyTorch Compiled (BF16) ---
    def pytorch_swiglu(x, w):
        y = torch.matmul(x, w)
        # Match QuACK's exact interleaving to keep the memory access pattern identical
        gate = y[..., ::2]
        up = y[..., 1::2]
        return F.silu(gate) * up

    print("Triggering torch.compile (this will take a minute for max-autotune)...")
    compiled_pytorch_swiglu = torch.compile(pytorch_swiglu, mode="max-autotune")
    # Force compilation to finish before the timing loop
    _ = compiled_pytorch_swiglu(x_bf16, w_concat_bf16)

    def pt_compiled_bf16():
        _ = compiled_pytorch_swiglu(x_bf16, w_concat_bf16)

    # --- 3. Setup QuACK Fused Gated GEMM (BF16) ---
    def quack_fused_bf16():
        _ = gemm_gated(
            A=x_bf16, 
            B=w_concat_bf16, 
            activation="swiglu", 
            postact_out=out_bf16,
            store_preact=False
        )

    # --- 4. Setup Custom Fused Dual GEMM (NVFP4) ---
    def custom_fp4_fused():
        _ = custom_kernel(data)

    # --- Timing Helper ---
    def time_fn(fn):
        for _ in range(warmups): fn()
        torch.cuda.synchronize()
        start = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
        end = [torch.cuda.Event(enable_timing=True) for _ in range(reps)]
        for i in range(reps):
            start[i].record()
            fn()
            end[i].record()
        torch.cuda.synchronize()
        return np.mean([s.elapsed_time(e) for s, e in zip(start, end)])

    # --- Execute and Measure ---
    print("\nBenchmarking PyTorch Compiled (BF16)...")
    pt_ms = time_fn(pt_compiled_bf16)

    print("Benchmarking QuACK Fused Gated GEMM (BF16)...")
    quack_ms = time_fn(quack_fused_bf16)
    
    print("Benchmarking Custom Fused Dual GEMM (NVFP4)...")
    custom_ms = time_fn(custom_fp4_fused)
    
    # Base floating-point operations: 2 * (A*Gate) + 2 * (A*Up) = 4 * M * N * K
    flops = 4 * M * N * K
    pt_tflops = (flops / (pt_ms / 1000.0)) / 1e12
    quack_tflops = (flops / (quack_ms / 1000.0)) / 1e12
    custom_tflops = (flops / (custom_ms / 1000.0)) / 1e12

    # --- Output ---
    print(f"\n📊 Fair Fused Compute Comparison (M={M}, N={N}, K={K}):")
    print(f"  {'Implementation':<25} | {'Latency':<10} | {'Throughput':<15}")
    print(f"  {'-'*25}-+-{'-'*10}-+-{'-'*15}")
    print(f"  {'PyTorch Compiled (BF16)':<25} | {pt_ms:7.4f} ms | {pt_tflops:10.2f} TFLOPs")
    print(f"  {'QuACK gemm_gated (BF16)':<25} | {quack_ms:7.4f} ms | {quack_tflops:10.2f} TFLOPs")
    print(f"  {'Custom Fused (NVFP4)':<25} | {custom_ms:7.4f} ms | {custom_tflops:10.2f} TFLOPs")
    print(f"\n⚡ Hardware Speedup (NVFP4 over PT BF16): {custom_tflops / pt_tflops:.2f}x")

check_implementation = make_match_reference(ref_kernel, rtol=1e-03, atol=1e-03)

if __name__ == "__main__":
    # Laguna XS.2 shapes: Hidden (K) = 2048, Intermediate (N) = 512
    # M = Sequence length / tokens processed. Shared expert processes everything.

    print("---------------Laguna_XS.2-----------------")
    m, n, k, l = 4096, 512, 2048, 1
    print(f"Generating NVFP4 inputs for Laguna XS.2 Shared Expert (M={m}, N={n}, K={k})...")
    data = generate_input(m, n, k, l, seed=42)
    
    print("Executing CuTe DualGEMM kernel (A*B1, A*B2) + SiLU fusion...")
    c_out = custom_kernel(data)
    
    print("Validating against PyTorch reference block-scaled GEMM...")
    check_implementation(data, c_out)
    print("Passed")

    print("Running production-grade benchmark...")
    benchmark_fair_fused_swiglu(custom_kernel, ref_kernel, data)

    print("---------------Laguna_XM.1-----------------")
    m, n, k, l = 4096, 512*8, 2048*8, 1
    print(f"Generating NVFP4 inputs for Laguna XM.1 Shared Expert (M={m}, N={n}, K={k})...")
    data = generate_input(m, n, k, l, seed=42)
    
    print("Executing CuTe DualGEMM kernel (A*B1, A*B2) + SiLU fusion...")
    c_out = custom_kernel(data)
    
    print("Validating against PyTorch reference block-scaled GEMM...")
    check_implementation(data, c_out)
    print("Passed")

    print("Running production-grade benchmark...")
    benchmark_fair_fused_swiglu(custom_kernel, ref_kernel, data)