# 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. from typing import List, Optional, Union import sympy import torch BYTES_PER_EL_FLOAT8 = 1 BYTES_PER_EL_BF16 = 2 gpu_name_to_specs = { "NVIDIA H100": { # https://www.nvidia.com/en-us/data-center/h100/, divide by 2 because no sparsity "bf16_peak_tops": 989e12, "fp8_peak_tops": 1979e12, # 2.4 TB per second, custom to Meta's H100 variant "peak_mem_bw_bytes_sec": 2.4e12, # based on experimental observation with sample large inputs "pct_achievable_gemm_tops": 0.78, # based on previous experience looking at pointwise triton kernels with large inputs, # which would hit about 2.2k GBPS on Meta's H100 variant "pct_achievable_mem_bw": 0.92, }, "NVIDIA B200": { # https://resources.nvidia.com/en-us-blackwell-architecture, page 19, # divide by 2 because no sparsity "bf16_peak_tops": 2.25e15, "fp8_peak_tops": 4.5e15, "fp4_peak_tops": 9.0e15, # https://resources.nvidia.com/en-us-blackwell-architecture, page 20 # 8.0 TB per second "peak_mem_bw_bytes_sec": 8.0e12, # for now, copy over from H100 # TODO(future): measure once we have the hardware "pct_achievable_gemm_tops": 0.78, # for now, copy over from H100 # TODO(future): measure once we have the hardware "pct_achievable_mem_bw": 0.92, }, "AMD Instinct MI300X": { # https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/data-sheets/amd-instinct-mi300x-data-sheet.pdf, page 1, "bf16_peak_tops": 1307e12, "fp8_peak_tops": 2614e12, # 5.3 TB per second "peak_mem_bw_bytes_sec": 5.3e12, # for now, copy over from H100 # TODO(future): run measurement on hardware "pct_achievable_gemm_tops": 0.78, # for now, copy over from H100 # TODO(future): run measurement on hardware "pct_achievable_mem_bw": 0.92, }, # TODO(future): more GPU names } def get_specs(): gpu_name = torch.cuda.get_device_name(0) return gpu_name_to_specs[gpu_name] # 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 KERNEL_LAUNCH_OVERHEAD_SEC = 0.002 * 0.001 def get_tensor_memory_traffic_ovhd_s( specs, dim0, dim1, tensor_role: str, float8_recipe_name: Optional[str], mx_recipe_name: Optional[str], fuse_with_prev=False, ) -> List[Union[sympy.Symbol, float]]: """ Calculates the roofline estimate of casting one of the gemm inputs (input, weight or grad_output) to float8 in fwd+bwd. Inputs: dim0 and dim1 (shape), tensor_role (input|weight|grad_output), recipe names Outputs: list of read/write traffic overhead in seconds, one for each kernel """ # assumes input bf16, output f8 numel = dim0 * dim1 res_bytes = None if float8_recipe_name == "tensorwise": if tensor_role == "weight": # x_bf16 = ... # kernel 1: x_bf16 -> max_abs_stage_1 -> tmp # kernel 2 (mem traffic not modeled): tmp -> max_abs_stage_2 -> max_abs # kernel 3 (fwd): x_bf16, max_abs -> to_float8 -> x_fp8_dim0 # kernel 4 (bwd): x_bf16, max_abs -> to_float8 -> x_fp8_dim1 if fuse_with_prev: kernel_1_rw = 0 else: # kernel 1: read numel, write 0 (assume size(tmp) ~ 0) kernel_1_rw = BYTES_PER_EL_BF16 * numel # kernel 3: read in bf16, write twice in float8 (row-major and col-major) kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel kernel_4_rw = kernel_3_rw res_bytes = [kernel_1_rw, 0, kernel_3_rw, kernel_4_rw] else: # x_bf16 = ... # kernel 1: x_bf16 -> max_abs_stage_1 -> tmp # kernel 2 (mem traffic not modeled): tmp -> max_abs_stage_2 -> max_abs # kernel 3: x_bf16, max_abs -> to_float8 -> x_fp8_dim0, x_fp8_dim1 if fuse_with_prev: kernel_1_rw = 0 else: # kernel 1: read numel, write 0 (assume size(tmp) ~ 0) kernel_1_rw = BYTES_PER_EL_BF16 * numel # kernel 3: read in bf16, write twice in float8 (row-major and col-major) kernel_3_rw = BYTES_PER_EL_BF16 * numel + 2 * BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw, 0, kernel_3_rw] elif float8_recipe_name == "rowwise": if tensor_role == "weight": # x_bf16 = ... # kernel 1 (fwd): x_bf16_dim0 -> x_float8_dim0 # kernel 2 (bwd): x_bf16_dim0 -> x_bf16_dim1 # kernel 3 (bwd): x_bf16_dim1 -> x_float8_dim1 # assume that we can't fuse 2 and 3 because that would require loading # the entire tensor to shared memory if fuse_with_prev: # assume we can fuse one of the reads with previous op kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel else: kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel kernel_2_rw = BYTES_PER_EL_BF16 * numel * 2 kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw, kernel_2_rw, kernel_3_rw] else: # x_bf16 = ... # kernel 1: x_bf16_dim0 -> x_float8_dim0, x_bf16_dim1 # kernel 2: x_bf16_dim1 -> x_float8_dim1 # assume that we can't fuse 1 and 2 because that would require loading # the entire tensor to shared memory if fuse_with_prev: # assume we can fuse one of the reads with previous op kernel_1_rw = ( 0 + BYTES_PER_EL_FLOAT8 * numel + BYTES_PER_EL_BF16 * numel ) else: kernel_1_rw = ( BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel + BYTES_PER_EL_BF16 * numel ) kernel_2_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw, kernel_2_rw] elif float8_recipe_name == "rowwise_with_gw_hp": if tensor_role in ("input", "grad_output"): # x_bf16 = ... # kernel 1 (fwd): x_bf16_dim0 -> x_float8_dim0 # bwd: no-op if fuse_with_prev: kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel else: kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw] elif tensor_role == "weight": # x_bf16 = ... # kernel 1 (fwd): w_bf16 -> w_float8_dim0, w_scale_dim0 # kernel 2 (bwd): w_scale_dim0 -> w_scale_tensorwise # kernel 3 (bwd): w_bf16, w_scale_tensorwise -> w_float8_dim1 kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel kernel_2_rw = 0 kernel_3_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw, kernel_2_rw, kernel_3_rw] else: assert False, "unsupported" else: assert mx_recipe_name in ( "mxfp8_emulated", "mxfp8_cutlass", "mxfp8_cublas", ), "unsupported" # For now, assume that we can't profitably fuse kernel 1 and kernel 2 # x_bf16 = ... # kernel 1: x_bf16 -> x_mxfp8_dim0 # kernel 2: x_bf16 -> x_mxfp8_dim1 if fuse_with_prev: kernel_1_rw = 0 + BYTES_PER_EL_FLOAT8 * numel else: kernel_1_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel kernel_2_rw = BYTES_PER_EL_BF16 * numel + BYTES_PER_EL_FLOAT8 * numel res_bytes = [kernel_1_rw, kernel_2_rw] # convert from bytes to seconds res_s = [ x / specs["peak_mem_bw_bytes_sec"] / specs["pct_achievable_mem_bw"] for x in res_bytes ] # take max of kernel_overhead, r/w time res_s = [sympy.Max(x, KERNEL_LAUNCH_OVERHEAD_SEC) for x in res_s] return res_s def get_individual_gemm_time_sympy( M: sympy.Symbol, K: sympy.Symbol, N: sympy.Symbol, dtype, mx_recipe_name ) -> sympy.Symbol: # compute bound specs = get_specs() gemm_ops = 2 * M * K * N if dtype is torch.bfloat16: peak_tops = specs["bf16_peak_tops"] elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2): peak_tops = specs["fp8_peak_tops"] else: assert False, "unsupported" compute_gemm_time_s = gemm_ops / peak_tops / specs["pct_achievable_gemm_tops"] # memory bound num_reads = M * K + K * N num_writes = M * N if mx_recipe_name is not None: assert mx_recipe_name in ( "mxfp8_emulated", "mxfp8_cutlass", "mxfp8_cublas", ), "unsupported" assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2), "unsupported" # adjust reads for MX scaling block_size = 32 num_scale_reads = num_reads // block_size # note: e8m0 bytes per element is the same as for e4m3|e5m2 num_reads = num_reads + num_scale_reads if dtype is torch.bfloat16: bytes_rw = num_reads * BYTES_PER_EL_BF16 + num_writes * BYTES_PER_EL_BF16 elif dtype in (torch.float8_e4m3fn, torch.float8_e5m2): # read in float8, output in bfloat16 bytes_rw = num_reads * BYTES_PER_EL_FLOAT8 + num_writes * BYTES_PER_EL_BF16 else: assert False, "unsupported" mem_gemm_time_s = ( bytes_rw / specs["peak_mem_bw_bytes_sec"] / specs["pct_achievable_mem_bw"] ) return sympy.Max(compute_gemm_time_s, mem_gemm_time_s, KERNEL_LAUNCH_OVERHEAD_SEC) def get_gemm_time_sympy( M: sympy.Symbol, K: sympy.Symbol, N: sympy.Symbol, dtype, float8_recipe_name: Optional[str], mx_recipe_name: Optional[str], ): # next: add rowwise_with_gw_hp here # note: this function is currently not super accurate for small shapes: # when M,K,N <= 1k,1k,1k it undercounts by around 2x gemm_dtype_input, gemm_dtype_grad_input, gemm_dtype_grad_weight = ( dtype, dtype, dtype, ) if float8_recipe_name == "rowwise_with_gw_hp": gemm_dtype_grad_weight = torch.bfloat16 gemm_output_time_s = get_individual_gemm_time_sympy( M, K, N, gemm_dtype_input, mx_recipe_name ) gemm_grad_input_time_s = get_individual_gemm_time_sympy( M, N, K, gemm_dtype_grad_input, mx_recipe_name ) gemm_grad_weight_time_s = get_individual_gemm_time_sympy( K, M, N, gemm_dtype_grad_weight, mx_recipe_name ) total = gemm_output_time_s + gemm_grad_input_time_s + gemm_grad_weight_time_s return total def get_float8_mem_sympy( M, K, N, float8_recipe_name: Optional[str], mx_recipe_name: Optional[str], enable_fusion_modeling: bool, ): specs = get_specs() # there are three gemms in the fwd/bwd of a linear: # # input @ weight_t = output # MxK @ KxN => MxN # # grad_output @ weight = grad_input # MxN @ NxK => MxK # # input_t @ grad_output = grad_weight # KxM @ MxN => KxN fwd_fp8_input_mem = get_tensor_memory_traffic_ovhd_s( specs, M, K, tensor_role="input", float8_recipe_name=float8_recipe_name, mx_recipe_name=mx_recipe_name, fuse_with_prev=enable_fusion_modeling, ) fwd_fp8_weight_mem = get_tensor_memory_traffic_ovhd_s( specs, K, N, tensor_role="weight", float8_recipe_name=float8_recipe_name, mx_recipe_name=mx_recipe_name, fuse_with_prev=False, ) gi_fp8_grad_output_mem = get_tensor_memory_traffic_ovhd_s( specs, M, N, tensor_role="grad_output", float8_recipe_name=float8_recipe_name, mx_recipe_name=mx_recipe_name, fuse_with_prev=enable_fusion_modeling, ) res = sum([*fwd_fp8_input_mem, *fwd_fp8_weight_mem, *gi_fp8_grad_output_mem]) return res