bench_flashinfer_2.py

import torch
import statistics
from task import input_t, output_t
from utils import make_match_reference

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


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)

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


# ---------------------------------------------------------------------------
# vLLM / Flashinfer baseline comparison for B200 Dual GEMM benchmarking
# ---------------------------------------------------------------------------

def _benchmark_fn(fn, warmup=10, iters=100):
    """Benchmark a callable using CUDA events. Returns median time in ms."""
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    times = []
    for _ in range(iters):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        fn()
        end.record()
        torch.cuda.synchronize()
        times.append(start.elapsed_time(end))
    return statistics.median(times)

def _prepare_vllm_dual_gemm_silu(a_fp4, b1_fp4, b2_fp4, sfa, sfb1, sfb2, m, n, k, l):
    """
    Direct vLLM CUTLASS baseline. 
    vLLM doesn't have a single 'dual gemm' Python API for FP4. 
    They literally launch two scaled_mm ops and run pointwise activations.
    """
    try:
        # FP4 block-scaled GEMM is cutlass_scaled_fp4_mm, NOT cutlass_scaled_mm
        # (the latter is the int8/fp8 path and rejects fp4 inputs -- that was the
        # source of the empty "()" failure).
        from vllm._custom_ops import cutlass_scaled_fp4_mm
    except (ImportError, AttributeError):
        return None

    # Use vLLM's OWN NVFP4 block-scale swizzle. cutlass_scaled_fp4_mm reads the
    # scales in vLLM's interleaved 128x4 layout -- which is NOT the same as the
    # torchao to_blocked layout that torch._scaled_mm wants (to_blocked adds an
    # extra inner 32x4 transpose). Hand-rolling it ran but fed mislaid scales;
    # swizzle_blockscale is the canonical vLLM helper, so the baseline is a
    # faithful NVFP4 path. Probe the known module locations across versions.
    swizzle_blockscale = None
    for _mod in ("vllm.model_executor.layers.quantization.utils.nvfp4_utils",
                 "vllm.model_executor.layers.quantization.utils.quant_utils"):
        try:
            import importlib
            swizzle_blockscale = getattr(importlib.import_module(_mod),
                                         "swizzle_blockscale")
            break
        except Exception:
            continue
    if swizzle_blockscale is None:
        raise ImportError(
            "vllm.cutlass_scaled_fp4_mm is present but swizzle_blockscale "
            "(NVFP4 block-scale swizzle) was not found -- cannot lay out scales "
            "correctly for the vLLM NVFP4 path."
        )

    # --- one-time preprocessing (NOT timed) ---
    # sfa/sfb are (mn, sf_k) float8_e4m3fn; swizzle_blockscale returns the same
    # logical shape in vLLM's interleaved layout, exactly what the kernel reads.
    scale_a = [swizzle_blockscale(sfa[:, :, i].contiguous()) for i in range(l)]
    scale_b1 = [swizzle_blockscale(sfb1[:, :, i].contiguous()) for i in range(l)]
    scale_b2 = [swizzle_blockscale(sfb2[:, :, i].contiguous()) for i in range(l)]

    # cutlass_scaled_fp4_mm computes a @ b.T with a:(m,k), b:(n,k) -- the weight
    # stays in (n, k) row-major (no transpose), matching res[m,n]=sum_k a[m,k]*b[n,k].
    # The op takes the packed FP4 as uint8 (2 nibbles/byte); passing the typed
    # torch.float4_e2m1fn_x2 trips the stable-ABI "ScalarType 45 not supported".
    a_slices = [a_fp4[:, :, i].contiguous().view(torch.uint8) for i in range(l)]
    b1_slices = [b1_fp4[:, :, i].contiguous().view(torch.uint8) for i in range(l)]
    b2_slices = [b2_fp4[:, :, i].contiguous().view(torch.uint8) for i in range(l)]

    # No global per-tensor scale in this data, so alpha = 1.0.
    alpha = torch.tensor(1.0, dtype=torch.float32, device="cuda")

    out = torch.empty((l, m, n), dtype=torch.float16, device="cuda").permute(1, 2, 0)

    def run():
        # Keep the timed region lean and fair: two FP4 GEMMs straight to fp16,
        # SiLU * mul fused in fp16. The previous version upcast both products to
        # fp32 and did an extra copy_, which at these tiny latency-bound shapes
        # cost more than the GEMMs themselves and inflated the speedup.
        for l_idx in range(l):
            r1 = cutlass_scaled_fp4_mm(
                a_slices[l_idx], b1_slices[l_idx],
                scale_a[l_idx], scale_b1[l_idx], alpha, torch.float16,
            )
            r2 = cutlass_scaled_fp4_mm(
                a_slices[l_idx], b2_slices[l_idx],
                scale_a[l_idx], scale_b2[l_idx], alpha, torch.float16,
            )
            out[:, :, l_idx] = torch.nn.functional.silu(r1) * r2
        return out

    return run

