diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py b/benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py index 8e5273d25ba..6303adf992b 100644 --- a/benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py +++ b/benchmarks/prototype/moe_training/mxfp8/bench_ep_pipeline.py @@ -40,6 +40,7 @@ from torchao.prototype.moe_training.ep.unpermute import _unpermute_bf16 from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, + set_mxfp8_grouped_mm_backend, ) device = torch.device("cuda") @@ -554,6 +555,7 @@ def print_results(experiments: List[Experiment]): def main(args: argparse.Namespace): """Main benchmark entry point.""" torch.random.manual_seed(123) + set_mxfp8_grouped_mm_backend(args.backend) # Set up process group setup_distributed() @@ -597,5 +599,11 @@ def setup_distributed(): action="store_true", help="Use torch.compile", ) + parser.add_argument( + "--backend", + choices=("legacy", "cutedsl"), + default="cutedsl", + help="MXFP8 grouped-mm backend", + ) args = parser.parse_args() main(args) diff --git a/benchmarks/prototype/moe_training/mxfp8/roofline_unified.py b/benchmarks/prototype/moe_training/mxfp8/roofline_unified.py index 1340be327e5..9cb896a9fb1 100644 --- a/benchmarks/prototype/moe_training/mxfp8/roofline_unified.py +++ b/benchmarks/prototype/moe_training/mxfp8/roofline_unified.py @@ -19,11 +19,17 @@ from triton.testing import do_bench from torchao.prototype.moe_training.kernels.mxfp8 import ( + mx_block_rearrange_2d_k_groups_cutedsl, mx_block_rearrange_2d_M_groups_cuda, + mxfp8_quantize_2d_1x32_cutedsl, + mxfp8_quantize_2d_32x1_cutedsl, torch_to_blocked_2d_M_groups, triton_mx_block_rearrange_2d_K_groups, triton_mx_block_rearrange_per_group_3d, ) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( + mx_block_rearrange_2d_m_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( mxfp8_quantize_cuda_3d, ) @@ -32,6 +38,7 @@ ) from torchao.prototype.moe_training.mxfp8_grouped_mm import ( _to_mxfp8_then_scaled_grouped_mm, + set_mxfp8_grouped_mm_backend, ) from torchao.prototype.moe_training.utils import generate_jagged_offs from torchao.prototype.mx_formats.config import ( @@ -42,6 +49,7 @@ from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper from torchao.quantization.quantize_.common import KernelPreference from torchao.testing.training.roofline_utils import ( + calibrate_specs, gpu_name_to_specs, ) @@ -49,7 +57,12 @@ class RooflineModel: """Roofline model for grouped GEMM on B200 GPU""" - def __init__(self, gpu_name="NVIDIA B200", power_limit_percent=100.0): + def __init__( + self, + gpu_name="NVIDIA B200", + power_limit_percent=100.0, + gpu_specs=None, + ): """ Args: gpu_name: GPU model name @@ -57,16 +70,30 @@ def __init__(self, gpu_name="NVIDIA B200", power_limit_percent=100.0): """ power_multiplier = power_limit_percent / 100.0 - if gpu_name in gpu_name_to_specs: + if gpu_specs is not None: + self.gpu_specs = gpu_specs + elif gpu_name in gpu_name_to_specs: self.gpu_specs = gpu_name_to_specs[gpu_name] + else: + raise ValueError(f"Unsupported GPU: {gpu_name}") + + if self.gpu_specs is not None: + bf16_gemm_pct = self.gpu_specs.get( + "pct_achievable_bf16_gemm_tops", + self.gpu_specs["pct_achievable_gemm_tops"], + ) + fp8_gemm_pct = self.gpu_specs.get( + "pct_achievable_fp8_gemm_tops", + self.gpu_specs["pct_achievable_gemm_tops"], + ) self.bf16_tflops = ( (self.gpu_specs["bf16_peak_tops"] / 1e12) - * self.gpu_specs["pct_achievable_gemm_tops"] + * bf16_gemm_pct * power_multiplier ) self.mxfp8_tflops = ( (self.gpu_specs["fp8_peak_tops"] / 1e12) - * self.gpu_specs["pct_achievable_gemm_tops"] + * fp8_gemm_pct * power_multiplier ) self.memory_bandwidth_gbs = ( @@ -74,8 +101,6 @@ def __init__(self, gpu_name="NVIDIA B200", power_limit_percent=100.0): * self.gpu_specs["pct_achievable_mem_bw"] * power_multiplier ) - else: - raise ValueError(f"Unsupported GPU: {gpu_name}") def compute_bf16_2d_3d_gemm_flops(self, M, K, N): """ @@ -427,31 +452,25 @@ def wrapper(): return time_ms -def benchmark_mxfp8_grouped_mm_fwd_bwd(x, w_t, offs, labels): +def benchmark_mxfp8_grouped_mm_fwd_bwd(x, w_t, offs, labels, backend: str): """Benchmark _to_mxfp8_then_scaled_grouped_mm forward + backward""" + set_mxfp8_grouped_mm_backend(backend) + torch._dynamo.reset() x_clone = x.clone().requires_grad_(True) w_t_clone = w_t.clone().requires_grad_(True) fn = torch.compile(_to_mxfp8_then_scaled_grouped_mm, fullgraph=True) - # Set all parameters explicitly as variables for positional args - A = x_clone - B_t = w_t_clone - offs_arg = offs - out_dtype = torch.bfloat16 - kernel_preference = KernelPreference.AUTO - wgrad_with_hp = False - scale_calculation_mode = MoEScaleCalculationMode.RCEIL - def wrapper(): out = fn( - A, - B_t, - offs_arg, - out_dtype, - kernel_preference, - wgrad_with_hp, - scale_calculation_mode, + x_clone, + w_t_clone, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.AUTO, + wgrad_with_hp=False, + scale_calculation_mode=MoEScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=False, ) loss = torch.nn.functional.mse_loss(out, labels) loss.backward() @@ -482,6 +501,22 @@ def benchmark_to_mxfp8_dim1_cuda(tensor, block_size=32): ) +def benchmark_mxfp8_quantize_2d_1x32_cutedsl(tensor): + return benchmark_cuda_function_in_microseconds( + lambda: mxfp8_quantize_2d_1x32_cutedsl(tensor) + ) + + +def benchmark_mxfp8_quantize_2d_32x1_cutedsl(tensor, block_size=32): + return benchmark_cuda_function_in_microseconds( + lambda: mxfp8_quantize_2d_32x1_cutedsl( + tensor, + block_size=block_size, + blocked_scale_output=False, + ) + ) + + def benchmark_mxfp8_quantize_cuda_3d(tensor, block_size=32): """Benchmark the 3D 32x1 quantizer on its input tensor.""" return benchmark_cuda_function_in_microseconds( @@ -534,12 +569,18 @@ def run( N: int = 4096, G: int = 8, breakdown_M: int = None, + backend: str = "cutedsl", outfile_speedup: str = "roofline_speedup_results.csv", outfile_quant_2d: str = "roofline_quant_2d_results.csv", outfile_quant_3d: str = "roofline_quant_3d_results.csv", plot_file: str = "roofline_unified.png", gpu_name: str = "NVIDIA B200", power_limit_percent: float = 100.0, + calibrate_roofline: bool = False, + calibration_m: int = 16384, + calibration_n: int = 8192, + calibration_k: int = 8192, + calibration_copy_numel: int = 256 * 1024 * 1024, ): """ Generate unified roofline analysis for MXFP8 grouped GEMM. @@ -549,19 +590,67 @@ def run( N: Output dimension per group (default: 4096) G: Number of groups (default: 8) breakdown_M: M value to use for kernel breakdown analysis (default: None, uses largest M from configs) + backend: MXFP8 grouped-mm backend to benchmark: legacy, cutedsl, or both outfile_speedup: CSV file for speedup results outfile_quant_2d: CSV file for 2D quantization results outfile_quant_3d: CSV file for 3D quantization results plot_file: PNG file to save unified plot gpu_name: GPU model (default: B200) power_limit_percent: Power limit as percentage (0-100, default: 100.0) + calibrate_roofline: Measure achievable GEMM and memory ceilings on this GPU + calibration_m: M dimension for GEMM calibration + calibration_n: N dimension for GEMM calibration + calibration_k: K dimension for GEMM calibration + calibration_copy_numel: Number of bf16 elements for memory calibration """ + K = 4096 if K == "" else int(K) + N = 4096 if N == "" else int(N) + G = 8 if G == "" else int(G) + breakdown_M = None if breakdown_M in (None, "") else int(breakdown_M) + power_limit_percent = ( + 100.0 if power_limit_percent == "" else float(power_limit_percent) + ) + if isinstance(calibrate_roofline, str): + calibrate_roofline = calibrate_roofline.lower() in ("1", "true", "yes") + calibration_m = int(calibration_m) + calibration_n = int(calibration_n) + calibration_k = int(calibration_k) + calibration_copy_numel = int(calibration_copy_numel) + print(f"GPU: {gpu_name}") print(f"Torch version: {torch.__version__}") print(f"\nFixed dimensions: K={K}, N={N}, G={G}") + assert backend in ("legacy", "cutedsl", "both") + mxfp8_backends = ("legacy", "cutedsl") if backend == "both" else (backend,) + print(f"MXFP8 backend: {backend}") print(f"Power limit: {power_limit_percent}%") - model = RooflineModel(gpu_name=gpu_name, power_limit_percent=power_limit_percent) + gpu_specs = None + if calibrate_roofline: + print("\nCalibrating roofline on current GPU...") + gpu_specs = calibrate_specs( + gpu_name=gpu_name, + mm_shape=(calibration_m, calibration_n, calibration_k), + copy_numel=calibration_copy_numel, + ) + print( + " copy bandwidth: " + f"{gpu_specs['calibrated_copy_bandwidth_bytes_sec'] / 1e9:.1f} GB/s" + ) + print(f" BF16 GEMM: {gpu_specs['calibrated_bf16_tops'] / 1e12:.1f} TFLOP/s") + print(f" FP8 GEMM: {gpu_specs['calibrated_fp8_tops'] / 1e12:.1f} TFLOP/s") + print( + " measured fractions: " + f"mem={gpu_specs['pct_achievable_mem_bw']:.3f}, " + f"bf16={gpu_specs['pct_achievable_bf16_gemm_tops']:.3f}, " + f"fp8={gpu_specs['pct_achievable_fp8_gemm_tops']:.3f}" + ) + + model = RooflineModel( + gpu_name=gpu_name, + power_limit_percent=power_limit_percent, + gpu_specs=gpu_specs, + ) print("\nGPU Specs:") print(f" BF16 TFLOPS: {model.bf16_tflops}") @@ -604,7 +693,7 @@ def run( x = torch.randn(M, K_val, dtype=torch.bfloat16, device="cuda") w = torch.randn(G_val, N_val, K_val, dtype=torch.bfloat16, device="cuda") w_t = w.contiguous().transpose(-2, -1) - offs = generate_jagged_offs(G_val, M) + offs = generate_jagged_offs(G_val, M, multiple_of=128) labels = torch.ones((M, N_val), device="cuda", dtype=torch.bfloat16) # Benchmark BF16 @@ -614,17 +703,26 @@ def run( f" BF16: Roofline={result['bf16_roofline_time_ms']:.3f}ms, Actual={bf16_actual_ms:.3f}ms" ) - # Benchmark MXFP8 - mxfp8_actual_ms = benchmark_mxfp8_grouped_mm_fwd_bwd(x, w_t, offs, labels) - result_dict["mxfp8_actual_time_ms"] = mxfp8_actual_ms - result_dict["actual_speedup"] = ( - bf16_actual_ms / mxfp8_actual_ms if bf16_actual_ms else None - ) - print( - f" MXFP8: Roofline={result['mxfp8_roofline_total_time_ms']:.3f}ms, Actual={mxfp8_actual_ms:.3f}ms" - ) - if result_dict["actual_speedup"]: - print(f" Actual Speedup: {result_dict['actual_speedup']:.3f}x") + for mxfp8_backend in mxfp8_backends: + mxfp8_actual_ms = benchmark_mxfp8_grouped_mm_fwd_bwd( + x, w_t, offs, labels, mxfp8_backend + ) + backend_time_key = f"mxfp8_actual_time_ms_{mxfp8_backend}" + backend_speedup_key = f"actual_speedup_{mxfp8_backend}" + result_dict[backend_time_key] = mxfp8_actual_ms + result_dict[backend_speedup_key] = ( + bf16_actual_ms / mxfp8_actual_ms if bf16_actual_ms else None + ) + if backend != "both": + result_dict["mxfp8_actual_time_ms"] = mxfp8_actual_ms + result_dict["actual_speedup"] = result_dict[backend_speedup_key] + print( + f" MXFP8 ({mxfp8_backend}): Roofline={result['mxfp8_roofline_total_time_ms']:.3f}ms, Actual={mxfp8_actual_ms:.3f}ms" + ) + if result_dict[backend_speedup_key]: + print( + f" Actual Speedup ({mxfp8_backend}): {result_dict[backend_speedup_key]:.3f}x" + ) speedup_results.append(result_dict) @@ -680,6 +778,17 @@ def run( f" triton_to_mxfp8_dim0: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={triton_dim0_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['triton_dim0_efficiency_pct']:.1f}%" ) + cutedsl_1x32_time_us = benchmark_mxfp8_quantize_2d_1x32_cutedsl(tensor) + cutedsl_1x32_bandwidth_gbs = total_gb / (cutedsl_1x32_time_us / 1e6) + result_dict["mxfp8_quantize_2d_1x32_cutedsl_us"] = cutedsl_1x32_time_us + result_dict["cutedsl_1x32_bandwidth_gbs"] = cutedsl_1x32_bandwidth_gbs + result_dict["cutedsl_1x32_efficiency_pct"] = ( + cutedsl_1x32_bandwidth_gbs / roofline_bandwidth_gbs + ) * 100 + print( + f" mxfp8_quantize_2d_1x32_cutedsl: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={cutedsl_1x32_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cutedsl_1x32_efficiency_pct']:.1f}%" + ) + # Benchmark triton_to_mxfp8_dim1 (CUDA) dim1_cuda_time_us = benchmark_to_mxfp8_dim1_cuda(tensor) dim1_cuda_bandwidth_gbs = total_gb / (dim1_cuda_time_us / 1e6) @@ -692,6 +801,17 @@ def run( f" to_mxfp8_dim1_cuda: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={dim1_cuda_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cuda_dim1_efficiency_pct']:.1f}%" ) + cutedsl_32x1_time_us = benchmark_mxfp8_quantize_2d_32x1_cutedsl(tensor) + cutedsl_32x1_bandwidth_gbs = total_gb / (cutedsl_32x1_time_us / 1e6) + result_dict["mxfp8_quantize_2d_32x1_cutedsl_us"] = cutedsl_32x1_time_us + result_dict["cutedsl_32x1_bandwidth_gbs"] = cutedsl_32x1_bandwidth_gbs + result_dict["cutedsl_32x1_efficiency_pct"] = ( + cutedsl_32x1_bandwidth_gbs / roofline_bandwidth_gbs + ) * 100 + print( + f" mxfp8_quantize_2d_32x1_cutedsl: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={cutedsl_32x1_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cutedsl_32x1_efficiency_pct']:.1f}%" + ) + quant_2d_results.append(result_dict) # Clean up tensors to free GPU memory @@ -808,7 +928,7 @@ def run( dtype=torch.uint8, device="cuda", ) - input_group_offsets = generate_jagged_offs(num_groups, M) + input_group_offsets = generate_jagged_offs(num_groups, M, multiple_of=128) # Benchmark CUDA kernel cuda_out = mx_block_rearrange_2d_M_groups_cuda( @@ -829,10 +949,30 @@ def run( f" mx_block_rearrange_2d_M_groups_cuda: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={cuda_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cuda_efficiency_pct']:.1f}%" ) + cutedsl_out = mx_block_rearrange_2d_m_groups_cutedsl( + input_tensor, + input_group_offsets, + ) + torch.testing.assert_close(cutedsl_out, cuda_out, rtol=0, atol=0) + cutedsl_time_us = benchmark_cuda_function_in_microseconds( + mx_block_rearrange_2d_m_groups_cutedsl, + input_tensor, + input_group_offsets, + ) + cutedsl_bandwidth_gbs = total_gb / (cutedsl_time_us / 1e6) + result_dict["mx_block_rearrange_2d_m_groups_cutedsl_us"] = cutedsl_time_us + result_dict["cutedsl_bandwidth_gbs"] = cutedsl_bandwidth_gbs + result_dict["cutedsl_efficiency_pct"] = ( + cutedsl_bandwidth_gbs / roofline_bandwidth_gbs + ) * 100 + print( + f" mx_block_rearrange_2d_m_groups_cutedsl: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={cutedsl_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cutedsl_efficiency_pct']:.1f}%" + ) + rearrange_results.append(result_dict) # Clean up tensors - del input_tensor, input_group_offsets, cuda_out + del input_tensor, input_group_offsets, cuda_out, cutedsl_out torch.cuda.empty_cache() # K-groups configurations (backward weight pass scales) @@ -873,7 +1013,9 @@ def run( dtype=torch.uint8, device="cuda", ) - scale_group_offsets = generate_jagged_offs(num_groups, M) // block_size + scale_group_offsets = ( + generate_jagged_offs(num_groups, M, multiple_of=128) // block_size + ) # Benchmark triton kernel triton_out = triton_mx_block_rearrange_2d_K_groups( @@ -894,10 +1036,30 @@ def run( f" triton_mx_block_rearrange_2d_K_groups: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={triton_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['triton_efficiency_pct']:.1f}%" ) + cutedsl_out = mx_block_rearrange_2d_k_groups_cutedsl( + input_tensor, + scale_group_offsets, + ) + torch.testing.assert_close(cutedsl_out, triton_out, rtol=0, atol=0) + cutedsl_time_us = benchmark_cuda_function_in_microseconds( + mx_block_rearrange_2d_k_groups_cutedsl, + input_tensor, + scale_group_offsets, + ) + cutedsl_bandwidth_gbs = total_gb / (cutedsl_time_us / 1e6) + result_dict["mx_block_rearrange_2d_k_groups_cutedsl_us"] = cutedsl_time_us + result_dict["cutedsl_bandwidth_gbs"] = cutedsl_bandwidth_gbs + result_dict["cutedsl_efficiency_pct"] = ( + cutedsl_bandwidth_gbs / roofline_bandwidth_gbs + ) * 100 + print( + f" mx_block_rearrange_2d_k_groups_cutedsl: Roofline={roofline_bandwidth_gbs:.1f} GB/s, Actual={cutedsl_bandwidth_gbs:.1f} GB/s, Efficiency={result_dict['cutedsl_efficiency_pct']:.1f}%" + ) + rearrange_results.append(result_dict) # Clean up tensors - del input_tensor, scale_group_offsets, triton_out + del input_tensor, scale_group_offsets, triton_out, cutedsl_out torch.cuda.empty_cache() df_rearrange = pd.DataFrame(rearrange_results) @@ -1003,7 +1165,7 @@ def run( x = torch.randn(M, K_val, dtype=torch.bfloat16, device="cuda") w = torch.randn(G_val, N_val, K_val, dtype=torch.bfloat16, device="cuda") w_t = w.contiguous().transpose(-2, -1) - offs = generate_jagged_offs(G_val, M) + offs = generate_jagged_offs(G_val, M, multiple_of=128) # Benchmark BF16 grouped GEMM bf16_gemm_time_us = benchmark_bf16_grouped_gemm(x, w_t, offs) @@ -1097,7 +1259,7 @@ def run( # We'll create grad_output and input, then quantize them grad_out = torch.randn(M, N_val, dtype=torch.bfloat16, device="cuda") x = torch.randn(M, K_val, dtype=torch.bfloat16, device="cuda") - offs = generate_jagged_offs(G_val, M) + offs = generate_jagged_offs(G_val, M, multiple_of=128) # Benchmark BF16 2D/2D grouped GEMM # For BF16, we need grad_out_t = grad_out.t().contiguous() to get (N, M) row-major @@ -1221,15 +1383,23 @@ def run( linestyle=":", label="Roofline Model", ) - ax1.plot( - df_speedup["M"], - df_speedup["actual_speedup"], - marker="s", - linewidth=2, - linestyle="-", - label="Actual Implementation", - color="purple", - ) + if "actual_speedup" in df_speedup: + actual_speedup_columns = [("actual_speedup", "Actual Implementation")] + else: + actual_speedup_columns = [ + (column, column.removeprefix("actual_speedup_")) + for column in df_speedup.columns + if column.startswith("actual_speedup_") + ] + for column, label in actual_speedup_columns: + ax1.plot( + df_speedup["M"], + df_speedup[column], + marker="s", + linewidth=2, + linestyle="-", + label=label, + ) ax1.axhline( y=1.0, color="red", @@ -1256,18 +1426,36 @@ def run( marker="s", linewidth=2, linestyle="-", - label="triton_to_mxfp8_dim0", + label="legacy dim0 quant (Triton)", color="blue", ) + ax2.plot( + df_quant_2d["M"], + df_quant_2d["cutedsl_1x32_efficiency_pct"], + marker="o", + linewidth=2, + linestyle="-", + label="CuTeDSL 1x32 quant", + color="green", + ) ax2.plot( df_quant_2d["M"], df_quant_2d["cuda_dim1_efficiency_pct"], marker="d", linewidth=2, linestyle="-", - label="to_mxfp8_dim1_cuda", + label="legacy dim1 quant (CUDA)", color="orange", ) + ax2.plot( + df_quant_2d["M"], + df_quant_2d["cutedsl_32x1_efficiency_pct"], + marker="x", + linewidth=2, + linestyle="-", + label="CuTeDSL 32x1 quant", + color="brown", + ) # 2D Rearrange kernels df_m_groups = df_rearrange[df_rearrange["kernel_type"] == "M_groups"] df_k_groups = df_rearrange[df_rearrange["kernel_type"] == "K_groups"] @@ -1277,18 +1465,36 @@ def run( marker="^", linewidth=2, linestyle="--", - label="CUDA M-groups scale blocked format", + label="legacy M-groups layout (CUDA)", color="purple", ) + ax2.plot( + df_m_groups["M"], + df_m_groups["cutedsl_efficiency_pct"], + marker="v", + linewidth=2, + linestyle="--", + label="CuTeDSL M-groups layout", + color="gray", + ) ax2.plot( df_k_groups["M"], df_k_groups["triton_efficiency_pct"], marker="d", linewidth=2, linestyle="--", - label="triton K-groups scale blocked format", + label="legacy K-groups layout (Triton)", color="red", ) + ax2.plot( + df_k_groups["M"], + df_k_groups["cutedsl_efficiency_pct"], + marker="p", + linewidth=2, + linestyle="--", + label="CuTeDSL K-groups layout", + color="black", + ) ax2.set_xlabel("Local Batch Size x Sequence Length (M)", fontsize=12) ax2.set_ylabel("Bandwidth Utilization (% of Peak)", fontsize=12) ax2.set_title(f"2D Quantization + Block Format Kernels (K={K}, N={N})", fontsize=13) @@ -1354,21 +1560,55 @@ def run( ax3.legend(fontsize=9) ax3.grid(True, alpha=0.3) - # Plot 4: Empty placeholder (reserved for future use) + # Plot 4: Backend runtime ratio ax4 = axes[0, 2] - ax4.text( - 0.5, - 0.5, - "Reserved for Future Analysis", - horizontalalignment="center", - verticalalignment="center", - transform=ax4.transAxes, - fontsize=14, - color="gray", - ) - ax4.set_xticks([]) - ax4.set_yticks([]) - ax4.grid(False) + if { + "mxfp8_actual_time_ms_legacy", + "mxfp8_actual_time_ms_cutedsl", + }.issubset(df_speedup.columns): + backend_runtime_ratio = ( + df_speedup["mxfp8_actual_time_ms_cutedsl"] + / df_speedup["mxfp8_actual_time_ms_legacy"] + ) + ax4.plot( + df_speedup["M"], + backend_runtime_ratio, + marker="o", + linewidth=2, + linestyle="-", + label="CuTeDSL / legacy runtime", + color="green", + ) + ax4.axhline( + y=1.0, + color="red", + linestyle=":", + linewidth=1.5, + label="parity", + ) + ax4.set_xscale("log", base=2) + ax4.set_xticks(df_speedup["M"]) + ax4.set_xticklabels([f"{int(m):,}" for m in df_speedup["M"]]) + ratio_min = backend_runtime_ratio.min() + ratio_max = backend_runtime_ratio.max() + ratio_margin = max(0.02, (ratio_max - ratio_min) * 0.2) + ax4.set_ylim(ratio_min - ratio_margin, ratio_max + ratio_margin) + else: + ax4.text( + 0.5, + 0.5, + "Run with --backend=both", + horizontalalignment="center", + verticalalignment="center", + transform=ax4.transAxes, + fontsize=12, + color="gray", + ) + ax4.set_xlabel("Local Batch Size x Sequence Length (M)", fontsize=12) + ax4.set_ylabel("Runtime Ratio", fontsize=12) + ax4.set_title("Backend Runtime Ratio", fontsize=13) + ax4.legend(fontsize=9) + ax4.grid(True, alpha=0.3) # Plot 5: 3D Quantization + Rearrange Kernels (Bandwidth %) ax5 = axes[1, 1] @@ -1380,7 +1620,7 @@ def run( marker="s", linewidth=2, linestyle="-", - label="mxfp8_quantize_cuda_3d", + label="CuTeDSL 3D quant + blocked scales", color="red", ) # 3D Rearrange kernel @@ -1390,7 +1630,7 @@ def run( marker="^", linewidth=2, linestyle="--", - label="triton per-group 3D scale blocked format", + label="legacy per-group 3D layout (Triton)", color="purple", ) ax5.set_xlabel("Local Batch Size x Sequence Length (M)", fontsize=12) @@ -1430,26 +1670,51 @@ def run( idx_large = df_quant_2d[df_quant_2d["M"] == M_large].index[0] # Forward pass kernel times (actual measurements in ms) - # Input quantization: use triton_to_mxfp8_dim0 for (M, K) - fwd_input_quant_ms = df_quant_2d.loc[idx_large, "triton_to_mxfp8_dim0_us"] / 1000 + fwd_input_quant_legacy_ms = ( + df_quant_2d.loc[idx_large, "triton_to_mxfp8_dim0_us"] / 1000 + ) + fwd_input_quant_cutedsl_ms = ( + df_quant_2d.loc[idx_large, "mxfp8_quantize_2d_1x32_cutedsl_us"] / 1000 + ) # Weight quantization: use mxfp8_quantize_cuda_3d directly on w_t, shape # (G, K, N), with no separate 3D scale rearrangement step. idx_3d_large = df_quant_3d[df_quant_3d["description"] == f"M={M_large}"].index[0] - fwd_weight_quant_ms = ( + fwd_weight_quant_cutedsl_ms = ( df_quant_3d.loc[idx_3d_large, "mxfp8_quantize_cuda_3d_us"] / 1000 ) + weight_tensor = torch.randn( + G_val, + N_val, + K_val, + dtype=torch.bfloat16, + device="cuda", + ) + fwd_weight_quant_legacy_ms = benchmark_triton_to_mxfp8_dim0(weight_tensor) / 1000 # Input scale rearrangement: M-groups for (M, K//32) idx_m_groups = df_rearrange[ (df_rearrange["kernel_type"] == "M_groups") & (df_rearrange["M"] == M_large) ].index[0] - fwd_input_scale_rearrange_ms = ( + fwd_input_scale_rearrange_legacy_ms = ( df_rearrange.loc[idx_m_groups, "mx_block_rearrange_2d_M_groups_cuda_us"] / 1000 ) + fwd_input_scale_rearrange_cutedsl_ms = ( + df_rearrange.loc[idx_m_groups, "mx_block_rearrange_2d_m_groups_cutedsl_us"] + / 1000 + ) - # Weight scales are emitted directly in blocked layout by the 3D quantizer. - fwd_weight_scale_rearrange_ms = 0.0 + idx_rearrange_3d_large = df_rearrange_3d[ + df_rearrange_3d["description"] == f"M={M_large}" + ].index[0] + fwd_weight_scale_rearrange_legacy_ms = ( + df_rearrange_3d.loc[ + idx_rearrange_3d_large, + "triton_mx_block_rearrange_per_group_3d_us", + ] + / 1000 + ) + fwd_weight_scale_rearrange_cutedsl_ms = 0.0 # GEMM: use actual MXFP8 2D/3D grouped GEMM time idx_gemm = df_grouped_gemm[df_grouped_gemm["M"] == M_large].index[0] @@ -1462,57 +1727,95 @@ def run( # Backward input: grad_output quantization (M, N) grad_out_tensor = torch.randn(M_large, N_val, dtype=torch.bfloat16, device="cuda") - bwd_input_grad_quant_ms = benchmark_triton_to_mxfp8_dim0(grad_out_tensor) / 1000 + bwd_input_grad_quant_legacy_ms = ( + benchmark_triton_to_mxfp8_dim0(grad_out_tensor) / 1000 + ) + bwd_input_grad_quant_cutedsl_ms = ( + benchmark_mxfp8_quantize_2d_1x32_cutedsl(grad_out_tensor) / 1000 + ) # Backward input: weight quantization is same as forward (reuse) - bwd_input_weight_quant_ms = fwd_weight_quant_ms + bwd_input_weight_quant_legacy_ms = fwd_weight_quant_legacy_ms + bwd_input_weight_quant_cutedsl_ms = fwd_weight_quant_cutedsl_ms # Backward input: grad scale rearrangement for (M, N//32) - M-groups grad_scales = torch.randint( 0, 256, size=(M_large, N_val // block_size), dtype=torch.uint8, device="cuda" ) - grad_offs = generate_jagged_offs(G_val, M_large) - bwd_input_grad_scale_rearrange_ms = ( + grad_offs = generate_jagged_offs(G_val, M_large, multiple_of=128) + bwd_input_grad_scale_rearrange_legacy_ms = ( benchmark_cuda_function_in_microseconds( mx_block_rearrange_2d_M_groups_cuda, grad_scales, grad_offs ) / 1000 ) + bwd_input_grad_scale_rearrange_cutedsl_ms = ( + benchmark_cuda_function_in_microseconds( + mx_block_rearrange_2d_m_groups_cutedsl, grad_scales, grad_offs + ) + / 1000 + ) # Backward input: weight scale rearrangement is same as forward (reuse) - bwd_input_weight_scale_rearrange_ms = fwd_weight_scale_rearrange_ms + bwd_input_weight_scale_rearrange_legacy_ms = fwd_weight_scale_rearrange_legacy_ms + bwd_input_weight_scale_rearrange_cutedsl_ms = fwd_weight_scale_rearrange_cutedsl_ms # Backward input: GEMM (same shape as forward, reuse) bwd_input_gemm_ms = fwd_gemm_ms # Backward weight pass - # Grad.T quantization (N, M) - grad_out_t_tensor = torch.randn(N_val, M_large, dtype=torch.bfloat16, device="cuda") - bwd_weight_grad_quant_ms = benchmark_triton_to_mxfp8_dim0(grad_out_t_tensor) / 1000 + bwd_weight_grad_quant_legacy_ms = ( + benchmark_to_mxfp8_dim1_cuda(grad_out_tensor) / 1000 + ) + bwd_weight_grad_quant_cutedsl_ms = ( + benchmark_mxfp8_quantize_2d_32x1_cutedsl(grad_out_tensor) / 1000 + ) - # Input quantization (M, K) - same as forward - bwd_weight_input_quant_ms = fwd_input_quant_ms + input_tensor_for_wgrad = torch.randn( + M_large, + K_val, + dtype=torch.bfloat16, + device="cuda", + ) + bwd_weight_input_quant_legacy_ms = ( + benchmark_to_mxfp8_dim1_cuda(input_tensor_for_wgrad) / 1000 + ) + bwd_weight_input_quant_cutedsl_ms = ( + benchmark_mxfp8_quantize_2d_32x1_cutedsl(input_tensor_for_wgrad) / 1000 + ) # Grad.T scale rearrangement (N, M//32) - K-groups idx_k_groups = df_rearrange[ (df_rearrange["kernel_type"] == "K_groups") & (df_rearrange["M"] == M_large) ].index[0] - bwd_weight_grad_scale_rearrange_ms = ( + bwd_weight_grad_scale_rearrange_legacy_ms = ( df_rearrange.loc[idx_k_groups, "triton_mx_block_rearrange_2d_K_groups_us"] / 1000 ) + bwd_weight_grad_scale_rearrange_cutedsl_ms = ( + df_rearrange.loc[idx_k_groups, "mx_block_rearrange_2d_k_groups_cutedsl_us"] + / 1000 + ) # Input scale rearrangement - need K-groups for (K, M//32) input_scales_k = torch.randint( 0, 256, size=(K_val, M_large // block_size), dtype=torch.uint8, device="cuda" ) - scale_group_offs = generate_jagged_offs(G_val, M_large) // block_size - bwd_weight_input_scale_rearrange_ms = ( + scale_group_offs = ( + generate_jagged_offs(G_val, M_large, multiple_of=128) // block_size + ) + bwd_weight_input_scale_rearrange_legacy_ms = ( benchmark_cuda_function_in_microseconds( triton_mx_block_rearrange_2d_K_groups, input_scales_k, scale_group_offs ) / 1000 ) + bwd_weight_input_scale_rearrange_cutedsl_ms = ( + benchmark_cuda_function_in_microseconds( + mx_block_rearrange_2d_k_groups_cutedsl, input_scales_k, scale_group_offs + ) + / 1000 + ) # GEMM: use actual MXFP8 2D/2D grouped GEMM time idx_gemm_2d2d = df_grouped_gemm_2d_2d[df_grouped_gemm_2d_2d["M"] == M_large].index[ @@ -1527,7 +1830,8 @@ def run( grad_out_tensor, grad_scales, grad_offs, - grad_out_t_tensor, + input_tensor_for_wgrad, + weight_tensor, input_scales_k, scale_group_offs, ) @@ -1536,76 +1840,114 @@ def run( # Data for stacked bars passes = ["Forward", "Backward\nInput", "Backward\nWeight"] - # Stack components (bottom to top) - quant_1 = [fwd_input_quant_ms, bwd_input_grad_quant_ms, bwd_weight_grad_quant_ms] - quant_2 = [ - fwd_weight_quant_ms, - bwd_input_weight_quant_ms, - bwd_weight_input_quant_ms, - ] - rearrange_1 = [ - fwd_input_scale_rearrange_ms, - bwd_input_grad_scale_rearrange_ms, - bwd_weight_grad_scale_rearrange_ms, - ] - rearrange_2 = [ - fwd_weight_scale_rearrange_ms, - bwd_input_weight_scale_rearrange_ms, - bwd_weight_input_scale_rearrange_ms, - ] + breakdown = { + "legacy": { + "quant_1": [ + fwd_input_quant_legacy_ms, + bwd_input_grad_quant_legacy_ms, + bwd_weight_grad_quant_legacy_ms, + ], + "quant_2": [ + fwd_weight_quant_legacy_ms, + bwd_input_weight_quant_legacy_ms, + bwd_weight_input_quant_legacy_ms, + ], + "rearrange_1": [ + fwd_input_scale_rearrange_legacy_ms, + bwd_input_grad_scale_rearrange_legacy_ms, + bwd_weight_grad_scale_rearrange_legacy_ms, + ], + "rearrange_2": [ + fwd_weight_scale_rearrange_legacy_ms, + bwd_input_weight_scale_rearrange_legacy_ms, + bwd_weight_input_scale_rearrange_legacy_ms, + ], + }, + "cutedsl": { + "quant_1": [ + fwd_input_quant_cutedsl_ms, + bwd_input_grad_quant_cutedsl_ms, + bwd_weight_grad_quant_cutedsl_ms, + ], + "quant_2": [ + fwd_weight_quant_cutedsl_ms, + bwd_input_weight_quant_cutedsl_ms, + bwd_weight_input_quant_cutedsl_ms, + ], + "rearrange_1": [ + fwd_input_scale_rearrange_cutedsl_ms, + bwd_input_grad_scale_rearrange_cutedsl_ms, + bwd_weight_grad_scale_rearrange_cutedsl_ms, + ], + "rearrange_2": [ + fwd_weight_scale_rearrange_cutedsl_ms, + bwd_input_weight_scale_rearrange_cutedsl_ms, + bwd_weight_input_scale_rearrange_cutedsl_ms, + ], + }, + } gemm = [fwd_gemm_ms, bwd_input_gemm_ms, bwd_weight_gemm_ms] - # Calculate cumulative bottoms for stacking - bottom_quant_2 = quant_1 - bottom_rearrange_1 = [quant_1[i] + quant_2[i] for i in range(3)] - bottom_rearrange_2 = [bottom_rearrange_1[i] + rearrange_1[i] for i in range(3)] - bottom_gemm = [bottom_rearrange_2[i] + rearrange_2[i] for i in range(3)] - - # Create stacked bars - x_pos = range(len(passes)) - bar_width = 0.6 - - ax6.bar(x_pos, quant_1, bar_width, label="Input/Grad Quant", color="#1f77b4") - ax6.bar( - x_pos, - quant_2, - bar_width, - bottom=bottom_quant_2, - label="Weight Quant", - color="#ff7f0e", - ) - ax6.bar( - x_pos, - rearrange_1, - bar_width, - bottom=bottom_rearrange_1, - label="Input/Grad Rearrange", - color="#2ca02c", - ) - ax6.bar( - x_pos, - rearrange_2, - bar_width, - bottom=bottom_rearrange_2, - label="Weight Rearrange", - color="#d62728", - ) - ax6.bar(x_pos, gemm, bar_width, bottom=bottom_gemm, label="GEMM", color="#9467bd") + # Create paired stacked bars + group_centers = list(range(len(passes))) + bar_width = 0.32 + backend_offsets = {"legacy": -bar_width / 2, "cutedsl": bar_width / 2} + component_specs = [ + ("quant_1", "Input/Grad Quant", "#1f77b4"), + ("quant_2", "Weight Quant", "#ff7f0e"), + ("rearrange_1", "Input/Grad Rearrange", "#2ca02c"), + ("rearrange_2", "Weight Rearrange", "#d62728"), + ("gemm", "GEMM", "#9467bd"), + ] + x_by_backend = { + backend_name: [center + offset for center in group_centers] + for backend_name, offset in backend_offsets.items() + } + totals_by_backend = {} + for backend_name, x_values in x_by_backend.items(): + bottoms = [0.0, 0.0, 0.0] + totals_by_backend[backend_name] = [0.0, 0.0, 0.0] + for component_name, label, color in component_specs: + values = ( + gemm + if component_name == "gemm" + else breakdown[backend_name][component_name] + ) + ax6.bar( + x_values, + values, + bar_width, + bottom=bottoms, + label=label if backend_name == "legacy" else None, + color=color, + alpha=0.70 if backend_name == "legacy" else 1.0, + ) + bottoms = [bottoms[i] + values[i] for i in range(3)] + totals_by_backend[backend_name] = bottoms # Formatting ax6.set_ylabel("Time (ms)", fontsize=12) ax6.set_title(f"Kernel Breakdown (M={M_large:,})", fontsize=13) - ax6.set_xticks(x_pos) + ax6.set_xticks(group_centers) ax6.set_xticklabels(passes, fontsize=10) ax6.grid(True, alpha=0.3, axis="y") + for backend_name, x_values in x_by_backend.items(): + for x_value in x_values: + ax6.text( + x_value, + -0.04, + backend_name, + ha="center", + va="top", + fontsize=8, + rotation=45, + transform=ax6.get_xaxis_transform(), + ) # Add total time labels on top of each bar - totals = [ - sum([quant_1[i], quant_2[i], rearrange_1[i], rearrange_2[i], gemm[i]]) - for i in range(3) - ] - for i, (pos, total) in enumerate(zip(x_pos, totals)): - ax6.text(pos, total, f"{total:.1f}", ha="center", va="bottom", fontsize=9) + for backend_name, x_values in x_by_backend.items(): + for pos, total in zip(x_values, totals_by_backend[backend_name]): + ax6.text(pos, total, f"{total:.1f}", ha="center", va="bottom", fontsize=8) # Add BF16 GEMM baseline reference lines # Forward and Backward Input use 2D/3D GEMM @@ -1616,29 +1958,19 @@ def run( ] # Draw horizontal lines for BF16 baseline at each bar position - bar_width_visual = 0.4 # Visual width for the line - ax6.plot( - [x_pos[0] - bar_width_visual, x_pos[0] + bar_width_visual], - [bf16_fwd_gemm_ms, bf16_fwd_gemm_ms], - color="red", - linestyle="--", - linewidth=2, - label="BF16 GEMM baseline", - ) - ax6.plot( - [x_pos[1] - bar_width_visual, x_pos[1] + bar_width_visual], - [bf16_fwd_gemm_ms, bf16_fwd_gemm_ms], - color="red", - linestyle="--", - linewidth=2, - ) - ax6.plot( - [x_pos[2] - bar_width_visual, x_pos[2] + bar_width_visual], - [bf16_bwd_weight_gemm_ms, bf16_bwd_weight_gemm_ms], - color="red", - linestyle="--", - linewidth=2, - ) + bar_width_visual = 0.42 + for i, baseline_ms in enumerate( + [bf16_fwd_gemm_ms, bf16_fwd_gemm_ms, bf16_bwd_weight_gemm_ms] + ): + label = "BF16 GEMM baseline" if i == 0 else None + ax6.plot( + [group_centers[i] - bar_width_visual, group_centers[i] + bar_width_visual], + [baseline_ms, baseline_ms], + color="red", + linestyle="--", + linewidth=2, + label=label, + ) # Add legend after all plot elements are added ax6.legend(loc="upper right", fontsize=8) @@ -1651,16 +1983,34 @@ def run( print("\n" + "=" * 80) print("SUMMARY STATISTICS") print("=" * 80) + if "actual_speedup" in df_speedup: + actual_speedup_summary = ( + f" Average actual speedup: {df_speedup['actual_speedup'].mean():.3f}x\n" + f" Median actual speedup: {df_speedup['actual_speedup'].median():.3f}x" + ) + else: + summary_lines = [] + for column in df_speedup.columns: + if column.startswith("actual_speedup_"): + backend_name = column.removeprefix("actual_speedup_") + summary_lines.append( + f" Average actual speedup ({backend_name}): {df_speedup[column].mean():.3f}x" + ) + summary_lines.append( + f" Median actual speedup ({backend_name}): {df_speedup[column].median():.3f}x" + ) + actual_speedup_summary = "\n".join(summary_lines) print( f""" Net Speedup Analysis: Average roofline speedup: {df_speedup["roofline_speedup"].mean():.3f}x - Average actual speedup: {df_speedup["actual_speedup"].mean():.3f}x - Median actual speedup: {df_speedup["actual_speedup"].median():.3f}x +{actual_speedup_summary} 2D Quantization Kernels: triton_to_mxfp8_dim0 avg efficiency: {df_quant_2d["triton_dim0_efficiency_pct"].mean():.1f}% + mxfp8_quantize_2d_1x32_cutedsl avg efficiency: {df_quant_2d["cutedsl_1x32_efficiency_pct"].mean():.1f}% to_mxfp8_dim1_cuda avg efficiency: {df_quant_2d["cuda_dim1_efficiency_pct"].mean():.1f}% + mxfp8_quantize_2d_32x1_cutedsl avg efficiency: {df_quant_2d["cutedsl_32x1_efficiency_pct"].mean():.1f}% 3D Quantization Kernels: mxfp8_quantize_cuda_3d avg efficiency: {df_quant_3d["cuda_3d_efficiency_pct"].mean():.1f}% @@ -1674,9 +2024,9 @@ def run( Configuration: K={K}, N={N}, G={G} Power Limit: {power_limit_percent}% - Peak BW: {model.memory_bandwidth_gbs:.1f} GB/s - Peak BF16 TFLOPS: {model.bf16_tflops:.1f} TFLOPS - Peak MXFP8 TFLOPS: {model.mxfp8_tflops:.1f} TFLOPS + Achievable BW roofline: {model.memory_bandwidth_gbs:.1f} GB/s + Achievable BF16 roofline: {model.bf16_tflops:.1f} TFLOPS + Achievable MXFP8 roofline: {model.mxfp8_tflops:.1f} TFLOPS """ ) print("=" * 80) diff --git a/scripts/validate_mxfp8_linear_cutedsl.py b/scripts/validate_mxfp8_linear_cutedsl.py new file mode 100644 index 00000000000..ca56afa85c6 --- /dev/null +++ b/scripts/validate_mxfp8_linear_cutedsl.py @@ -0,0 +1,212 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import argparse + +import torch +import torch.nn.functional as F + +from benchmarks.utils import benchmark_cuda_function_in_microseconds +from torchao.float8.float8_utils import compute_error +from torchao.prototype.moe_training.kernels.mxfp8.cute_utils import ( + _missing_cutedsl_runtime_packages, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + _mxfp8_cutedsl_kernels_available, +) +from torchao.prototype.moe_training.mxfp8_linear import ( + _to_mxfp8_then_scaled_mm, + set_mxfp8_linear_backend, +) +from torchao.prototype.mx_formats.config import ScaleCalculationMode +from torchao.quantization.quantize_.common import KernelPreference + + +def benchmark_cuda_graph_function_in_microseconds(f, *args, iters=1000, **kwargs): + for _ in range(10): + f(*args, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + f(*args, **kwargs) + for _ in range(10): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def benchmark_function(f, *args, use_cuda_graph=False, graph_iters=1000, **kwargs): + if use_cuda_graph: + return benchmark_cuda_graph_function_in_microseconds( + f, + *args, + iters=graph_iters, + **kwargs, + ) + return benchmark_cuda_function_in_microseconds(f, *args, **kwargs) + + +def run(x, weight, labels, scale_mode, wgrad_with_hp, backend): + set_mxfp8_linear_backend(backend) + if x.grad is not None: + x.grad.zero_() + if weight.grad is not None: + weight.grad.zero_() + out = _to_mxfp8_then_scaled_mm( + x, + weight, + kernel_preference=KernelPreference.AUTO, + scale_calculation_mode=scale_mode, + wgrad_with_hp=wgrad_with_hp, + backend=backend, + ) + F.mse_loss(out, labels).backward() + return out, x.grad, weight.grad + + +def benchmark_backend( + x, + weight, + labels, + scale_mode, + wgrad_with_hp, + backend, + use_cuda_graph=False, + graph_iters=1000, +): + x_bench = x.detach().clone().requires_grad_(True) + weight_bench = weight.detach().clone().requires_grad_(True) + return benchmark_function( + run, + x_bench, + weight_bench, + labels, + scale_mode, + wgrad_with_hp, + backend, + use_cuda_graph=use_cuda_graph, + graph_iters=graph_iters, + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--M", "--m", dest="m", type=int, default=4096) + parser.add_argument("--N", "--n", dest="n", type=int, default=4096) + parser.add_argument("--K", "--k", dest="k", type=int, default=2048) + parser.add_argument("--scale-mode", choices=("floor", "rceil"), default="rceil") + parser.add_argument("--wgrad-with-hp", action="store_true") + parser.add_argument( + "--bench-only", + choices=("legacy", "cutedsl"), + default=None, + ) + parser.add_argument("--deterministic", action="store_true") + parser.add_argument("--cuda-graph-bench", action="store_true") + parser.add_argument("--graph-iters", type=int, default=1000) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is not available") + if torch.cuda.get_device_capability()[0] != 10: + raise SystemExit(f"Requires SM 10.x; got {torch.cuda.get_device_capability()}") + if not _mxfp8_cutedsl_kernels_available: + missing = _missing_cutedsl_runtime_packages() + detail = ( + f"missing package(s): {', '.join(missing)}" + if missing + else "requires CUDA 12.8+ on SM 10.x" + ) + raise SystemExit(f"CuteDSL mxfp8 kernels unavailable ({detail})") + + scale_mode = ScaleCalculationMode(args.scale_mode) + if args.deterministic: + x = ( + torch.arange(args.m * args.k, device="cuda", dtype=torch.float32) + .remainder(251) + .reshape(args.m, args.k) + .to(torch.bfloat16) + ) + weight = ( + torch.arange(args.n * args.k, device="cuda", dtype=torch.float32) + .remainder(251) + .reshape(args.n, args.k) + .to(torch.bfloat16) + ) + else: + torch.manual_seed(0) + x = torch.randn(args.m, args.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(args.n, args.k, dtype=torch.bfloat16, device="cuda") + labels = torch.ones((args.m, args.n), dtype=torch.bfloat16, device="cuda") + + print( + f"shape=({args.m}, {args.n}, {args.k}) wgrad_with_hp={args.wgrad_with_hp} scale_mode={args.scale_mode}" + ) + if args.bench_only is None: + x_legacy = x.detach().clone().requires_grad_(True) + weight_legacy = weight.detach().clone().requires_grad_(True) + out_legacy, x_grad_legacy, weight_grad_legacy = run( + x_legacy, + weight_legacy, + labels, + scale_mode, + args.wgrad_with_hp, + "legacy", + ) + x_cutedsl = x.detach().clone().requires_grad_(True) + weight_cutedsl = weight.detach().clone().requires_grad_(True) + out_cutedsl, x_grad_cutedsl, weight_grad_cutedsl = run( + x_cutedsl, + weight_cutedsl, + labels, + scale_mode, + args.wgrad_with_hp, + "cutedsl", + ) + torch.cuda.synchronize() + + output_sqnr = compute_error(out_legacy, out_cutedsl) + input_grad_sqnr = compute_error(x_grad_legacy, x_grad_cutedsl) + weight_grad_sqnr = compute_error(weight_grad_legacy, weight_grad_cutedsl) + print(f"output_sqnr={output_sqnr:.2f}") + print(f"input_grad_sqnr={input_grad_sqnr:.2f}") + print(f"weight_grad_sqnr={weight_grad_sqnr:.2f}") + + if args.bench_only in (None, "legacy"): + legacy_us = benchmark_backend( + x, + weight, + labels, + scale_mode, + args.wgrad_with_hp, + "legacy", + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + print(f"legacy_us={legacy_us:.2f}") + if args.bench_only in (None, "cutedsl"): + cutedsl_us = benchmark_backend( + x, + weight, + labels, + scale_mode, + args.wgrad_with_hp, + "cutedsl", + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + print(f"cutedsl_us={cutedsl_us:.2f}") + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_mxfp8_grouped_mm.py b/test/prototype/moe_training/test_mxfp8_grouped_mm.py index 3a2cb82a623..9345669ea0e 100644 --- a/test/prototype/moe_training/test_mxfp8_grouped_mm.py +++ b/test/prototype/moe_training/test_mxfp8_grouped_mm.py @@ -35,6 +35,7 @@ _emulated_mxfp8_scaled_grouped_mm_2d_2d, _emulated_mxfp8_scaled_grouped_mm_2d_3d, _to_mxfp8_then_scaled_grouped_mm, + set_mxfp8_grouped_mm_backend, ) from torchao.prototype.moe_training.utils import ( _to_mxfp8_per_group_colwise, @@ -341,6 +342,190 @@ def test_mxfp8_grouped_gemm_with_dq_fwd_bwd( ) +@skip_if_rocm("ROCm not supported") +@pytest.mark.skipif( + not is_sm_version(10, 0), + reason="MXFP8 grouped GEMM requires SM100", +) +@pytest.mark.parametrize("backend", ("legacy", "cutedsl")) +def test_mxfp8_grouped_gemm_backend_fwd_bwd(backend): + M, K, N, num_experts = 4096, 1024, 2048, 8 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(num_experts, N, K, dtype=torch.bfloat16, device="cuda") + w_t = w.transpose(-2, -1).requires_grad_(True) + offs = generate_jagged_offs(num_experts, M, multiple_of=128) + + x_ref = x.detach().clone().requires_grad_(True) + w_t_ref = w_t.detach().clone().requires_grad_(True) + + set_mxfp8_grouped_mm_backend(backend) + try: + out = _to_mxfp8_then_scaled_grouped_mm( + x, + w_t, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.AUTO, + wgrad_with_hp=False, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=False, + ) + + ref_out = torch._grouped_mm(x_ref, w_t_ref, offs=offs, out_dtype=torch.bfloat16) + sqnr = compute_error(ref_out, out) + min_sqnr = 27.0 + assert sqnr >= min_sqnr, f"Output sqnr {sqnr} is too low, must be >= {min_sqnr}" + + labels = torch.ones_like(ref_out) + F.mse_loss(ref_out, labels).backward() + F.mse_loss(out, labels).backward() + + input_grad_sqnr = compute_error(x_ref.grad, x.grad) + min_input_grad_sqnr = 25.0 + assert input_grad_sqnr >= min_input_grad_sqnr, ( + f"Input grad sqnr {input_grad_sqnr} is too low, " + f"must be >= {min_input_grad_sqnr}" + ) + + weight_grad_sqnr = compute_error(w_t_ref.grad, w_t.grad) + min_weight_grad_sqnr = 24.0 + assert weight_grad_sqnr >= min_weight_grad_sqnr, ( + f"Weight grad sqnr {weight_grad_sqnr} is too low, " + f"must be >= {min_weight_grad_sqnr}" + ) + finally: + set_mxfp8_grouped_mm_backend("cutedsl") + + +@skip_if_rocm("ROCm not supported") +@pytest.mark.skipif( + not is_sm_version(10, 0), + reason="MXFP8 grouped GEMM requires SM100", +) +@pytest.mark.parametrize("backend", ("legacy", "cutedsl")) +def test_mxfp8_grouped_gemm_padded_backend_fwd_bwd(backend): + K, N, num_experts = 1024, 2048, 8 + group_sizes = [97, 113, 121, 99, 111, 123, 109, 123] + M = sum(group_sizes) + offs = torch.tensor(group_sizes, dtype=torch.int32, device="cuda").cumsum( + 0, dtype=torch.int32 + ) + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(num_experts, N, K, dtype=torch.bfloat16, device="cuda") + w_t = w.transpose(-2, -1).requires_grad_(True) + + x_ref = x.detach().clone().requires_grad_(True) + w_t_ref = w_t.detach().clone().requires_grad_(True) + + set_mxfp8_grouped_mm_backend(backend) + try: + out = _to_mxfp8_then_scaled_grouped_mm( + x, + w_t, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.AUTO, + wgrad_with_hp=True, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=True, + ) + + ref_out = _to_mxfp8_then_scaled_grouped_mm( + x_ref, + w_t_ref, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.EMULATED, + wgrad_with_hp=True, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=True, + ) + + output_sqnr = compute_error(ref_out, out) + min_output_sqnr = 27.0 + assert output_sqnr >= min_output_sqnr, ( + f"Output sqnr {output_sqnr} is too low, must be >= {min_output_sqnr}" + ) + + labels = torch.ones_like(ref_out) + F.mse_loss(ref_out, labels).backward() + F.mse_loss(out, labels).backward() + + input_grad_sqnr = compute_error(x_ref.grad, x.grad) + min_input_grad_sqnr = 25.0 + assert input_grad_sqnr >= min_input_grad_sqnr, ( + f"Input grad sqnr {input_grad_sqnr} is too low, " + f"must be >= {min_input_grad_sqnr}" + ) + + weight_grad_sqnr = compute_error(w_t_ref.grad, w_t.grad) + min_weight_grad_sqnr = 24.0 + assert weight_grad_sqnr >= min_weight_grad_sqnr, ( + f"Weight grad sqnr {weight_grad_sqnr} is too low, " + f"must be >= {min_weight_grad_sqnr}" + ) + finally: + set_mxfp8_grouped_mm_backend("cutedsl") + + +@skip_if_rocm("ROCm not supported") +@pytest.mark.skipif( + not is_sm_version(10, 0), + reason="MXFP8 grouped GEMM requires SM100", +) +@pytest.mark.parametrize("backend", ("legacy", "cutedsl")) +def test_mxfp8_grouped_gemm_mxtensor_backend_forward(backend): + block_size = 32 + M, K, N, num_experts = 4096, 1024, 2048, 8 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(num_experts, N, K, dtype=torch.bfloat16, device="cuda") + w_t = w.transpose(-2, -1) + offs = generate_jagged_offs(num_experts, M, multiple_of=128) + + x_scale, x_qdata = to_mx( + x, + elem_dtype=torch.float8_e4m3fn, + block_size=block_size, + scaling_mode=ScaleCalculationMode.RCEIL, + ) + x_mx = MXTensor.from_qdata_and_scales( + x_qdata, + x_scale, + orig_dtype=x.dtype, + block_size=block_size, + is_swizzled_scales=False, + ) + + set_mxfp8_grouped_mm_backend(backend) + try: + out = _to_mxfp8_then_scaled_grouped_mm( + x_mx, + w_t, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.AUTO, + wgrad_with_hp=True, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=False, + ) + finally: + set_mxfp8_grouped_mm_backend("cutedsl") + + ref_out = _to_mxfp8_then_scaled_grouped_mm( + x, + w_t, + offs=offs, + out_dtype=torch.bfloat16, + kernel_preference=KernelPreference.EMULATED, + wgrad_with_hp=True, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + pad_token_groups_for_grouped_mm=False, + ) + sqnr = compute_error(ref_out, out) + min_sqnr = 27.0 + assert sqnr >= min_sqnr, f"Output sqnr {sqnr} is too low, must be >= {min_sqnr}" + + @skip_if_rocm("ROCm not supported") def test_mxfp8_grouped_gemm_from_qdata_and_scales_matches_dynamic(): block_size = 32 diff --git a/test/prototype/moe_training/test_mxfp8_linear.py b/test/prototype/moe_training/test_mxfp8_linear.py index 4178b79bd9a..82ce67f503b 100644 --- a/test/prototype/moe_training/test_mxfp8_linear.py +++ b/test/prototype/moe_training/test_mxfp8_linear.py @@ -14,7 +14,10 @@ if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) -from torchao.prototype.moe_training.mxfp8_linear import MXFP8Linear +from torchao.prototype.moe_training.mxfp8_linear import ( + MXFP8Linear, + set_mxfp8_linear_backend, +) from torchao.quantization.quantize_.common.kernel_preference import KernelPreference @@ -53,18 +56,29 @@ def _build_linear_pair( @pytest.mark.parametrize( "kernel_preference", [KernelPreference.EMULATED, KernelPreference.AUTO] ) -def test_mxfp8_linear_fwd_bwd_sqnr(bias, wgrad_with_hp, kernel_preference): +@pytest.mark.parametrize("backend", ("legacy", "cutedsl")) +def test_mxfp8_linear_fwd_bwd_sqnr(bias, wgrad_with_hp, kernel_preference, backend): if kernel_preference == KernelPreference.AUTO and not is_sm_at_least_100(): pytest.skip("Real MXFP8 kernels require SM100+") M, K, N = 1024, 1024, 2048 - ref, mxfp8 = _build_linear_pair(K, N, bias, kernel_preference, wgrad_with_hp) + ref, mxfp8 = _build_linear_pair( + K, + N, + bias, + kernel_preference, + wgrad_with_hp, + ) x_ref = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) x_mxfp8 = x_ref.clone().detach().requires_grad_(True) out_ref = ref(x_ref) - out_mxfp8 = mxfp8(x_mxfp8) + set_mxfp8_linear_backend(backend) + try: + out_mxfp8 = mxfp8(x_mxfp8) + finally: + set_mxfp8_linear_backend("cutedsl") # Forward checks assert out_mxfp8.shape == out_ref.shape diff --git a/torchao/prototype/moe_training/mxfp8_grouped_mm.py b/torchao/prototype/moe_training/mxfp8_grouped_mm.py index eaf3a1c5afe..e306dbac442 100644 --- a/torchao/prototype/moe_training/mxfp8_grouped_mm.py +++ b/torchao/prototype/moe_training/mxfp8_grouped_mm.py @@ -13,16 +13,27 @@ _mxfp8_cuda_kernels_available as _mxfp8_cuda_kernels_available_quant, ) from torchao.prototype.moe_training.kernels.mxfp8 import ( + mx_block_rearrange_2d_k_groups_cutedsl, mx_block_rearrange_2d_M_groups_cuda, mxfp8_quantize_2d_1x32_cutedsl, + mxfp8_quantize_2d_32x1_cutedsl, mxfp8_quantize_cuda_3d, + pad_token_groups_cutedsl, triton_mx_block_rearrange_2d_K_groups, triton_mx_block_rearrange_per_group_3d, + unpad_token_groups_cutedsl, +) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( + mx_block_rearrange_2d_m_groups_cutedsl, ) from torchao.prototype.moe_training.utils import ( conditional_nostrict_trace, - pad_token_groups, - unpad_token_groups, +) +from torchao.prototype.moe_training.utils import ( + pad_token_groups as pad_token_groups_legacy, +) +from torchao.prototype.moe_training.utils import ( + unpad_token_groups as unpad_token_groups_legacy, ) from torchao.prototype.mx_formats.config import ( MXFP8Dim1CastKernelChoice, @@ -42,6 +53,75 @@ logger: logging.Logger = logging.getLogger(__name__) +_MXFP8_GROUPED_MM_BACKEND = "cutedsl" + + +def set_mxfp8_grouped_mm_backend(backend: str) -> None: + assert backend in ("legacy", "cutedsl") + global _MXFP8_GROUPED_MM_BACKEND + _MXFP8_GROUPED_MM_BACKEND = backend + + +def get_mxfp8_grouped_mm_backend() -> str: + return _MXFP8_GROUPED_MM_BACKEND + + +def _pad_token_groups( + input_act: torch.Tensor, + group_end_offsets: torch.Tensor, + alignment_size: int, + kernel_preference: KernelPreference, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if ( + _MXFP8_GROUPED_MM_BACKEND == "cutedsl" + and kernel_preference != KernelPreference.EMULATED + ): + alignment_size = max(alignment_size, 128) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + return pad_token_groups_cutedsl( + input_act, + group_end_offsets, + alignment_size=alignment_size, + ) + return pad_token_groups_legacy( + input_act, + group_end_offsets, + alignment_size=alignment_size, + kernel_preference=kernel_preference, + ) + + +def _unpad_token_groups( + padded_output: torch.Tensor, + original_group_end_offsets: torch.Tensor, + padded_group_start_offsets: torch.Tensor, + num_tokens: int, + alignment_size: int, + kernel_preference: KernelPreference, +) -> torch.Tensor: + if ( + _MXFP8_GROUPED_MM_BACKEND == "cutedsl" + and kernel_preference != KernelPreference.EMULATED + ): + alignment_size = max(alignment_size, 128) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + return unpad_token_groups_cutedsl( + padded_output, + original_group_end_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size=alignment_size, + ) + return unpad_token_groups_legacy( + padded_output, + original_group_end_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size=alignment_size, + kernel_preference=kernel_preference, + ) + + # Check if SM100 kernels are available # All SM100-dependent kernels are guarded at their definition sites _SM100_KERNELS_AVAILABLE = ( @@ -185,7 +265,7 @@ def forward( # Conditionally pad token groups if not aligned to block_size if pad_token_groups_for_grouped_mm: padded_input_act, padded_group_start_offsets, padded_group_end_offsets = ( - pad_token_groups( + _pad_token_groups( input_act, group_end_offsets, alignment_size=block_size, @@ -210,7 +290,7 @@ def forward( # Unpad output if padding was used if pad_token_groups_for_grouped_mm: - output = unpad_token_groups( + output = _unpad_token_groups( output, group_end_offsets, padded_group_start_offsets, @@ -271,7 +351,7 @@ def backward(ctx, grad_output: torch.Tensor): if pad_token_groups_for_grouped_mm: # padded start/end offsets same as what we saved from forward. # will be original offsets if we aren't using pad/unpad path - padded_grad_output, _, _ = pad_token_groups( + padded_grad_output, _, _ = _pad_token_groups( grad_output, original_group_end_offsets, alignment_size=block_size, @@ -306,7 +386,7 @@ def backward(ctx, grad_output: torch.Tensor): # Unpad grad_input if padding was used if pad_token_groups_for_grouped_mm: - grad_input = unpad_token_groups( + grad_input = _unpad_token_groups( grad_input, original_group_end_offsets, padded_group_start_offsets, @@ -521,9 +601,16 @@ def _compute_fwd_sm100( ) if not padded_input_act.is_swizzled_scales: # Convert scales to blocked layout for SM100 kernels - input_act_scales_blocked = mx_block_rearrange_2d_M_groups_cuda( - padded_input_act.scales, padded_group_end_offsets - ) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + input_act_scales_blocked = mx_block_rearrange_2d_m_groups_cutedsl( + padded_input_act.scale, + padded_group_end_offsets, + ) + else: + input_act_scales_blocked = mx_block_rearrange_2d_M_groups_cuda( + padded_input_act.scale, + padded_group_end_offsets, + ) else: input_act_e4m3, input_act_scales_blocked = mxfp8_quantize_2d_1x32_cutedsl( padded_input_act, @@ -531,16 +618,29 @@ def _compute_fwd_sm100( offs=padded_group_end_offsets, ) - # Quantize weights along dim0 (after transposing from (E, K, N) to (E, N, K)) - weight_e4m3, weight_scales = triton_to_mxfp8_dim0( - weight_t.transpose(-2, -1), block_size, scale_calculation_mode.value.lower() - ) - weight_scales_blocked = triton_mx_block_rearrange_per_group_3d(weight_scales) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + # Quantize weights directly to blocked tcgen05 scales. + weight_e4m3, weight_scales_blocked = mxfp8_quantize_cuda_3d( + weight_t._data if hasattr(weight_t, "_data") else weight_t, + block_size, + scale_block_dim1=block_size, + scale_block_dim2=1, + scaling_mode=scale_calculation_mode.value.lower(), + ) + weight_e4m3_for_gemm = weight_e4m3 + else: + weight_e4m3, weight_scales = triton_to_mxfp8_dim0( + weight_t.transpose(-2, -1), + block_size, + scale_calculation_mode.value.lower(), + ) + weight_scales_blocked = triton_mx_block_rearrange_per_group_3d(weight_scales) + weight_e4m3_for_gemm = weight_e4m3.transpose(-2, -1) # Compute output using SM100 kernel output = torch._scaled_grouped_mm( input_act_e4m3, - weight_e4m3.transpose(-2, -1), # Transpose back to (E, K, N) + weight_e4m3_for_gemm, input_act_scales_blocked, weight_scales_blocked, offs=padded_group_end_offsets, @@ -621,9 +721,16 @@ def _compute_dgrad_sm100( grad_out_e4m3, grad_output_scales_blocked = grad_output.qdata, grad_output.scale if not grad_output.is_swizzled_scales: # Convert scales to blocked layout for SM100 kernels - grad_output_scales_blocked = mx_block_rearrange_2d_M_groups_cuda( - grad_output.scales, group_end_offsets - ) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + grad_output_scales_blocked = mx_block_rearrange_2d_m_groups_cutedsl( + grad_output.scale, + group_end_offsets, + ) + else: + grad_output_scales_blocked = mx_block_rearrange_2d_M_groups_cuda( + grad_output.scale, + group_end_offsets, + ) else: grad_out_e4m3, grad_output_scales_blocked = mxfp8_quantize_2d_1x32_cutedsl( grad_output, @@ -746,41 +853,66 @@ def _compute_wgrad_sm100( ) return grad_weight.transpose(-2, -1) - # Use CUDA kernel for dim1 quant - grad_output_t_mx = _to_mxfp8_dim1_kernel_wrapper( - grad_output, - block_size, - elem_dtype=torch.float8_e4m3fn, - hp_dtype=grad_output.dtype, - kernel_preference=KernelPreference.AUTO, - cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, - scale_calculation_mode=scale_calculation_mode, - ) - grad_output_t_data = grad_output_t_mx.qdata - grad_output_t_scales = grad_output_t_mx.scale + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + grad_output_t_data, grad_output_t_scales = mxfp8_quantize_2d_32x1_cutedsl( + grad_output, + block_size=block_size, + scaling_mode=scale_calculation_mode.value.lower(), + blocked_scale_output=False, + ) + grad_output_t_data = grad_output_t_data.t() + input_act_t_data, input_act_t_scales = mxfp8_quantize_2d_32x1_cutedsl( + input_act, + block_size=block_size, + scaling_mode=scale_calculation_mode.value.lower(), + blocked_scale_output=False, + ) + input_act_t_data = input_act_t_data.t() + else: + grad_output_t_mx = _to_mxfp8_dim1_kernel_wrapper( + grad_output, + block_size, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=grad_output.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=scale_calculation_mode, + ) + grad_output_t_data = grad_output_t_mx.qdata + grad_output_t_scales = grad_output_t_mx.scale - input_act_t_mx = _to_mxfp8_dim1_kernel_wrapper( - input_act, - block_size, - elem_dtype=torch.float8_e4m3fn, - hp_dtype=input_act.dtype, - kernel_preference=KernelPreference.AUTO, - cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, - scale_calculation_mode=scale_calculation_mode, - ) - input_act_t_data = input_act_t_mx.qdata - input_act_t_scales = input_act_t_mx.scale + input_act_t_mx = _to_mxfp8_dim1_kernel_wrapper( + input_act, + block_size, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=input_act.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=scale_calculation_mode, + ) + input_act_t_data = input_act_t_mx.qdata + input_act_t_scales = input_act_t_mx.scale # Convert scales to blocked layout for SM100 kernels scale_group_offsets = group_end_offsets // block_size - grad_output_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( - grad_output_t_scales, - scale_group_offsets, - ) - input_act_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( - input_act_t_scales, - scale_group_offsets, - ) + if _MXFP8_GROUPED_MM_BACKEND == "cutedsl": + grad_output_t_scales_blocked = mx_block_rearrange_2d_k_groups_cutedsl( + grad_output_t_scales, + scale_group_offsets, + ) + input_act_t_scales_blocked = mx_block_rearrange_2d_k_groups_cutedsl( + input_act_t_scales, + scale_group_offsets, + ) + else: + grad_output_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + grad_output_t_scales, + scale_group_offsets, + ) + input_act_t_scales_blocked = triton_mx_block_rearrange_2d_K_groups( + input_act_t_scales, + scale_group_offsets, + ) # Compute grad_weight = grad_output_t @ input_act grad_weight = torch._scaled_grouped_mm( diff --git a/torchao/prototype/moe_training/mxfp8_linear.py b/torchao/prototype/moe_training/mxfp8_linear.py index 3e33e811cde..bb1125a584d 100644 --- a/torchao/prototype/moe_training/mxfp8_linear.py +++ b/torchao/prototype/moe_training/mxfp8_linear.py @@ -13,6 +13,9 @@ import torch from torch import nn +from torchao.prototype.moe_training.kernels.mxfp8 import ( + mxfp8_quantize_2d_1x32_cutedsl, +) from torchao.prototype.mx_formats.config import ( MXFP8Dim0CastKernelChoice, MXFP8Dim1CastKernelChoice, @@ -22,6 +25,54 @@ from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper from torchao.quantization.quantize_.common.kernel_preference import KernelPreference +_MXFP8_LINEAR_BACKEND = "cutedsl" + + +def set_mxfp8_linear_backend(backend: str) -> None: + assert backend in ("legacy", "cutedsl") + global _MXFP8_LINEAR_BACKEND + _MXFP8_LINEAR_BACKEND = backend + + +def get_mxfp8_linear_backend() -> str: + return _MXFP8_LINEAR_BACKEND + + +def _to_mxfp8_dim0( + data_hp: torch.Tensor, + elem_dtype: torch.dtype, + block_size: int, + scale_calculation_mode: ScaleCalculationMode, + kernel_preference: KernelPreference, + mxfp8_dim0_cast_kernel_choice: MXFP8Dim0CastKernelChoice, + backend: str, +) -> MXTensor: + if kernel_preference == KernelPreference.AUTO and backend == "cutedsl": + qdata, scale = mxfp8_quantize_2d_1x32_cutedsl( + data_hp, + block_size=block_size, + scaling_mode=scale_calculation_mode.value, + ) + return MXTensor( + qdata, + scale, + elem_dtype, + block_size, + data_hp.dtype, + kernel_preference, + None, + True, + ) + + return MXTensor.to_mx( + data_hp, + elem_dtype, + block_size, + scale_calculation_mode, + kernel_preference, + mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, + ) + # convenience wrapper def _to_mxfp8_then_scaled_mm( @@ -30,6 +81,7 @@ def _to_mxfp8_then_scaled_mm( kernel_preference: KernelPreference, scale_calculation_mode: ScaleCalculationMode, wgrad_with_hp: bool = False, + backend: str | None = None, ) -> torch.Tensor: """ Performs a matrix multiplication with MXFP8 quantization on both forward and backward passes. @@ -63,7 +115,13 @@ def _to_mxfp8_then_scaled_mm( grad_elem_dtype = torch.float8_e4m3fn block_size = 32 mxfp8_dim0_cast_kernel_choice = MXFP8Dim0CastKernelChoice.TRITON - mxfp8_dim1_cast_kernel_choice = MXFP8Dim1CastKernelChoice.CUDA + backend = get_mxfp8_linear_backend() if backend is None else backend + assert backend in ("legacy", "cutedsl") + mxfp8_dim1_cast_kernel_choice = ( + MXFP8Dim1CastKernelChoice.CUTEDSL + if backend == "cutedsl" + else MXFP8Dim1CastKernelChoice.CUDA + ) return mx_mm.apply( input_hp, @@ -77,6 +135,7 @@ def _to_mxfp8_then_scaled_mm( mxfp8_dim1_cast_kernel_choice, scale_calculation_mode, wgrad_with_hp, + backend, ) @@ -104,6 +163,7 @@ def forward( mxfp8_dim1_cast_kernel_choice: MXFP8Dim1CastKernelChoice, scale_calculation_mode: ScaleCalculationMode, wgrad_with_hp: bool, + backend: str, ): ctx.save_for_backward(input_hp, weight_hp) ctx.in_elem_dtype = in_elem_dtype @@ -115,26 +175,29 @@ def forward( ctx.mxfp8_dim0_cast_kernel_choice = mxfp8_dim0_cast_kernel_choice ctx.mxfp8_dim1_cast_kernel_choice = mxfp8_dim1_cast_kernel_choice ctx.scale_calculation_mode = scale_calculation_mode + ctx.backend = backend # input @ weight_t = output input_orig_shape = input_hp.shape input_hp_r = input_hp.reshape(-1, input_orig_shape[-1]) - input_mx_r_dim0 = MXTensor.to_mx( + input_mx_r_dim0 = _to_mxfp8_dim0( input_hp_r, in_elem_dtype, block_size, scale_calculation_mode, kernel_preference, - mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, + mxfp8_dim0_cast_kernel_choice, + backend, ) - weight_mx_dim0 = MXTensor.to_mx( + weight_mx_dim0 = _to_mxfp8_dim0( weight_hp, w_elem_dtype, block_size, scale_calculation_mode, kernel_preference, - mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, + mxfp8_dim0_cast_kernel_choice, + backend, ) output = torch.mm(input_mx_r_dim0, weight_mx_dim0.t()) output = output.reshape(*input_orig_shape[:-1], output.shape[-1]) @@ -153,6 +216,7 @@ def backward(ctx, grad_output_hp: torch.Tensor): mxfp8_dim1_cast_kernel_choice = ctx.mxfp8_dim1_cast_kernel_choice scale_calculation_mode = ctx.scale_calculation_mode wgrad_with_hp = ctx.wgrad_with_hp + backend = ctx.backend # grad_output may be non-contiguous (e.g. produced by a transposed or # otherwise strided downstream op). Both the dim0 cast (MXTensor.to_mx / @@ -167,13 +231,14 @@ def backward(ctx, grad_output_hp: torch.Tensor): input_hp_r = input_hp.reshape(-1, input_hp_orig_shape[-1]) # grad_output @ weight = grad_input - grad_output_mx_dim0 = MXTensor.to_mx( + grad_output_mx_dim0 = _to_mxfp8_dim0( grad_output_hp_r, grad_elem_dtype, block_size, scale_calculation_mode, kernel_preference, - mxfp8_dim0_cast_kernel_choice=mxfp8_dim0_cast_kernel_choice, + mxfp8_dim0_cast_kernel_choice, + backend, ) if ( @@ -266,6 +331,7 @@ def backward(ctx, grad_output_hp: torch.Tensor): None, None, None, + None, ) @@ -280,6 +346,7 @@ def __init__( kernel_preference: KernelPreference = KernelPreference.AUTO, scale_calculation_mode: ScaleCalculationMode = ScaleCalculationMode.RCEIL, wgrad_with_hp: bool = False, + backend: str | None = None, **kwargs, ): """ @@ -292,6 +359,7 @@ def __init__( self.kernel_preference = kernel_preference self.scale_calculation_mode = scale_calculation_mode self.wgrad_with_hp = wgrad_with_hp + self.backend = backend def forward(self, input: torch.Tensor) -> torch.Tensor: output = _to_mxfp8_then_scaled_mm( @@ -300,6 +368,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: kernel_preference=self.kernel_preference, scale_calculation_mode=self.scale_calculation_mode, wgrad_with_hp=self.wgrad_with_hp, + backend=self.backend, ) if self.bias is not None: output = output + self.bias.to(output.dtype) diff --git a/torchao/testing/training/roofline_utils.py b/torchao/testing/training/roofline_utils.py index 9390afd47b5..5c76fae8966 100644 --- a/torchao/testing/training/roofline_utils.py +++ b/torchao/testing/training/roofline_utils.py @@ -117,6 +117,83 @@ def get_specs(gpu_name: Optional[str] = None): return gpu_name_to_specs[gpu_name] +def get_device_peak_mem_bw_bytes_sec() -> Optional[float]: + props = torch.cuda.get_device_properties(0) + memory_clock_khz = getattr(props, "memory_clock_rate", 0) + memory_bus_width_bits = getattr(props, "memory_bus_width", 0) + if memory_clock_khz <= 0 or memory_bus_width_bits <= 0: + return None + return (memory_bus_width_bits / 8.0) * (memory_clock_khz * 1e3) * 2.0 + + +def _benchmark_cuda_ms(fn): + from triton.testing import do_bench + + return do_bench(fn, return_mode="median") + + +def _measure_copy_bandwidth_bytes_sec(numel: int) -> float: + src = torch.randn(numel, device="cuda", dtype=torch.bfloat16) + dst = torch.empty_like(src) + dst.copy_(src) + torch.cuda.synchronize() + ms = _benchmark_cuda_ms(lambda: dst.copy_(src)) + bytes_moved = 2 * src.numel() * src.element_size() + return bytes_moved / (ms / 1000.0) + + +def _measure_bf16_mm_tops(m: int, n: int, k: int) -> float: + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + b = torch.randn(k, n, device="cuda", dtype=torch.bfloat16) + torch.mm(a, b) + torch.cuda.synchronize() + ms = _benchmark_cuda_ms(lambda: torch.mm(a, b)) + return (2 * m * n * k) / (ms / 1000.0) + + +def _measure_fp8_scaled_mm_tops(m: int, n: int, k: int) -> float: + a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) + b_t = torch.randn(n, k, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) + b = b_t.contiguous().t() + scale_a = torch.ones(1, device="cuda") + scale_b = torch.ones(1, device="cuda") + torch._scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16) + torch.cuda.synchronize() + ms = _benchmark_cuda_ms( + lambda: torch._scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16) + ) + return (2 * m * n * k) / (ms / 1000.0) + + +def calibrate_specs( + gpu_name: Optional[str] = None, + mm_shape: tuple[int, int, int] = (16384, 8192, 8192), + copy_numel: int = 256 * 1024 * 1024, +): + specs = dict(get_specs(gpu_name)) + peak_mem_bw = get_device_peak_mem_bw_bytes_sec() + if peak_mem_bw is not None: + specs["peak_mem_bw_bytes_sec"] = peak_mem_bw + + copy_bw = _measure_copy_bandwidth_bytes_sec(copy_numel) + bf16_tops = _measure_bf16_mm_tops(*mm_shape) + fp8_tops = _measure_fp8_scaled_mm_tops(*mm_shape) + + specs["pct_achievable_mem_bw"] = copy_bw / specs["peak_mem_bw_bytes_sec"] + specs["pct_achievable_bf16_gemm_tops"] = bf16_tops / specs["bf16_peak_tops"] + specs["pct_achievable_fp8_gemm_tops"] = fp8_tops / specs["fp8_peak_tops"] + specs["pct_achievable_gemm_tops"] = min( + specs["pct_achievable_bf16_gemm_tops"], + specs["pct_achievable_fp8_gemm_tops"], + ) + specs["calibrated_copy_bandwidth_bytes_sec"] = copy_bw + specs["calibrated_bf16_tops"] = bf16_tops + specs["calibrated_fp8_tops"] = fp8_tops + specs["calibration_mm_shape"] = mm_shape + specs["calibration_copy_numel"] = copy_numel + return specs + + # Source: run a triton kernel with a single element read/write on an H100 and # measure GPU time from the trace # TODO(future): audit this across different hardware and triton/non-triton