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)