def _prepare_flashinfer_dual_gemm_silu(a_fp4, b1_fp4, b2_fp4, sfa, sfb1, sfb2, m, n, k, l):
    """
    Flashinfer baseline: uses flashinfer.gemm.bmm_fp4 if it is actually
    importable. The real symbol lives at flashinfer.gemm.bmm_fp4 (NOT
    flashinfer.bmm_fp4), so we probe for it correctly.

    If flashinfer (or its FP4 GEMM) is not available we return None so the
    caller can mark the row as skipped -- we deliberately do NOT silently fall
    back to the cuBLAS path, because reporting an identical cuBLAS run under a
    "Flashinfer" label is misleading (it was the cause of the two baselines
    being within noise of each other).

    Returns a zero-arg closure that performs only the timed compute, or None if
    flashinfer's FP4 GEMM is unavailable.
    """
    import importlib
    if importlib.util.find_spec("flashinfer") is None:
        # flashinfer genuinely not installed -> clean skip.
        return None
    # flashinfer IS present. mm_fp4 has lived at a few locations across versions
    # (top-level re-export, flashinfer.gemm, flashinfer.gemm.gemm_base). Probe
    # them in order; if none has it, let the real ImportError propagate so the
    # benchmark table shows the actual reason instead of a generic "unavailable".
    mm_fp4 = None
    _last_err = None
    for _modname, _attr in (("flashinfer", "mm_fp4"),
                            ("flashinfer.gemm", "mm_fp4"),
                            ("flashinfer.gemm.gemm_base", "mm_fp4")):
        try:
            mm_fp4 = getattr(importlib.import_module(_modname), _attr)
            break
        except Exception as _e:  # noqa: BLE001 -- record and try next location
            _last_err = _e
    if mm_fp4 is None:
        raise ImportError(f"flashinfer is installed but mm_fp4 not found: {_last_err}")

    # mm_fp4 expects 2D inputs: a is (m, k) fp4, b is (k, n) COLUMN-MAJOR fp4,
    # and the block scales must be in the 128x4 layout -- exactly the layout the
    # reference feeds to torch._scaled_mm via to_blocked(). Passing the raw
    # (mn, sf_k) scale slices (as before) feeds the kernel the wrong layout.
    # Hoist all slicing / re-blocking out of the timed region.
    # Mirrors the documented call mm_fp4(a_fp4, b_fp4.T, a_sf, b_sf.T, alpha, ...):
    #   a         -> (m, k) fp4
    #   b         -> (k, n) column-major fp4   (b_fp4.T)
    #   a_descale -> 2D 128x4-blocked scale
    #   b_descale -> 2D 128x4-blocked scale, transposed (b_sf.T)
    # This flashinfer build's internal check requires the packed FP4 as uint8
    # storage (it rejected the typed float4_e2m1fn_x2 with "mat1.dtype() ==
    # FLOAT4_E2M1X2 ... float4_e2m1fnx2 vs. uint8"). Same 1-byte reinterpret as
    # the vLLM path; view BEFORE transpose so the column-major b stays a clean view.
    a_slices = [a_fp4[:, :, i].contiguous().view(torch.uint8) for i in range(l)]
    b1_slices = [b1_fp4[:, :, i].view(torch.uint8).transpose(0, 1) for i in range(l)]
    b2_slices = [b2_fp4[:, :, i].view(torch.uint8).transpose(0, 1) for i in range(l)]

    def _blocked_2d(sf_slice):
        # 128x4-blocked scale kept as a 2D matrix (round_up(mn,128), round_up(sf_k,4)).
        # The flattened buffer triggered the "x.T on non-2D" warning and a wrong
        # b_descale orientation.
        mn, sfk = sf_slice.shape
        return to_blocked(sf_slice).reshape(ceil_div(mn, 128) * 128,
                                            ceil_div(sfk, 4) * 4).contiguous()

    sfa_b = [_blocked_2d(sfa[:, :, i]) for i in range(l)]
    sfb1_b = [_blocked_2d(sfb1[:, :, i]) for i in range(l)]
    sfb2_b = [_blocked_2d(sfb2[:, :, i]) for i in range(l)]

    # The cutlass/cudnn FP4 GEMM requires a real alpha tensor (global dequant
    # scalar); passing alpha=None failed with "argument #4 expected DLTensor*".
    # Our block scales already encode everything, so alpha = 1.0.
    alpha = torch.tensor(1.0, dtype=torch.float32, device="cuda")

    # Preallocate the GEMM outputs and pass out= so mm_fp4 does not allocate a
    # fresh result tensor on every timed call. (mm_fp4 still constructs a backend
    # runner and consults the AutoTuner per call -- that eager-dispatch cost is
    # inherent to this API and dominates at these tiny latency-bound shapes.)
    out1_buf = [torch.empty((m, n), dtype=torch.float16, device="cuda") for _ in range(l)]
    out2_buf = [torch.empty((m, n), dtype=torch.float16, device="cuda") for _ in range(l)]
    out = torch.empty((l, m, n), dtype=torch.float16, device="cuda").permute(1, 2, 0)

    def run():
        for i in range(l):
            mm_fp4(a_slices[i], b1_slices[i], sfa_b[i], sfb1_b[i].T,
                   alpha=alpha, out_dtype=torch.float16, out=out1_buf[i],
                   block_size=sf_vec_size)
            mm_fp4(a_slices[i], b2_slices[i], sfa_b[i], sfb2_b[i].T,
                   alpha=alpha, out_dtype=torch.float16, out=out2_buf[i],
                   block_size=sf_vec_size)
            out[:, :, i] = torch.nn.functional.silu(out1_buf[i]) * out2_buf[i]
        return out

    return run
def benchmark_comparison(m, n, k, l=1, seed=42, warmup=10, iters=100):
    """
    Run head-to-head benchmark of:
      1. CuTe DualGEMM kernel  (our fused CUTLASS kernel)
      2. vLLM baseline         (two cuBLAS scaled_mm + SiLU pointwise)
      3. Flashinfer baseline   (flashinfer.gemm.bmm_fp4 if importable)

    All methods share the SAME generated data tensors. Each method's one-time
    preprocessing (scale-factor reblocking, buffer allocation) is hoisted out
    of the timed region so that only the actual GEMM + SiLU compute is timed --
    matching the fact that the CuTe kernel is handed pre-permuted scale factors.

    Before timing, the CuTe kernel's output is verified against the PyTorch
    reference; if it does not match, its result is meaningless and the row is
    flagged rather than reported as a "win".

    Prints a comparison table with median kernel time and effective TFLOPS.
    """
    from laguna_dual_gemm import custom_kernel

    print(f"\n{'='*72}")
    print(f"  Dual GEMM + SiLU Benchmark  --  B200  (M={m}, N={n}, K={k}, L={l})")
    print(f"{'='*72}")

    data = generate_input(m, n, k, l, seed=seed)
    a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data

    # Theoretical FLOPs: 2*M*N*K per GEMM, two GEMMs
    flops = 2 * 2 * m * n * k * l

    # --- Correctness: the CuTe kernel must match the reference before any
    #     timing number is meaningful. ---
    custom_ok, custom_msg = None, ""
    try:
        custom_out = custom_kernel(data).clone()
        custom_ok, custom_msg = check_implementation(data, custom_out)
    except Exception as e:
        custom_ok, custom_msg = False, f"raised: {e}"
    status = "PASS" if custom_ok else "FAIL"
    print(f"\n  Correctness (CuTe vs PyTorch reference): {status}"
          + (f"  -- {custom_msg}" if not custom_ok else ""))

    results = {}

    # --- 1. CuTe DualGEMM (our kernel) ---
    try:
        custom_kernel(data)
        torch.cuda.synchronize()
        ms = _benchmark_fn(lambda: custom_kernel(data), warmup=warmup, iters=iters)
        tflops = flops / (ms * 1e-3) / 1e12
        results["CuTe DualGEMM (ours)"] = (ms, tflops)
    except Exception as e:
        results["CuTe DualGEMM (ours)"] = (None, None, str(e))

    try:
        vllm_run = _prepare_vllm_dual_gemm_silu(a, b1, b2, sfa, sfb1, sfb2, m, n, k, l)
        if vllm_run is None:
            results["vLLM (CUTLASS scaled_mm)"] = (None, None, "vllm._custom_ops.cutlass_scaled_fp4_mm unavailable")
        else:
            vllm_run()
            torch.cuda.synchronize()
            ms = _benchmark_fn(vllm_run, warmup=warmup, iters=iters)
            tflops = flops / (ms * 1e-3) / 1e12
            results["vLLM (CUTLASS scaled_mm)"] = (ms, tflops)
    except Exception as e:
        results["vLLM (CUTLASS scaled_mm)"] = (None, None, str(e))
    # --- 3. Flashinfer baseline (skipped, not faked, if unavailable) ---
    try:
        fi_run = _prepare_flashinfer_dual_gemm_silu(a, b1, b2, sfa, sfb1, sfb2, m, n, k, l)
        if fi_run is None:
            results["Flashinfer"] = (None, None, "flashinfer.mm_fp4 unavailable")
        else:
            fi_run()
            torch.cuda.synchronize()
            ms = _benchmark_fn(fi_run, warmup=warmup, iters=iters)
            tflops = flops / (ms * 1e-3) / 1e12
            results["Flashinfer"] = (ms, tflops)
    except Exception as e:
        results["Flashinfer"] = (None, None, str(e))

    # --- Print table ---
    print(f"\n{'Method':<30} {'Time (ms)':>12} {'TFLOPS':>10} {'Speedup':>10}")
    print("-" * 65)

    vllm_ms = None
    if "vLLM (CUTLASS scaled_mm)" in results and len(results["vLLM (CUTLASS scaled_mm)"]) == 2:
        vllm_ms = results["vLLM (CUTLASS scaled_mm)"][0]

    for name, vals in results.items():
        if len(vals) == 3:
            print(f"{name:<30} {'FAILED':>12} {'--':>10} {'--':>10}  ({vals[2]})")
            continue
        ms, tflops = vals
        speedup_str = "--"
        if vllm_ms and ms:
            speedup_str = f"{vllm_ms / ms:.2f}x"
        print(f"{name:<30} {ms:>11.4f}  {tflops:>9.2f}  {speedup_str:>10}")

    # Summary line -- only claim a speedup if the CuTe kernel is actually correct.
    cute_valid = (
        custom_ok
        and "CuTe DualGEMM (ours)" in results
        and len(results["CuTe DualGEMM (ours)"]) == 2
    )
    if not custom_ok:
        print("\n  >> CuTe result FAILED correctness -- speedup not reported "
              "(a wrong/no-op kernel is not a win).")
    else:
        if vllm_ms and cute_valid:
            cute_ms = results["CuTe DualGEMM (ours)"][0]
            print(f"\n  >> CuTe DualGEMM is {vllm_ms/cute_ms:.2f}x faster than vLLM cuBLAS baseline")
        if cute_valid and "Flashinfer" in results and len(results["Flashinfer"]) == 2:
            fi_ms = results["Flashinfer"][0]
            cute_ms = results["CuTe DualGEMM (ours)"][0]
            print(f"  >> CuTe DualGEMM is {fi_ms/cute_ms:.2f}x faster than Flashinfer baseline")

    # Reality check: at these shapes the GEMMs are tiny (a few GFLOP) and the
    # measurement is dominated by launch/dispatch latency, not throughput.
    # Treat the numbers as LATENCY, not as representative B200 FP4 TFLOPS.
    print("\n  note: M/N/K here are small -> latency-bound; TFLOPS are far below"
          " peak and\n        scale weakly with size. Use larger M for a throughput number.")

    print()
    return results


if __name__ == "__main__":
    # Laguna XS.2 shapes: Hidden (K) = 2048, Intermediate (N) = 512
    # M = Sequence length / tokens processed. Shared expert processes everything.
    shapes = [
        (4096, 512, 2048),  # SX2 prefill
        (4096, 512*2, 2048*2),  # M1 prefill
        (256, 512, 2048),   # SX2 decode
    ]

    for m, n, k in shapes:
        benchmark_comparison(m, n, k, l=1, seed=42, warmup=20, iters=200)