diff --git a/examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py b/examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py new file mode 100644 index 000000000..b159e14a0 --- /dev/null +++ b/examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py @@ -0,0 +1,1625 @@ +"""SM120 NVFP4 block-scaled GEMM benchmark. + +This example runs TileLang's SM120 NVFP4 block-scaled GEMM benchmark. +The user-facing surface is intentionally a single block-scaled GEMM entry; +SM120-specific scale layout, TMA, and persistent scheduling details stay inside +the implementation contract. + +It can optionally compare against the official CUTLASS GeForce NVFP4 example +79a. Run from the repository root: + + python examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py --m 8192 --n 8192 --k 8192 --run-cutlass +""" + +import argparse +from pathlib import Path +import re +import shutil +import subprocess +import sys + + +REPO_ROOT = Path(__file__).resolve().parents[2] + +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import torch +import tilelang +import tilelang.language as T +from tilelang.carver.arch import driver +from tilelang.profiler import do_bench +from tilelang.quantize import ( + blockscaled_chunk_kmajor_word_offset, + decode_ue4m3_scale_bytes, + swizzle_blockscaled_chunk_kmajor_scale_words, + tilelang_quantize_bf16_to_nvfp4_blockscaled, + unswizzle_blockscaled_chunk_kmajor_scale_words, +) + + +def _ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def _tflops(m: int, n: int, k: int, latency_ms: float) -> float: + return 2.0 * m * n * k / (latency_ms * 1.0e-3) / 1.0e12 + + +def _early_bench_options() -> argparse.Namespace: + parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) + parser.add_argument("--maxrregcount", type=int) + parser.add_argument("--ptxas-verbose", action="store_true") + args, _ = parser.parse_known_args() + return args + + +_EARLY_BENCH_OPTIONS = _early_bench_options() + + +_SM120_MICRO_PIPELINE = "sm120_backend_kblock_fulltile_package_pingpong" +_SM120_SCHEDULE = "pp_stream_output_tma_panel32_pipe2" +_SM120_SCALE_LAYOUT = "blockscaled_chunk_kmajor" +_SM120_SCALE_LOAD = "tma" +_SM120_AB_SHARED_STORAGE = "packed" +_SM120_THREADS = 384 + + +# Scale source contract for the SM120 optimized path: +# - semantic SFA/SFB shape is [M or N, K / 64] uint32; each word packs four +# scale bytes for four consecutive K/16 groups. +# - host storage uses CUTLASS BlockScaledBasicChunk K-major order: +# atom_shape=((32, 4), (16, 4)), atom_stride=((16, 4), (0, 1)). +# For one K=64 atom, the scale-byte offset is +# (row % 32) * 16 + (row // 32) * 4 + ((k // 16) % 4). +# The uint32 source layout below is the compressed form of that byte layout. +# - the benchmark always passes this swizzled storage to the kernel; reference +# checking keeps the semantic row-major copy separate. +def _device_compile_flags() -> list[str]: + flags = [ + "-DTL_SM120_FULLTILE_CUTLASS_SF_BASEPTR=1", + # A/B are TMA-loaded into the normal packed row-major shared view. + # The package-pingpong helper must use the same view when constructing + # its lane-local operand package. + "-DTL_SM120_FULLTILE_PACKAGE_ROWMAJOR_VIEW=1", + ] + if _EARLY_BENCH_OPTIONS.maxrregcount is not None: + if _EARLY_BENCH_OPTIONS.maxrregcount <= 0: + raise ValueError("--maxrregcount must be positive") + flags.append(f"--maxrregcount={_EARLY_BENCH_OPTIONS.maxrregcount}") + if _EARLY_BENCH_OPTIONS.ptxas_verbose: + flags.append("--ptxas-options=--verbose") + return flags + + +def _sm120_nvfp4_jit_pass_configs() -> dict: + pass_configs = {} + pass_configs[tilelang.PassConfigKey.TL_DEVICE_COMPILE_FLAGS] = _device_compile_flags() + return pass_configs + + +@tilelang.jit( + out_idx=None, + pass_configs=_sm120_nvfp4_jit_pass_configs(), +) +def tilelang_nvfp4_blockscaled_gemm( + M: int, + N: int, + K: int, + block_M: int, + block_N: int, + block_K: int, + num_stages: int, + warp_policy, + out_dtype, + producer_regs: int, + consumer_regs: int, +): + """SM120 NVFP4 block-scaled GEMM implementation. + + This function has a compact external contract. The selected SM120 + scheduling, scale-source layout, and MMA helper are fixed internally by + the SM120 implementation contract. + """ + + assert M % block_M == 0 + assert N % block_N == 0 + assert K % block_K == 0 + assert block_M % 2 == 0 + assert block_K % 64 == 0 + assert num_stages >= 2 + in_dtype = T.float4_e2m1fn + ab_shared_dtype = in_dtype + micro_pipeline = _SM120_MICRO_PIPELINE + sf_layout = _SM120_SCALE_LAYOUT + accum_dtype = T.float32 + half_block_N = block_N // 2 + sf_granularity_k = 16 + sf_words_per_block_k = block_K // 64 + + if _SM120_SCHEDULE == "pp_stream_output_tma_panel32_pipe2": + if sf_words_per_block_k < 4: + raise ValueError("The SM120 NVFP4 optimized path requires block_K >= 256 for scale TMA") + output_tma_panel32_pipe2 = True + output_tma_epi64x32 = False + tma_store_literal_copy = False + panel32_tma_store = False + panel64_tma_store = False + store_block_N = 32 + literal_panel_store = True + epilogue_store_block_M = block_M + epilogue_store_block_N = store_block_N + epilogue_store_slots = 2 + if block_N != 128: + raise ValueError("The SM120 NVFP4 optimized path requires block_N=128") + + sm_num = driver.get_num_sms() + n_blocks = T.ceildiv(N, block_N) + m_blocks = T.ceildiv(M, block_M) + total_tiles = n_blocks * m_blocks + k_tiles = K // block_K + stream_iters = T.ceildiv(total_tiles, sm_num) + owner_iters = T.ceildiv(stream_iters, 2) + + @T.macro + def copy_blockscaled_chunk_kmajor_scale_tile(SF, SF_shared, tile_row, block_rows, ko, stage, tx): + scale_cols = K // 64 + scale_tile_words = block_rows * sf_words_per_block_k + scale_lane = tx - 256 + for scale_iter in T.serial((scale_tile_words + 127) // 128): + scale_flat = scale_iter * 128 + scale_lane + if scale_flat < scale_tile_words: + global_flat = tile_row * block_rows * scale_cols + ko * sf_words_per_block_k * block_rows + scale_flat + SF_shared[ + stage, + scale_flat // sf_words_per_block_k, + scale_flat % sf_words_per_block_k, + ] = SF[global_flat // scale_cols, global_flat % scale_cols] + + @T.prim_func + def main( + A: T.Tensor((M, K), in_dtype), + B: T.Tensor((N, K), in_dtype), + SFA: T.Tensor((M, K // 64), T.uint32), + SFB: T.Tensor((N, K // 64), T.uint32), + C: T.Tensor((M, N), out_dtype), + ): + with T.Kernel(sm_num, threads=384) as block_id: + A_shared = T.alloc_shared((num_stages, block_M, block_K), ab_shared_dtype) + B_shared = T.alloc_shared((num_stages, block_N, block_K), ab_shared_dtype) + SFA_shared = T.alloc_shared((num_stages, block_M, sf_words_per_block_k), T.uint32) + SFB_shared = T.alloc_shared((num_stages, block_N, sf_words_per_block_k), T.uint32) + if literal_panel_store or panel32_tma_store or panel64_tma_store: + C_shared = T.alloc_shared((epilogue_store_slots, epilogue_store_block_M, epilogue_store_block_N), out_dtype) + + loaded = T.alloc_barrier([128] * (2 * num_stages)) + consumed = T.alloc_barrier([128] * (2 * num_stages)) + wg_order = T.alloc_barrier([128] * 2) + store_order = T.alloc_barrier([128] * 2) + + tx = T.get_thread_binding() + + if tx >= 256: + if producer_regs > 0: + T.set_max_nreg(producer_regs, 0) + for stream in T.serial(stream_iters): + tile_id = block_id + stream * sm_num + if tile_id < total_tiles: + tile_n = tile_id % n_blocks + tile_m = tile_id // n_blocks + owner = stream & 1 + owner_iter = stream // 2 + for ko in T.unroll(k_tiles, explicit=False, unroll_factor=1): + stage = ko % num_stages + phase = ((owner_iter * k_tiles + ko) // num_stages) & 1 + barrier_slot = owner * num_stages + stage + if stream > 0 and ko < num_stages: + prev_owner = (stream - 1) & 1 + prev_owner_iter = (stream - 1) // 2 + prev_barrier_slot = prev_owner * num_stages + stage + prev_last_ko = ((k_tiles - 1 - stage) // num_stages) * num_stages + stage + prev_phase = ((prev_owner_iter * k_tiles + prev_last_ko) // num_stages) & 1 + T.barrier_wait(consumed[prev_barrier_slot], prev_phase) + else: + T.barrier_wait(consumed[barrier_slot], phase ^ 1) + T.tma_copy( + A[ + tile_m * block_M : (tile_m + 1) * block_M, + ko * block_K : (ko + 1) * block_K, + ], + A_shared[stage, :, :], + barrier=loaded[barrier_slot], + leader_scope_threads=128, + ) + T.tma_copy( + B[ + tile_n * block_N : (tile_n + 1) * block_N, + ko * block_K : (ko + 1) * block_K, + ], + B_shared[stage, :, :], + barrier=loaded[barrier_slot], + leader_scope_threads=128, + ) + copy_blockscaled_chunk_kmajor_scale_tile(SFA, SFA_shared, tile_m, block_M, ko, stage, tx) + copy_blockscaled_chunk_kmajor_scale_tile(SFB, SFB_shared, tile_n, block_N, ko, stage, tx) + T.barrier_arrive(loaded[barrier_slot]) + + elif tx < 128: + if consumer_regs > 0: + T.set_max_nreg(consumer_regs, 1) + if panel32_tma_store: + C0_local_0 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C0_local_1 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C0_local_2 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C0_local_3 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + elif panel64_tma_store: + C0_local_0 = T.alloc_fragment((block_M, half_block_N), accum_dtype) + C0_local_1 = T.alloc_fragment((block_M, half_block_N), accum_dtype) + else: + C0_local = T.alloc_fragment((block_M, block_N), accum_dtype) + for owner_iter in T.serial(owner_iters): + stream = owner_iter * 2 + tile_id = block_id + stream * sm_num + if tile_id < total_tiles: + tile_n = tile_id % n_blocks + tile_m = tile_id // n_blocks + + if owner_iter > 0: + T.barrier_wait(wg_order[0], (owner_iter - 1) & 1) + + if panel32_tma_store: + T.clear(C0_local_0) + T.clear(C0_local_1) + T.clear(C0_local_2) + T.clear(C0_local_3) + elif panel64_tma_store: + T.clear(C0_local_0) + T.clear(C0_local_1) + else: + T.clear(C0_local) + for ko in T.unroll(k_tiles, explicit=False, unroll_factor=1): + stage = ko % num_stages + phase = ((owner_iter * k_tiles + ko) // num_stages) & 1 + barrier_slot = stage + T.barrier_wait(loaded[barrier_slot], phase) + if panel32_tma_store: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 0:store_block_N, :], + C0_local_0, + SFA_shared[stage, :, :], + SFB_shared[stage, 0:store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, store_block_N : 2 * store_block_N, :], + C0_local_1, + SFA_shared[stage, :, :], + SFB_shared[stage, store_block_N : 2 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 2 * store_block_N : 3 * store_block_N, :], + C0_local_2, + SFA_shared[stage, :, :], + SFB_shared[stage, 2 * store_block_N : 3 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 3 * store_block_N : 4 * store_block_N, :], + C0_local_3, + SFA_shared[stage, :, :], + SFB_shared[stage, 3 * store_block_N : 4 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + elif panel64_tma_store: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 0:half_block_N, :], + C0_local_0, + SFA_shared[stage, :, :], + SFB_shared[stage, 0:half_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, half_block_N:block_N, :], + C0_local_1, + SFA_shared[stage, :, :], + SFB_shared[stage, half_block_N:block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + else: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, :, :], + C0_local, + SFA_shared[stage, :, :], + SFB_shared[stage, :, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.barrier_arrive(consumed[barrier_slot]) + + T.barrier_arrive(wg_order[1]) + + if panel32_tma_store: + base_n0 = tile_n * block_N + T.copy(C0_local_0, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n1 = tile_n * block_N + store_block_N + T.copy(C0_local_1, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n1 : base_n1 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n2 = tile_n * block_N + 2 * store_block_N + T.copy(C0_local_2, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n3 = tile_n * block_N + 3 * store_block_N + T.copy(C0_local_3, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n3 : base_n3 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + elif panel64_tma_store: + if owner_iter > 0: + T.barrier_wait(store_order[0], (owner_iter - 1) & 1) + base_n0 = tile_n * block_N + T.copy(C0_local_0, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n2 = tile_n * block_N + half_block_N + T.copy(C0_local_1, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.barrier_arrive(store_order[1]) + elif literal_panel_store: + if owner_iter > 0: + T.barrier_wait(store_order[0], (owner_iter - 1) & 1) + if output_tma_panel32_pipe2: + for panel32 in T.serial(4): + epi_slot = panel32 & 1 + if panel32 >= 2: + T.tma_store_wait(1, True) + T.sync_threads(7, 128) + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + epi_slot * epilogue_store_block_M * epilogue_store_block_N, + panel32, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[epi_slot, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + tile_n * block_N + panel32 * epilogue_store_block_N : tile_n * block_N + + (panel32 + 1) * epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(7, 128) + elif output_tma_epi64x32: + for epi in T.serial(8): + epi_slot = epi % 4 + epi_m = epi // 4 + epi_n = epi % 4 + if epi >= 4: + T.tma_store_wait(3, True) + T.sync_threads(7, 128) + T.sm120_store_full_c_fragment_epi64x32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + epi_slot * epilogue_store_block_M * epilogue_store_block_N, + epi_m, + epi_n, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[epi_slot, :, :], + C[ + tile_m * block_M + epi_m * epilogue_store_block_M : tile_m * block_M + + (epi_m + 1) * epilogue_store_block_M, + tile_n * block_N + epi_n * epilogue_store_block_N : tile_n * block_N + + (epi_n + 1) * epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(7, 128) + elif tma_store_literal_copy: + base_n0 = tile_n * block_N + T.sm120_store_full_c_fragment_panel64_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 0, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + for linear in T.Parallel(block_M * epilogue_store_block_N): + store_i = linear // epilogue_store_block_N + store_j = linear % epilogue_store_block_N + C[tile_m * block_M + store_i, base_n0 + store_j] = C_shared[0, store_i, store_j] + + base_n2 = tile_n * block_N + half_block_N + T.sm120_store_full_c_fragment_panel64_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 1, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + for linear in T.Parallel(block_M * epilogue_store_block_N): + store_i = linear // epilogue_store_block_N + store_j = linear % epilogue_store_block_N + C[tile_m * block_M + store_i, base_n2 + store_j] = C_shared[0, store_i, store_j] + else: + base_n0 = tile_n * block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 0, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + # All consumer threads must wait before the reused C_shared panel is overwritten. + T.sync_threads(7, 128) + + base_n1 = tile_n * block_N + store_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 1, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n1 : base_n1 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(7, 128) + + base_n2 = tile_n * block_N + half_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 2, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(7, 128) + + base_n3 = tile_n * block_N + half_block_N + store_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C0_local.data, + 0, + C_shared.data, + 0, + 3, + ) + T.fence_proxy_async() + T.sync_threads(7, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n3 : base_n3 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.barrier_arrive(store_order[1]) + else: + T.copy(C0_local, C[tile_m * block_M, tile_n * block_N]) + + else: + if consumer_regs > 0: + T.set_max_nreg(consumer_regs, 1) + if panel32_tma_store: + C1_local_0 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C1_local_1 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C1_local_2 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + C1_local_3 = T.alloc_fragment((block_M, store_block_N), accum_dtype) + elif panel64_tma_store: + C1_local_0 = T.alloc_fragment((block_M, half_block_N), accum_dtype) + C1_local_1 = T.alloc_fragment((block_M, half_block_N), accum_dtype) + else: + C1_local = T.alloc_fragment((block_M, block_N), accum_dtype) + for owner_iter in T.serial(owner_iters): + stream = owner_iter * 2 + 1 + tile_id = block_id + stream * sm_num + if tile_id < total_tiles: + tile_n = tile_id % n_blocks + tile_m = tile_id // n_blocks + + T.barrier_wait(wg_order[1], owner_iter & 1) + + if panel32_tma_store: + T.clear(C1_local_0) + T.clear(C1_local_1) + T.clear(C1_local_2) + T.clear(C1_local_3) + elif panel64_tma_store: + T.clear(C1_local_0) + T.clear(C1_local_1) + else: + T.clear(C1_local) + for ko in T.unroll(k_tiles, explicit=False, unroll_factor=1): + stage = ko % num_stages + phase = ((owner_iter * k_tiles + ko) // num_stages) & 1 + barrier_slot = num_stages + stage + T.barrier_wait(loaded[barrier_slot], phase) + if panel32_tma_store: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 0:store_block_N, :], + C1_local_0, + SFA_shared[stage, :, :], + SFB_shared[stage, 0:store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, store_block_N : 2 * store_block_N, :], + C1_local_1, + SFA_shared[stage, :, :], + SFB_shared[stage, store_block_N : 2 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 2 * store_block_N : 3 * store_block_N, :], + C1_local_2, + SFA_shared[stage, :, :], + SFB_shared[stage, 2 * store_block_N : 3 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 3 * store_block_N : 4 * store_block_N, :], + C1_local_3, + SFA_shared[stage, :, :], + SFB_shared[stage, 3 * store_block_N : 4 * store_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + elif panel64_tma_store: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, 0:half_block_N, :], + C1_local_0, + SFA_shared[stage, :, :], + SFB_shared[stage, 0:half_block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, half_block_N:block_N, :], + C1_local_1, + SFA_shared[stage, :, :], + SFB_shared[stage, half_block_N:block_N, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + else: + T.mma_gemm_blockscaled( + A_shared[stage, :, :], + B_shared[stage, :, :], + C1_local, + SFA_shared[stage, :, :], + SFB_shared[stage, :, :], + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + T.barrier_arrive(consumed[barrier_slot]) + + T.barrier_arrive(wg_order[0]) + + if panel32_tma_store: + base_n0 = tile_n * block_N + T.copy(C1_local_0, C_shared[1, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[1, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n1 = tile_n * block_N + store_block_N + T.copy(C1_local_1, C_shared[1, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[1, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n1 : base_n1 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n2 = tile_n * block_N + 2 * store_block_N + T.copy(C1_local_2, C_shared[1, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[1, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n3 = tile_n * block_N + 3 * store_block_N + T.copy(C1_local_3, C_shared[1, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[1, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n3 : base_n3 + store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + elif panel64_tma_store: + T.barrier_wait(store_order[1], owner_iter & 1) + base_n0 = tile_n * block_N + T.copy(C1_local_0, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + + base_n2 = tile_n * block_N + half_block_N + T.copy(C1_local_1, C_shared[0, :, :]) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.barrier_arrive(store_order[0]) + elif literal_panel_store: + T.barrier_wait(store_order[1], owner_iter & 1) + if output_tma_panel32_pipe2: + for panel32 in T.serial(4): + epi_slot = panel32 & 1 + if panel32 >= 2: + T.tma_store_wait(1, True) + T.sync_threads(8, 128) + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + epi_slot * epilogue_store_block_M * epilogue_store_block_N, + panel32, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[epi_slot, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + tile_n * block_N + panel32 * epilogue_store_block_N : tile_n * block_N + + (panel32 + 1) * epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(8, 128) + elif output_tma_epi64x32: + for epi in T.serial(8): + epi_slot = epi % 4 + epi_m = epi // 4 + epi_n = epi % 4 + if epi >= 4: + T.tma_store_wait(3, True) + T.sync_threads(8, 128) + T.sm120_store_full_c_fragment_epi64x32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + epi_slot * epilogue_store_block_M * epilogue_store_block_N, + epi_m, + epi_n, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[epi_slot, :, :], + C[ + tile_m * block_M + epi_m * epilogue_store_block_M : tile_m * block_M + + (epi_m + 1) * epilogue_store_block_M, + tile_n * block_N + epi_n * epilogue_store_block_N : tile_n * block_N + + (epi_n + 1) * epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(8, 128) + elif tma_store_literal_copy: + base_n0 = tile_n * block_N + T.sm120_store_full_c_fragment_panel64_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 0, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + for linear in T.Parallel(block_M * epilogue_store_block_N): + store_i = linear // epilogue_store_block_N + store_j = linear % epilogue_store_block_N + C[tile_m * block_M + store_i, base_n0 + store_j] = C_shared[0, store_i, store_j] + + base_n2 = tile_n * block_N + half_block_N + T.sm120_store_full_c_fragment_panel64_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 1, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + for linear in T.Parallel(block_M * epilogue_store_block_N): + store_i = linear // epilogue_store_block_N + store_j = linear % epilogue_store_block_N + C[tile_m * block_M + store_i, base_n2 + store_j] = C_shared[0, store_i, store_j] + else: + base_n0 = tile_n * block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 0, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n0 : base_n0 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + # All consumer threads must wait before the reused C_shared panel is overwritten. + T.sync_threads(8, 128) + + base_n1 = tile_n * block_N + store_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 1, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n1 : base_n1 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(8, 128) + + base_n2 = tile_n * block_N + half_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 2, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n2 : base_n2 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.sync_threads(8, 128) + + base_n3 = tile_n * block_N + half_block_N + store_block_N + T.sm120_store_full_c_fragment_panel32_tma_bf16( + C1_local.data, + 0, + C_shared.data, + 0, + 3, + ) + T.fence_proxy_async() + T.sync_threads(8, 128) + T.tma_copy( + C_shared[0, :, :], + C[ + tile_m * block_M : (tile_m + 1) * block_M, + base_n3 : base_n3 + epilogue_store_block_N, + ], + leader_scope_threads=128, + ) + T.tma_store_wait(0, True) + T.barrier_arrive(store_order[0]) + else: + T.copy(C1_local, C[tile_m * block_M, tile_n * block_N]) + + return main + + +_FP4_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + + +def _make_packed_fp4(rows: int, cols: int, *, seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda") + generator.manual_seed(seed) + return torch.randint(-128, 128, (rows, cols // 2), device="cuda", dtype=torch.int8, generator=generator) + + +def _make_bf16_activation(rows: int, cols: int, *, seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda") + generator.manual_seed(seed) + return (torch.randn((rows, cols), device="cuda", dtype=torch.float32, generator=generator) * 2.0).to(torch.bfloat16) + + +def _make_ones_packed_fp4(rows: int, cols: int) -> torch.Tensor: + # Two FP4 e2m1 values with raw code 0x2, packed into one byte. + return torch.full((rows, cols // 2), 0x22, device="cuda", dtype=torch.int8) + + +def _make_constant_scale_words(rows: int, k: int, byte: int = 0x38) -> torch.Tensor: + word = byte | (byte << 8) | (byte << 16) | (byte << 24) + return torch.full((rows, k // 64), word, device="cuda", dtype=torch.uint32) + + +def _pack_scale_words(scale_bytes: torch.Tensor) -> torch.Tensor: + scale_i64 = scale_bytes.to(torch.int64).reshape(scale_bytes.shape[0], -1, 4) + words = scale_i64[:, :, 0] + words = words | (scale_i64[:, :, 1] << 8) + words = words | (scale_i64[:, :, 2] << 16) + words = words | (scale_i64[:, :, 3] << 24) + return words.to(torch.uint32).contiguous() + + +def _make_binary_scale_words(rows: int, k: int, *, seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda") + generator.manual_seed(seed) + scale_bytes = ( + torch.randint( + 0, + 2, + (rows, k // 16), + device="cuda", + dtype=torch.int64, + generator=generator, + ) + * 0x38 + ) + return _pack_scale_words(scale_bytes) + + +def _blockscaled_chunk_kmajor_word_offset(row: int, k64_word: int, block_rows: int = 128) -> tuple[int, int]: + """Return compressed word coordinates for CUTLASS BlockScaledBasicChunk. + + CUTLASS' element-level byte offset for one K=64 atom is: + (row % 32) * 16 + (row // 32) * 4 + ((k // 16) % 4) + + This benchmark stores the four adjacent ((k // 16) % 4) bytes in one + uint32, so the physical word coordinate is the byte offset divided by four. + """ + return blockscaled_chunk_kmajor_word_offset(row, k64_word, block_rows=block_rows) + + +def _swizzle_scale_words_blockscaled_chunk_kmajor(words: torch.Tensor, block_rows: int = 128, block_words: int = 4) -> torch.Tensor: + """Convert semantic row-major scale words to BlockScaledBasicChunk K-major. + + Input words[row, k_word] is the natural reference matrix: one uint32 packs + scale bytes for K groups [4*k_word + 0, ..., +3]. Output storage is the + uint32-compressed form of CUTLASS BlockScaledBasicChunk K-major order. TMA + copies this physical tile directly into shared memory. + """ + return swizzle_blockscaled_chunk_kmajor_scale_words(words, block_rows=block_rows, block_words=block_words) + + +def _check_blockscaled_chunk_kmajor_offsets() -> None: + expected = { + (0, 0): (0, 0), + (0, 1): (32, 0), + (0, 2): (64, 0), + (0, 3): (96, 0), + (1, 0): (1, 0), + (15, 0): (15, 0), + (16, 0): (16, 0), + (31, 0): (31, 0), + (32, 0): (0, 1), + (33, 0): (1, 1), + (64, 0): (0, 2), + (96, 0): (0, 3), + (127, 0): (31, 3), + } + for (row, k64_word), physical in expected.items(): + actual = _blockscaled_chunk_kmajor_word_offset(row, k64_word) + if actual != physical: + raise AssertionError( + f"BlockScaledBasicChunk K-major offset mismatch: row={row} k64_word={k64_word} expected={physical} actual={actual}" + ) + + +def _decode_binary_scale_words(words: torch.Tensor, k: int) -> torch.Tensor: + w = words.to(torch.int64) + scale_bytes = torch.empty((words.shape[0], k // 16), device=words.device, dtype=torch.int64) + scale_bytes[:, 0::4] = w & 0xFF + scale_bytes[:, 1::4] = (w >> 8) & 0xFF + scale_bytes[:, 2::4] = (w >> 16) & 0xFF + scale_bytes[:, 3::4] = (w >> 24) & 0xFF + return (scale_bytes != 0).to(torch.float32) + + +def _decode_ue4m3_scale_words(words: torch.Tensor, k: int) -> torch.Tensor: + w = words.to(torch.int64) + scale_bytes = torch.empty((words.shape[0], k // 16), device=words.device, dtype=torch.uint8) + scale_bytes[:, 0::4] = (w & 0xFF).to(torch.uint8) + scale_bytes[:, 1::4] = ((w >> 8) & 0xFF).to(torch.uint8) + scale_bytes[:, 2::4] = ((w >> 16) & 0xFF).to(torch.uint8) + scale_bytes[:, 3::4] = ((w >> 24) & 0xFF).to(torch.uint8) + return decode_ue4m3_scale_bytes(scale_bytes) + + +def _decode_rowmajor_fp4(packed: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + u = packed.contiguous().view(torch.uint8) + lut = torch.tensor(_FP4_E2M1_VALUES, device=packed.device, dtype=torch.float32) + out = torch.empty((rows, cols), device=packed.device, dtype=torch.float32) + out[:, 0::2] = lut[(u & 0x0F).long()] + out[:, 1::2] = lut[((u >> 4) & 0x0F).long()] + return out + + +def _verify_tilelang_output( + A: torch.Tensor, + B: torch.Tensor, + SFA: torch.Tensor, + SFB: torch.Tensor, + C: torch.Tensor, + out_dtype: torch.dtype, + scale_mode: str, + block_m: int = 128, + block_n: int = 128, + sm_num: int | None = None, +) -> None: + A_full = _decode_rowmajor_fp4(A, A.shape[0], A.shape[1] * 2) + B_full = _decode_rowmajor_fp4(B, B.shape[0], B.shape[1] * 2) + if scale_mode == "constant": + ref = (A_full @ B_full.T).to(out_dtype) + elif scale_mode in ("random_binary", "random_sfa", "random_sfb", "ue4m3"): + if scale_mode == "ue4m3": + sfa = _decode_ue4m3_scale_words(SFA, A_full.shape[1]) + sfb = _decode_ue4m3_scale_words(SFB, B_full.shape[1]) + else: + sfa = _decode_binary_scale_words(SFA, A_full.shape[1]) + sfb = _decode_binary_scale_words(SFB, B_full.shape[1]) + ref_f32 = torch.zeros((A_full.shape[0], B_full.shape[0]), device=C.device, dtype=torch.float32) + for k_sf in range(A_full.shape[1] // 16): + k0 = k_sf * 16 + k1 = k0 + 16 + a_chunk = A_full[:, k0:k1] * sfa[:, k_sf].unsqueeze(1) + b_chunk = B_full[:, k0:k1] * sfb[:, k_sf].unsqueeze(1) + ref_f32 += a_chunk @ b_chunk.T + ref = ref_f32.to(out_dtype) + else: + raise ValueError(f"Unsupported scale_mode={scale_mode!r}") + try: + torch.testing.assert_close(C, ref, rtol=0.0, atol=0.0) + except AssertionError: + diff = ref != C + mismatch = int(diff.sum().item()) + total = diff.numel() + print(f"TileLang mismatch summary: {mismatch}/{total} ({mismatch / total:.2%})") + if C.shape[0] % block_m == 0 and C.shape[1] % block_n == 0: + m_blocks = C.shape[0] // block_m + n_blocks = C.shape[1] // block_n + tile_counts = diff.reshape(m_blocks, block_m, n_blocks, block_n).sum(dim=(1, 3)) + flat = tile_counts.flatten() + top = torch.topk(flat, k=min(12, flat.numel())) + for value, index in zip(top.values.cpu().tolist(), top.indices.cpu().tolist()): + if int(value) == 0: + continue + tile_m = index // n_blocks + tile_n = index % n_blocks + tile_id = tile_m * n_blocks + tile_n + owner = "unknown" + if sm_num is not None and sm_num > 0: + owner = f"WG{(tile_id // sm_num) & 1}" + print(f" tile m={tile_m} n={tile_n} tile_id={tile_id} owner={owner} mismatch={int(value)}/{block_m * block_n}") + if block_n % 64 == 0: + panel_counts = diff.reshape(m_blocks, block_m, n_blocks, block_n // 64, 64).sum(dim=(1, 4)) + panel_flat = panel_counts.flatten() + panel_top = torch.topk(panel_flat, k=min(12, panel_flat.numel())) + for value, index in zip(panel_top.values.cpu().tolist(), panel_top.indices.cpu().tolist()): + if int(value) == 0: + continue + panel = index % (block_n // 64) + tile_index = index // (block_n // 64) + tile_m = tile_index // n_blocks + tile_n = tile_index % n_blocks + tile_id = tile_m * n_blocks + tile_n + owner = "unknown" + if sm_num is not None and sm_num > 0: + owner = f"WG{(tile_id // sm_num) & 1}" + print( + " panel64 " + f"m={tile_m} n={tile_n} panel={panel} tile_id={tile_id} owner={owner} " + f"mismatch={int(value)}/{block_m * 64}" + ) + if block_n % 32 == 0: + panel32_counts = diff.reshape(m_blocks, block_m, n_blocks, block_n // 32, 32).sum(dim=(1, 4)) + panel32_flat = panel32_counts.flatten() + panel32_top = torch.topk(panel32_flat, k=min(16, panel32_flat.numel())) + for value, index in zip(panel32_top.values.cpu().tolist(), panel32_top.indices.cpu().tolist()): + if int(value) == 0: + continue + panel = index % (block_n // 32) + tile_index = index // (block_n // 32) + tile_m = tile_index // n_blocks + tile_n = tile_index % n_blocks + tile_id = tile_m * n_blocks + tile_n + owner = "unknown" + if sm_num is not None and sm_num > 0: + owner = f"WG{(tile_id // sm_num) & 1}" + print( + " panel32 " + f"m={tile_m} n={tile_n} panel={panel} tile_id={tile_id} owner={owner} " + f"mismatch={int(value)}/{block_m * 32}" + ) + tile0 = diff[:block_m, :block_n] + for panel in range(block_n // 32): + panel_diff = tile0[:, panel * 32 : (panel + 1) * 32] + row32 = panel_diff.reshape(4, 32, 32).sum(dim=(1, 2)).cpu().tolist() + col8 = panel_diff.reshape(block_m, 4, 8).sum(dim=(0, 2)).cpu().tolist() + print(f" tile0-panel32-buckets panel={panel} row32={[int(x) for x in row32]} col8={[int(x) for x in col8]}") + raise + + +def run_tilelang(args: argparse.Namespace) -> tuple[float, float]: + _check_blockscaled_chunk_kmajor_offsets() + out_dtype = T.bfloat16 if args.out_dtype == "bfloat16" else T.float32 + out_torch_dtype = torch.bfloat16 if args.out_dtype == "bfloat16" else torch.float32 + warp_policy = getattr(T.GemmWarpPolicy, args.warp_policy) + + kernel = tilelang_nvfp4_blockscaled_gemm( + args.m, + args.n, + args.k, + args.block_m, + args.block_n, + args.block_k, + args.num_stages, + warp_policy, + out_dtype, + args.producer_regs, + args.consumer_regs, + ) + + source = kernel.get_kernel_source() + if args.dump_source: + dump_path = Path(args.dump_source) + dump_path.parent.mkdir(parents=True, exist_ok=True) + dump_path.write_text(source) + print(f"TileLang CUDA source: {dump_path}") + if args.dump_sass: + sass_path = Path(args.dump_sass) + sass_path.parent.mkdir(parents=True, exist_ok=True) + saved_compile_flags = kernel.compile_flags + device_compile_flags = _device_compile_flags() + if device_compile_flags: + merged_compile_flags = list(saved_compile_flags or []) + for flag in device_compile_flags: + if flag not in merged_compile_flags: + merged_compile_flags.append(flag) + kernel.compile_flags = merged_compile_flags + kernel.export_sass(str(sass_path)) + kernel.compile_flags = saved_compile_flags + print(f"TileLang SASS: {sass_path}") + if "sm120_mma_blockscaled_kblock_fulltile" not in source: + raise RuntimeError("TileLang source did not lower to the SM120 full-tile MMA helper") + + if args.input_mode == "ones": + A = _make_ones_packed_fp4(args.m, args.k) + B = _make_ones_packed_fp4(args.n, args.k) + elif args.input_mode == "bf16_quantized": + A_bf16 = _make_bf16_activation(args.m, args.k, seed=args.seed) + B_bf16 = _make_bf16_activation(args.n, args.k, seed=args.seed + 1) + A, SFA = tilelang_quantize_bf16_to_nvfp4_blockscaled(A_bf16) + B, SFB = tilelang_quantize_bf16_to_nvfp4_blockscaled(B_bf16) + SFA_semantic = unswizzle_blockscaled_chunk_kmajor_scale_words(SFA) + SFB_semantic = unswizzle_blockscaled_chunk_kmajor_scale_words(SFB) + else: + A = _make_packed_fp4(args.m, args.k, seed=args.seed) + B = _make_packed_fp4(args.n, args.k, seed=args.seed + 1) + if args.input_mode != "bf16_quantized": + if args.scale_mode == "constant": + SFA_semantic = _make_constant_scale_words(args.m, args.k) + SFB_semantic = _make_constant_scale_words(args.n, args.k) + elif args.scale_mode == "random_binary": + SFA_semantic = _make_binary_scale_words(args.m, args.k, seed=args.seed + 100) + SFB_semantic = _make_binary_scale_words(args.n, args.k, seed=args.seed + 200) + elif args.scale_mode == "random_sfa": + SFA_semantic = _make_binary_scale_words(args.m, args.k, seed=args.seed + 100) + SFB_semantic = _make_constant_scale_words(args.n, args.k) + elif args.scale_mode == "random_sfb": + SFA_semantic = _make_constant_scale_words(args.m, args.k) + SFB_semantic = _make_binary_scale_words(args.n, args.k, seed=args.seed + 200) + else: + raise ValueError(f"Unsupported scale_mode={args.scale_mode!r}") + SFA = _swizzle_scale_words_blockscaled_chunk_kmajor(SFA_semantic) + SFB = _swizzle_scale_words_blockscaled_chunk_kmajor(SFB_semantic) + C = torch.empty((args.m, args.n), device="cuda", dtype=out_torch_dtype) + + kernel(A, B, SFA, SFB, C) + torch.cuda.synchronize() + + if args.profile_only and args.profile_warmup_launches < 0: + raise ValueError("--profile-warmup-launches must be non-negative") + if args.profile_only: + for _ in range(args.profile_warmup_launches): + kernel(A, B, SFA, SFB, C) + kernel(A, B, SFA, SFB, C) + torch.cuda.synchronize() + + if args.verify: + _verify_tilelang_output( + A, + B, + SFA_semantic, + SFB_semantic, + C, + out_torch_dtype, + "ue4m3" if args.input_mode == "bf16_quantized" else args.scale_mode, + args.block_m, + args.block_n, + driver.get_num_sms(), + ) + print("TileLang correctness: passed") + + if args.profile_only: + launches = 2 + args.profile_warmup_launches + print(f"TileLang profile-only: ran {launches} kernel launch(es)") + return float("nan"), float("nan") + + latency_ms = do_bench( + lambda: kernel(A, B, SFA, SFB, C), + warmup=args.warmup_ms, + rep=args.rep_ms, + _n_warmup=args.n_warmup, + _n_repeat=args.n_repeat, + backend=args.backend, + return_mode=args.return_mode, + ) + return latency_ms, _tflops(args.m, args.n, args.k, latency_ms) + + +def _find_cmake(args: argparse.Namespace) -> str: + if args.cmake: + return args.cmake + cmake = shutil.which("cmake") + if cmake: + return cmake + env_cmake = Path(sys.executable).resolve().parent / "cmake" + if env_cmake.exists(): + return str(env_cmake) + raise FileNotFoundError("cmake was not found; pass --cmake") + + +def _find_nvcc(args: argparse.Namespace) -> str: + if args.nvcc: + return args.nvcc + nvcc = shutil.which("nvcc") + if nvcc: + return nvcc + default_nvcc = Path("/usr/local/cuda-12.8/bin/nvcc") + if default_nvcc.exists(): + return str(default_nvcc) + raise FileNotFoundError("nvcc was not found; pass --nvcc") + + +def _run_command(cmd: list[str], *, cwd: Path) -> str: + print("+ " + " ".join(cmd), flush=True) + completed = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + if completed.stdout: + print(completed.stdout, end="") + completed.check_returncode() + return completed.stdout + + +def _cutlass_binary_path(args: argparse.Namespace) -> Path: + if args.cutlass_binary: + return Path(args.cutlass_binary) + return Path(args.cutlass_build_dir) / "examples" / "79_blackwell_geforce_gemm" / "79a_blackwell_geforce_nvfp4_bf16_gemm" + + +def build_cutlass_79a(args: argparse.Namespace) -> Path: + binary = _cutlass_binary_path(args) + if binary.exists() and not args.rebuild_cutlass: + return binary + + cmake = _find_cmake(args) + nvcc = _find_nvcc(args) + build_dir = Path(args.cutlass_build_dir) + util_include = REPO_ROOT / "3rdparty" / "cutlass" / "tools" / "util" / "include" + + configure_cmd = [ + cmake, + "-S", + "3rdparty/cutlass", + "-B", + str(build_dir), + "-DCMAKE_BUILD_TYPE=Release", + "-DCUTLASS_NVCC_ARCHS=120a", + "-DCUTLASS_ENABLE_EXAMPLES=ON", + "-DCUTLASS_ENABLE_TOOLS=ON", + "-DCUTLASS_ENABLE_LIBRARY=OFF", + "-DCUTLASS_ENABLE_TESTS=OFF", + "-DCUTLASS_ENABLE_PROFILER=OFF", + f"-DCMAKE_CUDA_COMPILER={nvcc}", + f"-DCMAKE_CUDA_FLAGS=-I{util_include}", + f"-DCMAKE_CXX_FLAGS=-I{util_include}", + ] + _run_command(configure_cmd, cwd=REPO_ROOT) + + build_cmd = [ + cmake, + "--build", + str(build_dir), + "--target", + "79a_blackwell_geforce_nvfp4_bf16_gemm", + "-j", + str(args.cutlass_build_jobs), + ] + _run_command(build_cmd, cwd=REPO_ROOT) + + if not binary.exists(): + raise FileNotFoundError(f"CUTLASS 79a binary was not produced at {binary}") + return binary + + +def run_cutlass(args: argparse.Namespace) -> tuple[float, float]: + binary = build_cutlass_79a(args) + output = _run_command( + [ + str(binary), + f"--m={args.m}", + f"--n={args.n}", + f"--k={args.k}", + f"--iterations={args.cutlass_iterations}", + "--skip-verification", + ], + cwd=REPO_ROOT, + ) + + latency_match = re.search(r"Avg runtime:\s*([0-9.eE+-]+)\s*ms", output) + gflops_match = re.search(r"GFLOPS:\s*([0-9.eE+-]+)", output) + if latency_match is None or gflops_match is None: + raise RuntimeError("Could not parse CUTLASS latency/GFLOPS output") + latency_ms = float(latency_match.group(1)) + return latency_ms, float(gflops_match.group(1)) / 1000.0 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--m", type=int, default=2048) + parser.add_argument("--n", type=int, default=2048) + parser.add_argument("--k", type=int, default=2048) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-n", type=int, default=128) + parser.add_argument("--block-k", type=int, default=256) + parser.add_argument("--num-stages", type=int, default=2) + parser.add_argument("--warp-policy", choices=["Square"], default="Square") + parser.add_argument("--out-dtype", choices=["bfloat16", "float32"], default="bfloat16") + parser.add_argument("--input-mode", choices=["random", "ones", "bf16_quantized"], default="random") + parser.add_argument( + "--scale-mode", + choices=["constant", "random_binary", "random_sfa", "random_sfb"], + default="constant", + ) + parser.add_argument("--producer-regs", type=int, default=40) + parser.add_argument("--consumer-regs", type=int, default=224) + parser.add_argument("--maxrregcount", type=int) + parser.add_argument("--ptxas-verbose", action="store_true") + parser.add_argument("--dump-source") + parser.add_argument("--dump-sass") + parser.add_argument("--backend", choices=["event", "cupti", "cudagraph"], default="event") + parser.add_argument("--return-mode", choices=["min", "max", "mean", "median"], default="mean") + parser.add_argument("--warmup-ms", type=float, default=25) + parser.add_argument("--rep-ms", type=float, default=100) + parser.add_argument("--n-warmup", type=int, default=0) + parser.add_argument("--n-repeat", type=int, default=0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--verify", action="store_true") + parser.add_argument("--profile-only", action="store_true") + parser.add_argument("--profile-warmup-launches", type=int, default=0) + parser.add_argument("--run-cutlass", action="store_true") + parser.add_argument("--cutlass-iterations", type=int, default=20) + parser.add_argument("--cutlass-build-dir", default="build-cutlass-sm120") + parser.add_argument("--cutlass-build-jobs", type=int, default=8) + parser.add_argument("--cutlass-binary") + parser.add_argument("--rebuild-cutlass", action="store_true") + parser.add_argument("--cmake") + parser.add_argument("--nvcc") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + capability = torch.cuda.get_device_capability() + if capability < (12, 0): + raise RuntimeError(f"SM120 or newer is required, got compute capability {capability}") + + print(f"Shape: M={args.m}, N={args.n}, K={args.k}") + effective_scale_mode = "ue4m3_from_tilelang_bf16_quantizer" if args.input_mode == "bf16_quantized" else args.scale_mode + print( + f"TileLang tile: {args.block_m}x{args.block_n}x{args.block_k}, " + f"threads={_SM120_THREADS}, policy={args.warp_policy}, output={args.out_dtype}, " + f"input_mode={args.input_mode}, scale_mode={effective_scale_mode}" + ) + print( + "TileLang SM120 block-scaled path: " + f"scale_layout={_SM120_SCALE_LAYOUT}, producer_regs={args.producer_regs}, " + f"consumer_regs={args.consumer_regs}, maxrregcount={args.maxrregcount}" + ) + + tilelang_latency_ms, tilelang_tflops = run_tilelang(args) + print(f"TileLang latency: {tilelang_latency_ms:.4f} ms") + print(f"TileLang FLOPS: {tilelang_tflops:.2f} TFLOPS") + + if args.run_cutlass: + cutlass_latency_ms, cutlass_tflops = run_cutlass(args) + print(f"CUTLASS 79a latency: {cutlass_latency_ms:.4f} ms") + print(f"CUTLASS 79a FLOPS: {cutlass_tflops:.2f} TFLOPS") + print(f"TileLang / CUTLASS: {tilelang_tflops / cutlass_tflops:.3f}x") + + +if __name__ == "__main__": + main() diff --git a/examples/gemm_sm120/sm120_nvfp4_blockscaled_ws.py b/examples/gemm_sm120/sm120_nvfp4_blockscaled_ws.py new file mode 100644 index 000000000..4ed78def5 --- /dev/null +++ b/examples/gemm_sm120/sm120_nvfp4_blockscaled_ws.py @@ -0,0 +1,252 @@ +"""Experimental SM120 NVFP4 GEMM using TileLang warp-specialization blocks. + +This file is intentionally separate from the main benchmark. It explores the +existing ``T.ws`` abstraction for a compact producer/consumer shape: + +* ``T.ws(1)`` issues TMA copies for A, B, SFA, and SFB. +* ``T.ws(0)`` waits on the stage barrier and issues ``T.mma_gemm_blockscaled``. + +The scale tensors use the semantic row-major ``[M or N, K / 64]`` uint32 +contract. This example is for API/lowering study first; it is not the persistent +pingpong performance path. +""" + +import argparse +from pathlib import Path +import sys + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import torch +import tilelang +import tilelang.language as T +from tilelang.carver.arch import driver +from tilelang.profiler import do_bench + +from sm120_nvfp4_blockscaled_gemm import ( + _make_binary_scale_words, + _make_constant_scale_words, + _make_ones_packed_fp4, + _make_packed_fp4, + _tflops, + _verify_tilelang_output, +) + + +def _jit_pass_configs(ptxas_verbose: bool = False) -> dict: + pass_configs = {} + if ptxas_verbose: + pass_configs[tilelang.PassConfigKey.TL_ENABLE_PTXAS_VERBOSE_OUTPUT] = True + return pass_configs + + +@tilelang.jit(out_idx=None, pass_configs=_jit_pass_configs()) +def tilelang_nvfp4_blockscaled_ws( + M: int, + N: int, + K: int, + block_M: int, + block_N: int, + block_K: int, + num_stages: int, + threads: int, + warp_policy, + out_dtype, +): + """Compact T.ws producer/consumer NVFP4 block-scaled GEMM.""" + + assert M % block_M == 0 + assert N % block_N == 0 + assert K % block_K == 0 + assert block_K % 64 == 0 + assert num_stages >= 2 + + in_dtype = T.float4_e2m1fn + accum_dtype = T.float32 + sf_words_per_block_k = block_K // 64 + sf_granularity_k = 16 + + @T.prim_func + def main( + A: T.Tensor((M, K), in_dtype), + B: T.Tensor((N, K), in_dtype), + SFA: T.Tensor((M, K // 64), T.uint32), + SFB: T.Tensor((N, K // 64), T.uint32), + C: T.Tensor((M, N), out_dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by): + # Keep these allocations single-stage. T.Pipelined owns the + # multi-versioning for num_stages; adding an explicit leading + # num_stages dimension here would multiply the physical smem budget. + A_shared = T.alloc_shared((block_M, block_K), in_dtype) + B_shared = T.alloc_shared((block_N, block_K), in_dtype) + SFA_shared = T.alloc_shared((block_M, sf_words_per_block_k), T.uint32) + SFB_shared = T.alloc_shared((block_N, sf_words_per_block_k), T.uint32) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + data_ready = T.alloc_barrier(arrive_count=128) + compute_done = T.alloc_barrier(arrive_count=128) + + with T.ws(0): + T.clear(C_local) + + for ko in T.Pipelined(K // block_K, num_stages=num_stages): + with T.ws(1): + T.barrier_wait(compute_done, (ko + 1) % 2) + T.tma_copy( + A[by * block_M : (by + 1) * block_M, ko * block_K : (ko + 1) * block_K], + A_shared, + barrier=data_ready, + ) + T.tma_copy( + B[bx * block_N : (bx + 1) * block_N, ko * block_K : (ko + 1) * block_K], + B_shared, + barrier=data_ready, + ) + T.tma_copy( + SFA[ + by * block_M : (by + 1) * block_M, + ko * sf_words_per_block_k : (ko + 1) * sf_words_per_block_k, + ], + SFA_shared, + barrier=data_ready, + ) + T.tma_copy( + SFB[ + bx * block_N : (bx + 1) * block_N, + ko * sf_words_per_block_k : (ko + 1) * sf_words_per_block_k, + ], + SFB_shared, + barrier=data_ready, + ) + T.barrier_arrive(data_ready) + + with T.ws(0): + T.barrier_wait(data_ready, ko % 2) + T.mma_gemm_blockscaled( + A_shared, + B_shared, + C_local, + SFA_shared, + SFB_shared, + transpose_B=True, + policy=warp_policy, + clear_accum=False, + k_start=0, + sf_a_granularity_k=sf_granularity_k, + sf_b_granularity_k=sf_granularity_k, + ) + T.barrier_arrive(compute_done) + + with T.ws(0): + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--m", type=int, default=512) + parser.add_argument("--n", type=int, default=512) + parser.add_argument("--k", type=int, default=512) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-n", type=int, default=128) + parser.add_argument("--block-k", type=int, default=256) + parser.add_argument("--num-stages", type=int, default=2) + parser.add_argument("--threads", type=int, default=256) + parser.add_argument("--warp-policy", choices=["Square"], default="Square") + parser.add_argument("--out-dtype", choices=["bfloat16", "float32"], default="bfloat16") + parser.add_argument("--input-mode", choices=["random", "ones"], default="random") + parser.add_argument( + "--scale-mode", + choices=["constant", "random_binary", "random_sfa", "random_sfb"], + default="random_binary", + ) + parser.add_argument("--warmup-ms", type=float, default=1) + parser.add_argument("--rep-ms", type=float, default=1) + parser.add_argument("--backend", choices=["event", "cupti", "cudagraph"], default="event") + parser.add_argument("--return-mode", choices=["min", "max", "mean", "median"], default="min") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--verify", action="store_true") + parser.add_argument("--dump-source") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + capability = torch.cuda.get_device_capability() + if capability < (12, 0): + raise RuntimeError(f"SM120 or newer is required, got compute capability {capability}") + + out_dtype = T.bfloat16 if args.out_dtype == "bfloat16" else T.float32 + out_torch_dtype = torch.bfloat16 if args.out_dtype == "bfloat16" else torch.float32 + warp_policy = getattr(T.GemmWarpPolicy, args.warp_policy) + + kernel = tilelang_nvfp4_blockscaled_ws( + args.m, + args.n, + args.k, + args.block_m, + args.block_n, + args.block_k, + args.num_stages, + args.threads, + warp_policy, + out_dtype, + ) + source = kernel.get_kernel_source() + if args.dump_source: + dump_path = Path(args.dump_source) + dump_path.parent.mkdir(parents=True, exist_ok=True) + dump_path.write_text(source) + print(f"TileLang CUDA source: {dump_path}") + if "warp_specialize" not in source and "kWarpSpecializationScope" not in source: + print("warning: generated source does not show a textual warp-specialization marker") + + if args.input_mode == "ones": + A = _make_ones_packed_fp4(args.m, args.k) + B = _make_ones_packed_fp4(args.n, args.k) + else: + A = _make_packed_fp4(args.m, args.k, seed=args.seed) + B = _make_packed_fp4(args.n, args.k, seed=args.seed + 1) + + if args.scale_mode == "constant": + SFA = _make_constant_scale_words(args.m, args.k) + SFB = _make_constant_scale_words(args.n, args.k) + elif args.scale_mode == "random_binary": + SFA = _make_binary_scale_words(args.m, args.k, seed=args.seed + 100) + SFB = _make_binary_scale_words(args.n, args.k, seed=args.seed + 200) + elif args.scale_mode == "random_sfa": + SFA = _make_binary_scale_words(args.m, args.k, seed=args.seed + 100) + SFB = _make_constant_scale_words(args.n, args.k) + elif args.scale_mode == "random_sfb": + SFA = _make_constant_scale_words(args.m, args.k) + SFB = _make_binary_scale_words(args.n, args.k, seed=args.seed + 200) + else: + raise ValueError(f"Unsupported scale_mode={args.scale_mode!r}") + + C = torch.empty((args.m, args.n), device="cuda", dtype=out_torch_dtype) + kernel(A, B, SFA, SFB, C) + torch.cuda.synchronize() + + if args.verify: + _verify_tilelang_output(A, B, SFA, SFB, C, out_torch_dtype, args.scale_mode, args.block_m, args.block_n, driver.get_num_sms()) + print("TileLang WS correctness: passed") + + latency_ms = do_bench( + lambda: kernel(A, B, SFA, SFB, C), + warmup=args.warmup_ms, + rep=args.rep_ms, + backend=args.backend, + return_mode=args.return_mode, + ) + print(f"TileLang WS latency: {latency_ms:.4f} ms") + print(f"TileLang WS FLOPS: {_tflops(args.m, args.n, args.k, latency_ms):.2f} TFLOPS") + + +if __name__ == "__main__": + main() diff --git a/maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py b/maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py new file mode 100644 index 000000000..5eeac78ef --- /dev/null +++ b/maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py @@ -0,0 +1,387 @@ +import os +from pathlib import Path + +import tilelang +import tilelang.language as T +from tilelang import tvm +from tilelang.intrinsics import TensorCoreIntrinEmitterWithBlockScale, get_swizzle_layout +from tilelang.transform import simplify_prim_func + + +M = 128 +N = 128 +K = 256 +MICRO_K = 64 +SF_VEC_SIZE = 16 +SCALE_WORD_K = MICRO_K // SF_VEC_SIZE +FP4_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + + +def _make_swizzle_layout(shared_buf): + dtype = shared_buf.dtype + shape = shared_buf.shape + can_swizzle = shape[-1] * tvm.DataType(dtype).bits % 512 == 0 + if not can_swizzle: + return T.Layout(shape, lambda *args: args) + + def transform_func(i, j): + new_warp_i, new_warp_j = get_swizzle_layout(i, j, shape[-1], dtype) + return [new_warp_i, new_warp_j] + + return T.Layout(shape, transform_func) + + +@simplify_prim_func +def _make_tilelang_nvf4_kernel(): + in_dtype = T.float4_e2m1fn + out_dtype = T.float32 + accum_dtype = T.float32 + + micro_size_x = 16 + micro_size_y = 16 + micro_size_k = MICRO_K + + block_row_warps = 2 + block_col_warps = 2 + warp_row_tiles = 32 + warp_col_tiles = 32 + chunk = K + shared_scope = "shared.dyn" + + block_M = block_row_warps * warp_row_tiles + block_N = block_col_warps * warp_col_tiles + block_K = chunk + + A_shared_shape = (block_M, block_K) + B_shared_shape = (block_N, block_K) + SFA_shared_shape = (block_M, block_K // micro_size_k) + SFB_shared_shape = (block_N, block_K // micro_size_k) + C_shared_shape = ( + block_M // micro_size_x, + block_N // micro_size_y, + micro_size_x, + micro_size_y, + ) + + warp_size = 32 + threads = warp_size * (block_row_warps * block_col_warps) + local_size_a = (micro_size_x * micro_size_k) // warp_size + local_size_b = (micro_size_y * micro_size_k) // warp_size + local_size_c = (micro_size_x * micro_size_y) // warp_size + warp_rows = warp_row_tiles // micro_size_x + warp_cols = warp_col_tiles // micro_size_y + + mma_emitter = TensorCoreIntrinEmitterWithBlockScale( + a_dtype=in_dtype, + b_dtype=in_dtype, + accum_dtype=accum_dtype, + a_transposed=False, + b_transposed=True, + block_row_warps=block_row_warps, + block_col_warps=block_col_warps, + warp_row_tiles=warp_row_tiles, + warp_col_tiles=warp_col_tiles, + chunk=chunk, + kind="mxf4nvf4", + scale_vec_size=4, + stype="ue4m3", + ) + + @T.prim_func + def main( + A: T.Tensor((M, K), in_dtype), + B: T.Tensor((N, K), in_dtype), + SFA: T.Tensor((M, K // MICRO_K), T.uint32), + SFB: T.Tensor((N, K // MICRO_K), T.uint32), + C: T.Tensor((M, N), out_dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as ( + bx, + by, + ): + A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope) + B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope) + SFA_shared = T.alloc_shared(SFA_shared_shape, T.uint32, scope=shared_scope) + SFB_shared = T.alloc_shared(SFB_shared_shape, T.uint32, scope=shared_scope) + C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope) + A_local = T.alloc_local((warp_rows * local_size_a), in_dtype) + B_local = T.alloc_local((warp_cols * local_size_b), in_dtype) + C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype) + + T.annotate_layout( + { + A_shared: _make_swizzle_layout(A_shared), + B_shared: _make_swizzle_layout(B_shared), + } + ) + T.use_swizzle(panel_size=10) + + T.clear(C_local) + + for ko in T.Pipelined((K // block_K), num_stages=2): + for i, k in T.Parallel(block_M, block_K): + A_shared[i, k] = A[by * block_M + i, ko * block_K + k] + + for j, k in T.Parallel(block_N, block_K): + B_shared[j, k] = B[bx * block_N + j, ko * block_K + k] + + for i, k in T.Parallel(block_M, block_K // micro_size_k): + SFA_shared[i, k] = SFA[by * block_M + i, ko * (block_K // micro_size_k) + k] + + for j, k in T.Parallel(block_N, block_K // micro_size_k): + SFB_shared[j, k] = SFB[bx * block_N + j, ko * (block_K // micro_size_k) + k] + + for ki in T.serial(0, (block_K // micro_size_k)): + mma_emitter.ldmatrix_a(A_local, A_shared, ki) + mma_emitter.ldmatrix_b(B_local, B_shared, ki) + mma_emitter.mma( + A_local, + B_local, + C_local, + SFA_shared, + SFB_shared, + ki=ki, + k_start=0, + sf_a_granularity_k=SF_VEC_SIZE, + sf_b_granularity_k=SF_VEC_SIZE, + ) + + mma_emitter.stmatrix(C_local, C_shared) + + for i, j in T.Parallel(block_M, block_N): + C[by * block_M + i, bx * block_N + j] = C_shared[ + i // micro_size_x, + j // micro_size_y, + i % micro_size_x, + j % micro_size_y, + ] + + return main + + +def _pack_tilelang_sf_u32(sf_bytes): + import torch + + assert sf_bytes.dtype == torch.uint8 + mn, sf_blocks = sf_bytes.shape + assert sf_blocks % SCALE_WORD_K == 0 + words = sf_bytes.reshape(mn, sf_blocks // SCALE_WORD_K, SCALE_WORD_K).to(torch.int64) + packed = words[:, :, 0] | (words[:, :, 1] << 8) | (words[:, :, 2] << 16) | (words[:, :, 3] << 24) + return packed.to(torch.uint32).contiguous() + + +def _pack_cutlass_sf_bytes(sf_bytes): + """Pack logical (MN, K/16) UE4M3 bytes into CUTLASS Sm1xxBlockScaledConfig<16> order.""" + import torch + + assert sf_bytes.dtype == torch.uint8 + mn, sf_blocks = sf_bytes.shape + assert mn == 128 + assert sf_blocks == K // SF_VEC_SIZE + + out = torch.empty((mn * sf_blocks,), device=sf_bytes.device, dtype=torch.uint8) + for sf_idx in range(sf_blocks): + k64 = sf_idx // SCALE_WORD_K + k_word = sf_idx % SCALE_WORD_K + for row in range(mn): + offset = k64 * 512 + (row % 32) * 16 + (row // 32) * 4 + k_word + out[offset] = sf_bytes[row, sf_idx] + return out.contiguous() + + +def _pack_cutlass_colmajor_fp4_from_rowmajor(packed, rows: int, cols: int): + """Repack logical row-major FP4 bytes into CUTLASS ColumnMajor packed storage.""" + import torch + + assert packed.dtype == torch.int8 + logical_bytes = packed.contiguous().view(torch.uint8) + lo = logical_bytes & 0x0F + hi = (logical_bytes >> 4) & 0x0F + logical = torch.empty((rows, cols), device=packed.device, dtype=torch.uint8) + logical[:, 0::2] = lo + logical[:, 1::2] = hi + + col_major = logical.T.contiguous().reshape(-1) + packed_col_major = (col_major[0::2] | (col_major[1::2] << 4)).contiguous() + return packed_col_major.view(torch.int8) + + +def _swap_fp4_nibbles(packed): + import torch + + u = packed.contiguous().view(torch.uint8) + return (((u & 0x0F) << 4) | ((u >> 4) & 0x0F)).contiguous().view(torch.int8) + + +def _decode_rowmajor_fp4(packed, rows: int, cols: int): + import torch + + u = packed.contiguous().view(torch.uint8) + lut = torch.tensor(FP4_E2M1_VALUES, device=packed.device, dtype=torch.float32) + out = torch.empty((rows, cols), device=packed.device, dtype=torch.float32) + out[:, 0::2] = lut[(u & 0x0F).long()] + out[:, 1::2] = lut[((u >> 4) & 0x0F).long()] + return out + + +def _decode_ue4m3(sf_bytes): + import torch + + u = sf_bytes.to(torch.int32) + exponent = (u >> 3) & 0x0F + mantissa = u & 0x07 + normal = (1.0 + mantissa.to(torch.float32) / 8.0) * torch.pow(2.0, exponent.to(torch.float32) - 7.0) + subnormal = (mantissa.to(torch.float32) / 8.0) * torch.pow(torch.tensor(2.0, device=sf_bytes.device), -6.0) + return torch.where(exponent == 0, subnormal, normal) + + +def _blockscaled_fp4_reference(a, b, sfa_logical, sfb_logical): + a_f32 = _decode_rowmajor_fp4(a, M, K) + b_f32 = _decode_rowmajor_fp4(b, N, K) + sfa = _decode_ue4m3(sfa_logical).repeat_interleave(SF_VEC_SIZE, dim=1) + sfb = _decode_ue4m3(sfb_logical).repeat_interleave(SF_VEC_SIZE, dim=1) + return (a_f32 * sfa) @ (b_f32 * sfb).T + + +def _build_cutlass_extension(): + from torch.utils.cpp_extension import load + + repo = Path(__file__).resolve().parents[2] + cutlass_root_env = os.environ.get("CUTLASS_ROOT") + cutlass_root = Path(cutlass_root_env) if cutlass_root_env else repo / "3rdparty" / "cutlass" + include_paths = [ + cutlass_root / "include", + cutlass_root / "tools" / "util" / "include", + ] + if not all(path.exists() for path in include_paths): + raise RuntimeError( + "CUTLASS headers were not found. Set CUTLASS_ROOT to a CUTLASS checkout " + "or populate the repository's 3rdparty/cutlass submodule." + ) + + return load( + name="tilelang_cutlass_nvf4_ref", + sources=[str(repo / "maint" / "gemm" / "cutlass_nvf4_ref.cu")], + extra_include_paths=[str(path) for path in include_paths], + extra_cuda_cflags=[ + "-std=c++20", + "-arch=sm_120a", + "--expt-relaxed-constexpr", + "-DCUTLASS_ARCH_MMA_SM120_SUPPORTED", + ], + extra_cflags=["-std=c++20"], + verbose=True, + ) + + +def run_compare(): + import torch + + torch.manual_seed(0) + assert torch.cuda.is_available(), "CUDA is required" + + input_mode = os.environ.get("NVF4_INPUT_MODE", "random") + if input_mode == "constant": + byte = int(os.environ.get("NVF4_INPUT_BYTE", "0x22"), 0) + a = torch.full((M, K // 2), byte, device="cuda", dtype=torch.uint8).view(torch.int8) + b = torch.full((N, K // 2), byte, device="cuda", dtype=torch.uint8).view(torch.int8) + elif input_mode == "a_random_b_constant": + byte = int(os.environ.get("NVF4_INPUT_BYTE", "0x22"), 0) + a = torch.randint(-128, 128, (M, K // 2), device="cuda", dtype=torch.int8) + b = torch.full((N, K // 2), byte, device="cuda", dtype=torch.uint8).view(torch.int8) + elif input_mode == "a_constant_b_random": + byte = int(os.environ.get("NVF4_INPUT_BYTE", "0x22"), 0) + a = torch.full((M, K // 2), byte, device="cuda", dtype=torch.uint8).view(torch.int8) + b = torch.randint(-128, 128, (N, K // 2), device="cuda", dtype=torch.int8) + elif input_mode == "random": + a = torch.randint(-128, 128, (M, K // 2), device="cuda", dtype=torch.int8) + b = torch.randint(-128, 128, (N, K // 2), device="cuda", dtype=torch.int8) + else: + raise ValueError(f"Unsupported NVF4_INPUT_MODE={input_mode!r}") + + scale_mode = os.environ.get("NVF4_SCALE_MODE", "varying") + if scale_mode == "constant": + sfa_logical = torch.full((M, K // SF_VEC_SIZE), 0x38, device="cuda", dtype=torch.uint8) + sfb_logical = torch.full((N, K // SF_VEC_SIZE), 0x38, device="cuda", dtype=torch.uint8) + elif scale_mode in {"varying", "sfa_varying", "sfb_varying", "k_varying", "mn_varying"}: + # Avoid UE4M3 NaN encodings while still varying scales by row and K block. + sfa_logical = torch.empty((M, K // SF_VEC_SIZE), device="cuda", dtype=torch.uint8) + sfb_logical = torch.empty((N, K // SF_VEC_SIZE), device="cuda", dtype=torch.uint8) + row = torch.arange(M, device="cuda", dtype=torch.int32).reshape(M, 1) + col = torch.arange(K // SF_VEC_SIZE, device="cuda", dtype=torch.int32).reshape(1, K // SF_VEC_SIZE) + sfa_logical.copy_((0x20 + ((row * 3 + col * 5) % 32)).to(torch.uint8)) + sfb_logical.copy_((0x20 + ((row * 7 + col * 11) % 32)).to(torch.uint8)) + if scale_mode == "k_varying": + sfa_logical.copy_((0x20 + ((col * 5) % 32)).expand_as(sfa_logical).to(torch.uint8)) + sfb_logical.copy_((0x20 + ((col * 11) % 32)).expand_as(sfb_logical).to(torch.uint8)) + elif scale_mode == "mn_varying": + sfa_logical.copy_((0x20 + ((row * 3) % 32)).expand_as(sfa_logical).to(torch.uint8)) + sfb_logical.copy_((0x20 + ((row * 7) % 32)).expand_as(sfb_logical).to(torch.uint8)) + if scale_mode == "sfa_varying": + sfb_logical.fill_(0x38) + elif scale_mode == "sfb_varying": + sfa_logical.fill_(0x38) + else: + raise ValueError(f"Unsupported NVF4_SCALE_MODE={scale_mode!r}") + + sfa_tl = _pack_tilelang_sf_u32(sfa_logical) + sfb_tl = _pack_tilelang_sf_u32(sfb_logical) + sfa_cutlass = _pack_cutlass_sf_bytes(sfa_logical) + sfb_cutlass = _pack_cutlass_sf_bytes(sfb_logical) + + a_tl = _swap_fp4_nibbles(a) if os.environ.get("NVF4_SWAP_A_FOR_TILELANG", "0") == "1" else a + b_tl = _swap_fp4_nibbles(b) if os.environ.get("NVF4_SWAP_B_FOR_TILELANG", "0") == "1" else b + + kernel = tilelang.compile(_make_tilelang_nvf4_kernel(), target="cuda", out_idx=[4]) + c_tl = kernel(a_tl, b_tl, sfa_tl, sfb_tl) + + cutlass_ref = _build_cutlass_extension() + if os.environ.get("NVF4_CUTLASS_REPACK_B", "0") == "1": + b_cutlass = _pack_cutlass_colmajor_fp4_from_rowmajor(b, N, K) + else: + b_cutlass = b + c_in = torch.zeros((M, N), device="cuda", dtype=torch.float32) + c_cutlass = torch.empty((M, N), device="cuda", dtype=torch.float32) + cutlass_ref.cutlass_nvf4_gemm_128x128x256(a, b_cutlass, sfa_cutlass, sfb_cutlass, c_in, c_cutlass) + + torch.cuda.synchronize() + diff = (c_tl - c_cutlass).abs() + diff_t = (c_tl - c_cutlass.T).abs() + ref = _blockscaled_fp4_reference(a, b, sfa_logical, sfb_logical) + diff_tl_ref = (c_tl - ref).abs() + diff_cutlass_ref = (c_cutlass - ref).abs() + print("scale_mode:", scale_mode) + print("input_mode:", input_mode) + print("max_abs_diff:", diff.max().item()) + print("mean_abs_diff:", diff.mean().item()) + print("max_abs_diff_transposed:", diff_t.max().item()) + print("mean_abs_diff_transposed:", diff_t.mean().item()) + print("max_abs_diff_tilelang_ref:", diff_tl_ref.max().item()) + print("mean_abs_diff_tilelang_ref:", diff_tl_ref.mean().item()) + print("max_abs_diff_cutlass_ref:", diff_cutlass_ref.max().item()) + print("mean_abs_diff_cutlass_ref:", diff_cutlass_ref.mean().item()) + if os.environ.get("NVF4_SKIP_ASSERT", "0") == "1": + return + torch.testing.assert_close(c_tl, c_cutlass, rtol=0.0, atol=0.0) + print("TileLang NVF4 block-scale output matches CUTLASS exactly.") + + +if __name__ == "__main__": + run_compare() diff --git a/maint/gemm/cutlass_nvf4_ref.cu b/maint/gemm/cutlass_nvf4_ref.cu new file mode 100644 index 000000000..d9735e97b --- /dev/null +++ b/maint/gemm/cutlass_nvf4_ref.cu @@ -0,0 +1,160 @@ +#include + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/packed_stride.hpp" + +namespace { + +using namespace cute; + +using ElementA = cutlass::float_e2m1_t; +using ElementB = cutlass::float_e2m1_t; +using ElementC = float; +using ElementD = float; +using ElementAccumulator = float; +using ElementCompute = float; + +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::RowMajor; +using LayoutD = cutlass::layout::RowMajor; + +using ElementPairA = cutlass::nv_float4_t; +using ElementPairB = cutlass::nv_float4_t; + +static constexpr int AlignmentA = + 16 * 8 / cutlass::sizeof_bits::value; +static constexpr int AlignmentB = + 16 * 8 / cutlass::sizeof_bits::value; +static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; +static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + +using TileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; + +using CollectiveEpilogue = + typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape, + ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC, + ElementD, LayoutD, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + +using CollectiveMainloop = + typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp, + ElementPairA, LayoutA, AlignmentA, ElementPairB, LayoutB, AlignmentB, + ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedPingpong>::CollectiveOp; + +template struct Dummy { + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +using GemmKernel = typename Dummy::GemmKernel; +using Gemm = typename Dummy::Gemm; +using StrideA = typename GemmKernel::StrideA; +using StrideB = typename GemmKernel::StrideB; +using StrideC = typename GemmKernel::StrideC; +using StrideD = typename GemmKernel::StrideD; + +void check_tensor(torch::Tensor const &tensor, char const *name, + torch::DeviceType device) { + TORCH_CHECK(tensor.device().type() == device, name, " must be on CUDA"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +void cutlass_nvf4_gemm_128x128x256(torch::Tensor A, torch::Tensor B, + torch::Tensor SFA, torch::Tensor SFB, + torch::Tensor C, torch::Tensor D) { +#if !(defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || \ + defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)) + TORCH_CHECK( + false, + "CUTLASS was not compiled with SM120/SM121 block-scale MMA support"); +#else + constexpr int M = 128; + constexpr int N = 128; + constexpr int K = 256; + + check_tensor(A, "A", torch::kCUDA); + check_tensor(B, "B", torch::kCUDA); + check_tensor(SFA, "SFA", torch::kCUDA); + check_tensor(SFB, "SFB", torch::kCUDA); + check_tensor(C, "C", torch::kCUDA); + check_tensor(D, "D", torch::kCUDA); + + TORCH_CHECK(A.numel() == M * K / 2, + "A must contain packed NVF4 bytes for 128x256"); + TORCH_CHECK(B.numel() == N * K / 2, + "B must contain packed NVF4 bytes for 128x256"); + TORCH_CHECK(SFA.numel() == M * (K / 16), + "SFA must be CUTLASS-layout UE4M3 bytes"); + TORCH_CHECK(SFB.numel() == N * (K / 16), + "SFB must be CUTLASS-layout UE4M3 bytes"); + TORCH_CHECK(C.numel() == M * N, "C must be 128x128 f32"); + TORCH_CHECK(D.numel() == M * N, "D must be 128x128 f32"); + + auto problem = cute::make_shape(M, N, K, 1); + auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); + auto stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); + auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + + using Sm1xxBlkScaledConfig = + typename GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem); + auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem); + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + problem, + {reinterpret_cast(A.data_ptr()), stride_A, + reinterpret_cast(B.data_ptr()), stride_B, + reinterpret_cast(SFA.data_ptr()), + layout_SFA, + reinterpret_cast(SFB.data_ptr()), + layout_SFB}, + {{1.0f, 0.0f}, + reinterpret_cast(C.data_ptr()), + stride_C, + reinterpret_cast(D.data_ptr()), + stride_D}}; + + Gemm gemm; + auto status = gemm.can_implement(arguments); + TORCH_CHECK(status == cutlass::Status::kSuccess, + "CUTLASS can_implement failed: ", cutlassGetStatusString(status)); + + size_t workspace_size = Gemm::get_workspace_size(arguments); + auto workspace = torch::empty( + {static_cast(workspace_size)}, + torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA)); + + status = gemm.initialize(arguments, workspace.data_ptr()); + TORCH_CHECK(status == cutlass::Status::kSuccess, + "CUTLASS initialize failed: ", cutlassGetStatusString(status)); + + status = gemm.run(); + TORCH_CHECK(status == cutlass::Status::kSuccess, + "CUTLASS run failed: ", cutlassGetStatusString(status)); + TORCH_CHECK(cudaDeviceSynchronize() == cudaSuccess, "CUTLASS kernel failed"); +#endif +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("cutlass_nvf4_gemm_128x128x256", &cutlass_nvf4_gemm_128x128x256, + "CUTLASS SM120 NVF4 block-scaled GEMM 128x128x256"); +} diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc index cfdb27f5e..9672ece57 100644 --- a/src/cuda/codegen/codegen_cuda.cc +++ b/src/cuda/codegen/codegen_cuda.cc @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -33,6 +34,18 @@ using namespace ffi; namespace { +bool EnableSM120CompactUnpackedFP4Shared() { + const char *value = std::getenv("TL_SM120_COMPACT_UNPACKED_FP4_SHARED"); + return value != nullptr && value[0] != '\0' && + !(value[0] == '0' && value[1] == '\0'); +} + +bool EnableSM120MmaRawUnpackedFP4AccessPtr() { + const char *value = std::getenv("TL_SM120_MMA_RAW_UNPACKED_FP4_ACCESS_PTR"); + return value != nullptr && value[0] != '\0' && + !(value[0] == '0' && value[1] == '\0'); +} + bool IsValidCPAsyncTransferBytes(int64_t bytes) { return bytes == 4 || bytes == 8 || bytes == 16; } @@ -604,6 +617,9 @@ std::string CodeGenTileLangCUDA::Finish() { if (need_mma_instruction_h_) { decl_stream << "#include \n"; } + if (need_gemm_sm120_h_) { + decl_stream << "#include \n"; + } if (need_wgmma_instruction_h_) { decl_stream << "#include \n"; } @@ -2184,7 +2200,12 @@ std::string CodeGenTileLangCUDA::GetBufferRef(DataType t, } if (pointed_to.is_float4_e2m1_unpacked() && (scope == "shared" || scope == "shared.dyn")) { - ptr_os << "uint8_t"; + if (EnableSM120CompactUnpackedFP4Shared()) { + enable_fp4_ = true; + ptr_os << "fp4_e2_t"; + } else { + ptr_os << "uint8_t"; + } } else { PrintType(pointed_to, ptr_os); } @@ -2205,6 +2226,12 @@ std::string CodeGenTileLangCUDA::GetBufferRef(DataType t, return os.str(); } std::string index_str = PrintExpr(index); + bool is_compact_scalar_fp4 = + buffer_element_dtype.lanes() == 1 && + (buffer_element_dtype.is_float4_e2m1fn() || + (EnableSM120CompactUnpackedFP4Shared() && + buffer_element_dtype.is_float4_e2m1_unpacked() && + (scope == "shared" || scope == "shared.dyn"))); if ((t.bits() == 4 && !t.is_float4()) || (t.bits() == 1 && t.is_int())) { // Scalar int4/uint4 storage is byte-packed (2 logical elements per byte). // Vector int4 loads/stores reinterpret the underlying packed bytes as the @@ -2220,8 +2247,7 @@ std::string CodeGenTileLangCUDA::GetBufferRef(DataType t, << " + " << index_str << ")"; } else if (t == buffer_element_dtype) { int div_factor = 1; - if (buffer_element_dtype.is_float4_e2m1fn() && - buffer_element_dtype.lanes() == 1) { + if (is_compact_scalar_fp4) { div_factor = 2; } index_str = @@ -2231,8 +2257,7 @@ std::string CodeGenTileLangCUDA::GetBufferRef(DataType t, // Fix fp4 pointer arithmetic: fp4 elements are 4-bit packed 2 per byte. // fp4* + n incorrectly advances n bytes (skipping 2n elements). int div_factor = 1; - if (buffer_element_dtype.is_float4_e2m1fn() && - buffer_element_dtype.lanes() == 1) { + if (is_compact_scalar_fp4) { div_factor = 2; } index_str = @@ -2652,12 +2677,63 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { } print_extern_call_stmt(ss.str(), 0, 1); } else if (op->op.same_as(tl::ptx_ldmatrix())) { + ICHECK_EQ(op->args.size(), 4U); int trans = Downcast(op->args[0])->value; int num = Downcast(op->args[1])->value; std::string func_name = "tl::ptx_ldmatrix_x" + std::to_string(num); if (trans == 1) func_name += "_trans"; - print_extern_call_stmt(func_name, 2); + arith::Analyzer analyzer; + auto print_fp4_packed_local_ptr = + [&](const PrimExpr &expr) -> std::optional { + const VarNode *buffer_var = nullptr; + PrimExpr offset; + + if (const auto *call = expr.as()) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + ICHECK_GE(call->args.size(), 3U); + buffer_var = call->args[1].as(); + offset = call->args[2]; + } else if (call->op.same_as(tl::access_ptr())) { + ICHECK_EQ(call->args.size(), 3U) + << "tl.access_ptr expects 3 args: (BufferLoad, extent, rw_mask)"; + const auto *load = call->args[0].as(); + ICHECK(load) << "tl.access_ptr arg0 must be BufferLoad"; + if (load->indices.size() != 1) { + return std::nullopt; + } + buffer_var = load->buffer->data.get(); + offset = load->indices[0]; + } else if (call->op.same_as(builtin::address_of())) { + const auto *load = call->args[0].as(); + ICHECK(load) << "address_of arg must be BufferLoad"; + if (load->indices.size() != 1) { + return std::nullopt; + } + buffer_var = load->buffer->data.get(); + offset = load->indices[0]; + } + } + + if (buffer_var == nullptr || !offset.defined()) { + return std::nullopt; + } + auto it = fp4_packed_buffers_.find(buffer_var); + if (it == fp4_packed_buffers_.end()) { + return std::nullopt; + } + PrimExpr packed_offset = analyzer.Simplify(truncdiv(offset, 2)); + return "(" + it->second + " + (" + this->PrintExpr(packed_offset) + "))"; + }; + + std::string src_ptr = this->PrintExpr(op->args[2]); + std::optional packed_dst_ptr = + print_fp4_packed_local_ptr(op->args[3]); + std::string dst_ptr = packed_dst_ptr.has_value() + ? packed_dst_ptr.value() + : this->PrintExpr(op->args[3]); + this->PrintIndent(); + this->stream << func_name << "(" << src_ptr << ", " << dst_ptr << ");\n"; } else if (op->op.same_as(tl::ptx_stmatrix())) { int trans = Downcast(op->args[0])->value; int num = Downcast(op->args[1])->value; @@ -3041,6 +3117,423 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { replacer.register_rule("(C_ptr)", c_ref); replacer.register_rule("(C_offset)", c_bias); this->stream << replacer.rewrite(mma_call); + } else if (op->op.same_as(tl::ptx_mma_block_scale())) { + need_gemm_sm120_h_ = true; + // arg 0: accum_dtype + // arg 1: shape: mXnXkX + // arg 2: A layout: row/col + // arg 3: B layout: row/col + // arg 4: kind: mxf4nvf4 + // arg 5: scale_vec_size: 4 + // arg 6: A dtype: e2m1 + // arg 7: B dtype: e2m1 + // arg 8: scale_type: ue4m3 + // arg 9: A_data pointer + // arg 10: A_offset + // arg 11: B_data pointer + // arg 12: B_offset + // arg 13: C_data pointer + // arg 14: C_offset + // arg 15: scale_a_data + // arg 16: scale_b_data + // arg 17: scale_a_byte_id + // arg 18: scale_a_thread_id + // arg 19: scale_b_byte_id + // arg 20: scale_b_thread_id + ICHECK_EQ(op->args.size(), 21U); + std::string accum_dtype = Downcast(op->args[0])->value; + std::string shape = Downcast(op->args[1])->value; + std::string A_layout = Downcast(op->args[2])->value; + std::string B_layout = Downcast(op->args[3])->value; + std::string kind = Downcast(op->args[4])->value; + int scale_vec_size = static_cast(Downcast(op->args[5])->value); + std::string A_dtype = Downcast(op->args[6])->value; + std::string B_dtype = Downcast(op->args[7])->value; + std::string scale_type = Downcast(op->args[8])->value; + std::string a_ref = this->PrintExpr(op->args[9]); + std::string a_offset = this->PrintExpr(op->args[10]); + std::string b_ref = this->PrintExpr(op->args[11]); + std::string b_offset = this->PrintExpr(op->args[12]); + std::string c_ref = this->PrintExpr(op->args[13]); + std::string c_offset = this->PrintExpr(op->args[14]); + std::string scale_a = this->PrintExpr(op->args[15]); + std::string scale_b = this->PrintExpr(op->args[16]); + std::string scale_a_byte_id = this->PrintExpr(op->args[17]); + std::string scale_a_thread_id = this->PrintExpr(op->args[18]); + std::string scale_b_byte_id = this->PrintExpr(op->args[19]); + std::string scale_b_thread_id = this->PrintExpr(op->args[20]); + + bool supported_mxf4nvf4_4x_ue4m3 = + accum_dtype == "float32" && shape == "m16n8k64" && A_layout == "row" && + B_layout == "col" && kind == "mxf4nvf4" && scale_vec_size == 4 && + A_dtype == "e2m1" && B_dtype == "e2m1" && scale_type == "ue4m3"; + ICHECK(supported_mxf4nvf4_4x_ue4m3) + << "Unsupported ptx_mma_block_scale configuration: accum_dtype=" + << accum_dtype << ", shape=" << shape << ", A_layout=" << A_layout + << ", B_layout=" << B_layout << ", kind=" << kind + << ", scale_vec_size=" << scale_vec_size << ", A_dtype=" << A_dtype + << ", B_dtype=" << B_dtype << ", scale_type=" << scale_type + << ". Currently supported: f32 m16n8k64 row.col kind::mxf4nvf4 " + "scale_vec::4X e2m1.e2m1 f32 ue4m3."; + + auto resolve_fp4_packed_buffer = + [&](const PrimExpr &var_expr, std::string &ref, std::string &offset) { + if (const VarNode *var = var_expr.as()) { + auto it = fp4_packed_buffers_.find(var); + if (it != fp4_packed_buffers_.end()) { + ref = it->second; + offset = "(" + offset + ") / 2"; + } + } + }; + resolve_fp4_packed_buffer(op->args[9], a_ref, a_offset); + resolve_fp4_packed_buffer(op->args[11], b_ref, b_offset); + + this->PrintIndent(); + + std::string mma_kind_enum = "tl::SM120MmaBlockScaledKind::kMxf4nvf4"; + std::string scale_type_enum = "tl::SM120MmaScaleType::kUE4M3"; + std::string mma_call = + "tl::sm120_mma_sync_blockscaled<(MMA_KIND), (SCALE_VEC_SIZE), " + "(SCALE_TYPE)>(" + "reinterpret_cast((C_ptr) + (C_offset)), " + "reinterpret_cast((A_ptr) + (A_offset)), " + "reinterpret_cast((B_ptr) + (B_offset)), " + "reinterpret_cast((C_ptr) + (C_offset)), " + "(*reinterpret_cast((scale_a))), " + "(*reinterpret_cast((scale_b))), " + "static_cast((scale_a_byte_id)), " + "static_cast((scale_a_thread_id)), " + "static_cast((scale_b_byte_id)), " + "static_cast((scale_b_thread_id)));\n"; + + tl::codegen::Replacer replacer; + replacer.register_rule("(MMA_KIND)", mma_kind_enum); + replacer.register_rule("(SCALE_VEC_SIZE)", std::to_string(scale_vec_size)); + replacer.register_rule("(SCALE_TYPE)", scale_type_enum); + replacer.register_rule("(A_ptr)", a_ref); + replacer.register_rule("(A_offset)", a_offset); + replacer.register_rule("(B_ptr)", b_ref); + replacer.register_rule("(B_offset)", b_offset); + replacer.register_rule("(C_ptr)", c_ref); + replacer.register_rule("(C_offset)", c_offset); + replacer.register_rule("(scale_a)", scale_a); + replacer.register_rule("(scale_b)", scale_b); + replacer.register_rule("(scale_a_byte_id)", scale_a_byte_id); + replacer.register_rule("(scale_a_thread_id)", scale_a_thread_id); + replacer.register_rule("(scale_b_byte_id)", scale_b_byte_id); + replacer.register_rule("(scale_b_thread_id)", scale_b_thread_id); + this->stream << replacer.rewrite(mma_call); + } else if (op->op.same_as(tl::sm120_mma_blockscaled_cute_consumer_bridge())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // arg 2: A shared tile base pointer + // arg 3: B shared tile base pointer + // arg 4: SFA shared tile base pointer + // arg 5: SFB shared tile base pointer + // arg 6: static K64 index inside the block_K tile + ICHECK_EQ(op->args.size(), 7U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + auto void_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + this->PrintExpr(op->args[idx]) + + ")"; + }; + auto mma_operand_void_ptr_arg = [&](size_t idx) { + if (EnableSM120MmaRawUnpackedFP4AccessPtr()) { + std::optional element_type = + GetAccessPtrElementType(op->args[idx]); + const auto *call = op->args[idx].as(); + if (element_type.has_value() && + element_type.value().is_float4_e2m1_unpacked() && call != nullptr) { + if (call->op.same_as(builtin::tvm_access_ptr())) { + ICHECK_GE(call->args.size(), 3U); + return "reinterpret_cast(reinterpret_cast(" + + this->PrintExpr(call->args[1]) + ") + (" + + this->PrintExpr(call->args[2]) + "))"; + } + if (call->op.same_as(tl::access_ptr())) { + ICHECK_EQ(call->args.size(), 3U) << "tl.access_ptr expects 3 args: " + "(BufferLoad, extent, rw_mask)"; + const auto *load = call->args[0].as(); + ICHECK(load) << "tl.access_ptr arg0 must be BufferLoad"; + ICHECK_EQ(load->indices.size(), 1U) + << "SM120 raw FP4 bridge access expects 1D shared operand"; + return "reinterpret_cast(reinterpret_cast(" + + GetVarID(load->buffer->data.get()) + ") + (" + + this->PrintExpr(load->indices[0]) + "))"; + } + if (call->op.same_as(builtin::address_of())) { + ICHECK_EQ(call->args.size(), 1U); + const auto *load = call->args[0].as(); + ICHECK(load) << "address_of arg must be BufferLoad"; + ICHECK_EQ(load->indices.size(), 1U) + << "SM120 raw FP4 bridge address_of expects 1D shared operand"; + return "reinterpret_cast(reinterpret_cast(" + + GetVarID(load->buffer->data.get()) + ") + (" + + this->PrintExpr(load->indices[0]) + "))"; + } + } + } + return void_ptr_arg(idx); + }; + auto uint32_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + + this->PrintExpr(op->args[idx]) + ")"; + }; + + this->PrintIndent(); + this->stream << "tl::sm120_mma_blockscaled_cute_consumer_bridge(" + << "reinterpret_cast((" << c_ref << ") + (" << c_offset + << ")), " << mma_operand_void_ptr_arg(2) << ", " + << mma_operand_void_ptr_arg(3) << ", " << uint32_ptr_arg(4) + << ", " << uint32_ptr_arg(5) << ", " + << this->PrintExpr(op->args[6]) << ");\n"; + } else if (op->op.same_as(tl::sm120_store_full_c_fragment_panel64_bf16())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // arg 2: C shared output slice pointer + // arg 3: C shared output slice offset + // arg 4: panel id inside the full 128x128 accumulator + ICHECK_EQ(op->args.size(), 5U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + std::string c_shared_ref = this->PrintExpr(op->args[2]); + std::string c_shared_offset = this->PrintExpr(op->args[3]); + + this->PrintIndent(); + this->stream << "tl::detail::sm120_store_full_c_fragment_panel64_bf16(" + << "reinterpret_cast((" << c_ref << ") + (" + << c_offset << ")), " + << "(reinterpret_cast(" << c_shared_ref << ") + (" + << c_shared_offset << ")), " << this->PrintExpr(op->args[4]) + << ");\n"; + } else if (op->op.same_as( + tl::sm120_store_full_c_fragment_panel32_tma_bf16())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // arg 2: C shared output slice pointer + // arg 3: C shared output slice offset + // arg 4: 32-column panel id inside the full 128x128 accumulator + ICHECK_EQ(op->args.size(), 5U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + std::string c_shared_ref = this->PrintExpr(op->args[2]); + std::string c_shared_offset = this->PrintExpr(op->args[3]); + + this->PrintIndent(); + this->stream << "tl::detail::sm120_store_full_c_fragment_panel32_tma_bf16(" + << "reinterpret_cast((" << c_ref << ") + (" + << c_offset << ")), " + << "(reinterpret_cast(" << c_shared_ref << ") + (" + << c_shared_offset << ")), " << this->PrintExpr(op->args[4]) + << ");\n"; + } else if (op->op.same_as( + tl::sm120_store_full_c_fragment_epi64x32_tma_bf16())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // arg 2: C shared output slice pointer + // arg 3: C shared output slice offset + // arg 4: 64-row epilogue tile id inside the full 128x128 accumulator + // arg 5: 32-column epilogue tile id inside the full 128x128 accumulator + ICHECK_EQ(op->args.size(), 6U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + std::string c_shared_ref = this->PrintExpr(op->args[2]); + std::string c_shared_offset = this->PrintExpr(op->args[3]); + + this->PrintIndent(); + this->stream << "tl::detail::sm120_store_full_c_fragment_epi64x32_tma_bf16(" + << "reinterpret_cast((" << c_ref << ") + (" + << c_offset << ")), " + << "(reinterpret_cast(" << c_shared_ref << ") + (" + << c_shared_offset << ")), " << this->PrintExpr(op->args[4]) + << ", " << this->PrintExpr(op->args[5]) << ");\n"; + } else if ( + op->op.same_as( + tl::sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // args 2..5: A ldmatrix source pointers + // args 6..9: B ldmatrix source pointers + // arg 10: SFA shared base pointer + // arg 11: SFB shared base pointer + // arg 12: static K64 index inside the block_K tile + ICHECK_EQ(op->args.size(), 13U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + auto void_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + this->PrintExpr(op->args[idx]) + + ")"; + }; + auto uint32_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + + this->PrintExpr(op->args[idx]) + ")"; + }; + + this->PrintIndent(); + this->stream + << "tl::sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide(" + << "reinterpret_cast((" << c_ref << ") + (" << c_offset + << ")), " << void_ptr_arg(2) << ", " << void_ptr_arg(3) << ", " + << void_ptr_arg(4) << ", " << void_ptr_arg(5) << ", " << void_ptr_arg(6) + << ", " << void_ptr_arg(7) << ", " << void_ptr_arg(8) << ", " + << void_ptr_arg(9) << ", " << uint32_ptr_arg(10) << ", " + << uint32_ptr_arg(11) << ", " << this->PrintExpr(op->args[12]) + << ");\n"; + } else if ( + op->op.same_as( + tl::sm120_mma_blockscaled_kblock_fulltile_package_pingpong())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // arg 2: A shared K-stage base + // arg 3: B shared K-stage base + // arg 4: SFA shared K-stage base + // arg 5: SFB shared K-stage base + ICHECK_EQ(op->args.size(), 6U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + auto void_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + this->PrintExpr(op->args[idx]) + + ")"; + }; + auto uint32_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + + this->PrintExpr(op->args[idx]) + ")"; + }; + + this->PrintIndent(); + this->stream + << "tl::sm120_mma_blockscaled_kblock_fulltile_package_pingpong(" + << "reinterpret_cast((" << c_ref << ") + (" << c_offset + << ")), " << void_ptr_arg(2) << ", " << void_ptr_arg(3) << ", " + << uint32_ptr_arg(4) << ", " << uint32_ptr_arg(5) << ");\n"; + } else if (op->op.same_as( + tl::sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // args 2..5: A ldmatrix source pointers + // args 6..9: B ldmatrix source pointers + // arg 10: SFA shared base pointer + // arg 11: SFB shared base pointer + // arg 12: static K64 index inside the block_K tile + ICHECK_EQ(op->args.size(), 13U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + auto void_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + this->PrintExpr(op->args[idx]) + + ")"; + }; + auto uint32_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + + this->PrintExpr(op->args[idx]) + ")"; + }; + + this->PrintIndent(); + this->stream << "tl::sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide(" + << "reinterpret_cast((" << c_ref << ") + (" << c_offset + << ")), " << void_ptr_arg(2) << ", " << void_ptr_arg(3) << ", " + << void_ptr_arg(4) << ", " << void_ptr_arg(5) << ", " + << void_ptr_arg(6) << ", " << void_ptr_arg(7) << ", " + << void_ptr_arg(8) << ", " << void_ptr_arg(9) << ", " + << uint32_ptr_arg(10) << ", " << uint32_ptr_arg(11) << ", " + << this->PrintExpr(op->args[12]) << ");\n"; + } else if (op->op.same_as(tl::sm120_mma_blockscaled_kblock_fulltile())) { + need_gemm_sm120_h_ = true; + // arg 0: C fragment pointer + // arg 1: C fragment offset + // args 2..5: A ldmatrix source pointers + // args 6..9: B ldmatrix source pointers + // args 10..13: SFA scale-word pointers + // args 14..17: SFB scale-word pointers + // args 18..21: SFB replicate scale-word pointers + // arg 22: static K64 index inside the block_K tile + ICHECK_EQ(op->args.size(), 23U); + std::string c_ref = this->PrintExpr(op->args[0]); + std::string c_offset = this->PrintExpr(op->args[1]); + auto void_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + this->PrintExpr(op->args[idx]) + + ")"; + }; + auto mma_operand_void_ptr_arg = [&](size_t idx) { + if (EnableSM120MmaRawUnpackedFP4AccessPtr()) { + std::optional element_type = + GetAccessPtrElementType(op->args[idx]); + const auto *call = op->args[idx].as(); + if (element_type.has_value() && + element_type.value().is_float4_e2m1_unpacked() && call != nullptr) { + // The SM120 FP4 LDSM helper consumes byte row-start addresses. The + // compact shared path still uses logical FP4 offsets for normal + // scalar accesses, but A/B MMA operand pointers must keep the + // CUTLASS-compatible byte rowstart instead of applying GetBufferRef's + // generic /2 compacting. + if (call->op.same_as(builtin::tvm_access_ptr())) { + ICHECK_GE(call->args.size(), 3U); + return "reinterpret_cast(reinterpret_cast(" + + this->PrintExpr(call->args[1]) + ") + (" + + this->PrintExpr(call->args[2]) + "))"; + } + if (call->op.same_as(tl::access_ptr())) { + ICHECK_EQ(call->args.size(), 3U) << "tl.access_ptr expects 3 args: " + "(BufferLoad, extent, rw_mask)"; + const auto *load = call->args[0].as(); + ICHECK(load) << "tl.access_ptr arg0 must be BufferLoad"; + ICHECK_EQ(load->indices.size(), 1U) + << "SM120 raw FP4 MMA access expects 1D shared operand"; + return "reinterpret_cast(reinterpret_cast(" + + GetVarID(load->buffer->data.get()) + ") + (" + + this->PrintExpr(load->indices[0]) + "))"; + } + if (call->op.same_as(builtin::address_of())) { + ICHECK_EQ(call->args.size(), 1U); + const auto *load = call->args[0].as(); + ICHECK(load) << "address_of arg must be BufferLoad"; + ICHECK_EQ(load->indices.size(), 1U) + << "SM120 raw FP4 MMA address_of expects 1D shared operand"; + return "reinterpret_cast(reinterpret_cast(" + + GetVarID(load->buffer->data.get()) + ") + (" + + this->PrintExpr(load->indices[0]) + "))"; + } + } + } + return void_ptr_arg(idx); + }; + auto uint32_ptr_arg = [&](size_t idx) { + return "reinterpret_cast(" + + this->PrintExpr(op->args[idx]) + ")"; + }; + + this->PrintIndent(); + this->stream << "tl::sm120_mma_blockscaled_kblock_fulltile(" + << "reinterpret_cast((" << c_ref << ") + (" << c_offset + << ")), " << mma_operand_void_ptr_arg(2) << ", " + << mma_operand_void_ptr_arg(3) << ", " + << mma_operand_void_ptr_arg(4) << ", " + << mma_operand_void_ptr_arg(5) << ", " + << mma_operand_void_ptr_arg(6) << ", " + << mma_operand_void_ptr_arg(7) << ", " + << mma_operand_void_ptr_arg(8) << ", " + << mma_operand_void_ptr_arg(9) << ", " << uint32_ptr_arg(10) + << ", " << uint32_ptr_arg(11) << ", " << uint32_ptr_arg(12) + << ", " << uint32_ptr_arg(13) << ", " << uint32_ptr_arg(14) + << ", " << uint32_ptr_arg(15) << ", " << uint32_ptr_arg(16) + << ", " << uint32_ptr_arg(17) << ", " << uint32_ptr_arg(18) + << ", " << uint32_ptr_arg(19) << ", " << uint32_ptr_arg(20) + << ", " << uint32_ptr_arg(21) << ", " + << this->PrintExpr(op->args[22]) << ");\n"; } else if (op->op.same_as(builtin::ptx_mma_sp())) { // arg 0: shape: mXnXkX // arg 1: A layout: row/col @@ -4785,6 +5278,10 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { std::string scope = GetPtrStorageScope(op->buffer->data); const VarNode *buffer = op->buffer->data.as(); DataType alloc_dtype = op->buffer->dtype; + bool is_float4_unpacked_shared = alloc_dtype.is_float4_e2m1_unpacked() && + (scope == "shared" || scope == "shared.dyn"); + bool is_compact_unpacked_fp4_shared = + is_float4_unpacked_shared && EnableSM120CompactUnpackedFP4Shared(); if (scope.find("wmma.") == 0) { if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") { ICHECK( @@ -4809,9 +5306,6 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { } else if (scope == "local.descriptor.tcgen05_instr") { stream << "tl::Tcgen05InstrDescriptor " << vid << ";\n"; } else { - bool is_float4_unpacked_shared = - alloc_dtype.is_float4_e2m1_unpacked() && - (scope == "shared" || scope == "shared.dyn"); bool is_fp4_scalar_local = alloc_dtype.is_float4_e2m1fn() && alloc_dtype.is_scalar() && scope == "local"; bool is_int4_scalar_local = @@ -4820,7 +5314,12 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { if (!is_fp4_scalar_local && !is_int4_scalar_local) { PrintStorageScope(scope, stream); if (is_float4_unpacked_shared) { - stream << "uint8_t"; + if (is_compact_unpacked_fp4_shared) { + enable_fp4_ = true; + stream << "fp4_e2_t"; + } else { + stream << "uint8_t"; + } } else { PrintType(alloc_dtype, stream); } @@ -4839,9 +5338,13 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocBufferNode *op) { if (scope.find("wmma.") == 0) { constant_size = GetWmmaFragmentSize(scope, buffer, constant_size); } - if ((alloc_dtype == DataType::Int(4) || alloc_dtype == DataType::UInt(4)) && + if (((alloc_dtype == DataType::Int(4) || + alloc_dtype == DataType::UInt(4)) || + (alloc_dtype.is_float4_e2m1fn() && alloc_dtype.is_scalar())) && scope == "shared") { constant_size = (constant_size + 1) / 2; + } else if (is_compact_unpacked_fp4_shared) { + constant_size = (constant_size + 1) / 2; } else if (alloc_dtype == DataType::Int(1) && scope == "shared") { constant_size = constant_size / 32; } diff --git a/src/cuda/codegen/codegen_cuda.h b/src/cuda/codegen/codegen_cuda.h index 653af6264..7f8a27353 100644 --- a/src/cuda/codegen/codegen_cuda.h +++ b/src/cuda/codegen/codegen_cuda.h @@ -130,6 +130,8 @@ private: bool need_mma_h_{false}; // whether need tl mma instruction header bool need_mma_instruction_h_{false}; + // whether need tl SM120 block-scaled GEMM helper header + bool need_gemm_sm120_h_{false}; // whether need tl wgmma instruction header bool need_wgmma_instruction_h_{false}; // whether need tl tcgen05mma instruction header diff --git a/src/cuda/op/copy.cc b/src/cuda/op/copy.cc index 48c14438d..aca8e9353 100644 --- a/src/cuda/op/copy.cc +++ b/src/cuda/op/copy.cc @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -45,6 +46,12 @@ PrimExpr MakeTmaLeaderCondition(PrimExpr thread_extent) { return Call(DataType::Bool(), tl_shuffle_elect(), {std::move(thread_extent)}); } +bool EnableSM120CompactUnpackedFP4Shared() { + const char *value = std::getenv("TL_SM120_COMPACT_UNPACKED_FP4_SHARED"); + return value != nullptr && value[0] != '\0' && + !(value[0] == '0' && value[1] == '\0'); +} + int TMAPayloadElementBits(DataType dtype) { // IR elements are 8-bit unpacked SMEM slots, but TMA/mbarrier transactions // count the packed FP4 payload loaded from global memory. @@ -92,6 +99,17 @@ int64_t TMATransactionBytesFromElements(int64_t elements, DataType dtype) { return TMABytesFromElements(elements, TMAPayloadElementBits(dtype)); } +PrimExpr CompactUnpackedFP4SharedOffset(PrimExpr offset, DataType dtype) { + if (!EnableSM120CompactUnpackedFP4Shared() || + !dtype.is_float4_e2m1_unpacked()) { + return offset; + } + // Keep TMA offsets in logical FP4 element units. access_ptr lowers through + // CUDA BufferRef, which applies the compact shared logical-to-physical /2 + // conversion exactly once for unpacked FP4 shared buffers. + return offset; +} + int64_t TMAElementsForBytes(int64_t bytes, DataType dtype) { ICHECK_EQ((bytes * 8) % dtype.bits(), 0) << bytes << " bytes cannot be represented as whole elements of " << dtype; @@ -1608,12 +1626,21 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &lower_args, "to normal copy"; return fallback_to_normal("undecodable shared swizzle layout"); } - // Recast element-space layout into byte-address space. - // Because CuTe swizzle are based on byte addresses. + // Recast the swizzle atom into byte-address space because CuTe/TMA + // swizzles are based on byte addresses. For scalar NVFP4, only the swizzle + // atom can be recast cleanly from 4-bit logical elements; the plain layout + // may contain half-byte contiguous strides and remains in element space for + // offset/coordinate analysis below. int elem_bits = shared_tensor->dtype.bits(); - // E.g., composed_bytes = Sw<3,4,3> o 0 o (64,64,8):(128,2,8192) - auto composed_bytes = composed.value().Recast(elem_bits, /*new_bits=*/8); - const auto *sw = composed_bytes->swizzle.get(); + cute::Swizzle swizzle_bytes; + if (shared_tensor->dtype.is_float4_e2m1fn() && + shared_tensor->dtype.is_scalar()) { + swizzle_bytes = composed.value()->swizzle.Recast(elem_bits, /*new_bits=*/8); + } else { + // E.g., composed_bytes = Sw<3,4,3> o 0 o (64,64,8):(128,2,8192) + swizzle_bytes = composed.value().Recast(elem_bits, /*new_bits=*/8)->swizzle; + } + const auto *sw = swizzle_bytes.get(); int b_bits = sw->b_bits, m_base = sw->m_base, s_shift = sw->s_shift; if (!sw->IsSwizzled()) { desc.swizzle = static_cast(CU_TENSOR_MAP_SWIZZLE_NONE); @@ -2324,8 +2351,10 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &lower_args, } PrimExpr elements = analyzer->Simplify(shared_elements); + PrimExpr tma_shared_offset = + CompactUnpackedFP4SharedOffset(shared_offset, shared_tensor->dtype); PrimExpr shared_addr = shared_tensor.access_ptr( - is_load ? 2 : 1, DataType::Handle(), 1, shared_offset, elements); + is_load ? 2 : 1, DataType::Handle(), 1, tma_shared_offset, elements); PrimExpr global_addr = global_tensor.access_ptr( is_load ? 1 : 2, DataType::Handle(), 1, global_offset, elements); diff --git a/src/op/builtin.cc b/src/op/builtin.cc index 249545a96..b7a887677 100644 --- a/src/op/builtin.cc +++ b/src/op/builtin.cc @@ -61,6 +61,12 @@ DataType CuTensorMapType() { return DataType::UInt(8, 128); } TVM_REGISTER_OP("tl." #OpName) \ .set_attr("TScriptPrinterName", #OpName) +#define TIR_DEFINE_TL_OPAQUE_BUILTIN(OpName, NumInputs) \ + TIR_DEFINE_TL_BUILTIN(OpName) \ + .set_num_inputs(NumInputs) \ + .set_attr("TCallEffectKind", \ + Integer(CallEffectKind::kOpaque)) + // Pointer access metadata op (frontend-only, lowered later). TIR_DEFINE_TL_BUILTIN(access_ptr) .set_num_inputs(3) @@ -250,6 +256,27 @@ TIR_DEFINE_TL_BUILTIN(ptx_wgmma_sp_rs) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_TL_OPAQUE_BUILTIN(ptx_mma_block_scale, 21); + +TIR_DEFINE_TL_OPAQUE_BUILTIN(sm120_mma_blockscaled_kblock_fulltile, 23); + +TIR_DEFINE_TL_OPAQUE_BUILTIN( + sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide, 13); + +TIR_DEFINE_TL_OPAQUE_BUILTIN( + sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide, 13); + +TIR_DEFINE_TL_OPAQUE_BUILTIN( + sm120_mma_blockscaled_kblock_fulltile_package_pingpong, 6); + +TIR_DEFINE_TL_OPAQUE_BUILTIN(sm120_mma_blockscaled_cute_consumer_bridge, 7); + +TIR_DEFINE_TL_OPAQUE_BUILTIN(sm120_store_full_c_fragment_panel64_bf16, 5); + +TIR_DEFINE_TL_OPAQUE_BUILTIN(sm120_store_full_c_fragment_panel32_tma_bf16, 5); + +TIR_DEFINE_TL_OPAQUE_BUILTIN(sm120_store_full_c_fragment_epi64x32_tma_bf16, 6); + TIR_DEFINE_TL_BUILTIN(ptx_tcgen05_mma_ss) .set_num_inputs(14) .set_attr("TCallEffectKind", diff --git a/src/op/builtin.h b/src/op/builtin.h index 562f06c4a..30b233e53 100644 --- a/src/op/builtin.h +++ b/src/op/builtin.h @@ -437,6 +437,55 @@ TVM_DLL const Op &ptx_wgmma_sp_ss(); */ TVM_DLL const Op &ptx_wgmma_sp_rs(); +/*! + * \brief tvm intrinsic for ptx tensor core mma with block scaling on SM120a. + */ +TVM_DLL const Op &ptx_mma_block_scale(); + +/*! + * \brief Private SM120a NVF4 full-tile K-block MMA helper. + */ +TVM_DLL const Op &sm120_mma_blockscaled_kblock_fulltile(); + +/*! + * \brief Private SM120a NVF4 full-tile K-block MMA helper with backend-owned + * scale-owner words. + */ +TVM_DLL const Op &sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide(); + +/*! + * \brief Private SM120a NVF4 full-tile helper with A-full/B-panel schedule and + * backend-owned scale-owner words. + */ +TVM_DLL const Op & +sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide(); + +/*! + * \brief Private SM120a NVF4 full-tile package-lifecycle probe. + */ +TVM_DLL const Op &sm120_mma_blockscaled_kblock_fulltile_package_pingpong(); + +/*! + * \brief Private SM120a NVF4 full shared-tile consumer bridge. + */ +TVM_DLL const Op &sm120_mma_blockscaled_cute_consumer_bridge(); + +/*! + * \brief Private SM120a NVF4 full accumulator fragment panel store helper. + */ +TVM_DLL const Op &sm120_store_full_c_fragment_panel64_bf16(); + +/*! + * \brief Private SM120a NVF4 full accumulator fragment TMA-store panel helper. + */ +TVM_DLL const Op &sm120_store_full_c_fragment_panel32_tma_bf16(); + +/*! + * \brief Private SM120a NVF4 full accumulator fragment 64x32 epilogue TMA-store + * helper. + */ +TVM_DLL const Op &sm120_store_full_c_fragment_epi64x32_tma_bf16(); + /*! * \brief tvm intrinsic for tcgen05 mma shared-shared instructions. */ diff --git a/src/tl_templates/cuda/gemm_sm120.h b/src/tl_templates/cuda/gemm_sm120.h new file mode 100644 index 000000000..cc4764732 --- /dev/null +++ b/src/tl_templates/cuda/gemm_sm120.h @@ -0,0 +1,1758 @@ +#pragma once + +#include "common.h" +#include "ldsm.h" + +#ifndef __CUDACC_RTC__ +#include +#endif + +namespace tl { + +enum class SM120MmaBlockScaledKind : int { + kMxf4nvf4 = 0, +}; + +enum class SM120MmaScaleType : int { + kUE4M3 = 0, +}; + +template +struct SM120MmaBlockScaledConfig { + static constexpr bool kSupported = false; +}; + +template <> +struct SM120MmaBlockScaledConfig { + static constexpr bool kSupported = true; +}; + +namespace detail { + +// CUTLASS BlockScaledBasicChunk K-major scale layout, compressed to uint32 +// words. The scale-byte offset for one 128x64 atom is: +// +// (idx % 32) * 16 + (idx / 32) * 4 + k16 +// +// where k16 is the scale byte inside the four K/16 groups. TileLang packs those +// four adjacent scale bytes in one uint32, so this helper returns the flattened +// uint32 word offset: +// +// ki * 128 + (idx % 32) * 4 + (idx / 32) +// +// This is the byte offset divided by four, not the byte-offset formula with +// only the final +k16 term removed. +TL_DEVICE uint32_t sm120_blockscaled_chunk_kmajor_sf_word(uint32_t idx, + uint32_t ki) { + return ki * 128u + (idx & 31u) * 4u + (idx >> 5); +} + +// First-class scale package for the SM120 blockscaled TV copy-view work. The +// current implementation intentionally preserves the compact-selector semantic +// rows used by production. The important contract boundary is that scale +// loading is now isolated from the A/B operand package, so future +// get_layoutSFA_TV / get_layoutSFB_TV lowering can change the producer-lane +// copy view without rewriting the OMMA.SF issue loop. +struct SM120ScaleTVPackage { + uint32_t sa0, sa1; + uint32_t sb0, sb1; +}; + +template +TL_DEVICE void sm120_load_scale_tv_package(ScalePkg &pkg, + const uint32_t *sfa_base, + const uint32_t *sfb_base, + uint32_t lane, uint32_t warp_m, + uint32_t warp_n, uint32_t scale_k) { + uint32_t const qlane = lane & 3u; + uint32_t const sfa_row = 8u * (lane & 1u) + (lane >> 2); + uint32_t const sfb_col = lane >> 2; + uint32_t const a_owner_in_pair = qlane >> 1; + uint32_t const scale_m0 = warp_m * 64u + a_owner_in_pair * 16u + sfa_row; + uint32_t const scale_n0 = warp_n * 64u + qlane * 8u + sfb_col; + pkg.sa0 = sfa_base[sm120_blockscaled_chunk_kmajor_sf_word(scale_m0, scale_k)]; + pkg.sa1 = + sfa_base[sm120_blockscaled_chunk_kmajor_sf_word(scale_m0 + 32u, scale_k)]; + pkg.sb0 = sfb_base[sm120_blockscaled_chunk_kmajor_sf_word(scale_n0, scale_k)]; + pkg.sb1 = + sfb_base[sm120_blockscaled_chunk_kmajor_sf_word(scale_n0 + 32u, scale_k)]; +} + +TL_DEVICE void sm120_copy_scale_tv_package(SM120ScaleTVPackage &pkg, + const uint32_t *sfa_base, + const uint32_t *sfb_base, + int k_block_idx) { + uint32_t const tx = uint32_t(int(threadIdx.x) & 127); + uint32_t const lane = tx & 31u; + uint32_t const warp = tx >> 5; + uint32_t const warp_m = warp & 1u; + uint32_t const warp_n = warp >> 1; + sm120_load_scale_tv_package(pkg, sfa_base, sfb_base, lane, warp_m, warp_n, + uint32_t(k_block_idx)); +} + +// SM120 blockscaled fulltile kernels precompute some shared-memory addresses +// and keep operand fragments in named scalar registers. The generic +// ptx_ldmatrix_x4(void*, void*) helper takes a C++ pointer and a contiguous +// local buffer, so keep this address-register/scalar-register form local to the +// SM120 operand-package helpers. +TL_DEVICE void sm120_ldmatrix_x4_u32_addr(uint32_t smem_int_ptr, uint32_t &d0, + uint32_t &d1, uint32_t &d2, + uint32_t &d3) { + asm volatile( + "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(smem_int_ptr)); +} + +TL_DEVICE void sm120_ldmatrix_x4_u32(void const *const smem_ptr, uint32_t &d0, + uint32_t &d1, uint32_t &d2, uint32_t &d3) { + sm120_ldmatrix_x4_u32_addr(smem_ptr_to_uint(smem_ptr), d0, d1, d2, d3); +} + +// This copy atom consumes a padded/unpacked 4-bit shared layout. It matches the +// F8F6F4/MXF-style path where FP4 nibbles are placed in 8-bit slots with p64 +// padding. Dense SM120 NVFP4/mxf4nvf4 uses packed U4 shared storage and the +// ordinary x4.m8n8.shared.b16 ldmatrix path above. +TL_DEVICE void sm120_ldmatrix_x4_fp4_u32_addr(uint32_t smem_int_ptr, + uint32_t &d0, uint32_t &d1, + uint32_t &d2, uint32_t &d3) { + asm volatile("ldmatrix.sync.aligned.m8n16.x4.shared.b8x16.b4x16_p64 " + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(smem_int_ptr)); +#if !defined(TL_SM120_FULLTILE_FP4_LDSM_NO_SHIFT) + d0 <<= 2; + d1 <<= 2; + d2 <<= 2; + d3 <<= 2; +#endif +} + +TL_DEVICE void sm120_ldmatrix_x4_fp4_u32(void const *const smem_ptr, + uint32_t &d0, uint32_t &d1, + uint32_t &d2, uint32_t &d3) { + sm120_ldmatrix_x4_fp4_u32_addr(smem_ptr_to_uint(smem_ptr), d0, d1, d2, d3); +} + +TL_DEVICE void sm120_ldmatrix_x4_blockscaled_operand_addr(uint32_t smem_int_ptr, + uint32_t &d0, + uint32_t &d1, + uint32_t &d2, + uint32_t &d3) { +#if defined(TL_SM120_FULLTILE_FP4_LDSM) + sm120_ldmatrix_x4_fp4_u32_addr(smem_int_ptr, d0, d1, d2, d3); +#else + sm120_ldmatrix_x4_u32_addr(smem_int_ptr, d0, d1, d2, d3); +#if defined(TL_SM120_FULLTILE_FP4_REG_SHIFT) + d0 <<= 2; + d1 <<= 2; + d2 <<= 2; + d3 <<= 2; +#endif +#endif +} + +TL_DEVICE void sm120_ldmatrix_x4_blockscaled_operand(void const *const smem_ptr, + uint32_t &d0, uint32_t &d1, + uint32_t &d2, + uint32_t &d3) { +#if defined(TL_SM120_FULLTILE_FP4_LDSM) + sm120_ldmatrix_x4_fp4_u32(smem_ptr, d0, d1, d2, d3); +#else + sm120_ldmatrix_x4_u32(smem_ptr, d0, d1, d2, d3); +#if defined(TL_SM120_FULLTILE_FP4_REG_SHIFT) + d0 <<= 2; + d1 <<= 2; + d2 <<= 2; + d3 <<= 2; +#endif +#endif +} + +TL_DEVICE uint32_t sm120_cute_accum_pair_offset_from_tl_pair(uint32_t tl_pair) { + uint32_t const i = tl_pair >> 5; + uint32_t const rem_i = tl_pair & 31u; + uint32_t const j = rem_i >> 3; + uint32_t const rem_j = rem_i & 7u; + uint32_t const half = rem_j >> 2; + uint32_t const reg_pair = rem_j & 3u; + return i * 4u + j * 32u + half * 16u + reg_pair; +} + +TL_DEVICE void +sm120_store_full_c_fragment_panel64_bf16(float const *const c_fragment, + bfloat16_t *const c_shared, + int const panel) { + uint32_t const tx = uint32_t(threadIdx.x) & 127u; + if (int(tx >> 6) != panel) { + return; + } + +#pragma unroll + for (int i = 0; i < 64; ++i) { + uint32_t const tl_pair = uint32_t(i) * 2u; +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) + uint32_t const src_pair = + sm120_cute_accum_pair_offset_from_tl_pair(tl_pair); +#else + uint32_t const src_pair = tl_pair; +#endif + float2 const value = + *reinterpret_cast(c_fragment + src_pair); + uint1 packed; + reinterpret_cast<__nv_bfloat162 *>(&packed)[0] = + __float22bfloat162_rn(value); + uint32_t const row = ((tx & 63u) >> 5) * 64u + uint32_t(i >> 4) * 16u + + uint32_t(i & 1) * 8u + ((tx & 31u) >> 2); + uint32_t const col = uint32_t((i & 15) >> 1) * 8u + (tx & 3u) * 2u; + *reinterpret_cast(c_shared + row * 64u + col) = packed; + } +} + +TL_DEVICE uint32_t sm120_tma_store_128x32_bf16_swizzle_offset(uint32_t row, + uint32_t col8) { + uint32_t const tx = ((row & 31u) << 2) | (col8 >> 3); + return ((row >> 5) * 1024u) + ((tx >> 2) * 32u) + + (((((tx & 31u) >> 4) + ((tx & 3u) >> 1)) & 1u) * 16u) + + (((((tx & 15u) >> 3) + (tx & 1u)) & 1u) * 8u); +} + +TL_DEVICE void +sm120_store_full_c_fragment_panel32_tma_bf16(float const *const c_fragment, + bfloat16_t *const c_shared, + int const panel32) { + uint32_t const tx = uint32_t(threadIdx.x) & 127u; + uint32_t const panel64 = uint32_t(panel32) >> 1; + uint32_t const sub32 = uint32_t(panel32) & 1u; + if ((tx >> 6) != panel64) { + return; + } + +#pragma unroll + for (int i = 0; i < 64; ++i) { + uint32_t const col64 = uint32_t((i & 15) >> 1) * 8u + (tx & 3u) * 2u; + if ((col64 >> 5) != sub32) { + continue; + } + + uint32_t const tl_pair = uint32_t(i) * 2u; +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) + uint32_t const src_pair = + sm120_cute_accum_pair_offset_from_tl_pair(tl_pair); +#else + uint32_t const src_pair = tl_pair; +#endif + float2 const value = + *reinterpret_cast(c_fragment + src_pair); + uint1 packed; + reinterpret_cast<__nv_bfloat162 *>(&packed)[0] = + __float22bfloat162_rn(value); + + uint32_t const row = ((tx & 63u) >> 5) * 64u + uint32_t(i >> 4) * 16u + + uint32_t(i & 1) * 8u + ((tx & 31u) >> 2); + uint32_t const col32 = col64 & 31u; + uint32_t const swizzled = + sm120_tma_store_128x32_bf16_swizzle_offset(row, col32 & ~7u) + + (col32 & 7u); + *reinterpret_cast(c_shared + swizzled) = packed; + } +} + +TL_DEVICE void sm120_store_full_c_fragment_epi64x32_tma_bf16( + float const *const c_fragment, bfloat16_t *const c_shared, int const row64, + int const col32_panel) { + uint32_t const tx = uint32_t(threadIdx.x) & 127u; + uint32_t const col64_panel = uint32_t(col32_panel) >> 1; + uint32_t const sub32 = uint32_t(col32_panel) & 1u; + if ((tx >> 6) != col64_panel) { + return; + } + +#pragma unroll + for (int i = 0; i < 64; ++i) { + uint32_t const col64 = uint32_t((i & 15) >> 1) * 8u + (tx & 3u) * 2u; + if ((col64 >> 5) != sub32) { + continue; + } + + uint32_t const row = ((tx & 63u) >> 5) * 64u + uint32_t(i >> 4) * 16u + + uint32_t(i & 1) * 8u + ((tx & 31u) >> 2); + if ((row >> 6) != uint32_t(row64)) { + continue; + } + + uint32_t const tl_pair = uint32_t(i) * 2u; +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) + uint32_t const src_pair = + sm120_cute_accum_pair_offset_from_tl_pair(tl_pair); +#else + uint32_t const src_pair = tl_pair; +#endif + float2 const value = + *reinterpret_cast(c_fragment + src_pair); + uint1 packed; + reinterpret_cast<__nv_bfloat162 *>(&packed)[0] = + __float22bfloat162_rn(value); + + uint32_t const local_row = row & 63u; + uint32_t const col32 = col64 & 31u; + uint32_t const swizzled = + sm120_tma_store_128x32_bf16_swizzle_offset(local_row, col32 & ~7u) + + (col32 & 7u); + *reinterpret_cast(c_shared + swizzled) = packed; + } +} + +// SM120a NVF4 block-scaled warp MMA: +// mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X +TL_DEVICE void sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( + float *d, uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, + uint32_t b1, const float *c, uint32_t scale_a, uint32_t scale_b, + uint16_t scale_a_byte_id = 0, uint16_t scale_a_thread_id = 0, + uint16_t scale_b_byte_id = 0, uint16_t scale_b_thread_id = 0) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::" + "4X.f32.e2m1.e2m1.f32.ue4m3 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13}, " + "{%14}, {%15, %16}, " + "{%17}, {%18, %19};\n" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(c[0]), + "f"(c[1]), "f"(c[2]), "f"(c[3]), "r"(scale_a), "h"(scale_a_byte_id), + "h"(scale_a_thread_id), "r"(scale_b), "h"(scale_b_byte_id), + "h"(scale_b_thread_id)); +} + +TL_DEVICE void sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_accum_regs( + float *d, uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, + uint32_t b1, uint32_t scale_a, uint32_t scale_b, + uint16_t scale_a_byte_id = 0, uint16_t scale_a_thread_id = 0, + uint16_t scale_b_byte_id = 0, uint16_t scale_b_thread_id = 0) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::" + "4X.f32.e2m1.e2m1.f32.ue4m3 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%0, %1, %2, %3}, " + "{%10}, {%11, %12}, " + "{%13}, {%14, %15};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(scale_a), + "h"(scale_a_byte_id), "h"(scale_a_thread_id), "r"(scale_b), + "h"(scale_b_byte_id), "h"(scale_b_thread_id)); +} + +TL_DEVICE void sm120_mma2_m16n8k64_mxf4nvf4_4x_ue4m3_same_b_regs( + float *d0, uint32_t a00, uint32_t a01, uint32_t a02, uint32_t a03, + const float *c0, uint32_t scale_a0, float *d1, uint32_t a10, uint32_t a11, + uint32_t a12, uint32_t a13, const float *c1, uint32_t scale_a1, uint32_t b0, + uint32_t b1, uint32_t scale_b) { + uint16_t const zero = 0; + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::" + "4X.f32.e2m1.e2m1.f32.ue4m3 " + "{%0, %1, %2, %3}, " + "{%8, %9, %10, %11}, " + "{%16, %17}, " + "{%18, %19, %20, %21}, " + "{%26}, {%29, %29}, " + "{%28}, {%29, %29};\n" + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::" + "4X.f32.e2m1.e2m1.f32.ue4m3 " + "{%4, %5, %6, %7}, " + "{%12, %13, %14, %15}, " + "{%16, %17}, " + "{%22, %23, %24, %25}, " + "{%27}, {%29, %29}, " + "{%28}, {%29, %29};\n" + : "=&f"(d0[0]), "=&f"(d0[1]), "=&f"(d0[2]), "=&f"(d0[3]), "=&f"(d1[0]), + "=&f"(d1[1]), "=&f"(d1[2]), "=&f"(d1[3]) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), "r"(a10), "r"(a11), "r"(a12), + "r"(a13), "r"(b0), "r"(b1), "f"(c0[0]), "f"(c0[1]), "f"(c0[2]), + "f"(c0[3]), "f"(c1[0]), "f"(c1[1]), "f"(c1[2]), "f"(c1[3]), + "r"(scale_a0), "r"(scale_a1), "r"(scale_b), "h"(zero)); +} + +TL_DEVICE void sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3( + float *d, const uint32_t *a, const uint32_t *b, const float *c, + uint32_t scale_a, uint32_t scale_b, uint16_t scale_a_byte_id = 0, + uint16_t scale_a_thread_id = 0, uint16_t scale_b_byte_id = 0, + uint16_t scale_b_thread_id = 0) { + sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( + d, a[0], a[1], a[2], a[3], b[0], b[1], c, scale_a, scale_b, + scale_a_byte_id, scale_a_thread_id, scale_b_byte_id, scale_b_thread_id); +} + +TL_DEVICE uint32_t sm120_fulltile_compact_k_swizzle_offset(uint32_t tx, + int k_block_idx) { + switch (k_block_idx & 3) { + case 0: + return ((tx & 7u) >> 1) * 32u; + case 1: + return ((tx & 7u) >> 2) * 64u + (((((tx & 3u) >> 1) + 1u) & 1u) * 32u); + case 2: + return (((((tx & 7u) >> 2) + 1u) & 1u) * 64u) + (((tx & 3u) >> 1) * 32u); + default: + return (((((tx & 7u) >> 2) + 1u) & 1u) * 64u) + + (((((tx & 3u) >> 1) + 1u) & 1u) * 32u); + } +} + +TL_DEVICE uint32_t sm120_fulltile_compact_a_offset(uint32_t tx, int k_block_idx, + int row_idx) { + return (((tx & 63u) >> 5) * 2048u) + ((tx & 15u) * 128u) + + sm120_fulltile_compact_k_swizzle_offset(tx, k_block_idx) + + (((((tx & 31u) >> 4) + (tx & 1u)) & 1u) * 16u) + + uint32_t(row_idx) * 4096u; +} + +TL_DEVICE uint32_t sm120_fulltile_compact_b_offset(uint32_t tx, int k_block_idx, + int panel_idx) { + return ((tx >> 6) * 2048u) + (((tx & 31u) >> 4) * 1024u) + + ((tx & 7u) * 128u) + + sm120_fulltile_compact_k_swizzle_offset(tx, k_block_idx) + + (((((tx & 15u) >> 3) + (tx & 1u)) & 1u) * 16u) + + uint32_t(panel_idx) * 4096u; +} + +TL_DEVICE uint32_t sm120_fulltile_rowmajor_a_offset(uint32_t tx, + int k_block_idx, + int row_idx) { + return (((tx & 63u) >> 5) * 8192u) + ((tx & 15u) * 128u) + + sm120_fulltile_compact_k_swizzle_offset(tx, k_block_idx) + + (((((tx & 31u) >> 4) + (tx & 1u)) & 1u) * 16u) + + uint32_t(row_idx) * 2048u; +} + +TL_DEVICE uint32_t sm120_fulltile_rowmajor_b_offset(uint32_t tx, + int k_block_idx, + int panel_idx) { + return ((tx >> 6) * 8192u) + (((tx & 31u) >> 4) * 1024u) + + ((tx & 7u) * 128u) + + sm120_fulltile_compact_k_swizzle_offset(tx, k_block_idx) + + (((((tx & 15u) >> 3) + (tx & 1u)) & 1u) * 16u) + + uint32_t(panel_idx) * 2048u; +} + +TL_DEVICE uint32_t sm120_fulltile_package_a_offset(uint32_t tx, int k_block_idx, + int row_idx) { +#if defined(TL_SM120_FULLTILE_PACKAGE_ROWMAJOR_VIEW) + return sm120_fulltile_rowmajor_a_offset(tx, k_block_idx, row_idx); +#else + return sm120_fulltile_compact_a_offset(tx, k_block_idx, row_idx); +#endif +} + +TL_DEVICE uint32_t sm120_fulltile_package_b_offset(uint32_t tx, int k_block_idx, + int panel_idx) { +#if defined(TL_SM120_FULLTILE_PACKAGE_ROWMAJOR_VIEW) + return sm120_fulltile_rowmajor_b_offset(tx, k_block_idx, panel_idx); +#else + return sm120_fulltile_compact_b_offset(tx, k_block_idx, panel_idx); +#endif +} + +} // namespace detail + +template +TL_DEVICE void sm120_mma_sync_blockscaled(float *d, const uint32_t *a, + const uint32_t *b, const float *c, + uint32_t scale_a, uint32_t scale_b, + uint16_t scale_a_byte_id = 0, + uint16_t scale_a_thread_id = 0, + uint16_t scale_b_byte_id = 0, + uint16_t scale_b_thread_id = 0) { + static_assert(Kind == SM120MmaBlockScaledKind::kMxf4nvf4, + "Only kind::mxf4nvf4 is supported"); + static_assert(ScaleVecSize == 4, "Only scale_vec::4X is supported"); + static_assert(SType == SM120MmaScaleType::kUE4M3, + "kind::mxf4nvf4 only supports ue4m3 scale factors"); + static_assert( + SM120MmaBlockScaledConfig::kSupported, + "Unsupported sm120 mma.block_scale configuration"); + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3( + d, a, b, c, scale_a, scale_b, scale_a_byte_id, scale_a_thread_id, + scale_b_byte_id, scale_b_thread_id); +} + +TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile( + float *c, const void *a_smem_0, const void *a_smem_1, const void *a_smem_2, + const void *a_smem_3, const void *b_smem_0, const void *b_smem_1, + const void *b_smem_2, const void *b_smem_3, const uint32_t *sfa_0, + const uint32_t *sfa_1, const uint32_t *sfa_2, const uint32_t *sfa_3, + const uint32_t *sfb_0, const uint32_t *sfb_1, const uint32_t *sfb_2, + const uint32_t *sfb_3, const uint32_t *sfb_rep_0, const uint32_t *sfb_rep_1, + const uint32_t *sfb_rep_2, const uint32_t *sfb_rep_3, int k_block_idx); + +TL_DEVICE void sm120_mma_blockscaled_cute_consumer_bridge( + float *c, const void *a_smem_base, const void *b_smem_base, + const uint32_t *sfa_smem_base, const uint32_t *sfb_smem_base, + int k_block_idx) { +#if defined(TL_SM120_FULLTILE_CUTLASS_SF_BASEPTR) + uint32_t const tx = uint32_t(int(threadIdx.x) & 127); + char const *const a_base = static_cast(a_smem_base); + char const *const b_base = static_cast(b_smem_base); + sm120_mma_blockscaled_kblock_fulltile( + c, a_base + detail::sm120_fulltile_compact_a_offset(tx, k_block_idx, 0), + a_base + detail::sm120_fulltile_compact_a_offset(tx, k_block_idx, 1), + a_base + detail::sm120_fulltile_compact_a_offset(tx, k_block_idx, 2), + a_base + detail::sm120_fulltile_compact_a_offset(tx, k_block_idx, 3), + b_base + detail::sm120_fulltile_compact_b_offset(tx, k_block_idx, 0), + b_base + detail::sm120_fulltile_compact_b_offset(tx, k_block_idx, 1), + b_base + detail::sm120_fulltile_compact_b_offset(tx, k_block_idx, 2), + b_base + detail::sm120_fulltile_compact_b_offset(tx, k_block_idx, 3), + sfa_smem_base, sfa_smem_base, sfa_smem_base, sfa_smem_base, sfb_smem_base, + sfb_smem_base, sfb_smem_base, sfb_smem_base, sfb_smem_base, sfb_smem_base, + sfb_smem_base, sfb_smem_base, k_block_idx); +#else + asm volatile("" ::"l"(c), "l"(a_smem_base), "l"(b_smem_base), + "l"(sfa_smem_base), "l"(sfb_smem_base), "r"(k_block_idx) + : "memory"); +#endif +} + +TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile( + float *c, const void *a_smem_0, const void *a_smem_1, const void *a_smem_2, + const void *a_smem_3, const void *b_smem_0, const void *b_smem_1, + const void *b_smem_2, const void *b_smem_3, const uint32_t *sfa_0, + const uint32_t *sfa_1, const uint32_t *sfa_2, const uint32_t *sfa_3, + const uint32_t *sfb_0, const uint32_t *sfb_1, const uint32_t *sfb_2, + const uint32_t *sfb_3, const uint32_t *sfb_rep_0, const uint32_t *sfb_rep_1, + const uint32_t *sfb_rep_2, const uint32_t *sfb_rep_3, int k_block_idx) { + uint32_t b00, b01, b02, b03; + uint32_t b10, b11, b12, b13; + uint32_t b20, b21, b22, b23; + uint32_t b30, b31, b32, b33; + +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) + int const warp128 = (int(threadIdx.x) & 127) >> 5; + int const m_group = warp128 & 1; + int const n_group = warp128 >> 1; + constexpr int kCuteRowstartBackstep = +#if defined(TL_SM120_FULLTILE_RAW_UNPACKED_FP4_ACCESS_PTR) + 6144; +#elif defined(TL_SM120_FULLTILE_COMPACT_UNPACKED_FP4_SHARED) + 3072; +#else + 6144; +#endif + constexpr int kCuteKAtomStride = +#if defined(TL_SM120_FULLTILE_RAW_UNPACKED_FP4_ACCESS_PTR) + 4096; +#elif defined(TL_SM120_FULLTILE_COMPACT_UNPACKED_FP4_SHARED) + 2048; +#else + 4096; +#endif + char const *const a_cute_m0 = + static_cast(a_smem_0) - kCuteRowstartBackstep * m_group; + char const *const b_cute_n0 = + static_cast(b_smem_0) - kCuteRowstartBackstep * n_group; +#endif + +#if defined(TL_SM120_FULLTILE_SINGLE_OMMA) +#if !defined(TL_SM120_FULLTILE_SINGLE_KI) +#define TL_SM120_FULLTILE_SINGLE_KI 0 +#endif +#if !defined(TL_SM120_FULLTILE_SINGLE_I) +#define TL_SM120_FULLTILE_SINGLE_I 0 +#endif +#if !defined(TL_SM120_FULLTILE_SINGLE_J) +#define TL_SM120_FULLTILE_SINGLE_J 0 +#endif +#if !defined(TL_SM120_FULLTILE_SINGLE_HALF) +#define TL_SM120_FULLTILE_SINGLE_HALF 0 +#endif + if (k_block_idx != TL_SM120_FULLTILE_SINGLE_KI) { + return; + } + { + uint32_t single_a0, single_a1, single_a2, single_a3; + uint32_t single_b0, single_b1, single_b2, single_b3; +#if defined(TL_SM120_FULLTILE_SINGLE_SCRATCH_ONES) +#if TL_SM120_FULLTILE_SINGLE_I != 0 || TL_SM120_FULLTILE_SINGLE_J != 0 +#error \ + "TL_SM120_FULLTILE_SINGLE_SCRATCH_ONES currently supports site I=0,J=0 only" +#endif +#if !defined(TL_SM120_FULLTILE_SINGLE_SCRATCH_BYTE) +#define TL_SM120_FULLTILE_SINGLE_SCRATCH_BYTE 0x22 +#endif + __shared__ __align__(128) unsigned char tl_sm120_single_a_scratch[4 * 512]; + __shared__ __align__(128) unsigned char tl_sm120_single_b_scratch[4 * 512]; + int const single_warp = (int(threadIdx.x) & 127) >> 5; + int const lane = int(threadIdx.x) & 31; + unsigned char *const a_scratch = + tl_sm120_single_a_scratch + single_warp * 512; + unsigned char *const b_scratch = + tl_sm120_single_b_scratch + single_warp * 512; +#pragma unroll + for (int scratch_i = 0; scratch_i < 16; ++scratch_i) { + int const byte_i = lane + scratch_i * 32; + a_scratch[byte_i] = TL_SM120_FULLTILE_SINGLE_SCRATCH_BYTE; + b_scratch[byte_i] = TL_SM120_FULLTILE_SINGLE_SCRATCH_BYTE; + } + __syncwarp(); + detail::sm120_ldmatrix_x4_fp4_u32(a_scratch + lane * 16, single_a0, + single_a1, single_a2, single_a3); + detail::sm120_ldmatrix_x4_fp4_u32(b_scratch + lane * 16, single_b0, + single_b1, single_b2, single_b3); +#elif defined(TL_SM120_FULLTILE_SINGLE_REAL_SCRATCH) +#if TL_SM120_FULLTILE_SINGLE_I != 0 || TL_SM120_FULLTILE_SINGLE_J != 0 +#error \ + "TL_SM120_FULLTILE_SINGLE_REAL_SCRATCH currently supports site I=0,J=0 only" +#endif + __shared__ __align__( + 128) unsigned char tl_sm120_single_a_real_scratch[4 * 512]; + __shared__ __align__( + 128) unsigned char tl_sm120_single_b_real_scratch[4 * 512]; + int const single_warp = (int(threadIdx.x) & 127) >> 5; + int const lane = int(threadIdx.x) & 31; + unsigned char *const a_scratch = + tl_sm120_single_a_real_scratch + single_warp * 512; + unsigned char *const b_scratch = + tl_sm120_single_b_real_scratch + single_warp * 512; + unsigned char const *const a_src = + static_cast(a_smem_0); + unsigned char const *const b_src = + static_cast(b_smem_0); +#pragma unroll + for (int byte_i = 0; byte_i < 8; ++byte_i) { + a_scratch[lane * 16 + byte_i] = a_src[byte_i]; + b_scratch[lane * 16 + byte_i] = b_src[byte_i]; + a_scratch[lane * 16 + byte_i + 8] = 0; + b_scratch[lane * 16 + byte_i + 8] = 0; + } + __syncwarp(); + detail::sm120_ldmatrix_x4_fp4_u32(a_scratch + lane * 16, single_a0, + single_a1, single_a2, single_a3); + detail::sm120_ldmatrix_x4_fp4_u32(b_scratch + lane * 16, single_b0, + single_b1, single_b2, single_b3); +#else +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) + detail::sm120_ldmatrix_x4_fp4_u32( + a_cute_m0 + kCuteKAtomStride * TL_SM120_FULLTILE_SINGLE_I, single_a0, + single_a1, single_a2, single_a3); +#else +#if TL_SM120_FULLTILE_SINGLE_I == 0 + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_smem_0, single_a0, single_a1, single_a2, single_a3); +#elif TL_SM120_FULLTILE_SINGLE_I == 1 + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_smem_1, single_a0, single_a1, single_a2, single_a3); +#elif TL_SM120_FULLTILE_SINGLE_I == 2 + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_smem_2, single_a0, single_a1, single_a2, single_a3); +#elif TL_SM120_FULLTILE_SINGLE_I == 3 + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_smem_3, single_a0, single_a1, single_a2, single_a3); +#else +#error "TL_SM120_FULLTILE_SINGLE_I must be in [0, 3]" +#endif +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) + detail::sm120_ldmatrix_x4_fp4_u32( + b_cute_n0 + kCuteKAtomStride * TL_SM120_FULLTILE_SINGLE_J, single_b0, + single_b1, single_b2, single_b3); +#else +#if TL_SM120_FULLTILE_SINGLE_J == 0 + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_smem_0, single_b0, single_b1, single_b2, single_b3); +#elif TL_SM120_FULLTILE_SINGLE_J == 1 + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_smem_1, single_b0, single_b1, single_b2, single_b3); +#elif TL_SM120_FULLTILE_SINGLE_J == 2 + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_smem_2, single_b0, single_b1, single_b2, single_b3); +#elif TL_SM120_FULLTILE_SINGLE_J == 3 + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_smem_3, single_b0, single_b1, single_b2, single_b3); +#else +#error "TL_SM120_FULLTILE_SINGLE_J must be in [0, 3]" +#endif +#endif +#endif + +#if TL_SM120_FULLTILE_SINGLE_I == 0 + uint32_t const single_sa = *sfa_0; +#elif TL_SM120_FULLTILE_SINGLE_I == 1 + uint32_t const single_sa = *sfa_1; +#elif TL_SM120_FULLTILE_SINGLE_I == 2 + uint32_t const single_sa = *sfa_2; +#elif TL_SM120_FULLTILE_SINGLE_I == 3 + uint32_t const single_sa = *sfa_3; +#else +#error "TL_SM120_FULLTILE_SINGLE_I must be in [0, 3]" +#endif + +#if TL_SM120_FULLTILE_SINGLE_J == 0 +#if TL_SM120_FULLTILE_SINGLE_HALF == 0 + uint32_t const single_sb = *sfb_0; + uint32_t const single_b_lo = single_b0; + uint32_t const single_b_hi = single_b1; +#else + uint32_t const single_sb = *sfb_rep_0; + uint32_t const single_b_lo = single_b2; + uint32_t const single_b_hi = single_b3; +#endif +#elif TL_SM120_FULLTILE_SINGLE_J == 1 +#if TL_SM120_FULLTILE_SINGLE_HALF == 0 + uint32_t const single_sb = *sfb_1; + uint32_t const single_b_lo = single_b0; + uint32_t const single_b_hi = single_b1; +#else + uint32_t const single_sb = *sfb_rep_1; + uint32_t const single_b_lo = single_b2; + uint32_t const single_b_hi = single_b3; +#endif +#elif TL_SM120_FULLTILE_SINGLE_J == 2 +#if TL_SM120_FULLTILE_SINGLE_HALF == 0 + uint32_t const single_sb = *sfb_2; + uint32_t const single_b_lo = single_b0; + uint32_t const single_b_hi = single_b1; +#else + uint32_t const single_sb = *sfb_rep_2; + uint32_t const single_b_lo = single_b2; + uint32_t const single_b_hi = single_b3; +#endif +#elif TL_SM120_FULLTILE_SINGLE_J == 3 +#if TL_SM120_FULLTILE_SINGLE_HALF == 0 + uint32_t const single_sb = *sfb_3; + uint32_t const single_b_lo = single_b0; + uint32_t const single_b_hi = single_b1; +#else + uint32_t const single_sb = *sfb_rep_3; + uint32_t const single_b_lo = single_b2; + uint32_t const single_b_hi = single_b3; +#endif +#else +#error "TL_SM120_FULLTILE_SINGLE_J must be in [0, 3]" +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) + int const single_c_offset = TL_SM120_FULLTILE_SINGLE_I * 4 + + TL_SM120_FULLTILE_SINGLE_J * 32 + + TL_SM120_FULLTILE_SINGLE_HALF * 16; +#else + int const single_c_offset = TL_SM120_FULLTILE_SINGLE_I * 32 + + TL_SM120_FULLTILE_SINGLE_J * 8 + + TL_SM120_FULLTILE_SINGLE_HALF * 4; +#endif + float *d = c + single_c_offset; +#if defined(TL_SM120_FULLTILE_REG_DEBUG) +#if defined(TL_SM120_FULLTILE_REG_TAG_DEBUG) + { + uint32_t const debug_warp = (uint32_t(int(threadIdx.x) & 127) >> 5); + uint32_t const debug_lane = uint32_t(int(threadIdx.x) & 31); + uint32_t const debug_base = + 0xa5000000u | (debug_warp << 16) | (debug_lane << 8); +#define TL_SM120_REG_DEBUG_TAG(SLOT) \ + d[SLOT] = __uint_as_float(debug_base | uint32_t(SLOT)) + TL_SM120_REG_DEBUG_TAG(0); + TL_SM120_REG_DEBUG_TAG(1); + TL_SM120_REG_DEBUG_TAG(2); + TL_SM120_REG_DEBUG_TAG(3); + TL_SM120_REG_DEBUG_TAG(4); + TL_SM120_REG_DEBUG_TAG(5); + TL_SM120_REG_DEBUG_TAG(6); + TL_SM120_REG_DEBUG_TAG(7); + TL_SM120_REG_DEBUG_TAG(8); + TL_SM120_REG_DEBUG_TAG(9); + TL_SM120_REG_DEBUG_TAG(10); + TL_SM120_REG_DEBUG_TAG(11); +#undef TL_SM120_REG_DEBUG_TAG + } +#else + d[0] = __uint_as_float(single_a0); + d[1] = __uint_as_float(single_a1); + d[2] = __uint_as_float(single_a2); + d[3] = __uint_as_float(single_a3); + d[4] = __uint_as_float(single_b0); + d[5] = __uint_as_float(single_b1); + d[6] = __uint_as_float(single_b2); + d[7] = __uint_as_float(single_b3); + d[8] = __uint_as_float(single_b_lo); + d[9] = __uint_as_float(single_b_hi); + d[10] = __uint_as_float(single_sa); + d[11] = __uint_as_float(single_sb); +#endif + return; +#endif + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( + d, single_a0, single_a1, single_a2, single_a3, single_b_lo, single_b_hi, + d, single_sa, single_sb); + return; + } +#endif + +#if !defined(TL_SM120_FULLTILE_BOUNDED_2X2_PACKAGE) && \ + !defined(TL_SM120_FULLTILE_AFULL_B_PANEL_STREAM) +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 0, b00, b01, b02, b03); + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + kCuteKAtomStride, b10, b11, b12, + b13); + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 2 * kCuteKAtomStride, b20, b21, + b22, b23); + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 3 * kCuteKAtomStride, b30, b31, + b32, b33); +#else + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_0, b00, b01, b02, b03); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_1, b10, b11, b12, b13); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_2, b20, b21, b22, b23); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_3, b30, b31, b32, b33); +#endif + +#if defined(TL_SM120_FULLTILE_CUTLASS_SF_BASEPTR) + uint32_t const scale_warp = (uint32_t(int(threadIdx.x) & 127) >> 5); + uint32_t const scale_warp_m = scale_warp & 1u; + uint32_t const scale_warp_n = scale_warp >> 1; + uint32_t const scale_lane = uint32_t(int(threadIdx.x) & 31); + uint32_t const scale_sfa_row = 8u * (scale_lane & 1u) + (scale_lane >> 2); + uint32_t const scale_sfb_col = scale_lane >> 2; + uint32_t const scale_k = uint32_t(k_block_idx); + uint32_t const scale_m0 = scale_warp_m * 64u + scale_sfa_row; + uint32_t const scale_m1 = scale_m0 + 16u; + uint32_t const scale_m2 = scale_m0 + 32u; + uint32_t const scale_m3 = scale_m0 + 48u; + uint32_t const scale_n0 = scale_warp_n * 64u + scale_sfb_col; + uint32_t const scale_n1 = scale_n0 + 16u; + uint32_t const scale_n2 = scale_n0 + 32u; + uint32_t const scale_n3 = scale_n0 + 48u; + uint32_t const *sfa_base = sfa_0; + uint32_t const *sfb_base = sfb_0; + uint32_t const sa0 = sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_m0, scale_k)]; + uint32_t const sa1 = sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_m1, scale_k)]; + uint32_t const sa2 = sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_m2, scale_k)]; + uint32_t const sa3 = sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_m3, scale_k)]; + uint32_t const sb0 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n0, scale_k)]; + uint32_t const sb1 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n1, scale_k)]; + uint32_t const sb2 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n2, scale_k)]; + uint32_t const sb3 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n3, scale_k)]; + uint32_t const sbr0 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n0 + 8u, scale_k)]; + uint32_t const sbr1 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n1 + 8u, scale_k)]; + uint32_t const sbr2 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n2 + 8u, scale_k)]; + uint32_t const sbr3 = sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word( + scale_n3 + 8u, scale_k)]; +#define TL_SM120_SFA0 sa0 +#define TL_SM120_SFA1 sa1 +#define TL_SM120_SFA2 sa2 +#define TL_SM120_SFA3 sa3 +#else + uint32_t const sb0 = *sfb_0; + uint32_t const sb1 = *sfb_1; + uint32_t const sb2 = *sfb_2; + uint32_t const sb3 = *sfb_3; + uint32_t const sbr0 = *sfb_rep_0; + uint32_t const sbr1 = *sfb_rep_1; + uint32_t const sbr2 = *sfb_rep_2; + uint32_t const sbr3 = *sfb_rep_3; +#define TL_SM120_SFA0 (*sfa_0) +#define TL_SM120_SFA1 (*sfa_1) +#define TL_SM120_SFA2 (*sfa_2) +#define TL_SM120_SFA3 (*sfa_3) +#endif +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) && \ + !defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) + float c_cute[128]; +#pragma unroll + for (int mi = 0; mi < 4; ++mi) { +#pragma unroll + for (int nj = 0; nj < 4; ++nj) { +#pragma unroll + for (int half = 0; half < 2; ++half) { + int const tl_offset = mi * 32 + nj * 8 + half * 4; + int const cute_offset = mi * 4 + nj * 32 + half * 16; + *reinterpret_cast(c_cute + cute_offset) = + *reinterpret_cast(c + tl_offset); + } + } + } + float *c_mma = c_cute; +#else + float *c_mma = c; +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) || \ + defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) +#define TL_SM120_C_OFFSET(I, J, HALF) ((I) * 4 + (J) * 32 + (HALF) * 16) +#else +#define TL_SM120_C_OFFSET(I, J, HALF) ((I) * 32 + (J) * 8 + (HALF) * 4) +#endif + +#define TL_SM120_MMA_N8(I, J, HALF, A0, A1, A2, A3, B0, B1, SA, SB) \ + do { \ + float *d = c_mma + TL_SM120_C_OFFSET(I, J, HALF); \ + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs(d, A0, A1, A2, A3, B0, \ + B1, d, SA, SB); \ + } while (0) + +#define TL_SM120_MMA_ROW(I, A0, A1, A2, A3, SA) \ + do { \ + TL_SM120_MMA_N8(I, 0, 0, A0, A1, A2, A3, b00, b01, SA, sb0); \ + TL_SM120_MMA_N8(I, 0, 1, A0, A1, A2, A3, b02, b03, SA, sbr0); \ + TL_SM120_MMA_N8(I, 1, 0, A0, A1, A2, A3, b10, b11, SA, sb1); \ + TL_SM120_MMA_N8(I, 1, 1, A0, A1, A2, A3, b12, b13, SA, sbr1); \ + TL_SM120_MMA_N8(I, 2, 0, A0, A1, A2, A3, b20, b21, SA, sb2); \ + TL_SM120_MMA_N8(I, 2, 1, A0, A1, A2, A3, b22, b23, SA, sbr2); \ + TL_SM120_MMA_N8(I, 3, 0, A0, A1, A2, A3, b30, b31, SA, sb3); \ + TL_SM120_MMA_N8(I, 3, 1, A0, A1, A2, A3, b32, b33, SA, sbr3); \ + } while (0) + + uint32_t a0, a1, a2, a3; +#if defined(TL_SM120_FULLTILE_AFULL_B_PANEL_STREAM) + uint32_t a4, a5, a6, a7; + uint32_t a8, a9, a10, a11; + uint32_t a12, a13, a14, a15; +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 0, a0, a1, a2, a3); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + kCuteKAtomStride, a4, a5, a6, + a7); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 2 * kCuteKAtomStride, a8, a9, + a10, a11); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 3 * kCuteKAtomStride, a12, a13, + a14, a15); +#else + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a4, a5, a6, a7); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a8, a9, a10, a11); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a12, a13, a14, a15); +#endif + +#if defined(TL_SM120_FULLTILE_AFULL_SCALE_REGS) + uint32_t const sa0 = *sfa_0; + uint32_t const sa1 = *sfa_1; + uint32_t const sa2 = *sfa_2; + uint32_t const sa3 = *sfa_3; + uint32_t const sb0 = *sfb_0; + uint32_t const sb1 = *sfb_1; + uint32_t const sb2 = *sfb_2; + uint32_t const sb3 = *sfb_3; + uint32_t const sbr0 = *sfb_rep_0; + uint32_t const sbr1 = *sfb_rep_1; + uint32_t const sbr2 = *sfb_rep_2; + uint32_t const sbr3 = *sfb_rep_3; +#define TL_SM120_SFA0 sa0 +#define TL_SM120_SFA1 sa1 +#define TL_SM120_SFA2 sa2 +#define TL_SM120_SFA3 sa3 +#define TL_SM120_SFB0 sb0 +#define TL_SM120_SFB1 sb1 +#define TL_SM120_SFB2 sb2 +#define TL_SM120_SFB3 sb3 +#define TL_SM120_SFBR0 sbr0 +#define TL_SM120_SFBR1 sbr1 +#define TL_SM120_SFBR2 sbr2 +#define TL_SM120_SFBR3 sbr3 +#else +#define TL_SM120_SFA0 (*sfa_0) +#define TL_SM120_SFA1 (*sfa_1) +#define TL_SM120_SFA2 (*sfa_2) +#define TL_SM120_SFA3 (*sfa_3) +#define TL_SM120_SFB0 (*sfb_0) +#define TL_SM120_SFB1 (*sfb_1) +#define TL_SM120_SFB2 (*sfb_2) +#define TL_SM120_SFB3 (*sfb_3) +#define TL_SM120_SFBR0 (*sfb_rep_0) +#define TL_SM120_SFBR1 (*sfb_rep_1) +#define TL_SM120_SFBR2 (*sfb_rep_2) +#define TL_SM120_SFBR3 (*sfb_rep_3) +#endif + +#if defined(TL_SM120_FULLTILE_AFULL_B_PANEL_ADDRS) +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) + uint32_t const b_panel_0_addr = smem_ptr_to_uint(b_cute_n0 + 0); + uint32_t const b_panel_1_addr = + smem_ptr_to_uint(b_cute_n0 + kCuteKAtomStride); + uint32_t const b_panel_2_addr = + smem_ptr_to_uint(b_cute_n0 + 2 * kCuteKAtomStride); + uint32_t const b_panel_3_addr = + smem_ptr_to_uint(b_cute_n0 + 3 * kCuteKAtomStride); +#else + uint32_t const b_panel_0_addr = smem_ptr_to_uint(b_smem_0); + uint32_t const b_panel_1_addr = smem_ptr_to_uint(b_smem_1); + uint32_t const b_panel_2_addr = smem_ptr_to_uint(b_smem_2); + uint32_t const b_panel_3_addr = smem_ptr_to_uint(b_smem_3); +#endif +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) +#if defined(TL_SM120_FULLTILE_AFULL_B_PANEL_ADDRS) +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_0_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_1_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_2_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_3_addr, b00, b01, \ + b02, b03) +#else +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 0, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + kCuteKAtomStride, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 2 * kCuteKAtomStride, b00, \ + b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 3 * kCuteKAtomStride, b00, \ + b01, b02, b03) +#endif +#else +#if defined(TL_SM120_FULLTILE_AFULL_B_PANEL_ADDRS) +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_0_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_1_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_2_addr, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_blockscaled_operand_addr(b_panel_3_addr, b00, b01, \ + b02, b03) +#else +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_0, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_1, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_2, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_3, b00, b01, b02, b03) +#endif +#endif + +#define TL_SM120_MMA_B_PANEL(J, LOAD_B, SFB, SFB_REP) \ + do { \ + LOAD_B(); \ + uint32_t const sb = (SFB); \ + uint32_t const sbr = (SFB_REP); \ + TL_SM120_MMA_N8(0, J, 0, a0, a1, a2, a3, b00, b01, TL_SM120_SFA0, sb); \ + TL_SM120_MMA_N8(1, J, 0, a4, a5, a6, a7, b00, b01, TL_SM120_SFA1, sb); \ + TL_SM120_MMA_N8(2, J, 0, a8, a9, a10, a11, b00, b01, TL_SM120_SFA2, sb); \ + TL_SM120_MMA_N8(3, J, 0, a12, a13, a14, a15, b00, b01, TL_SM120_SFA3, sb); \ + TL_SM120_MMA_N8(0, J, 1, a0, a1, a2, a3, b02, b03, TL_SM120_SFA0, sbr); \ + TL_SM120_MMA_N8(1, J, 1, a4, a5, a6, a7, b02, b03, TL_SM120_SFA1, sbr); \ + TL_SM120_MMA_N8(2, J, 1, a8, a9, a10, a11, b02, b03, TL_SM120_SFA2, sbr); \ + TL_SM120_MMA_N8(3, J, 1, a12, a13, a14, a15, b02, b03, TL_SM120_SFA3, \ + sbr); \ + } while (0) + + TL_SM120_MMA_B_PANEL(0, TL_SM120_LOAD_B_PANEL_0, TL_SM120_SFB0, + TL_SM120_SFBR0); + TL_SM120_MMA_B_PANEL(1, TL_SM120_LOAD_B_PANEL_1, TL_SM120_SFB1, + TL_SM120_SFBR1); + TL_SM120_MMA_B_PANEL(2, TL_SM120_LOAD_B_PANEL_2, TL_SM120_SFB2, + TL_SM120_SFBR2); + TL_SM120_MMA_B_PANEL(3, TL_SM120_LOAD_B_PANEL_3, TL_SM120_SFB3, + TL_SM120_SFBR3); + +#undef TL_SM120_MMA_B_PANEL +#undef TL_SM120_LOAD_B_PANEL_3 +#undef TL_SM120_LOAD_B_PANEL_2 +#undef TL_SM120_LOAD_B_PANEL_1 +#undef TL_SM120_LOAD_B_PANEL_0 +#undef TL_SM120_SFBR3 +#undef TL_SM120_SFBR2 +#undef TL_SM120_SFBR1 +#undef TL_SM120_SFBR0 +#undef TL_SM120_SFB3 +#undef TL_SM120_SFB2 +#undef TL_SM120_SFB1 +#undef TL_SM120_SFB0 +#undef TL_SM120_SFA3 +#undef TL_SM120_SFA2 +#undef TL_SM120_SFA1 +#undef TL_SM120_SFA0 +#elif defined(TL_SM120_FULLTILE_PAIR_ASM_B) + uint32_t a4, a5, a6, a7; +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + kCuteKAtomStride, a4, a5, \ + a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 2 * kCuteKAtomStride, a0, \ + a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 3 * kCuteKAtomStride, a4, \ + a5, a6, a7); \ + } while (0) +#else +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a4, a5, a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a4, a5, a6, a7); \ + } while (0) +#endif + +#define TL_SM120_MMA2_SAME_B(I0, I1, J, HALF, B0, B1, SA0, SA1, SB) \ + do { \ + float *d0 = c_mma + TL_SM120_C_OFFSET(I0, J, HALF); \ + float *d1 = c_mma + TL_SM120_C_OFFSET(I1, J, HALF); \ + detail::sm120_mma2_m16n8k64_mxf4nvf4_4x_ue4m3_same_b_regs( \ + d0, a0, a1, a2, a3, d0, SA0, d1, a4, a5, a6, a7, d1, SA1, B0, B1, SB); \ + } while (0) + +#define TL_SM120_MMA2_PAIR_ROWS(I0, I1, SA0, SA1) \ + do { \ + TL_SM120_MMA2_SAME_B(I0, I1, 0, 0, b00, b01, SA0, SA1, sb0); \ + TL_SM120_MMA2_SAME_B(I0, I1, 0, 1, b02, b03, SA0, SA1, sbr0); \ + TL_SM120_MMA2_SAME_B(I0, I1, 1, 0, b10, b11, SA0, SA1, sb1); \ + TL_SM120_MMA2_SAME_B(I0, I1, 1, 1, b12, b13, SA0, SA1, sbr1); \ + TL_SM120_MMA2_SAME_B(I0, I1, 2, 0, b20, b21, SA0, SA1, sb2); \ + TL_SM120_MMA2_SAME_B(I0, I1, 2, 1, b22, b23, SA0, SA1, sbr2); \ + TL_SM120_MMA2_SAME_B(I0, I1, 3, 0, b30, b31, SA0, SA1, sb3); \ + TL_SM120_MMA2_SAME_B(I0, I1, 3, 1, b32, b33, SA0, SA1, sbr3); \ + } while (0) + + TL_SM120_LOAD_A_PAIR_01(); + TL_SM120_MMA2_PAIR_ROWS(0, 1, *sfa_0, *sfa_1); + TL_SM120_LOAD_A_PAIR_23(); + TL_SM120_MMA2_PAIR_ROWS(2, 3, *sfa_2, *sfa_3); + +#undef TL_SM120_MMA2_PAIR_ROWS +#undef TL_SM120_MMA2_SAME_B +#undef TL_SM120_LOAD_A_PAIR_23 +#undef TL_SM120_LOAD_A_PAIR_01 +#elif defined(TL_SM120_FULLTILE_BOUNDED_2X2_PACKAGE) + uint32_t a4, a5, a6, a7; +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + kCuteKAtomStride, a4, a5, \ + a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 2 * kCuteKAtomStride, a0, \ + a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 3 * kCuteKAtomStride, a4, \ + a5, a6, a7); \ + } while (0) +#else +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a4, a5, a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a4, a5, a6, a7); \ + } while (0) +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_B) +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 0, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + kCuteKAtomStride, b00, b01, \ + b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 2 * kCuteKAtomStride, b00, \ + b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_fp4_u32(b_cute_n0 + 3 * kCuteKAtomStride, b00, \ + b01, b02, b03) +#else +#define TL_SM120_LOAD_B_PANEL_0() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_0, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_1() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_1, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_2() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_2, b00, b01, b02, b03) +#define TL_SM120_LOAD_B_PANEL_3() \ + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_3, b00, b01, b02, b03) +#endif + +#define TL_SM120_MMA_PAIR_FWD(I0, I1, J, HALF, B0, B1, SA0, SA1, SB) \ + do { \ + TL_SM120_MMA_N8(I0, J, HALF, a0, a1, a2, a3, B0, B1, SA0, SB); \ + TL_SM120_MMA_N8(I1, J, HALF, a4, a5, a6, a7, B0, B1, SA1, SB); \ + } while (0) + +#define TL_SM120_MMA_PAIR_REV(I0, I1, J, HALF, B0, B1, SA0, SA1, SB) \ + do { \ + TL_SM120_MMA_N8(I1, J, HALF, a4, a5, a6, a7, B0, B1, SA1, SB); \ + TL_SM120_MMA_N8(I0, J, HALF, a0, a1, a2, a3, B0, B1, SA0, SB); \ + } while (0) + +#define TL_SM120_MMA_B_PANEL(I0, I1, J, LOAD_B, SFB, SFB_REP, SA0, SA1) \ + do { \ + LOAD_B(); \ + uint32_t const sb = *(SFB); \ + uint32_t const sbr = *(SFB_REP); \ + TL_SM120_MMA_PAIR_FWD(I0, I1, J, 0, b00, b01, SA0, SA1, sb); \ + TL_SM120_MMA_PAIR_REV(I0, I1, J, 1, b02, b03, SA0, SA1, sbr); \ + } while (0) + +#define TL_SM120_MMA_A_PAIR(I0, I1, LOAD_A, SA0, SA1) \ + do { \ + LOAD_A(); \ + TL_SM120_MMA_B_PANEL(I0, I1, 0, TL_SM120_LOAD_B_PANEL_0, sfb_0, sfb_rep_0, \ + SA0, SA1); \ + TL_SM120_MMA_B_PANEL(I0, I1, 1, TL_SM120_LOAD_B_PANEL_1, sfb_1, sfb_rep_1, \ + SA0, SA1); \ + TL_SM120_MMA_B_PANEL(I0, I1, 2, TL_SM120_LOAD_B_PANEL_2, sfb_2, sfb_rep_2, \ + SA0, SA1); \ + TL_SM120_MMA_B_PANEL(I0, I1, 3, TL_SM120_LOAD_B_PANEL_3, sfb_3, sfb_rep_3, \ + SA0, SA1); \ + } while (0) + + TL_SM120_MMA_A_PAIR(0, 1, TL_SM120_LOAD_A_PAIR_01, *sfa_0, *sfa_1); + TL_SM120_MMA_A_PAIR(2, 3, TL_SM120_LOAD_A_PAIR_23, *sfa_2, *sfa_3); + +#undef TL_SM120_MMA_A_PAIR +#undef TL_SM120_MMA_B_PANEL +#undef TL_SM120_MMA_PAIR_REV +#undef TL_SM120_MMA_PAIR_FWD +#undef TL_SM120_LOAD_B_PANEL_3 +#undef TL_SM120_LOAD_B_PANEL_2 +#undef TL_SM120_LOAD_B_PANEL_1 +#undef TL_SM120_LOAD_B_PANEL_0 +#undef TL_SM120_LOAD_A_PAIR_23 +#undef TL_SM120_LOAD_A_PAIR_01 +#elif defined(TL_SM120_FULLTILE_A2_BPAIR_ORDER) + uint32_t a4, a5, a6, a7; +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + kCuteKAtomStride, a4, a5, \ + a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 2 * kCuteKAtomStride, a0, \ + a1, a2, a3); \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 3 * kCuteKAtomStride, a4, \ + a5, a6, a7); \ + } while (0) +#else +#define TL_SM120_LOAD_A_PAIR_01() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a4, a5, a6, a7); \ + } while (0) +#define TL_SM120_LOAD_A_PAIR_23() \ + do { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a4, a5, a6, a7); \ + } while (0) +#endif + +#define TL_SM120_MMA_PAIR_FWD(I0, I1, J, HALF, B0, B1, SA0, SA1, SB) \ + do { \ + TL_SM120_MMA_N8(I0, J, HALF, a0, a1, a2, a3, B0, B1, SA0, SB); \ + TL_SM120_MMA_N8(I1, J, HALF, a4, a5, a6, a7, B0, B1, SA1, SB); \ + } while (0) + +#define TL_SM120_MMA_PAIR_REV(I0, I1, J, HALF, B0, B1, SA0, SA1, SB) \ + do { \ + TL_SM120_MMA_N8(I1, J, HALF, a4, a5, a6, a7, B0, B1, SA1, SB); \ + TL_SM120_MMA_N8(I0, J, HALF, a0, a1, a2, a3, B0, B1, SA0, SB); \ + } while (0) + +#define TL_SM120_MMA_PAIR_ROWS(I0, I1, SA0, SA1) \ + do { \ + TL_SM120_MMA_PAIR_FWD(I0, I1, 0, 0, b00, b01, SA0, SA1, sb0); \ + TL_SM120_MMA_PAIR_REV(I0, I1, 0, 1, b02, b03, SA0, SA1, sbr0); \ + TL_SM120_MMA_PAIR_FWD(I0, I1, 1, 0, b10, b11, SA0, SA1, sb1); \ + TL_SM120_MMA_PAIR_REV(I0, I1, 1, 1, b12, b13, SA0, SA1, sbr1); \ + TL_SM120_MMA_PAIR_FWD(I0, I1, 2, 0, b20, b21, SA0, SA1, sb2); \ + TL_SM120_MMA_PAIR_REV(I0, I1, 2, 1, b22, b23, SA0, SA1, sbr2); \ + TL_SM120_MMA_PAIR_FWD(I0, I1, 3, 0, b30, b31, SA0, SA1, sb3); \ + TL_SM120_MMA_PAIR_REV(I0, I1, 3, 1, b32, b33, SA0, SA1, sbr3); \ + } while (0) + + TL_SM120_LOAD_A_PAIR_01(); + TL_SM120_MMA_PAIR_ROWS(0, 1, TL_SM120_SFA0, TL_SM120_SFA1); + TL_SM120_LOAD_A_PAIR_23(); + TL_SM120_MMA_PAIR_ROWS(2, 3, TL_SM120_SFA2, TL_SM120_SFA3); + +#undef TL_SM120_MMA_PAIR_ROWS +#undef TL_SM120_MMA_PAIR_REV +#undef TL_SM120_MMA_PAIR_FWD +#undef TL_SM120_LOAD_A_PAIR_23 +#undef TL_SM120_LOAD_A_PAIR_01 +#elif defined(TL_SM120_FULLTILE_ROLLING_A_ORDER) +#if defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) +#define TL_SM120_LOAD_A_ROW(I) \ + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + (I) * kCuteKAtomStride, a0, \ + a1, a2, a3) +#else +#define TL_SM120_LOAD_A_ROW(I) \ + do { \ + if ((I) == 0) { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); \ + } else if ((I) == 1) { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a0, a1, a2, a3); \ + } else if ((I) == 2) { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); \ + } else { \ + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a0, a1, a2, a3); \ + } \ + } while (0) +#endif + +#define TL_SM120_MMA_N8_ROLL(I, J, HALF, B0, B1, SA, SB) \ + do { \ + TL_SM120_LOAD_A_ROW(I); \ + TL_SM120_MMA_N8(I, J, HALF, a0, a1, a2, a3, B0, B1, SA, SB); \ + } while (0) + +#define TL_SM120_MMA_N8_ROLL_FWD(J, HALF, B0, B1, SB) \ + do { \ + TL_SM120_MMA_N8_ROLL(0, J, HALF, B0, B1, *sfa_0, SB); \ + TL_SM120_MMA_N8_ROLL(1, J, HALF, B0, B1, *sfa_1, SB); \ + TL_SM120_MMA_N8_ROLL(2, J, HALF, B0, B1, *sfa_2, SB); \ + TL_SM120_MMA_N8_ROLL(3, J, HALF, B0, B1, *sfa_3, SB); \ + } while (0) + +#define TL_SM120_MMA_N8_ROLL_REV(J, HALF, B0, B1, SB) \ + do { \ + TL_SM120_MMA_N8_ROLL(3, J, HALF, B0, B1, *sfa_3, SB); \ + TL_SM120_MMA_N8_ROLL(2, J, HALF, B0, B1, *sfa_2, SB); \ + TL_SM120_MMA_N8_ROLL(1, J, HALF, B0, B1, *sfa_1, SB); \ + TL_SM120_MMA_N8_ROLL(0, J, HALF, B0, B1, *sfa_0, SB); \ + } while (0) + + TL_SM120_MMA_N8_ROLL_FWD(0, 0, b00, b01, sb0); + TL_SM120_MMA_N8_ROLL_REV(0, 1, b02, b03, sbr0); + TL_SM120_MMA_N8_ROLL_FWD(1, 0, b10, b11, sb1); + TL_SM120_MMA_N8_ROLL_REV(1, 1, b12, b13, sbr1); + TL_SM120_MMA_N8_ROLL_FWD(2, 0, b20, b21, sb2); + TL_SM120_MMA_N8_ROLL_REV(2, 1, b22, b23, sbr2); + TL_SM120_MMA_N8_ROLL_FWD(3, 0, b30, b31, sb3); + TL_SM120_MMA_N8_ROLL_REV(3, 1, b32, b33, sbr3); + +#undef TL_SM120_MMA_N8_ROLL_REV +#undef TL_SM120_MMA_N8_ROLL_FWD +#undef TL_SM120_MMA_N8_ROLL +#undef TL_SM120_LOAD_A_ROW +#elif defined(TL_SM120_FULLTILE_CUTE_ROWSTART) || \ + defined(TL_SM120_FULLTILE_CUTE_ROWSTART_A) + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 0, a0, a1, a2, a3); + TL_SM120_MMA_ROW(0, a0, a1, a2, a3, TL_SM120_SFA0); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + kCuteKAtomStride, a0, a1, a2, + a3); + TL_SM120_MMA_ROW(1, a0, a1, a2, a3, TL_SM120_SFA1); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 2 * kCuteKAtomStride, a0, a1, + a2, a3); + TL_SM120_MMA_ROW(2, a0, a1, a2, a3, TL_SM120_SFA2); + detail::sm120_ldmatrix_x4_fp4_u32(a_cute_m0 + 3 * kCuteKAtomStride, a0, a1, + a2, a3); + TL_SM120_MMA_ROW(3, a0, a1, a2, a3, TL_SM120_SFA3); +#else + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); + TL_SM120_MMA_ROW(0, a0, a1, a2, a3, TL_SM120_SFA0); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a0, a1, a2, a3); + TL_SM120_MMA_ROW(1, a0, a1, a2, a3, TL_SM120_SFA1); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); + TL_SM120_MMA_ROW(2, a0, a1, a2, a3, TL_SM120_SFA2); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a0, a1, a2, a3); + TL_SM120_MMA_ROW(3, a0, a1, a2, a3, TL_SM120_SFA3); +#endif + +#if !defined(TL_SM120_FULLTILE_BOUNDED_2X2_PACKAGE) && \ + !defined(TL_SM120_FULLTILE_AFULL_B_PANEL_STREAM) +#undef TL_SM120_SFA3 +#undef TL_SM120_SFA2 +#undef TL_SM120_SFA1 +#undef TL_SM120_SFA0 +#endif + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) && \ + !defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) +#pragma unroll + for (int mi = 0; mi < 4; ++mi) { +#pragma unroll + for (int nj = 0; nj < 4; ++nj) { +#pragma unroll + for (int half = 0; half < 2; ++half) { + int const tl_offset = mi * 32 + nj * 8 + half * 4; + int const cute_offset = mi * 4 + nj * 32 + half * 16; + *reinterpret_cast(c + tl_offset) = + *reinterpret_cast(c_cute + cute_offset); + } + } + } +#endif + +#undef TL_SM120_MMA_ROW +#undef TL_SM120_MMA_N8 +#undef TL_SM120_C_OFFSET +} + +TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide( + float *c, const void *a_smem_0, const void *a_smem_1, const void *a_smem_2, + const void *a_smem_3, const void *b_smem_0, const void *b_smem_1, + const void *b_smem_2, const void *b_smem_3, const uint32_t *sfa_base, + const uint32_t *sfb_base, int k_block_idx) { + uint32_t b00, b01, b02, b03; + uint32_t b10, b11, b12, b13; + uint32_t b20, b21, b22, b23; + uint32_t b30, b31, b32, b33; + + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_0, b00, b01, b02, b03); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_1, b10, b11, b12, b13); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_2, b20, b21, b22, b23); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_3, b30, b31, b32, b33); + + uint32_t const tx = uint32_t(int(threadIdx.x) & 127); + uint32_t const lane = tx & 31u; + uint32_t const qlane = lane & 3u; + uint32_t const warp = tx >> 5; + uint32_t const warp_m = warp & 1u; + uint32_t const warp_n = warp >> 1; + uint32_t const sfa_row = 8u * (lane & 1u) + (lane >> 2); + uint32_t const sfb_col = lane >> 2; + uint32_t const a_owner_in_pair = qlane >> 1; + uint32_t const scale_k = uint32_t(k_block_idx); + + uint32_t const scale_m0 = warp_m * 64u + a_owner_in_pair * 16u + sfa_row; + uint32_t const scale_m1 = scale_m0 + 32u; + uint32_t const scale_n0 = warp_n * 64u + qlane * 8u + sfb_col; + uint32_t const scale_n1 = scale_n0 + 32u; + uint32_t const sa_owner0 = + sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_m0, + scale_k)]; + uint32_t const sa_owner1 = + sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_m1, + scale_k)]; + uint32_t const sb_owner0 = + sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_n0, + scale_k)]; + uint32_t const sb_owner1 = + sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_n1, + scale_k)]; + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) || \ + defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) +#define TL_SM120_OWNER_C_OFFSET(I, J, HALF) ((I) * 4 + (J) * 32 + (HALF) * 16) +#else +#define TL_SM120_OWNER_C_OFFSET(I, J, HALF) ((I) * 32 + (J) * 8 + (HALF) * 4) +#endif + +#define TL_SM120_OWNER_MMA_N8(I, J, HALF, A0, A1, A2, A3, B0, B1, SA, SB, \ + SA_TID, SB_TID) \ + do { \ + float *d = c + TL_SM120_OWNER_C_OFFSET(I, J, HALF); \ + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( \ + d, A0, A1, A2, A3, B0, B1, d, SA, SB, 0, uint16_t(SA_TID), 0, \ + uint16_t(SB_TID)); \ + } while (0) + +#define TL_SM120_OWNER_MMA_ROW(I, A0, A1, A2, A3, SA, SA_TID) \ + do { \ + TL_SM120_OWNER_MMA_N8(I, 0, 0, A0, A1, A2, A3, b00, b01, SA, sb_owner0, \ + SA_TID, 0); \ + TL_SM120_OWNER_MMA_N8(I, 0, 1, A0, A1, A2, A3, b02, b03, SA, sb_owner0, \ + SA_TID, 1); \ + TL_SM120_OWNER_MMA_N8(I, 1, 0, A0, A1, A2, A3, b10, b11, SA, sb_owner0, \ + SA_TID, 2); \ + TL_SM120_OWNER_MMA_N8(I, 1, 1, A0, A1, A2, A3, b12, b13, SA, sb_owner0, \ + SA_TID, 3); \ + TL_SM120_OWNER_MMA_N8(I, 2, 0, A0, A1, A2, A3, b20, b21, SA, sb_owner1, \ + SA_TID, 0); \ + TL_SM120_OWNER_MMA_N8(I, 2, 1, A0, A1, A2, A3, b22, b23, SA, sb_owner1, \ + SA_TID, 1); \ + TL_SM120_OWNER_MMA_N8(I, 3, 0, A0, A1, A2, A3, b30, b31, SA, sb_owner1, \ + SA_TID, 2); \ + TL_SM120_OWNER_MMA_N8(I, 3, 1, A0, A1, A2, A3, b32, b33, SA, sb_owner1, \ + SA_TID, 3); \ + } while (0) + + uint32_t a0, a1, a2, a3; + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a0, a1, a2, a3); + TL_SM120_OWNER_MMA_ROW(0, a0, a1, a2, a3, sa_owner0, 0); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a0, a1, a2, a3); + TL_SM120_OWNER_MMA_ROW(1, a0, a1, a2, a3, sa_owner0, 1); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a0, a1, a2, a3); + TL_SM120_OWNER_MMA_ROW(2, a0, a1, a2, a3, sa_owner1, 0); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a0, a1, a2, a3); + TL_SM120_OWNER_MMA_ROW(3, a0, a1, a2, a3, sa_owner1, 1); + +#undef TL_SM120_OWNER_MMA_ROW +#undef TL_SM120_OWNER_MMA_N8 +#undef TL_SM120_OWNER_C_OFFSET +} + +TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide( + float *c, const void *a_smem_0, const void *a_smem_1, const void *a_smem_2, + const void *a_smem_3, const void *b_smem_0, const void *b_smem_1, + const void *b_smem_2, const void *b_smem_3, const uint32_t *sfa_base, + const uint32_t *sfb_base, int k_block_idx) { + uint32_t b00, b01, b02, b03; + uint32_t b10, b11, b12, b13; + + // Load B0 first so the following A/scale setup gives it real distance before + // its first OMMA use. + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_0, b00, b01, b02, b03); + + uint32_t const tx = uint32_t(int(threadIdx.x) & 127); + uint32_t const lane = tx & 31u; + uint32_t const qlane = lane & 3u; + uint32_t const warp = tx >> 5; + uint32_t const warp_m = warp & 1u; + uint32_t const warp_n = warp >> 1; + uint32_t const sfa_row = 8u * (lane & 1u) + (lane >> 2); + uint32_t const sfb_col = lane >> 2; + uint32_t const a_owner_in_pair = qlane >> 1; + uint32_t const scale_k = uint32_t(k_block_idx); + + uint32_t const scale_m0 = warp_m * 64u + a_owner_in_pair * 16u + sfa_row; + uint32_t const scale_m1 = scale_m0 + 32u; + uint32_t const scale_n0 = warp_n * 64u + qlane * 8u + sfb_col; + uint32_t const scale_n1 = scale_n0 + 32u; + uint32_t const sa_owner0 = + sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_m0, + scale_k)]; + uint32_t const sa_owner1 = + sfa_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_m1, + scale_k)]; + uint32_t const sb_owner0 = + sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_n0, + scale_k)]; + uint32_t const sb_owner1 = + sfb_base[detail::sm120_blockscaled_chunk_kmajor_sf_word(scale_n1, + scale_k)]; + + uint32_t a00, a01, a02, a03; + uint32_t a10, a11, a12, a13; + uint32_t a20, a21, a22, a23; + uint32_t a30, a31, a32, a33; + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_0, a00, a01, a02, a03); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_1, a10, a11, a12, a13); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_2, a20, a21, a22, a23); + detail::sm120_ldmatrix_x4_blockscaled_operand(a_smem_3, a30, a31, a32, a33); + +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) || \ + defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) +#define TL_SM120_AFULL_C_OFFSET(I, J, HALF) ((I) * 4 + (J) * 32 + (HALF) * 16) +#else +#define TL_SM120_AFULL_C_OFFSET(I, J, HALF) ((I) * 32 + (J) * 8 + (HALF) * 4) +#endif + +#define TL_SM120_AFULL_MMA_N8(I, J, HALF, A0, A1, A2, A3, B0, B1, SA, SB, \ + SA_TID, SB_TID) \ + do { \ + float *d = c + TL_SM120_AFULL_C_OFFSET(I, J, HALF); \ + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( \ + d, A0, A1, A2, A3, B0, B1, d, SA, SB, 0, uint16_t(SA_TID), 0, \ + uint16_t(SB_TID)); \ + } while (0) + +#define TL_SM120_AFULL_MMA_PANEL_LO(J, B0, B1, SB, SB_TID) \ + do { \ + TL_SM120_AFULL_MMA_N8(0, J, 0, a00, a01, a02, a03, B0, B1, sa_owner0, SB, \ + 0, SB_TID); \ + TL_SM120_AFULL_MMA_N8(1, J, 0, a10, a11, a12, a13, B0, B1, sa_owner0, SB, \ + 1, SB_TID); \ + TL_SM120_AFULL_MMA_N8(2, J, 0, a20, a21, a22, a23, B0, B1, sa_owner1, SB, \ + 0, SB_TID); \ + TL_SM120_AFULL_MMA_N8(3, J, 0, a30, a31, a32, a33, B0, B1, sa_owner1, SB, \ + 1, SB_TID); \ + } while (0) + +#define TL_SM120_AFULL_MMA_PANEL_HI(J, B2, B3, SB, SB_TID) \ + do { \ + TL_SM120_AFULL_MMA_N8(3, J, 1, a30, a31, a32, a33, B2, B3, sa_owner1, SB, \ + 1, SB_TID); \ + TL_SM120_AFULL_MMA_N8(2, J, 1, a20, a21, a22, a23, B2, B3, sa_owner1, SB, \ + 0, SB_TID); \ + TL_SM120_AFULL_MMA_N8(1, J, 1, a10, a11, a12, a13, B2, B3, sa_owner0, SB, \ + 1, SB_TID); \ + TL_SM120_AFULL_MMA_N8(0, J, 1, a00, a01, a02, a03, B2, B3, sa_owner0, SB, \ + 0, SB_TID); \ + } while (0) + + TL_SM120_AFULL_MMA_PANEL_LO(0, b00, b01, sb_owner0, 0); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_1, b10, b11, b12, b13); + TL_SM120_AFULL_MMA_PANEL_HI(0, b02, b03, sb_owner0, 1); + TL_SM120_AFULL_MMA_PANEL_LO(1, b10, b11, sb_owner0, 2); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_2, b00, b01, b02, b03); + TL_SM120_AFULL_MMA_PANEL_HI(1, b12, b13, sb_owner0, 3); + TL_SM120_AFULL_MMA_PANEL_LO(2, b00, b01, sb_owner1, 0); + detail::sm120_ldmatrix_x4_blockscaled_operand(b_smem_3, b10, b11, b12, b13); + TL_SM120_AFULL_MMA_PANEL_HI(2, b02, b03, sb_owner1, 1); + TL_SM120_AFULL_MMA_PANEL_LO(3, b10, b11, sb_owner1, 2); + TL_SM120_AFULL_MMA_PANEL_HI(3, b12, b13, sb_owner1, 3); + +#undef TL_SM120_AFULL_MMA_PANEL_HI +#undef TL_SM120_AFULL_MMA_PANEL_LO +#undef TL_SM120_AFULL_MMA_N8 +#undef TL_SM120_AFULL_C_OFFSET +} + +struct SM120FulltileABOwnerWidePackage { + uint32_t a00, a01, a02, a03; + uint32_t a10, a11, a12, a13; + uint32_t a20, a21, a22, a23; + uint32_t a30, a31, a32, a33; + uint32_t b00, b01, b02, b03; + uint32_t b10, b11, b12, b13; + uint32_t b20, b21, b22, b23; + uint32_t b30, b31, b32, b33; +}; + +TL_DEVICE void +sm120_copy_fulltile_ab_owner_wide_package(SM120FulltileABOwnerWidePackage &pkg, + const char *a_base, + const char *b_base, int k_block_idx) { + uint32_t const tx = uint32_t(int(threadIdx.x) & 127); + + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_base + detail::sm120_fulltile_package_a_offset(tx, k_block_idx, 0), + pkg.a00, pkg.a01, pkg.a02, pkg.a03); + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_base + detail::sm120_fulltile_package_a_offset(tx, k_block_idx, 1), + pkg.a10, pkg.a11, pkg.a12, pkg.a13); + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_base + detail::sm120_fulltile_package_a_offset(tx, k_block_idx, 2), + pkg.a20, pkg.a21, pkg.a22, pkg.a23); + detail::sm120_ldmatrix_x4_blockscaled_operand( + a_base + detail::sm120_fulltile_package_a_offset(tx, k_block_idx, 3), + pkg.a30, pkg.a31, pkg.a32, pkg.a33); + + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_base + detail::sm120_fulltile_package_b_offset(tx, k_block_idx, 0), + pkg.b00, pkg.b01, pkg.b02, pkg.b03); + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_base + detail::sm120_fulltile_package_b_offset(tx, k_block_idx, 1), + pkg.b10, pkg.b11, pkg.b12, pkg.b13); + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_base + detail::sm120_fulltile_package_b_offset(tx, k_block_idx, 2), + pkg.b20, pkg.b21, pkg.b22, pkg.b23); + detail::sm120_ldmatrix_x4_blockscaled_operand( + b_base + detail::sm120_fulltile_package_b_offset(tx, k_block_idx, 3), + pkg.b30, pkg.b31, pkg.b32, pkg.b33); +} + +TL_DEVICE void sm120_gemm_fulltile_ab_owner_wide_package( + float *c, const SM120FulltileABOwnerWidePackage &pkg, + const detail::SM120ScaleTVPackage &scale_pkg) { +#if defined(TL_SM120_FULLTILE_CUTE_ACCUM_LAYOUT) || \ + defined(TL_SM120_FULLTILE_CUTE_ACCUM_DIRECT) +#define TL_SM120_PKG_C_OFFSET(I, J, HALF) ((I) * 4 + (J) * 32 + (HALF) * 16) +#else +#define TL_SM120_PKG_C_OFFSET(I, J, HALF) ((I) * 32 + (J) * 8 + (HALF) * 4) +#endif + +#define TL_SM120_PKG_MMA_N8(I, J, HALF, A0, A1, A2, A3, B0, B1, SA, SB, \ + SA_TID, SB_TID) \ + do { \ + float *d = c + TL_SM120_PKG_C_OFFSET(I, J, HALF); \ + detail::sm120_mma_m16n8k64_mxf4nvf4_4x_ue4m3_regs( \ + d, A0, A1, A2, A3, B0, B1, d, SA, SB, 0, uint16_t(SA_TID), 0, \ + uint16_t(SB_TID)); \ + } while (0) + +#define TL_SM120_PKG_MMA_ROW(I, A0, A1, A2, A3, SA, SA_TID) \ + do { \ + TL_SM120_PKG_MMA_N8(I, 0, 0, A0, A1, A2, A3, pkg.b00, pkg.b01, SA, \ + scale_pkg.sb0, SA_TID, 0); \ + TL_SM120_PKG_MMA_N8(I, 0, 1, A0, A1, A2, A3, pkg.b02, pkg.b03, SA, \ + scale_pkg.sb0, SA_TID, 1); \ + TL_SM120_PKG_MMA_N8(I, 1, 0, A0, A1, A2, A3, pkg.b10, pkg.b11, SA, \ + scale_pkg.sb0, SA_TID, 2); \ + TL_SM120_PKG_MMA_N8(I, 1, 1, A0, A1, A2, A3, pkg.b12, pkg.b13, SA, \ + scale_pkg.sb0, SA_TID, 3); \ + TL_SM120_PKG_MMA_N8(I, 2, 0, A0, A1, A2, A3, pkg.b20, pkg.b21, SA, \ + scale_pkg.sb1, SA_TID, 0); \ + TL_SM120_PKG_MMA_N8(I, 2, 1, A0, A1, A2, A3, pkg.b22, pkg.b23, SA, \ + scale_pkg.sb1, SA_TID, 1); \ + TL_SM120_PKG_MMA_N8(I, 3, 0, A0, A1, A2, A3, pkg.b30, pkg.b31, SA, \ + scale_pkg.sb1, SA_TID, 2); \ + TL_SM120_PKG_MMA_N8(I, 3, 1, A0, A1, A2, A3, pkg.b32, pkg.b33, SA, \ + scale_pkg.sb1, SA_TID, 3); \ + } while (0) + + TL_SM120_PKG_MMA_ROW(0, pkg.a00, pkg.a01, pkg.a02, pkg.a03, scale_pkg.sa0, 0); + TL_SM120_PKG_MMA_ROW(1, pkg.a10, pkg.a11, pkg.a12, pkg.a13, scale_pkg.sa0, 1); + TL_SM120_PKG_MMA_ROW(2, pkg.a20, pkg.a21, pkg.a22, pkg.a23, scale_pkg.sa1, 0); + TL_SM120_PKG_MMA_ROW(3, pkg.a30, pkg.a31, pkg.a32, pkg.a33, scale_pkg.sa1, 1); + +#undef TL_SM120_PKG_MMA_ROW +#undef TL_SM120_PKG_MMA_N8 +#undef TL_SM120_PKG_C_OFFSET +} + +TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile_package_pingpong( + float *c, const void *a_smem_base, const void *b_smem_base, + const uint32_t *sfa_smem_base, const uint32_t *sfb_smem_base) { + const char *a_base = static_cast(a_smem_base); + const char *b_base = static_cast(b_smem_base); + SM120FulltileABOwnerWidePackage pkg0; + SM120FulltileABOwnerWidePackage pkg1; + detail::SM120ScaleTVPackage scale_pkg0; + detail::SM120ScaleTVPackage scale_pkg1; + + sm120_copy_fulltile_ab_owner_wide_package(pkg0, a_base, b_base, 0); + detail::sm120_copy_scale_tv_package(scale_pkg0, sfa_smem_base, sfb_smem_base, + 0); + sm120_copy_fulltile_ab_owner_wide_package(pkg1, a_base, b_base, 1); + detail::sm120_copy_scale_tv_package(scale_pkg1, sfa_smem_base, sfb_smem_base, + 1); + sm120_gemm_fulltile_ab_owner_wide_package(c, pkg0, scale_pkg0); + sm120_copy_fulltile_ab_owner_wide_package(pkg0, a_base, b_base, 2); + detail::sm120_copy_scale_tv_package(scale_pkg0, sfa_smem_base, sfb_smem_base, + 2); + sm120_gemm_fulltile_ab_owner_wide_package(c, pkg1, scale_pkg1); + sm120_copy_fulltile_ab_owner_wide_package(pkg1, a_base, b_base, 3); + detail::sm120_copy_scale_tv_package(scale_pkg1, sfa_smem_base, sfb_smem_base, + 3); + sm120_gemm_fulltile_ab_owner_wide_package(c, pkg0, scale_pkg0); + sm120_gemm_fulltile_ab_owner_wide_package(c, pkg1, scale_pkg1); +} + +} // namespace tl diff --git a/src/transform/merge_shared_memory_allocations.cc b/src/transform/merge_shared_memory_allocations.cc index 0370e1b42..caed6e32b 100644 --- a/src/transform/merge_shared_memory_allocations.cc +++ b/src/transform/merge_shared_memory_allocations.cc @@ -34,6 +34,7 @@ #include #include +#include #include #include #include @@ -72,6 +73,68 @@ static bool IsStaticSharedMemory(Var buffer_var) { storage_scope.tag.empty(); } +static DataType GetStorageSizeExprDType(const Buffer &buffer) { + DataType size_dtype = DataType::Int(32); + if (!buffer->shape.empty()) { + size_dtype = buffer->shape[0].dtype(); + } + if (!size_dtype.is_int() && !size_dtype.is_uint()) { + size_dtype = DataType::Int(32); + } + return size_dtype; +} + +static bool EnableSM120CompactUnpackedFP4Shared() { + const char *value = std::getenv("TL_SM120_COMPACT_UNPACKED_FP4_SHARED"); + return value != nullptr && value[0] != '\0' && + !(value[0] == '0' && value[1] == '\0'); +} + +static int64_t GetSharedStorageBitsPerLogicalElement(DataType dtype) { + if (EnableSM120CompactUnpackedFP4Shared() && + dtype.is_float4_e2m1_unpacked()) { + return 4 * dtype.lanes(); + } + return static_cast(dtype.bits()) * dtype.lanes(); +} + +static PrimExpr GetSharedStorageSizeBytes(const Buffer &buffer) { + DataType size_dtype = GetStorageSizeExprDType(buffer); + int64_t element_bits = GetSharedStorageBitsPerLogicalElement(buffer->dtype); + + // Buffer shapes are expressed in logical elements. Do not use + // DataType::bytes() here: it rounds sub-byte scalar types up per element, so + // scalar packed FP4 would be charged as one byte per value instead of two + // values per byte. Compute total bits first, then round the whole allocation + // up to bytes. + PrimExpr size_bits = make_const(size_dtype, element_bits); + for (const PrimExpr &extent : buffer->shape) { + PrimExpr e = extent; + if (e.dtype() != size_dtype) { + e = cast(size_dtype, e); + } + size_bits = size_bits * e; + } + + if (element_bits % 8 == 0) { + return indexdiv(size_bits, make_const(size_dtype, 8)); + } + return indexdiv(size_bits + make_const(size_dtype, 7), + make_const(size_dtype, 8)); +} + +static PrimExpr SharedByteOffsetToLogicalIndexOffset(PrimExpr byte_offset, + DataType dtype) { + // buffer_byte_offsets_ stores byte offsets into the merged uint8 arena, but + // the non-alias rewrite still indexes the original typed buffer. Convert the + // byte offset back to that buffer's logical element index. Packed scalar + // NVFP4 has two logical elements per byte. + if (dtype.is_float4_e2m1fn() && dtype.is_scalar()) { + return byte_offset * make_const(byte_offset.dtype(), 2); + } + return indexdiv(byte_offset, dtype.bytes() * dtype.lanes()); +} + /*! * \brief collect the mapping from the buffer var to its allocate */ @@ -570,25 +633,7 @@ private: for (const VarNode *var : sorted_vars) { const AllocBufferNode *alloc = shmem_allocs_.at(var); - int64_t bytes_per_elem = static_cast( - alloc->buffer->dtype.bytes() * alloc->buffer->dtype.lanes()); - - DataType size_dtype = DataType::Int(32); - if (!alloc->buffer->shape.empty()) { - size_dtype = alloc->buffer->shape[0].dtype(); - } - if (!size_dtype.is_int() && !size_dtype.is_uint()) { - size_dtype = DataType::Int(32); - } - - PrimExpr size_expr = make_const(size_dtype, bytes_per_elem); - for (const PrimExpr &extent : alloc->buffer->shape) { - PrimExpr e = extent; - if (e.dtype() != size_dtype) { - e = cast(size_dtype, e); - } - size_expr = size_expr * e; - } + PrimExpr size_expr = GetSharedStorageSizeBytes(alloc->buffer); int alignment = align_bytes_; auto align_it = shmem_alignment_map_.find(var); @@ -627,9 +672,8 @@ private: auto alloc_it = shmem_allocs_.find(buffer_var_node); if (alloc_it != shmem_allocs_.end()) { const AllocBufferNode *alloc = alloc_it->second; - PrimExpr buffer_size_bytes = alloc->buffer->shape[0] * - alloc->buffer->dtype.bytes() * - alloc->buffer->dtype.lanes(); + PrimExpr buffer_size_bytes = + GetSharedStorageSizeBytes(alloc->buffer); LOG(DEBUG) << " Buffer: " << buffer_var_node->name_hint << " (Type: " << alloc->buffer->dtype << ")" << ", Start Offset: " << byte_offset @@ -812,7 +856,7 @@ private: auto it = buffer_byte_offsets_.find(buffer_var.get()); ICHECK(it != buffer_byte_offsets_.end()) << "buffer_var = " << buffer_var->name_hint << ", dtype = " << dtype; - return indexdiv(it->second, dtype.bytes() * dtype.lanes()); + return SharedByteOffsetToLogicalIndexOffset(it->second, dtype); } bool HasBufferOffset(const Var &buffer_var) { @@ -1406,30 +1450,16 @@ private: } const AllocBufferNode *alloc = shmem_allocs_.at(var); - int64_t bytes_per_elem = static_cast( - alloc->buffer->dtype.bytes() * alloc->buffer->dtype.lanes()); - DataType size_dtype = DataType::Int(32); - if (!alloc->buffer->shape.empty()) { - size_dtype = alloc->buffer->shape[0].dtype(); - } - if (!size_dtype.is_int() && !size_dtype.is_uint()) { - size_dtype = DataType::Int(32); - } - - PrimExpr size_expr = make_const(size_dtype, bytes_per_elem); - for (const PrimExpr &extent : alloc->buffer->shape) { - PrimExpr e = extent; - if (e.dtype() != size_dtype) { - e = cast(size_dtype, e); - } - size_expr = size_expr * e; - } + PrimExpr size_expr = GetSharedStorageSizeBytes(alloc->buffer); + DataType size_dtype = GetStorageSizeExprDType(alloc->buffer); info.size_dtype = size_dtype; info.size_expr = size_expr; auto const_size = GetRef(alloc).ConstantAllocationSize(); if (const_size.has_value()) { - info.const_size_bytes = const_size.value() * bytes_per_elem; + int64_t element_bits = + GetSharedStorageBitsPerLogicalElement(alloc->buffer->dtype); + info.const_size_bytes = (const_size.value() * element_bits + 7) / 8; } buf_infos.push_back(std::move(info)); diff --git a/testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py b/testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py new file mode 100644 index 000000000..71894e35f --- /dev/null +++ b/testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py @@ -0,0 +1,705 @@ +import pytest + +import tilelang +import tilelang.language as T +import tilelang.testing +from tilelang import tvm +from tilelang.cuda.intrinsics.layout.mma_layout import mma_load_a_32x32_to_shared_16x64_layout +from tilelang.cuda.intrinsics.macro.mma_macro_generator import SM120BlockScaledFullTilePackageContract +from tilelang.intrinsics import TensorCoreIntrinEmitterWithBlockScale, get_swizzle_layout +from tilelang.quantize import blockscaled_chunk_kmajor_word_offset +from tilelang.transform import simplify_prim_func + + +_FP4_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + + +def _make_sm120_fulltile_contract(): + emitter = TensorCoreIntrinEmitterWithBlockScale( + a_dtype=T.float4_e2m1fn, + b_dtype=T.float4_e2m1fn, + accum_dtype=T.float32, + a_transposed=False, + b_transposed=True, + block_row_warps=2, + block_col_warps=2, + warp_row_tiles=64, + warp_col_tiles=64, + chunk=256, + reduce_k=1, + num_elems_per_byte=2, + kind="mxf4nvf4", + scale_vec_size=4, + stype="ue4m3", + ) + return SM120BlockScaledFullTilePackageContract.for_package_pingpong( + emitter, + sf_layout="blockscaled_chunk_kmajor", + ) + + +def _oracle_blockscaled_chunk_kmajor_flat_word(row: int, kblock: int) -> int: + physical_row, physical_word = blockscaled_chunk_kmajor_word_offset(row, kblock) + return physical_row * 4 + physical_word + + +def _make_swizzle_layout(shared_buf): + dtype = shared_buf.dtype + shape = shared_buf.shape + + can_swizzle = shape[-1] * tvm.DataType(dtype).bits % 512 == 0 + if not can_swizzle: + return T.Layout(shape, lambda *args: args) + + def transform_func(i, j): + new_warp_i, new_warp_j = get_swizzle_layout(i, j, shape[-1], dtype) + return [new_warp_i, new_warp_j] + + return T.Layout(shape, transform_func) + + +@simplify_prim_func +def _make_nvf4_matmul_codegen_kernel( + M, + N, + K, + num_stages=2, + *, + block_row_warps=2, + block_col_warps=2, + warp_row_tiles=32, + warp_col_tiles=32, + micro_pipeline=None, + sf_layout=None, +): + assert K % 64 == 0 + in_dtype = T.float4_e2m1fn + out_dtype = T.float32 + accum_dtype = T.float32 + + micro_size_k = 64 + + chunk = K + shared_scope = "shared.dyn" + + block_M = block_row_warps * warp_row_tiles + block_N = block_col_warps * warp_col_tiles + block_K = chunk + + A_shape = (M, K) + B_shape = (N, K) + SFA_shape = (M, K // micro_size_k) + SFB_shape = (N, K // micro_size_k) + A_shared_shape = (block_M, block_K) + B_shared_shape = (block_N, block_K) + SFA_shared_shape = (block_M, block_K // micro_size_k) + SFB_shared_shape = (block_N, block_K // micro_size_k) + + warp_size = 32 + threads = warp_size * (block_row_warps * block_col_warps) + + @T.prim_func + def main( + A: T.Tensor(A_shape, in_dtype), + B: T.Tensor(B_shape, in_dtype), + SFA: T.Tensor(SFA_shape, T.uint32), + SFB: T.Tensor(SFB_shape, T.uint32), + C: T.Tensor((M, N), out_dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as ( + bx, + by, + ): + A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope) + B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope) + SFA_shared = T.alloc_shared(SFA_shared_shape, T.uint32, scope=shared_scope) + SFB_shared = T.alloc_shared(SFB_shared_shape, T.uint32, scope=shared_scope) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + T.use_swizzle(panel_size=10) + + for ko in T.Pipelined((K // block_K), num_stages=num_stages): + for i, k in T.Parallel(block_M, block_K): + A_shared[i, k] = A[by * block_M + i, ko * block_K + k] + + for j, k in T.Parallel(block_N, block_K): + B_shared[j, k] = B[bx * block_N + j, ko * block_K + k] + + for i, k in T.Parallel(block_M, block_K // micro_size_k): + SFA_shared[i, k] = SFA[by * block_M + i, ko * (block_K // micro_size_k) + k] + + for j, k in T.Parallel(block_N, block_K // micro_size_k): + SFB_shared[j, k] = SFB[bx * block_N + j, ko * (block_K // micro_size_k) + k] + + T.mma_gemm_blockscaled( + A_shared, + B_shared, + C_local, + SFA_shared, + SFB_shared, + transpose_B=True, + clear_accum=True, + k_start=ko * block_K, + sf_a_granularity_k=16, + sf_b_granularity_k=16, + micro_pipeline=micro_pipeline, + sf_layout=sf_layout, + ) + + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + + +def _decode_rowmajor_fp4(packed, rows: int, cols: int): + import torch + + u = packed.contiguous().view(torch.uint8) + lut = torch.tensor(_FP4_E2M1_VALUES, device=packed.device, dtype=torch.float32) + out = torch.empty((rows, cols), device=packed.device, dtype=torch.float32) + out[:, 0::2] = lut[(u & 0x0F).long()] + out[:, 1::2] = lut[((u >> 4) & 0x0F).long()] + return out + + +def _make_packed_fp4_inputs(M: int, N: int, K: int, input_mode: str): + import torch + + if input_mode == "constant": + a = torch.full((M, K // 2), 0x22, device="cuda", dtype=torch.uint8).view(torch.int8) + b = torch.full((N, K // 2), 0x22, device="cuda", dtype=torch.uint8).view(torch.int8) + elif input_mode == "random": + a = torch.randint(-128, 128, (M, K // 2), device="cuda", dtype=torch.int8) + b = torch.randint(-128, 128, (N, K // 2), device="cuda", dtype=torch.int8) + elif input_mode == "a_random_b_alternating": + a = torch.randint(-128, 128, (M, K // 2), device="cuda", dtype=torch.int8) + b = torch.full((N, K // 2), 0x21, device="cuda", dtype=torch.uint8).view(torch.int8) + elif input_mode == "a_constant_b_random": + a = torch.full((M, K // 2), 0x22, device="cuda", dtype=torch.uint8).view(torch.int8) + b = torch.randint(-128, 128, (N, K // 2), device="cuda", dtype=torch.int8) + else: + raise ValueError(f"Unsupported input_mode={input_mode!r}") + return a, b + + +def _make_constant_scale_words(rows: int, K: int, byte: int = 0x38): + import torch + + word = byte | (byte << 8) | (byte << 16) | (byte << 24) + return torch.full((rows, K // 64), word, device="cuda", dtype=torch.uint32) + + +def _pack_scale_words(scale_bytes): + import torch + + scale_i64 = scale_bytes.to(torch.int64).reshape(scale_bytes.shape[0], -1, 4) + word = scale_i64[:, :, 0] + word = word | (scale_i64[:, :, 1] << 8) + word = word | (scale_i64[:, :, 2] << 16) + word = word | (scale_i64[:, :, 3] << 24) + return word.to(torch.uint32).contiguous() + + +def _make_varying_power_of_two_scale_words(rows: int, K: int): + import torch + + scale_choices = torch.tensor([0x30, 0x38, 0x40], device="cuda", dtype=torch.uint8) + row = torch.arange(rows, device="cuda", dtype=torch.int64)[:, None] + col = torch.arange(K // 16, device="cuda", dtype=torch.int64)[None, :] + scale_bytes = scale_choices[(row + 2 * col) % scale_choices.numel()] + return _pack_scale_words(scale_bytes), scale_bytes + + +def _decode_ue4m3_scale_bytes(scale_bytes): + import torch + + u = scale_bytes.to(torch.int32) + exponent = (u >> 3) & 0x0F + mantissa = u & 0x07 + normal = (1.0 + mantissa.to(torch.float32) / 8.0) * torch.pow(2.0, exponent.to(torch.float32) - 7.0) + subnormal = (mantissa.to(torch.float32) / 8.0) * torch.pow(torch.tensor(2.0, device=scale_bytes.device), -6.0) + return torch.where(exponent == 0, subnormal, normal) + + +def _reference_constant_scale_gemm(A, B, M: int, N: int, K: int): + a_f32 = _decode_rowmajor_fp4(A, M, K) + b_f32 = _decode_rowmajor_fp4(B, N, K) + return a_f32 @ b_f32.T + + +def _reference_blockscaled_gemm(A, B, SFA, SFB, M: int, N: int, K: int): + a_f32 = _decode_rowmajor_fp4(A, M, K) + b_f32 = _decode_rowmajor_fp4(B, N, K) + sfa = _decode_ue4m3_scale_bytes(SFA).repeat_interleave(16, dim=1) + sfb = _decode_ue4m3_scale_bytes(SFB).repeat_interleave(16, dim=1) + return (a_f32 * sfa) @ (b_f32 * sfb).T + + +def test_nvf4_mma_block_scale_fragment_layouts_match_cute(): + # CUTLASS/CuTe SM120 ALayout: + # Layout, Shape<_8,_2,_2>>, + # Stride, Stride<_16,_8,_512>>>. + seen = set() + for thread_id in range(32): + for local_id in range(32): + coord = mma_load_a_32x32_to_shared_16x64_layout(thread_id, local_id) + assert 0 <= coord[0] < 16 + assert 0 <= coord[1] < 64 + seen.add(coord) + assert len(seen) == 16 * 64 + + assert mma_load_a_32x32_to_shared_16x64_layout(0, 0) == (0, 0) + assert mma_load_a_32x32_to_shared_16x64_layout(0, 8) == (8, 0) + assert mma_load_a_32x32_to_shared_16x64_layout(31, 31) == (15, 63) + + +def test_nvf4_mma_block_scale_lane_scale_mapping_matches_cute(): + sfa_rows = [TensorCoreIntrinEmitterWithBlockScale._sfa_row_in_atom(tx) for tx in range(32)] + sfb_cols = [TensorCoreIntrinEmitterWithBlockScale._sfb_col_in_atom(tx) for tx in range(32)] + + assert sfa_rows == [ + 0, + 8, + 0, + 8, + 1, + 9, + 1, + 9, + 2, + 10, + 2, + 10, + 3, + 11, + 3, + 11, + 4, + 12, + 4, + 12, + 5, + 13, + 5, + 13, + 6, + 14, + 6, + 14, + 7, + 15, + 7, + 15, + ] + assert sfb_cols == [ + 0, + 0, + 0, + 0, + 1, + 1, + 1, + 1, + 2, + 2, + 2, + 2, + 3, + 3, + 3, + 3, + 4, + 4, + 4, + 4, + 5, + 5, + 5, + 5, + 6, + 6, + 6, + 6, + 7, + 7, + 7, + 7, + ] + + +@pytest.mark.parametrize( + "kwargs", + [ + {"kind": "mxf4nvf4", "scale_vec_size": 2, "stype": "ue4m3"}, + {"kind": "mxf4nvf4", "scale_vec_size": 4, "stype": "ue8m0"}, + {"kind": "mxf4", "scale_vec_size": 4, "stype": "ue4m3"}, + ], +) +def test_nvf4_mma_block_scale_rejects_unsupported_configs(kwargs): + with pytest.raises(ValueError, match="Unsupported SM120 block-scale MMA config"): + TensorCoreIntrinEmitterWithBlockScale( + a_dtype=T.float4_e2m1fn, + b_dtype=T.float4_e2m1fn, + accum_dtype=T.float32, + a_transposed=False, + b_transposed=True, + block_row_warps=2, + block_col_warps=2, + warp_row_tiles=32, + warp_col_tiles=32, + chunk=256, + **kwargs, + ) + + +@pytest.mark.parametrize( + "dtype_kwargs", + [ + {"a_dtype": T.float16, "b_dtype": T.float4_e2m1fn, "accum_dtype": T.float32}, + {"a_dtype": T.float4_e2m1fn, "b_dtype": T.float16, "accum_dtype": T.float32}, + {"a_dtype": T.float4_e2m1fn, "b_dtype": T.float4_e2m1fn, "accum_dtype": T.float16}, + ], +) +def test_nvf4_mma_block_scale_rejects_incompatible_dtypes(dtype_kwargs): + with pytest.raises(ValueError, match="mxf4nvf4 expects"): + TensorCoreIntrinEmitterWithBlockScale( + a_transposed=False, + b_transposed=True, + block_row_warps=2, + block_col_warps=2, + warp_row_tiles=32, + warp_col_tiles=32, + chunk=256, + kind="mxf4nvf4", + scale_vec_size=4, + stype="ue4m3", + **dtype_kwargs, + ) + + +def test_sm120_fulltile_package_contract_matches_current_pingpong_path(): + contract = _make_sm120_fulltile_contract() + + assert (contract.tile_m, contract.tile_n, contract.tile_k) == (128, 128, 256) + assert (contract.block_row_warps, contract.block_col_warps) == (2, 2) + assert (contract.warp_rows, contract.warp_cols) == (4, 4) + assert (contract.warp_row_tiles, contract.warp_col_tiles) == (64, 64) + assert (contract.kblocks, contract.micro_size_k) == (4, 64) + assert (contract.scale_package_words_sfa, contract.scale_package_words_sfb) == (2, 2) + assert (contract.issue_count_per_warp, contract.issue_count_per_warpgroup) == (32, 128) + + +def test_sm120_fulltile_package_contract_describes_compact_selector_copy_view(): + contract = _make_sm120_fulltile_contract() + + assert contract.compact_selector_scale_rows(lane=0, warp_m=0, warp_n=0) == ((0, 32), (0, 32)) + assert contract.compact_selector_scale_rows(lane=1, warp_m=0, warp_n=0) == ((8, 40), (8, 40)) + assert contract.compact_selector_scale_rows(lane=2, warp_m=1, warp_n=1) == ((80, 112), (80, 112)) + assert contract.compact_selector_scale_rows(lane=31, warp_m=1, warp_n=1) == ((95, 127), (95, 127)) + + +@pytest.mark.parametrize("lane", [0, 1, 2, 17, 31]) +@pytest.mark.parametrize("warp_m", [0, 1]) +@pytest.mark.parametrize("warp_n", [0, 1]) +@pytest.mark.parametrize("kblock", [0, 2, 3]) +def test_sm120_fulltile_package_contract_word_offsets_match_source_layout_oracle(lane, warp_m, warp_n, kblock): + contract = _make_sm120_fulltile_contract() + + sfa_rows, sfb_rows = contract.compact_selector_scale_rows(lane=lane, warp_m=warp_m, warp_n=warp_n) + expected_sfa = tuple(_oracle_blockscaled_chunk_kmajor_flat_word(row, kblock) for row in sfa_rows) + expected_sfb = tuple(_oracle_blockscaled_chunk_kmajor_flat_word(row, kblock) for row in sfb_rows) + + assert contract.compact_selector_scale_word_offsets(lane=lane, warp_m=warp_m, warp_n=warp_n, kblock=kblock) == ( + expected_sfa, + expected_sfb, + ) + + +def test_sm120_fulltile_package_contract_describes_pingpong_lifecycle(): + contract = _make_sm120_fulltile_contract() + + assert contract.package_pingpong_lifecycle() == ( + ("copy", 0, 0), + ("copy", 1, 1), + ("gemm", 0, 0), + ("copy", 0, 2), + ("gemm", 1, 1), + ("copy", 1, 3), + ("gemm", 0, 2), + ("gemm", 1, 3), + ) + + +def test_sm120_fulltile_package_contract_pingpong_lifecycle_is_data_ready(): + contract = _make_sm120_fulltile_contract() + + package_kblock = {} + gemmed_kblocks = [] + for op, package_id, kblock in contract.package_pingpong_lifecycle(): + if op == "copy": + package_kblock[package_id] = kblock + elif op == "gemm": + assert package_kblock[package_id] == kblock + gemmed_kblocks.append(kblock) + else: + raise AssertionError(f"unexpected package lifecycle op: {op}") + + assert gemmed_kblocks == [0, 1, 2, 3] + + +def test_sm120_fulltile_package_contract_describes_omma_sf_issue_schedule(): + contract = _make_sm120_fulltile_contract() + schedule = contract.omma_sf_issue_schedule_per_warp() + + assert len(schedule) == contract.issue_count_per_warp + assert schedule[0] == (0, 0, 0, 0, 0, 0, 0) + assert schedule[1] == (0, 0, 1, 0, 0, 0, 1) + assert schedule[2] == (0, 1, 0, 0, 0, 0, 2) + assert schedule[3] == (0, 1, 1, 0, 0, 0, 3) + assert schedule[16] == (2, 0, 0, 1, 0, 0, 0) + assert schedule[-1] == (3, 3, 1, 1, 1, 1, 3) + + for mma_i, mma_j, n8_half, sfa_word, sfb_word, scale_a_tid, scale_b_tid in schedule: + assert sfa_word == (0 if mma_i < 2 else 1) + assert sfb_word == (0 if mma_j < 2 else 1) + assert scale_a_tid == (mma_i & 1) + assert scale_b_tid == (mma_j & 1) * 2 + n8_half + + assert sorted({issue[5] for issue in schedule}) == [0, 1] + assert sorted({issue[6] for issue in schedule}) == [0, 1, 2, 3] + assert sum(1 for issue in schedule if issue[3] == 0) == 16 + assert sum(1 for issue in schedule if issue[3] == 1) == 16 + assert sum(1 for issue in schedule if issue[4] == 0) == 16 + assert sum(1 for issue in schedule if issue[4] == 1) == 16 + + +def test_sm120_fulltile_package_contract_rejects_shape_drift(): + emitter = TensorCoreIntrinEmitterWithBlockScale( + a_dtype=T.float4_e2m1fn, + b_dtype=T.float4_e2m1fn, + accum_dtype=T.float32, + a_transposed=False, + b_transposed=True, + block_row_warps=2, + block_col_warps=4, + warp_row_tiles=64, + warp_col_tiles=64, + chunk=256, + reduce_k=1, + num_elems_per_byte=2, + kind="mxf4nvf4", + scale_vec_size=4, + stype="ue4m3", + ) + + with pytest.raises(ValueError, match="128x128x256"): + SM120BlockScaledFullTilePackageContract.for_package_pingpong( + emitter, + sf_layout="blockscaled_chunk_kmajor", + ) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +@pytest.mark.parametrize("K", [64, 128, 256]) +def test_nvf4_mma_block_scale_codegen(K): + kernel = tilelang.compile( + _make_nvf4_matmul_codegen_kernel(128, 128, K), + target="cuda", + out_idx=[4], + ) + src = kernel.get_kernel_source() + assert "#include " in src + assert "#include " not in src + assert "sm120_mma_sync_blockscaled" in src + assert "SFA_shared" in src + assert "SFB_shared" in src + assert "scale_a_local" not in src + assert "scale_b_local" not in src + assert "SM120MmaBlockScaledKind::kMxf4nvf4" in src + assert "SM120MmaScaleType::kUE4M3" in src + fp4_tile_bytes = 128 * K // 2 + sf_tile_bytes = 128 * (K // 64) * 4 + assert f"void* B_shared = ((void*)((char*)buf_dyn_shmem + {fp4_tile_bytes}));" in src + assert f"void* SFA_shared = ((void*)((char*)buf_dyn_shmem + {2 * fp4_tile_bytes}));" in src + assert f"void* SFB_shared = ((void*)((char*)buf_dyn_shmem + {2 * fp4_tile_bytes + sf_tile_bytes}));" in src + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +def test_nvf4_mma_block_scale_rejects_unknown_micro_pipeline(): + with pytest.raises(ValueError, match="Unsupported block-scaled micro_pipeline"): + tilelang.compile( + _make_nvf4_matmul_codegen_kernel( + 128, + 128, + 256, + warp_row_tiles=64, + warp_col_tiles=64, + micro_pipeline="sm120_backend_kblock_fulltile_unknown", + sf_layout="blockscaled_chunk_kmajor", + ), + target="cuda", + out_idx=[4], + ) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +def test_nvf4_mma_block_scale_rejects_legacy_cutlass_128x4_layout_alias(): + with pytest.raises(ValueError, match="Unsupported SM120 scale layout: cutlass_128x4"): + tilelang.compile( + _make_nvf4_matmul_codegen_kernel( + 128, + 128, + 256, + warp_row_tiles=64, + warp_col_tiles=64, + micro_pipeline="sm120_backend_kblock_fulltile_package_pingpong", + sf_layout="cutlass_128x4", + ), + target="cuda", + out_idx=[4], + ) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +def test_nvf4_mma_block_scale_package_pingpong_contract_lowers_fulltile(): + kernel = tilelang.compile( + _make_nvf4_matmul_codegen_kernel( + 128, + 128, + 256, + warp_row_tiles=64, + warp_col_tiles=64, + micro_pipeline="sm120_backend_kblock_fulltile_package_pingpong", + sf_layout="blockscaled_chunk_kmajor", + ), + target="cuda", + out_idx=[4], + ) + + src = kernel.get_kernel_source() + assert "sm120_mma_blockscaled_kblock_fulltile_package_pingpong" in src + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +def test_nvf4_mma_block_scale_packed_smem_offsets(): + kernel = tilelang.compile( + _make_nvf4_matmul_codegen_kernel(256, 256, 256, num_stages=3), + target="cuda", + out_idx=[4], + ) + src = kernel.get_kernel_source() + assert "void* A_shared = ((void*)((char*)buf_dyn_shmem + 0));" in src + assert "void* B_shared = ((void*)((char*)buf_dyn_shmem + 24576));" in src + assert "void* SFA_shared = ((void*)((char*)buf_dyn_shmem + 49152));" in src + assert "void* SFB_shared = ((void*)((char*)buf_dyn_shmem + 52224));" in src + + +def test_nvf4_mma_block_scale_packed_smem_non_alias_offset_units(): + @T.prim_func + def before( + A: T.Tensor((16,), T.float4_e2m1fn), + B: T.Tensor((16,), T.float4_e2m1fn), + ): + with T.Kernel(1, threads=32): + a = T.alloc_shared((3,), T.float4_e2m1fn) + b = T.alloc_shared((4,), T.float4_e2m1fn) + a[0] = A[0] + b[0] = A[1] + B[0] = a[0] + B[1] = b[0] + + mod = tvm.IRModule.from_expr( + before.with_attr("global_symbol", "main").with_attr( + "target", + tvm.target.Target("webgpu"), + ) + ) + mod = tilelang.transform.LowerOpaqueBlock()(mod) + mod = tilelang.transform.FlattenBuffer()(mod) + mod = tilelang.transform.Simplify()(mod) + mod = tilelang.transform.MergeSharedMemoryAllocations()(mod) + + src = mod.script() + # Three logical FP4 values use two physical bytes. After 16-byte alignment, + # the next buffer starts at logical FP4 offset 32, not byte offset 16. + assert "b[32]" in src + assert "b[16]" not in src + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +@pytest.mark.parametrize( + "K,input_mode", + [ + (64, "random"), + (128, "a_random_b_alternating"), + (256, "a_constant_b_random"), + (256, "random"), + ], +) +def test_nvf4_mma_block_scale_constant_scale_correctness(K, input_mode): + import torch + + torch.manual_seed(0) + M = N = 128 + kernel = tilelang.compile( + _make_nvf4_matmul_codegen_kernel(M, N, K), + target="cuda", + out_idx=[4], + ) + + A, B = _make_packed_fp4_inputs(M, N, K, input_mode) + SFA = _make_constant_scale_words(M, K) + SFB = _make_constant_scale_words(N, K) + + C = kernel(A, B, SFA, SFB) + ref = _reference_constant_scale_gemm(A, B, M, N, K) + torch.testing.assert_close(C, ref, rtol=0.0, atol=0.0) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(12, 0) +def test_nvf4_mma_block_scale_varying_scale_correctness(): + import torch + + torch.manual_seed(0) + M = N = 128 + K = 128 + kernel = tilelang.compile( + _make_nvf4_matmul_codegen_kernel(M, N, K), + target="cuda", + out_idx=[4], + ) + + A, B = _make_packed_fp4_inputs(M, N, K, "random") + SFA, sfa_bytes = _make_varying_power_of_two_scale_words(M, K) + SFB, sfb_bytes = _make_varying_power_of_two_scale_words(N, K) + + C = kernel(A, B, SFA, SFB) + ref = _reference_blockscaled_gemm(A, B, sfa_bytes, sfb_bytes, M, N, K) + torch.testing.assert_close(C, ref, rtol=0.0, atol=0.0) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py b/testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py new file mode 100644 index 000000000..672d32c0d --- /dev/null +++ b/testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py @@ -0,0 +1,47 @@ +import importlib.util +import sys +from pathlib import Path + + +def _load_sm120_example(monkeypatch): + repo_root = Path(__file__).resolve().parents[3] + example = repo_root / "examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py" + monkeypatch.setattr(sys, "argv", [str(example)]) + spec = importlib.util.spec_from_file_location("sm120_nvfp4_blockscaled_gemm_example", example) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_sm120_nvfp4_example_cli_keeps_internal_strategy_fixed(monkeypatch): + module = _load_sm120_example(monkeypatch) + args = module.parse_args() + + assert not hasattr(args, "micro_pipeline") + assert not hasattr(args, "sf_layout") + assert not hasattr(args, "scale_storage_layout") + assert not hasattr(args, "manual_ws2") + + assert module._SM120_MICRO_PIPELINE == "sm120_backend_kblock_fulltile_package_pingpong" + assert module._SM120_SCALE_LAYOUT == "blockscaled_chunk_kmajor" + + +def test_sm120_nvfp4_example_source_has_no_legacy_strategy_flags(): + repo_root = Path(__file__).resolve().parents[3] + source = (repo_root / "examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py").read_text() + + forbidden_flags = [ + "--micro-pipeline", + "--manual-ws2", + "--manual-ws2-split", + "--manual-ws2-sf-load", + "--manual-ws2-sf-layout", + "--manual-ws2-ab-shared-storage", + "--manual-ws2-ab-copy-view", + "--scale-storage-layout", + "--sf-layout", + "--path", + ] + for flag in forbidden_flags: + assert flag not in source diff --git a/testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py b/testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py new file mode 100644 index 000000000..08514ae73 --- /dev/null +++ b/testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py @@ -0,0 +1,745 @@ +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +tilelang_testing = pytest.importorskip("tilelang.testing") + +from tilelang.quantize import nvfp4 as nvfp4_utils +from tilelang.quantize import ( + blockscaled_chunk_kmajor_word_offset, + decode_packed_fp4_e2m1, + decode_ue4m3_scale_bytes, + encode_fp4_e2m1_values, + encode_ue4m3_scale_bytes, + pack_blockscaled_chunk_kmajor_scale_bytes, + pack_nvfp4_scale_bytes, + tilelang_quantize_bf16_to_nvfp4_blockscaled, + quantize_bf16_to_nvfp4_blockscaled, + quantize_nvfp4_blockscaled, + swizzle_blockscaled_chunk_kmajor_scale_words, + unswizzle_blockscaled_chunk_kmajor_scale_words, +) + + +def _pack_semantic_scale_words(scale_bytes): + scale_i64 = scale_bytes.to(torch.int64).reshape(scale_bytes.shape[0], scale_bytes.shape[1] // 4, 4) + words = scale_i64[:, :, 0] + words = words | (scale_i64[:, :, 1] << 8) + words = words | (scale_i64[:, :, 2] << 16) + words = words | (scale_i64[:, :, 3] << 24) + return words.to(torch.uint32).contiguous() + + +def _expected_blockscaled_chunk_kmajor_byte_location(row: int, k16_idx: int, k64_cols: int = 4) -> tuple[int, int, int]: + k64_word = k16_idx // 4 + byte_lane = k16_idx % 4 + row_block = row // 128 + row_in_block = row % 128 + flat_word = row_block * 128 * k64_cols + k64_word * 128 + (row_in_block % 32) * 4 + (row_in_block // 32) + physical_row = flat_word // k64_cols + physical_word = flat_word % k64_cols + return physical_row, physical_word, byte_lane + + +def _flat_word_offset(row: int, k64_word: int) -> int: + physical_row, physical_word = blockscaled_chunk_kmajor_word_offset(row, k64_word) + return physical_row * 4 + physical_word + + +def _packed_byte(packed, row: int, word: int, byte_lane: int) -> int: + return (int(packed[row, word].item()) >> (8 * byte_lane)) & 0xFF + + +def _packed_word(packed, logical_row: int, k64_word: int) -> int: + k64_cols = packed.shape[1] + row_in_block = logical_row % 128 + flat_word = (logical_row // 128) * 128 * k64_cols + k64_word * 128 + (row_in_block % 32) * 4 + (row_in_block // 32) + return int(packed.reshape(-1)[flat_word].item()) + + +def _sfa_row_in_lane(lane: int) -> int: + return nvfp4_utils._sm120_sfa_row_in_lane(lane) + + +def _sfb_col_in_lane(lane: int) -> int: + return nvfp4_utils._sm120_sfb_col_in_lane(lane) + + +def _sfa_selector_source_lane(lane: int, scale_a_thread_id: int) -> int: + return nvfp4_utils._sm120_sfa_selector_source_lane(lane, scale_a_thread_id) + + +def _sfb_selector_source_lane(lane: int, scale_b_thread_id: int) -> int: + return nvfp4_utils._sm120_sfb_selector_source_lane(lane, scale_b_thread_id) + + +def _current_compact_sfa_row(lane: int, warp_m: int, mma_i: int) -> int: + return nvfp4_utils._sm120_compact_sfa_effective_row(lane, warp_m, mma_i) + + +def _current_compact_sfb_row(lane: int, warp_n: int, mma_j: int, half: int) -> int: + return nvfp4_utils._sm120_compact_sfb_effective_row(lane, warp_n, mma_j, half) + + +def _compact_issue_contract(lane: int, warp_m: int, warp_n: int): + return nvfp4_utils._sm120_compact_scale_issue_contract(lane, warp_m, warp_n) + + +def _compact_copy_view_contract(lane: int, warp_m: int, warp_n: int, k64_word: int): + return nvfp4_utils._sm120_compact_scale_copy_view_contract(lane, warp_m, warp_n, k64_word) + + +def _compact_producer_contract(warp_m: int, warp_n: int): + return nvfp4_utils._sm120_compact_scale_producer_contract(warp_m, warp_n) + + +def _compact_package_consumer_byte_offsets(lane: int, warp_m: int, warp_n: int, k64_word: int): + return nvfp4_utils._sm120_compact_package_consumer_byte_offsets(lane, warp_m, warp_n, k64_word) + + +def _cutlass_issue_site_contract(site: int): + return nvfp4_utils._sm120_cutlass_issue_site_contract(site) + + +def _tilelang_current_issue_site_contract(site: int): + return nvfp4_utils._sm120_tilelang_current_issue_site_contract(site) + + +def _unpack_with_oracle(packed, rows: int, k16_cols: int): + out = torch.empty((rows, k16_cols), dtype=torch.uint8) + k64_cols = k16_cols // 4 + for row in range(rows): + for k16_idx in range(k16_cols): + physical_row, physical_word, byte_lane = _expected_blockscaled_chunk_kmajor_byte_location(row, k16_idx, k64_cols) + out[row, k16_idx] = _packed_byte(packed, physical_row, physical_word, byte_lane) + return out + + +def test_blockscaled_chunk_kmajor_word_offset_fixed_cases(): + expected = { + (0, 0): (0, 0), + (0, 1): (32, 0), + (0, 2): (64, 0), + (0, 3): (96, 0), + (31, 0): (31, 0), + (32, 0): (0, 1), + (63, 1): (63, 1), + (64, 2): (64, 2), + (96, 3): (96, 3), + (127, 3): (127, 3), + } + for (row, k64_word), physical in expected.items(): + assert blockscaled_chunk_kmajor_word_offset(row, k64_word) == physical + + +def test_sm120_contract_helpers_stay_internal_to_nvfp4_module(): + import tilelang.quantize as quantize + + assert hasattr(nvfp4_utils, "_sm120_compact_scale_issue_contract") + assert hasattr(nvfp4_utils, "_sm120_compact_scale_copy_view_contract") + assert not hasattr(quantize, "_sm120_compact_scale_issue_contract") + assert not hasattr(quantize, "_sm120_compact_scale_copy_view_contract") + + +def test_blockscaled_chunk_kmajor_byte_word_and_flat_offsets(): + # CUTLASS documents BlockScaledBasicChunk at scale-byte granularity. TileLang + # stores four K/16 scale bytes in one uint32, so the CUDA helper consumes the + # corresponding flat word offset. + byte_offsets_for_k16_zero = { + 0: 0, + 32: 4, + 64: 8, + 96: 12, + 127: 508, + } + for row, byte_offset in byte_offsets_for_k16_zero.items(): + physical_row, physical_word, byte_lane = _expected_blockscaled_chunk_kmajor_byte_location(row, 0) + assert byte_lane == 0 + assert physical_row * 16 + physical_word * 4 + byte_lane == byte_offset + assert _flat_word_offset(row, 0) == byte_offset // 4 + + assert [_flat_word_offset(row, 0) for row in (0, 32, 64, 96)] == [0, 1, 2, 3] + assert [_flat_word_offset(row, 0) * 4 for row in (0, 32, 64, 96)] == [0, 4, 8, 12] + + +def test_blockscaled_chunk_kmajor_explains_compact_source_address_examples(): + # These are source-layout examples only: they show why BlockScaledBasicChunk + # K-major can produce compact scale-word addresses such as the CuTe-style + # examples A=[0,4,8,12] and B=[0,4,32,36]. A real OMMA.SF scale package must + # additionally prove that the selected source lanes name the right semantic + # SFA/SFB rows for every issue. + sfa_source_rows = [0, 32, 64, 96] + sfb_source_rows = [0, 32, 2, 34] + + assert [_flat_word_offset(row, 0) * 4 for row in sfa_source_rows] == [0, 4, 8, 12] + assert [_flat_word_offset(row, 0) * 4 for row in sfb_source_rows] == [0, 4, 32, 36] + + # The production compact-selector package still names the rows selected by + # OMMA.SF's thread-id network. For lane 0 this is a different access pattern; + # the source layout is already swizzled, but a CuTe-style scale slot package is + # still a separate lowering contract. + current_lane0_sfa_rows = [_current_compact_sfa_row(0, 0, mma_i) for mma_i in range(4)] + current_lane0_sfb_rows = [_current_compact_sfb_row(0, 0, mma_j, 0) for mma_j in range(4)] + + assert current_lane0_sfa_rows == [0, 16, 32, 48] + assert [_flat_word_offset(row, 0) * 4 for row in current_lane0_sfa_rows] == [0, 256, 4, 260] + assert current_lane0_sfb_rows == [0, 16, 32, 48] + assert [_flat_word_offset(row, 0) * 4 for row in current_lane0_sfb_rows] == [0, 256, 4, 260] + + +def test_compact_package_producer_contract_inverts_omma_sf_source_lane_selectors(): + producer_rows = _compact_producer_contract(warp_m=0, warp_n=0) + + sfa_issue0 = [row for row in producer_rows if row["kind"] == "SFA" and row["mma_i"] == 0 and row["producer_lane"] < 4] + sfa_issue1 = [row for row in producer_rows if row["kind"] == "SFA" and row["mma_i"] == 1 and row["producer_lane"] < 4] + assert [(row["producer_lane"], row["consumers"], row["row"]) for row in sfa_issue0] == [ + (0, (0, 2), 0), + (1, (1, 3), 8), + ] + assert [(row["producer_lane"], row["consumers"], row["row"]) for row in sfa_issue1] == [ + (2, (0, 2), 16), + (3, (1, 3), 24), + ] + + sfb_issue0 = [row for row in producer_rows if row["kind"] == "SFB" and row["mma_j"] == 0 and row["producer_lane"] < 4] + assert [(row["producer_lane"], row["consumers"], row["row"]) for row in sfb_issue0] == [ + (0, (0, 1, 2, 3), 0), + (1, (0, 1, 2, 3), 8), + ] + + # The inverse contract is the semantic constraint a future scale package must + # keep. Compact source addresses alone are insufficient unless they are tied + # back to these producer-lane rows. + for row in producer_rows: + if row["kind"] == "SFA": + for consumer_lane in row["consumers"]: + assert _current_compact_sfa_row(consumer_lane, 0, row["mma_i"]) == row["row"] + else: + for consumer_lane in row["consumers"]: + assert _current_compact_sfb_row(consumer_lane, 0, row["mma_j"], row["half"]) == row["row"] + + +def test_current_register_package_shape_cannot_be_relabelled_as_cute_compact_slots(): + # With the current issue order and selector ids, lane 0 consumes scale rows + # through the existing sa0/sa1 and sb0/sb1 slots. These are the true offsets + # seen by the current OMMA.SF package after applying the source-lane selector. + current_offsets = _compact_package_consumer_byte_offsets(lane=0, warp_m=0, warp_n=0, k64_word=0) + + assert current_offsets["SFA"] == (0, 256, 4, 260) + assert current_offsets["SFB"] == (0, 128, 256, 384, 4, 132, 260, 388) + + # The compact source-address examples are real properties of the swizzled + # source layout, but they are not reachable by merely changing the current + # smem load offsets while preserving this package's semantic rows. + assert current_offsets["SFA"] != (0, 4, 8, 12) + assert current_offsets["SFB"][:4] != (0, 4, 32, 36) + + +def test_cutlass_fulltile_issue_site_contract_matches_known_static_dump_points(): + expected = { + 0: (0, 0, 0, 0, 0, 0, 0, 0), + 1: (0, 1, 0, 0, 4, 0, 0, 4), + 3: (0, 3, 0, 0, 12, 0, 0, 12), + 4: (0, 3, 0, 1, 12, 16, 64, 28), + 7: (0, 0, 0, 1, 0, 16, 64, 16), + 8: (0, 0, 1, 0, 0, 4, 128, 32), + 31: (0, 0, 3, 1, 0, 28, 448, 112), + 32: (1, 0, 0, 0, 16, 32, 16, 0), + 63: (1, 0, 3, 1, 16, 60, 464, 112), + 96: (3, 0, 0, 0, 48, 96, 48, 0), + 127: (3, 0, 3, 1, 48, 124, 496, 112), + } + for site, values in expected.items(): + issue = _cutlass_issue_site_contract(site) + assert ( + issue["k_block"], + issue["m_atom"], + issue["n16_col"], + issue["n8_half"], + issue["SFA_slot0"], + issue["SFB_slot0"], + issue["B_slot0"], + issue["C_slot0"], + ) == values + + +def test_current_tilelang_issue_order_differs_from_cutlass_serpentine_order(): + # Sites 0..3 agree on m order. Starting at site 4 CUTLASS reverses m_atom for + # the odd N8 half, while the current helper continues row-major over I. + assert [_cutlass_issue_site_contract(site)["m_atom"] for site in range(8)] == [0, 1, 2, 3, 3, 2, 1, 0] + assert [_tilelang_current_issue_site_contract(site)["mma_i"] for site in range(8)] == [0] * 8 + + assert [_cutlass_issue_site_contract(site)["SFA_slot0"] for site in range(8)] == [0, 4, 8, 12, 12, 8, 4, 0] + assert [_tilelang_current_issue_site_contract(site)["sa_tid"] for site in range(8)] == [0] * 8 + + +def test_sfa_compact_source_example_spans_warp_m_groups_not_one_consumer_package(): + # The SFA [0,4,8,12] byte-address example corresponds to rows + # [0,32,64,96]. For the current TileLang warp decomposition, lane 0 reaches + # those rows only by grouping the same producer-lane/reg-group slots across + # warp_m=0 and warp_m=1. A single warp_m consumer issue sequence remains + # [0,16,32,48] and therefore has the non-compact offsets checked above. + compact_rows = [ + _current_compact_sfa_row(lane=0, warp_m=0, mma_i=0), + _current_compact_sfa_row(lane=0, warp_m=0, mma_i=2), + _current_compact_sfa_row(lane=0, warp_m=1, mma_i=0), + _current_compact_sfa_row(lane=0, warp_m=1, mma_i=2), + ] + assert compact_rows == [0, 32, 64, 96] + assert [_flat_word_offset(row, 0) * 4 for row in compact_rows] == [0, 4, 8, 12] + + one_warp_rows = [_current_compact_sfa_row(lane=0, warp_m=0, mma_i=mma_i) for mma_i in range(4)] + assert one_warp_rows == [0, 16, 32, 48] + assert [_flat_word_offset(row, 0) * 4 for row in one_warp_rows] == [0, 256, 4, 260] + + +def test_sfb_compact_source_example_spans_b_operand_lane_groups(): + # The SFB [0,4,32,36] byte-address example corresponds to rows + # [0,32,2,34]. In the current semantic model those are the two SFB reg groups + # owned by producer lanes 0 and 8. This matches the B operand's lane-grouped + # view: lanes separated by 8 advance the B column-in-lane coordinate while + # keeping qlane fixed. + compact_rows = [ + _current_compact_sfb_row(lane=0, warp_n=0, mma_j=0, half=0), + _current_compact_sfb_row(lane=0, warp_n=0, mma_j=2, half=0), + _current_compact_sfb_row(lane=8, warp_n=0, mma_j=0, half=0), + _current_compact_sfb_row(lane=8, warp_n=0, mma_j=2, half=0), + ] + assert compact_rows == [0, 32, 2, 34] + assert [_flat_word_offset(row, 0) * 4 for row in compact_rows] == [0, 4, 32, 36] + + one_consumer_rows = [_current_compact_sfb_row(lane=0, warp_n=0, mma_j=mma_j, half=0) for mma_j in range(4)] + assert one_consumer_rows == [0, 16, 32, 48] + assert [_flat_word_offset(row, 0) * 4 for row in one_consumer_rows] == [0, 256, 4, 260] + + +def test_pack_blockscaled_chunk_kmajor_scale_bytes_fixed_byte_offsets(): + rows = 256 + k16_cols = 32 + scale_bytes = torch.zeros((rows, k16_cols), dtype=torch.uint8) + cases = [ + (0, 0, 0x11), + (32, 0, 0x22), + (64, 8, 0x33), + (96, 12, 0x44), + (127, 15, 0x55), + (128, 16, 0x66), + (159, 31, 0x77), + (255, 27, 0x88), + ] + for row, k16_idx, value in cases: + scale_bytes[row, k16_idx] = value + + packed = pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes) + + assert packed.shape == (rows, k16_cols // 4) + assert packed.dtype == torch.uint32 + for row, k16_idx, value in cases: + physical_row, physical_word, byte_lane = _expected_blockscaled_chunk_kmajor_byte_location(row, k16_idx, packed.shape[1]) + assert _packed_byte(packed, physical_row, physical_word, byte_lane) == value + + +def test_pack_blockscaled_chunk_kmajor_scale_bytes_random_binary_512x32_matches_oracle(): + rows = 512 + k16_cols = 512 // 16 + generator = torch.Generator(device="cpu").manual_seed(17) + scale_bytes = torch.randint(0, 2, (rows, k16_cols), generator=generator, dtype=torch.uint8) * 0x38 + + packed = pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes) + + assert packed.shape == (rows, k16_cols // 4) + assert packed.dtype == torch.uint32 + assert torch.equal(_unpack_with_oracle(packed, rows, k16_cols), scale_bytes) + + +def test_pack_blockscaled_chunk_kmajor_scale_bytes_matches_word_swizzle(): + rows = 512 + k16_cols = 512 // 16 + scale_bytes = torch.arange(rows * k16_cols, dtype=torch.uint8).reshape(rows, k16_cols) + + semantic_words = _pack_semantic_scale_words(scale_bytes) + assert torch.equal(pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes), swizzle_blockscaled_chunk_kmajor_scale_words(semantic_words)) + + +def test_pack_nvfp4_scale_bytes_default_matches_blockscaled_chunk_kmajor_layout(): + rows = 512 + k16_cols = 512 // 16 + scale_bytes = torch.arange(rows * k16_cols, dtype=torch.uint8).reshape(rows, k16_cols) + + expected = pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes) + actual = pack_nvfp4_scale_bytes(scale_bytes) + + assert torch.equal(actual, expected) + + +def test_blockscaled_chunk_kmajor_matches_current_omma_sf_selector_contract(): + rows = 128 + k = 256 + k64_words = k // 64 + semantic_words = torch.empty((rows, k64_words), dtype=torch.uint32) + for row in range(rows): + for k64_word in range(k64_words): + semantic_words[row, k64_word] = row * 1000 + k64_word + + packed = swizzle_blockscaled_chunk_kmajor_scale_words(semantic_words) + + # Current TileLang package keeps two SFA words and two SFB words per lane. + # OMMA.SF's *_thread_id selects another lane in the same quad, so the + # effective rows are derived from the selector id, not from the local lane's + # qlane alone. This is the compact-selector contract used by the production + # package path in gemm_sm120.h. + for lane in range(32): + qlane = lane & 3 + sfa_base_row = _sfa_row_in_lane(lane) + sfb_base_col = _sfb_col_in_lane(lane) + for warp_m in range(2): + for warp_n in range(2): + for k64_word in range(k64_words): + local_sfa_owner = qlane >> 1 + local_sfb_owner = qlane + local_sa0_row = warp_m * 64 + local_sfa_owner * 16 + sfa_base_row + local_sa1_row = local_sa0_row + 32 + local_sb0_row = warp_n * 64 + local_sfb_owner * 8 + sfb_base_col + local_sb1_row = local_sb0_row + 32 + + assert _packed_word(packed, local_sa0_row, k64_word) == int(semantic_words[local_sa0_row, k64_word]) + assert _packed_word(packed, local_sa1_row, k64_word) == int(semantic_words[local_sa1_row, k64_word]) + assert _packed_word(packed, local_sb0_row, k64_word) == int(semantic_words[local_sb0_row, k64_word]) + assert _packed_word(packed, local_sb1_row, k64_word) == int(semantic_words[local_sb1_row, k64_word]) + + for mma_i in range(4): + effective_sfa_row = _current_compact_sfa_row(lane, warp_m, mma_i) + assert _packed_word(packed, effective_sfa_row, k64_word) == int(semantic_words[effective_sfa_row, k64_word]) + + for mma_j in range(4): + for half in range(2): + effective_sfb_row = _current_compact_sfb_row(lane, warp_n, mma_j, half) + assert _packed_word(packed, effective_sfb_row, k64_word) == int(semantic_words[effective_sfb_row, k64_word]) + + +def test_compact_package_issue_table_models_omma_sf_source_lane_selectors(): + issues = _compact_issue_contract(lane=0, warp_m=0, warp_n=0) + + assert len(issues) == 4 * 4 * 2 + assert issues[:4] == [ + { + "mma_i": 0, + "mma_j": 0, + "half": 0, + "sa_reg": "sa0", + "sa_tid": 0, + "sa_source_lane": 0, + "sfa_row": 0, + "sb_reg": "sb0", + "sb_tid": 0, + "sb_source_lane": 0, + "sfb_row": 0, + }, + { + "mma_i": 0, + "mma_j": 0, + "half": 1, + "sa_reg": "sa0", + "sa_tid": 0, + "sa_source_lane": 0, + "sfa_row": 0, + "sb_reg": "sb0", + "sb_tid": 1, + "sb_source_lane": 1, + "sfb_row": 8, + }, + { + "mma_i": 0, + "mma_j": 1, + "half": 0, + "sa_reg": "sa0", + "sa_tid": 0, + "sa_source_lane": 0, + "sfa_row": 0, + "sb_reg": "sb0", + "sb_tid": 2, + "sb_source_lane": 2, + "sfb_row": 16, + }, + { + "mma_i": 0, + "mma_j": 1, + "half": 1, + "sa_reg": "sa0", + "sa_tid": 0, + "sa_source_lane": 0, + "sfa_row": 0, + "sb_reg": "sb0", + "sb_tid": 3, + "sb_source_lane": 3, + "sfb_row": 24, + }, + ] + + assert [issue["sfa_row"] for issue in issues[0::8]] == [0, 16, 32, 48] + assert [issue["sfb_row"] for issue in issues[:8]] == [0, 8, 16, 24, 32, 40, 48, 56] + + +def test_compact_package_issue_table_matches_cpp_macro_issue_order(): + issues = _compact_issue_contract(lane=0, warp_m=0, warp_n=0) + assert [ + ( + issue["mma_i"], + issue["mma_j"], + issue["half"], + issue["sa_reg"], + issue["sa_tid"], + issue["sb_reg"], + issue["sb_tid"], + ) + for issue in issues + ] == [ + (0, 0, 0, "sa0", 0, "sb0", 0), + (0, 0, 1, "sa0", 0, "sb0", 1), + (0, 1, 0, "sa0", 0, "sb0", 2), + (0, 1, 1, "sa0", 0, "sb0", 3), + (0, 2, 0, "sa0", 0, "sb1", 0), + (0, 2, 1, "sa0", 0, "sb1", 1), + (0, 3, 0, "sa0", 0, "sb1", 2), + (0, 3, 1, "sa0", 0, "sb1", 3), + (1, 0, 0, "sa0", 1, "sb0", 0), + (1, 0, 1, "sa0", 1, "sb0", 1), + (1, 1, 0, "sa0", 1, "sb0", 2), + (1, 1, 1, "sa0", 1, "sb0", 3), + (1, 2, 0, "sa0", 1, "sb1", 0), + (1, 2, 1, "sa0", 1, "sb1", 1), + (1, 3, 0, "sa0", 1, "sb1", 2), + (1, 3, 1, "sa0", 1, "sb1", 3), + (2, 0, 0, "sa1", 0, "sb0", 0), + (2, 0, 1, "sa1", 0, "sb0", 1), + (2, 1, 0, "sa1", 0, "sb0", 2), + (2, 1, 1, "sa1", 0, "sb0", 3), + (2, 2, 0, "sa1", 0, "sb1", 0), + (2, 2, 1, "sa1", 0, "sb1", 1), + (2, 3, 0, "sa1", 0, "sb1", 2), + (2, 3, 1, "sa1", 0, "sb1", 3), + (3, 0, 0, "sa1", 1, "sb0", 0), + (3, 0, 1, "sa1", 1, "sb0", 1), + (3, 1, 0, "sa1", 1, "sb0", 2), + (3, 1, 1, "sa1", 1, "sb0", 3), + (3, 2, 0, "sa1", 1, "sb1", 0), + (3, 2, 1, "sa1", 1, "sb1", 1), + (3, 3, 0, "sa1", 1, "sb1", 2), + (3, 3, 1, "sa1", 1, "sb1", 3), + ] + + +def test_cpp_package_macro_keeps_compact_selector_issue_order(): + header = Path(__file__).resolve().parents[3] / "src/tl_templates/cuda/gemm_sm120.h" + source = header.read_text() + row_macro = source[source.index("#define TL_SM120_PKG_MMA_ROW") : source.index(" TL_SM120_PKG_MMA_ROW(0")] + + position = 0 + for mma_j, half, sb_reg, sb_tid in [ + (0, 0, "scale_pkg.sb0", 0), + (0, 1, "scale_pkg.sb0", 1), + (1, 0, "scale_pkg.sb0", 2), + (1, 1, "scale_pkg.sb0", 3), + (2, 0, "scale_pkg.sb1", 0), + (2, 1, "scale_pkg.sb1", 1), + (3, 0, "scale_pkg.sb1", 2), + (3, 1, "scale_pkg.sb1", 3), + ]: + call = f"TL_SM120_PKG_MMA_N8(I, {mma_j}, {half}," + position = row_macro.index(call, position) + selector = f"{sb_reg}, SA_TID, {sb_tid});" + assert selector in row_macro[position:] + + row_calls = source[source.index(" TL_SM120_PKG_MMA_ROW(0") : source.index("#undef TL_SM120_PKG_MMA_ROW")] + position = 0 + for mma_i, sa_reg, sa_tid in [ + (0, "scale_pkg.sa0", 0), + (1, "scale_pkg.sa0", 1), + (2, "scale_pkg.sa1", 0), + (3, "scale_pkg.sa1", 1), + ]: + call = f"TL_SM120_PKG_MMA_ROW({mma_i}," + position = row_calls.index(call, position) + assert sa_reg in row_calls[position:] + assert f"{sa_tid});" in row_calls[position:] + + +def test_compact_package_issue_table_matches_effective_row_helpers(): + for lane in range(32): + for warp_m in range(2): + for warp_n in range(2): + for issue in _compact_issue_contract(lane, warp_m, warp_n): + assert issue["sfa_row"] == _current_compact_sfa_row(lane, warp_m, issue["mma_i"]) + assert issue["sfb_row"] == _current_compact_sfb_row( + lane, + warp_n, + issue["mma_j"], + issue["half"], + ) + + +def test_compact_package_copy_view_contract_maps_issue_rows_to_kmajor_words(): + issues = _compact_copy_view_contract(lane=0, warp_m=0, warp_n=0, k64_word=0) + + assert [ + ( + issue["mma_i"], + issue["mma_j"], + issue["half"], + issue["sfa_row"], + issue["sfa_word_coord"], + issue["sfa_word_offset"], + issue["sfb_row"], + issue["sfb_word_coord"], + issue["sfb_word_offset"], + ) + for issue in issues[:8] + ] == [ + (0, 0, 0, 0, (0, 0), 0, 0, (0, 0), 0), + (0, 0, 1, 0, (0, 0), 0, 8, (8, 0), 32), + (0, 1, 0, 0, (0, 0), 0, 16, (16, 0), 64), + (0, 1, 1, 0, (0, 0), 0, 24, (24, 0), 96), + (0, 2, 0, 0, (0, 0), 0, 32, (0, 1), 1), + (0, 2, 1, 0, (0, 0), 0, 40, (8, 1), 33), + (0, 3, 0, 0, (0, 0), 0, 48, (16, 1), 65), + (0, 3, 1, 0, (0, 0), 0, 56, (24, 1), 97), + ] + + assert [issue["sfa_word_offset"] for issue in issues[0::8]] == [0, 64, 1, 65] + + +def test_sm120_cuda_helper_keeps_scale_tv_package_boundary(): + header = Path(__file__).resolve().parents[3] / "src/tl_templates/cuda/gemm_sm120.h" + source = header.read_text() + + assert "struct SM120ScaleTVPackage" in source + assert "template " in source + assert "sm120_load_scale_tv_package" in source + assert "sm120_copy_scale_tv_package(SM120ScaleTVPackage &pkg," in source + assert "sm120_gemm_fulltile_ab_owner_wide_package(" in source + assert "const detail::SM120ScaleTVPackage &scale_pkg" in source + assert ( + "uint32_t sa0, sa1;" + not in source[source.index("struct SM120FulltileABOwnerWidePackage") : source.index("sm120_copy_fulltile_ab_owner_wide_package")] + ) + assert "detail::SM120ScaleTVPackage scale_pkg0;" in source + assert "detail::SM120ScaleTVPackage scale_pkg1;" in source + assert source.index("detail::sm120_copy_scale_tv_package(scale_pkg0") < source.index( + "sm120_gemm_fulltile_ab_owner_wide_package(c, pkg0, scale_pkg0)" + ) + assert "pkg.sa0 =" in source + assert "pkg.sb1 =" in source + + +def test_blockscaled_chunk_kmajor_scale_packer_rejects_invalid_shapes(): + with pytest.raises(ValueError, match="rows multiple of 128"): + pack_blockscaled_chunk_kmajor_scale_bytes(torch.zeros((127, 16), dtype=torch.uint8)) + + with pytest.raises(ValueError, match="K/16 columns multiple of 16"): + pack_blockscaled_chunk_kmajor_scale_bytes(torch.zeros((128, 12), dtype=torch.uint8)) + + with pytest.raises(TypeError, match="torch.uint8"): + pack_blockscaled_chunk_kmajor_scale_bytes(torch.zeros((128, 16), dtype=torch.int32)) + + +def test_encode_fp4_e2m1_values_and_pack_order(): + values = torch.tensor([[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.5, -1.0]], dtype=torch.float32) + codes = encode_fp4_e2m1_values(values) + assert torch.equal(codes, torch.tensor([[0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x9, 0xA]], dtype=torch.uint8)) + + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).contiguous().view(torch.int8) + assert torch.equal(decode_packed_fp4_e2m1(packed), values) + + +def test_encode_ue4m3_scale_bytes_known_values(): + values = torch.tensor([0.0, 2.0**-9, 2.0**-6, 1.0, 2.0, 448.0], dtype=torch.float32) + encoded = encode_ue4m3_scale_bytes(values, rounding="nearest") + assert torch.equal(encoded, torch.tensor([0x00, 0x01, 0x08, 0x38, 0x40, 0x7E], dtype=torch.uint8)) + torch.testing.assert_close(decode_ue4m3_scale_bytes(encoded), values) + assert torch.isnan(decode_ue4m3_scale_bytes(torch.tensor([0x7F], dtype=torch.uint8))).all() + + +def test_quantize_nvfp4_blockscaled_bf16_activation_contract(): + rows = 128 + cols = 256 + x = torch.zeros((rows, cols), dtype=torch.bfloat16) + pattern = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, 0.0]) + x[0, :16] = pattern.to(torch.bfloat16) + + packed_fp4, packed_scales, scale_bytes = quantize_bf16_to_nvfp4_blockscaled(x, return_scale_bytes=True) + + assert packed_fp4.shape == (rows, cols // 2) + assert packed_fp4.dtype == torch.int8 + assert packed_scales.shape == (rows, cols // 64) + assert packed_scales.dtype == torch.uint32 + assert scale_bytes.shape == (rows, cols // 16) + assert scale_bytes.dtype == torch.uint8 + assert scale_bytes[0, 0].item() == 0x38 + assert torch.equal(packed_scales, pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes)) + + decoded = decode_packed_fp4_e2m1(packed_fp4) * decode_ue4m3_scale_bytes(scale_bytes).repeat_interleave(16, dim=1) + torch.testing.assert_close(decoded[0, :16], pattern, rtol=0.0, atol=0.0) + + +def test_quantize_nvfp4_blockscaled_random_bf16_has_bounded_error(): + rows = 128 + cols = 256 + generator = torch.Generator(device="cpu").manual_seed(19) + x = (torch.randn((rows, cols), generator=generator, dtype=torch.float32) * 2.0).to(torch.bfloat16) + + packed_fp4, scale_source, scale_bytes = quantize_bf16_to_nvfp4_blockscaled(x, return_scale_bytes=True) + decoded = decode_packed_fp4_e2m1(packed_fp4) * decode_ue4m3_scale_bytes(scale_bytes).repeat_interleave(16, dim=1) + + scale = decode_ue4m3_scale_bytes(scale_bytes).repeat_interleave(16, dim=1) + error = (decoded - x.to(torch.float32)).abs() + assert torch.isfinite(decoded).all() + assert torch.all(error <= scale + 1e-6) + assert torch.equal(scale_source, pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes)) + + +def test_quantize_nvfp4_blockscaled_explicit_layout_matches_default(): + rows = 128 + cols = 256 + generator = torch.Generator(device="cpu").manual_seed(23) + x = (torch.randn((rows, cols), generator=generator, dtype=torch.float32) * 2.0).to(torch.bfloat16) + + default = quantize_bf16_to_nvfp4_blockscaled(x) + explicit = quantize_bf16_to_nvfp4_blockscaled(x, scale_layout="blockscaled_chunk_kmajor") + + assert torch.equal(explicit[0], default[0]) + assert torch.equal(explicit[1], default[1]) + + +@tilelang_testing.requires_cuda +@tilelang_testing.requires_cuda_compute_version_ge(10, 0) +@pytest.mark.parametrize("rows, cols", [(128, 256), (256, 512)]) +def test_tilelang_quantize_nvfp4_blockscaled_matches_reference_layout_and_error_bound(rows, cols): + generator = torch.Generator(device="cuda").manual_seed(rows + cols) + x = (torch.randn((rows, cols), generator=generator, device="cuda", dtype=torch.float32) * 2.0).to(torch.bfloat16) + + packed_tl, scale_source_tl = tilelang_quantize_bf16_to_nvfp4_blockscaled(x) + _, scale_source_ref, scale_bytes_ref = quantize_bf16_to_nvfp4_blockscaled(x, return_scale_bytes=True) + + assert packed_tl.shape == (rows, cols // 2) + assert packed_tl.dtype == torch.int8 + assert scale_source_tl.shape == (rows, cols // 64) + assert scale_source_tl.dtype == torch.uint32 + assert torch.equal(scale_source_tl.cpu(), scale_source_ref.cpu()) + + semantic_words = unswizzle_blockscaled_chunk_kmajor_scale_words(scale_source_tl) + assert torch.equal(swizzle_blockscaled_chunk_kmajor_scale_words(semantic_words).cpu(), scale_source_tl.cpu()) + + scale = decode_ue4m3_scale_bytes(scale_bytes_ref).repeat_interleave(16, dim=1) + decoded = decode_packed_fp4_e2m1(packed_tl) * scale + error = (decoded - x.to(torch.float32)).abs() + assert torch.isfinite(decoded).all() + assert torch.all(error <= scale + 1e-6) + + +def test_quantize_nvfp4_blockscaled_alias(): + x = torch.zeros((128, 256), dtype=torch.bfloat16) + expected = quantize_bf16_to_nvfp4_blockscaled(x) + actual = quantize_nvfp4_blockscaled(x) + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[1], expected[1]) diff --git a/tilelang/cuda/intrinsics/__init__.py b/tilelang/cuda/intrinsics/__init__.py index 8601d9342..d140e6142 100644 --- a/tilelang/cuda/intrinsics/__init__.py +++ b/tilelang/cuda/intrinsics/__init__.py @@ -7,6 +7,7 @@ from .layout.utils import ( # noqa: F401 from .macro.mma_macro_generator import ( # noqa: F401 TensorCoreIntrinEmitter, TensorCoreIntrinEmitterWithLadderTransform, + TensorCoreIntrinEmitterWithBlockScale, ) from .layout.mma_layout import get_swizzle_layout # noqa: F401 diff --git a/tilelang/cuda/intrinsics/layout/mma_layout.py b/tilelang/cuda/intrinsics/layout/mma_layout.py index 205ce92cb..ebf5809c1 100644 --- a/tilelang/cuda/intrinsics/layout/mma_layout.py +++ b/tilelang/cuda/intrinsics/layout/mma_layout.py @@ -188,6 +188,81 @@ def mma_load_b_32x8_to_shared_16x16_layout(thread_id, local_id): return row, col +def shared_16x64_to_mma_a_32x32_layout(i, j): + """A fragment layout for m16n8k64 e2m1 row-major operand.""" + thread_id = 4 * (i % 8) + (j % 32) // 8 + local_id = 16 * (i // 8) + 8 * (j // 32) + j % 8 + return thread_id, local_id + + +def shared_64x16_to_mma_a_32x32_layout_trans(i, j): + return shared_16x64_to_mma_a_32x32_layout(j, i) + + +def shared_8x64_to_mma_b_32x16_layout(i, j): + """B fragment layout for m16n8k64 e2m1 col-major operand.""" + thread_id = 4 * i + (j % 32) // 8 + local_id = 8 * (j // 32) + j % 8 + return thread_id, local_id + + +def shared_64x8_to_mma_b_32x16_layout_trans(i, j): + return shared_8x64_to_mma_b_32x16_layout(j, i) + + +shared_16x64_to_mma_32x32_layout_sr_a = shared_16x64_to_mma_a_32x32_layout +shared_16x64_to_mma_32x32_layout_rs_a = shared_64x16_to_mma_a_32x32_layout_trans +shared_8x64_to_mma_32x16_layout_sr_b = shared_8x64_to_mma_b_32x16_layout +shared_8x64_to_mma_32x16_layout_rs_b = shared_64x8_to_mma_b_32x16_layout_trans + + +def ldmatrix_32x32_to_shared_16x64_layout_a(thread_id): + """Row-start addresses for ldmatrix.x4 loading A m16k64 e2m1.""" + row = thread_id % 16 + col = (thread_id // 16) * 32 + return row, col + + +def ldmatrix_32x32_to_shared_16x64_layout_b(thread_id): + """Row-start addresses for ldmatrix.x4 loading B n16k64 e2m1.""" + row = (thread_id // 16) * 8 + (thread_id % 8) + col = ((thread_id % 16) // 8) * 32 + return row, col + + +def ldmatrix_32x16_to_shared_8x64_layout_b(thread_id): + """Row-start addresses for ldmatrix.x2 loading B n8k64 e2m1.""" + row = thread_id % 8 + col = ((thread_id // 8) % 2) * 32 + return row, col + + +def mma_load_a_32x32_to_shared_16x64_layout(thread_id, local_id): + """Inverse: (thread_id, local_id) -> (m, k) for A m16k64 e2m1. + + Matches CUTLASS/CuTe SM120 ALayout: + Layout, Shape<_8,_2,_2>>, + Stride, Stride<_16,_8,_512>>>. + """ + tid_m = thread_id % 4 + tid_k = thread_id // 4 + val_k8 = local_id % 8 + val_m8 = (local_id // 8) % 2 + val_k32 = local_id // 16 + row = tid_k + 8 * val_m8 + col = tid_m * 8 + val_k8 + 32 * val_k32 + return row, col + + +def mma_load_b_32x16_to_shared_8x64_layout(thread_id, local_id): + """Inverse: (thread_id, local_id) -> (n, k) for B n8k64 e2m1.""" + group_id = thread_id // 4 + tid_in_group = thread_id % 4 + row = group_id + col = tid_in_group * 8 + (local_id % 8) + 32 * (local_id // 8) + return row, col + + def shared_16x16_to_mma_32x8_smoothlayout(i, j): return (i * 2 + j // 8, j % 8) diff --git a/tilelang/cuda/intrinsics/macro/__init__.py b/tilelang/cuda/intrinsics/macro/__init__.py index 658f79122..6247e5d94 100644 --- a/tilelang/cuda/intrinsics/macro/__init__.py +++ b/tilelang/cuda/intrinsics/macro/__init__.py @@ -1,6 +1,7 @@ from .mma_macro_generator import ( # noqa: F401 TensorCoreIntrinEmitter, TensorCoreIntrinEmitterWithLadderTransform, + TensorCoreIntrinEmitterWithBlockScale, ) from .mma_sp_macro_generator import SparseTensorCoreIntrinEmitter # noqa: F401 diff --git a/tilelang/cuda/intrinsics/macro/mma_macro_generator.py b/tilelang/cuda/intrinsics/macro/mma_macro_generator.py index 19ca70eda..44fbde7f4 100644 --- a/tilelang/cuda/intrinsics/macro/mma_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/mma_macro_generator.py @@ -1,5 +1,6 @@ from __future__ import annotations import tilelang.language as T +from dataclasses import dataclass from typing import Literal from collections.abc import Callable from tilelang.common import TransformKind @@ -28,11 +29,256 @@ from tilelang.cuda.intrinsics.layout.mma_layout import ( mma_load_a_32x16_to_shared_16x32_layout, mma_load_b_32x16_to_shared_16x32_layout, mma_load_a_32x8_to_shared_16x16_layout, + shared_16x64_to_mma_32x32_layout_sr_a, + shared_8x64_to_mma_32x16_layout_sr_b, + ldmatrix_32x32_to_shared_16x64_layout_a, + ldmatrix_32x32_to_shared_16x64_layout_b, + ldmatrix_32x16_to_shared_8x64_layout_b, ) lift = convert +@dataclass(frozen=True) +class BlockScaleMmaConfig: + """Static SM120 warp-level block-scale MMA configuration.""" + + kind: str + mma_prefix: str + atom_k: int + scale_vec_size: int + sf_vec_size: int + scale_type: str + a_dtype_abbrv: str + b_dtype_abbrv: str + accum_dtype: str = T.float32 + active_sfa_threads: int = 16 + active_sfb_threads: int = 8 + + +@dataclass(frozen=True) +class SM120BlockScaledFullTilePackageContract: + """Internal contract for the SM120 full-tile blockscaled package path. + + This is deliberately a lowering-level contract, not a user option. The CUDA + helper is still the instruction backend today, but the emitter must name the + fragment/package shape it expects before it can route a tile to that helper. + """ + + tile_m: int + tile_n: int + tile_k: int + block_row_warps: int + block_col_warps: int + warp_rows: int + warp_cols: int + warp_row_tiles: int + warp_col_tiles: int + kblocks: int + micro_size_k: int + sf_layout: str + scale_package_words_sfa: int + scale_package_words_sfb: int + issue_count_per_warp: int + issue_count_per_warpgroup: int + + @classmethod + def for_package_pingpong( + cls, + emitter: TensorCoreIntrinEmitterWithBlockScale, + *, + sf_layout: str, + kblocks: int = 4, + ) -> SM120BlockScaledFullTilePackageContract: + contract = cls( + tile_m=int(emitter.block_row_warps) * int(emitter.warp_row_tiles), + tile_n=int(emitter.block_col_warps) * int(emitter.warp_col_tiles), + tile_k=kblocks * int(emitter.micro_size_k), + block_row_warps=int(emitter.block_row_warps), + block_col_warps=int(emitter.block_col_warps), + warp_rows=int(emitter.warp_rows), + warp_cols=int(emitter.warp_cols), + warp_row_tiles=int(emitter.warp_row_tiles), + warp_col_tiles=int(emitter.warp_col_tiles), + kblocks=kblocks, + micro_size_k=int(emitter.micro_size_k), + sf_layout=sf_layout, + scale_package_words_sfa=2, + scale_package_words_sfb=2, + issue_count_per_warp=32, + issue_count_per_warpgroup=128, + ) + contract.validate() + return contract + + def validate(self) -> None: + if (self.tile_m, self.tile_n, self.tile_k) != (128, 128, 256): + raise ValueError( + f"SM120 full-tile package contract currently requires tile shape 128x128x256, got {self.tile_m}x{self.tile_n}x{self.tile_k}" + ) + if (self.warp_rows, self.warp_cols) != (4, 4): + raise ValueError( + "SM120 full-tile package contract currently requires each consumer " + f"warp to issue a 4x4 MMA atom grid, got {self.warp_rows}x{self.warp_cols}" + ) + if (self.block_row_warps, self.block_col_warps) != (2, 2): + raise ValueError( + "SM120 full-tile package contract currently requires a 2x2 consumer " + f"warp layout, got {self.block_row_warps}x{self.block_col_warps}" + ) + if (self.warp_row_tiles, self.warp_col_tiles) != (64, 64): + raise ValueError( + "SM120 full-tile package contract currently requires each consumer " + f"warp tile to be 64x64, got {self.warp_row_tiles}x{self.warp_col_tiles}" + ) + if self.kblocks != 4 or self.micro_size_k != 64: + raise ValueError( + "SM120 full-tile package contract currently requires four K=64 " + f"blocks, got kblocks={self.kblocks}, micro_size_k={self.micro_size_k}" + ) + if self.sf_layout != "blockscaled_chunk_kmajor": + raise ValueError("SM120 full-tile package contract requires sf_layout='blockscaled_chunk_kmajor'") + if (self.scale_package_words_sfa, self.scale_package_words_sfb) != (2, 2): + raise ValueError( + "SM120 compact scale TV package requires 2 SFA and 2 SFB " + f"words, got {self.scale_package_words_sfa}, {self.scale_package_words_sfb}" + ) + if self.issue_count_per_warp != 32 or self.issue_count_per_warpgroup != 128: + raise ValueError( + "SM120 full-tile package contract expects 32 OMMA.SF issues per " + f"warp and 128 per warpgroup, got {self.issue_count_per_warp}, " + f"{self.issue_count_per_warpgroup}" + ) + + @staticmethod + def blockscaled_chunk_kmajor_sf_word(row: int, kblock: int) -> int: + """Return the uint32 scale-word offset for the source/smem scale layout.""" + + if row < 0 or row >= 128: + raise ValueError(f"scale row must be in [0, 128), got {row}") + if kblock < 0 or kblock >= 4: + raise ValueError(f"kblock must be in [0, 4), got {kblock}") + return kblock * 128 + (row & 31) * 4 + (row >> 5) + + def compact_selector_scale_rows(self, lane: int, warp_m: int, warp_n: int) -> tuple[tuple[int, int], tuple[int, int]]: + """Return SFA/SFB semantic rows loaded by the current compact TV package.""" + + if lane < 0 or lane >= 32: + raise ValueError(f"lane must be in [0, 32), got {lane}") + if warp_m < 0 or warp_m >= self.block_row_warps: + raise ValueError(f"warp_m must be in [0, {self.block_row_warps}), got {warp_m}") + if warp_n < 0 or warp_n >= self.block_col_warps: + raise ValueError(f"warp_n must be in [0, {self.block_col_warps}), got {warp_n}") + + qlane = lane & 3 + sfa_row = 8 * (lane & 1) + (lane >> 2) + sfb_col = lane >> 2 + a_owner_in_pair = qlane >> 1 + scale_m0 = warp_m * 64 + a_owner_in_pair * 16 + sfa_row + scale_n0 = warp_n * 64 + qlane * 8 + sfb_col + return (scale_m0, scale_m0 + 32), (scale_n0, scale_n0 + 32) + + def compact_selector_scale_word_offsets( + self, lane: int, warp_m: int, warp_n: int, kblock: int + ) -> tuple[tuple[int, int], tuple[int, int]]: + sfa_rows, sfb_rows = self.compact_selector_scale_rows(lane, warp_m, warp_n) + return ( + tuple(self.blockscaled_chunk_kmajor_sf_word(row, kblock) for row in sfa_rows), + tuple(self.blockscaled_chunk_kmajor_sf_word(row, kblock) for row in sfb_rows), + ) + + @staticmethod + def sfa_selector_source_lane(lane: int, scale_a_thread_id: int) -> int: + return (lane & ~3) | ((scale_a_thread_id & 1) << 1) | (lane & 1) + + @staticmethod + def sfb_selector_source_lane(lane: int, scale_b_thread_id: int) -> int: + return (lane & ~3) | (scale_b_thread_id & 3) + + def compact_selector_effective_rows( + self, lane: int, warp_m: int, warp_n: int, issue: tuple[int, int, int, int, int, int, int] + ) -> tuple[int, int]: + """Return semantic SFA/SFB rows consumed by one compact-selector issue.""" + + mma_i, mma_j, n8_half, sfa_word, sfb_word, scale_a_tid, scale_b_tid = issue + del mma_i, mma_j, n8_half + sfa_source_lane = self.sfa_selector_source_lane(lane, scale_a_tid) + sfb_source_lane = self.sfb_selector_source_lane(lane, scale_b_tid) + sfa_rows, _ = self.compact_selector_scale_rows(sfa_source_lane, warp_m, warp_n) + _, sfb_rows = self.compact_selector_scale_rows(sfb_source_lane, warp_m, warp_n) + return sfa_rows[sfa_word], sfb_rows[sfb_word] + + def package_pingpong_lifecycle(self) -> tuple[tuple[str, int, int], ...]: + """Return the current copy/gemm package lifecycle from the CUDA helper. + + Each tuple is ``(op, register_package, kblock)``. Scale and A/B packages + share the same register package id in the current implementation. + """ + + return ( + ("copy", 0, 0), + ("copy", 1, 1), + ("gemm", 0, 0), + ("copy", 0, 2), + ("gemm", 1, 1), + ("copy", 1, 3), + ("gemm", 0, 2), + ("gemm", 1, 3), + ) + + def omma_sf_issue_schedule_per_warp(self) -> tuple[tuple[int, int, int, int, int, int, int], ...]: + """Return the current per-warp OMMA.SF issue schedule. + + Each tuple is ``(mma_i, mma_j, n8_half, sfa_word, sfb_word, + scale_a_thread_id, scale_b_thread_id)``. ``sfa_word`` and ``sfb_word`` + are indices inside the current compact-selector scale package, not + source-memory offsets. + """ + + issues = [] + for mma_i in range(4): + sfa_word = 0 if mma_i < 2 else 1 + scale_a_thread_id = mma_i & 1 + for mma_j in range(4): + sfb_word = 0 if mma_j < 2 else 1 + for n8_half in range(2): + scale_b_thread_id = (mma_j & 1) * 2 + n8_half + issues.append( + ( + mma_i, + mma_j, + n8_half, + sfa_word, + sfb_word, + scale_a_thread_id, + scale_b_thread_id, + ) + ) + return tuple(issues) + + +_SUPPORTED_BLOCK_SCALE_MMA_CONFIGS = { + ("mxf4nvf4", 4, "ue4m3"): BlockScaleMmaConfig( + kind="mxf4nvf4", + mma_prefix="m16n8k64", + atom_k=64, + scale_vec_size=4, + sf_vec_size=16, + scale_type="ue4m3", + a_dtype_abbrv="e2m1", + b_dtype_abbrv="e2m1", + ), +} + + +def _get_block_scale_mma_config(kind: str, scale_vec_size: int, scale_type: str) -> BlockScaleMmaConfig: + key = (kind, scale_vec_size, scale_type) + if key not in _SUPPORTED_BLOCK_SCALE_MMA_CONFIGS: + supported = ", ".join(str(k) for k in sorted(_SUPPORTED_BLOCK_SCALE_MMA_CONFIGS)) + raise ValueError(f"Unsupported SM120 block-scale MMA config {key}; supported: {supported}") + return _SUPPORTED_BLOCK_SCALE_MMA_CONFIGS[key] + + class TensorCoreIntrinEmitter: """ To eliminate Python syntax within TIR Macro. @@ -352,7 +598,7 @@ class TensorCoreIntrinEmitter: tx, _, warp_m = self.extract_thread_binding(thread_binding) trans = self.a_transposed - for i in T.serial(warp_rows): + for i in T.unroll(warp_rows): wi, wk = warp_m * warp_row_tiles + i * micro_size_x, rk * chunk + ki * micro_size_k if ldmatrix_available: @@ -464,7 +710,7 @@ class TensorCoreIntrinEmitter: tx, warp_n, _ = self.extract_thread_binding(thread_binding) trans = not b_transposed - for i in T.serial(warp_cols): + for i in T.unroll(warp_cols): # Assign B_shared_elem wi, wk = ( warp_n * warp_col_tiles + i * micro_size_y, @@ -715,6 +961,9 @@ class TensorCoreIntrinEmitter: elif dtype_bits == 8: transform_func_sr_a = shared_16x32_to_mma_32x16_layout_sr_a transform_func_sr_b = shared_16x32_to_mma_32x16_layout_sr_b + elif dtype_bits == 4: + transform_func_sr_a = shared_16x64_to_mma_32x32_layout_sr_a + transform_func_sr_b = shared_8x64_to_mma_32x16_layout_sr_b else: raise ValueError(f"Unsupported dtype {dtype}") @@ -1223,3 +1472,2504 @@ class TensorCoreIntrinEmitterWithLadderTransform(TensorCoreIntrinEmitter): ) return _warp_mma(A_local_buf, B_local_buf, C_local_buf) + + +class TensorCoreIntrinEmitterWithBlockScale(TensorCoreIntrinEmitter): + """SM120 warp-level block-scale MMA emitter. + + The emitter keeps scale-factor storage explicit, matching TileLang's + TCGEN05 block-scaled style while targeting warp-level ``mma.sync``. + """ + + def __init__( + self, + a_dtype: str = T.float4_e2m1fn, + b_dtype: str = T.float4_e2m1fn, + accum_dtype: str = T.float32, + a_transposed: bool = False, + b_transposed: bool = False, + block_row_warps: int = 2, + block_col_warps: int = 2, + warp_row_tiles: int = 32, + warp_col_tiles: int = 32, + chunk: int = 64, + reduce_k: int = 1, + num_elems_per_byte: int = 1, + is_m_first: bool | None = False, + thread_var: Var | None = None, + kind: str = "mxf4nvf4", + scale_vec_size: int = 4, + stype: str = "ue4m3", + ): + self.block_scale_config = _get_block_scale_mma_config(kind, scale_vec_size, stype) + a_dtype_abbrv = self._get_dtype_abbrv(str(a_dtype)) + b_dtype_abbrv = self._get_dtype_abbrv(str(b_dtype)) + if ( + a_dtype_abbrv != self.block_scale_config.a_dtype_abbrv + or b_dtype_abbrv != self.block_scale_config.b_dtype_abbrv + or str(accum_dtype) != self.block_scale_config.accum_dtype + ): + raise ValueError( + f"{self.block_scale_config.kind} expects a_dtype={self.block_scale_config.a_dtype_abbrv}, " + f"b_dtype={self.block_scale_config.b_dtype_abbrv}, " + f"accum_dtype={self.block_scale_config.accum_dtype}; " + f"got a_dtype={a_dtype}, b_dtype={b_dtype}, accum_dtype={accum_dtype}" + ) + self.kind = self.block_scale_config.kind + self.scale_vec_size = self.block_scale_config.scale_vec_size + self.stype = self.block_scale_config.scale_type + self.sf_vec_size = self.block_scale_config.sf_vec_size + super().__init__( + a_dtype=a_dtype, + b_dtype=b_dtype, + accum_dtype=accum_dtype, + a_transposed=a_transposed, + b_transposed=b_transposed, + block_row_warps=block_row_warps, + block_col_warps=block_col_warps, + warp_row_tiles=warp_row_tiles, + warp_col_tiles=warp_col_tiles, + chunk=chunk, + reduce_k=reduce_k, + num_elems_per_byte=num_elems_per_byte, + is_m_first=is_m_first, + thread_var=thread_var, + ) + + def _initialize_k_dim(self, a_dtype=T.float16): + self.k_dim = self.block_scale_config.atom_k + + def _initialize_abbrev(self, a_dtype, b_dtype, accum_dtype): + self.a_dtype_abbrv = self.block_scale_config.a_dtype_abbrv + self.b_dtype_abbrv = self.block_scale_config.b_dtype_abbrv + self.accum_dtype_abbrv = self._get_dtype_abbrv(accum_dtype) + + def _initialize_mma_prefix(self, k_dim: int = 16): + self.mma_prefix = self.block_scale_config.mma_prefix + + def ldmatrix_a(self, A_local_buf: Buffer, A_shared_buf: Buffer | BufferRegion, ki: PrimExpr, rk: PrimExpr | None = 0): + warp_row_tiles = self.warp_row_tiles + warp_rows = self.warp_rows + chunk = self.chunk + micro_size_x = self.micro_size_x + micro_size_k = self.micro_size_k + local_size_a = self.local_size_a + a_transposed = self.a_transposed + + thread_binding = self.get_thread_binding() + A_region = self._legalize_to_buffer_region(A_shared_buf) + A_buf = A_region.buffer + A_base0 = A_region.region[-2].min + A_base1 = A_region.region[-1].min + A_other = [r.min for r in A_region.region[:-2]] + + @T.macro + def _warp_ld_a_e2m1(A_local_buf, A_shared_buf, ki, thread_binding, rk=0): + tx, _, warp_m = self.extract_thread_binding(thread_binding) + for i in T.unroll(warp_rows): + wi = warp_m * warp_row_tiles + i * micro_size_x + wk = rk * chunk + ki * micro_size_k + row_off, col_off = ldmatrix_32x32_to_shared_16x64_layout_a(tx) + if a_transposed: + T.ptx_ldmatrix( + T.bool(False), + 4, + T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + row_off, A_base1 + wi + col_off)], + "r", + extent=local_size_a, + ), + T.access_ptr(A_local_buf[i * local_size_a], "w", extent=local_size_a), + ) + else: + T.ptx_ldmatrix( + T.bool(False), + 4, + T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wi + row_off, A_base1 + wk + col_off)], + "r", + extent=local_size_a, + ), + T.access_ptr(A_local_buf[i * local_size_a], "w", extent=local_size_a), + ) + + return _warp_ld_a_e2m1(A_local_buf, A_region, ki, thread_binding, rk) + + def ldmatrix_b(self, B_local_buf: Buffer, B_shared_buf: Buffer | BufferRegion, ki: PrimExpr, rk: PrimExpr | None = 0): + warp_col_tiles = self.warp_col_tiles + warp_cols = self.warp_cols + chunk = self.chunk + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_b = self.local_size_b + b_transposed = self.b_transposed + replicate_b = self.n_dim == 16 + + thread_binding = self.get_thread_binding() + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + @T.macro + def _warp_ld_b_e2m1(B_local_buf, B_shared_buf, ki, thread_binding, rk=0): + tx, warp_n, _ = self.extract_thread_binding(thread_binding) + for i in T.unroll(warp_cols): + wi = warp_n * warp_col_tiles + i * micro_size_y + wk = rk * chunk + ki * micro_size_k + if replicate_b: + row_off, col_off = ldmatrix_32x32_to_shared_16x64_layout_b(tx) + else: + row_off, col_off = ldmatrix_32x16_to_shared_8x64_layout_b(tx) + if b_transposed: + T.ptx_ldmatrix( + T.bool(False), + 4 if replicate_b else 2, + T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + wi + row_off, B_base1 + wk + col_off)], + "r", + extent=local_size_b, + ), + T.access_ptr(B_local_buf[i * local_size_b], "w", extent=local_size_b), + ) + else: + T.ptx_ldmatrix( + T.bool(True), + 4 if replicate_b else 2, + T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + wk + row_off, B_base1 + wi + col_off)], + "r", + extent=local_size_b, + ), + T.access_ptr(B_local_buf[i * local_size_b], "w", extent=local_size_b), + ) + + return _warp_ld_b_e2m1(B_local_buf, B_region, ki, thread_binding, rk) + + def ldmatrix_b_atom( + self, + B_local_buf: Buffer, + B_shared_buf: Buffer | BufferRegion, + ki: PrimExpr, + inst_n_idx: PrimExpr | int, + rk: PrimExpr | None = 0, + ): + warp_col_tiles = self.warp_col_tiles + chunk = self.chunk + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_b = self.local_size_b + b_transposed = self.b_transposed + replicate_b = self.n_dim == 16 + + thread_binding = self.get_thread_binding() + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + @T.macro + def _warp_ld_b_atom_e2m1(B_local_buf, B_shared_buf, ki, thread_binding, rk=0): + tx, warp_n, _ = self.extract_thread_binding(thread_binding) + wi = warp_n * warp_col_tiles + inst_n_idx * micro_size_y + wk = rk * chunk + ki * micro_size_k + if replicate_b: + row_off, col_off = ldmatrix_32x32_to_shared_16x64_layout_b(tx) + else: + row_off, col_off = ldmatrix_32x16_to_shared_8x64_layout_b(tx) + if b_transposed: + T.ptx_ldmatrix( + T.bool(False), + 4 if replicate_b else 2, + T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + wi + row_off, B_base1 + wk + col_off)], + "r", + extent=local_size_b, + ), + T.access_ptr(B_local_buf[0], "w", extent=local_size_b), + ) + else: + T.ptx_ldmatrix( + T.bool(True), + 4 if replicate_b else 2, + T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + wk + row_off, B_base1 + wi + col_off)], + "r", + extent=local_size_b, + ), + T.access_ptr(B_local_buf[0], "w", extent=local_size_b), + ) + + return _warp_ld_b_atom_e2m1(B_local_buf, B_region, ki, thread_binding, rk) + + def _scale_region_parts(self, scale_buf: Buffer | BufferRegion): + if isinstance(scale_buf, BufferRegion): + scale_region = scale_buf + elif isinstance(scale_buf, Buffer): + scale_region = self._legalize_to_buffer_region(scale_buf) + else: + raise ValueError(f"Unsupported scale buffer type: {type(scale_buf)}") + return ( + scale_region.buffer, + [r.min for r in scale_region.region[:-2]], + scale_region.region[-2].min, + scale_region.region[-1].min, + ) + + @staticmethod + def _sfa_row_in_atom(tx: PrimExpr): + # CUTLASS SFALayout for k64 uses ((2,2,8),64), stride ((8,0,1),16). + # With K-major flattening, the M coordinate is 8 * (lane % 2) + lane // 4. + return 8 * (tx % 2) + (tx // 4) + + @staticmethod + def _sfb_col_in_atom(tx: PrimExpr): + # CUTLASS SFBLayout for k64 uses ((4,8),64), stride ((0,1),8), so the + # logical N coordinate is lane // 4 with broadcast across four groups. + return tx // 4 + + def _scale_word_k(self, k_start: PrimExpr, ki: PrimExpr, sf_granularity_k: int): + packed_word_k = int(sf_granularity_k) * 4 + if packed_word_k != self.sf_vec_size * 4: + raise ValueError( + f"{self.kind} expects packed scale words covering {self.sf_vec_size * 4} K elements, " + f"got sf_granularity_k={sf_granularity_k}" + ) + _k_start = tvm.tirx.const(k_start, "int32") if isinstance(k_start, int) else k_start + return (_k_start + self.micro_size_k * ki) // packed_word_k + + def mma( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + replicate_b = self.n_dim == 16 + if sf_layout not in ("rowmajor", "blockscaled_chunk_kmajor"): + raise ValueError(f"Unsupported SM120 scale layout: {sf_layout}") + + def _cutlass_sf_word(idx, word_k): + return T.call_pure_extern("int32", "tl::detail::sm120_blockscaled_chunk_kmajor_sf_word", idx, word_k) + + @T.macro + def _warp_mma_block_scale(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + for i, j in T.grid(warp_rows, warp_cols): + scale_m = warp_m * warp_row_tiles + i * micro_size_x + sfa_row + scale_n = warp_n * warp_col_tiles + j * micro_size_y + sfb_col + if sf_layout == "blockscaled_chunk_kmajor": + scale_a_word = _cutlass_sf_word(scale_m, scale_a_word_k) + scale_b_word = _cutlass_sf_word(scale_n, scale_b_word_k) + scale_a_ptr = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_a_word // 4, SFA_base_k + scale_a_word % 4)], + "r", + ) + scale_b_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_b_word // 4, SFB_base_k + scale_b_word % 4)], + "r", + ) + else: + scale_a_ptr = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)], + "r", + ) + scale_b_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_k)], + "r", + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + if sf_layout == "blockscaled_chunk_kmajor": + scale_b_rep_n = scale_n + 8 + scale_b_rep_word = _cutlass_sf_word(scale_b_rep_n, scale_b_word_k) + scale_b_rep_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_b_rep_word // 4, SFB_base_k + scale_b_rep_word % 4)], + "r", + ) + else: + scale_b_rep_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_block_scale(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding) + + def ldscale( + self, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + replicate_b = self.n_dim == 16 + if sf_layout not in ("rowmajor", "blockscaled_chunk_kmajor"): + raise ValueError(f"Unsupported SM120 scale layout: {sf_layout}") + + def _cutlass_sf_word(idx, word_k): + return T.call_pure_extern("int32", "tl::detail::sm120_blockscaled_chunk_kmajor_sf_word", idx, word_k) + + @T.macro + def _warp_ldscale_block_scale(SFA_local_buf, SFB_local_buf, SFB_rep_local_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + for i in T.unroll(warp_rows): + scale_m = warp_m * warp_row_tiles + i * micro_size_x + sfa_row + if sf_layout == "blockscaled_chunk_kmajor": + scale_a_word = _cutlass_sf_word(scale_m, scale_a_word_k) + SFA_local_buf[i] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_a_word // 4, SFA_base_k + scale_a_word % 4)] + else: + SFA_local_buf[i] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)] + for j in T.unroll(warp_cols): + scale_n = warp_n * warp_col_tiles + j * micro_size_y + sfb_col + if sf_layout == "blockscaled_chunk_kmajor": + scale_b_word = _cutlass_sf_word(scale_n, scale_b_word_k) + SFB_local_buf[j] = SFB_data[tuple(SFB_other) + (SFB_base_n + scale_b_word // 4, SFB_base_k + scale_b_word % 4)] + else: + SFB_local_buf[j] = SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_k)] + if replicate_b: + if sf_layout == "blockscaled_chunk_kmajor": + scale_b_rep_n = scale_n + 8 + scale_b_rep_word = _cutlass_sf_word(scale_b_rep_n, scale_b_word_k) + SFB_rep_local_buf[j] = SFB_data[ + tuple(SFB_other) + (SFB_base_n + scale_b_rep_word // 4, SFB_base_k + scale_b_rep_word % 4) + ] + else: + SFB_rep_local_buf[j] = SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n + 8, SFB_base_k + scale_b_word_k)] + + return _warp_ldscale_block_scale( + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def ldscale_fragment( + self, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + """Load SM120 block-scale fragments into local registers. + + This is currently a thin wrapper over the existing scale-word load. + The separate name gives the SM120 MMA lowering a stable hook for a + CUTLASS-like scale-fragment copy path. + """ + return self.ldscale( + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + SFA_buf, + SFB_buf, + ki=ki, + k_start=k_start, + sf_a_granularity_k=sf_a_granularity_k, + sf_b_granularity_k=sf_b_granularity_k, + sf_layout=sf_layout, + ) + + def ldscale_fragment_b_owner( + self, + SFA_fragment_buf, + SFB_owner_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + """Load SFA normally and pack four B scale owners across a lane quad.""" + + warp_rows = self.warp_rows + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + if sf_layout != "rowmajor": + raise ValueError("ldscale_fragment_b_owner currently supports rowmajor scale layout only") + if self.n_dim != 16 or int(self.warp_cols) != 2: + raise ValueError("ldscale_fragment_b_owner currently requires replicated B with warp_cols=2") + + @T.macro + def _warp_ldscale_b_owner(SFA_fragment_buf, SFB_owner_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + qlane = tx & 3 + for i in T.unroll(warp_rows): + scale_m = warp_m * warp_row_tiles + i * micro_size_x + sfa_row + SFA_fragment_buf[i] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)] + SFB_owner_buf[0] = SFB_data[ + tuple(SFB_other) + (SFB_base_n + warp_n * warp_col_tiles + sfb_col + qlane * 8, SFB_base_k + scale_b_word_k) + ] + + return _warp_ldscale_b_owner( + SFA_fragment_buf, + SFB_owner_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def ldscale_fragment_ab_owner( + self, + SFA_owner_buf, + SFB_owner_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + """Load two A scale owners and one B scale owner for full-tile SM120.""" + + warp_rows = self.warp_rows + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + if sf_layout != "rowmajor": + raise ValueError("ldscale_fragment_ab_owner currently supports rowmajor scale layout only") + if int(warp_rows) != 4 or self.n_dim != 16 or int(self.warp_cols) != 2: + raise ValueError("ldscale_fragment_ab_owner currently requires fulltile warp_rows=4, warp_cols=2") + + @T.macro + def _warp_ldscale_ab_owner(SFA_owner_buf, SFB_owner_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + qlane = tx & 3 + a_owner_in_pair = qlane >> 1 + for g in T.unroll(2): + scale_m = warp_m * warp_row_tiles + g * (2 * micro_size_x) + a_owner_in_pair * micro_size_x + sfa_row + SFA_owner_buf[g] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)] + SFB_owner_buf[0] = SFB_data[ + tuple(SFB_other) + (SFB_base_n + warp_n * warp_col_tiles + sfb_col + qlane * 8, SFB_base_k + scale_b_word_k) + ] + + return _warp_ldscale_ab_owner( + SFA_owner_buf, + SFB_owner_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def ldscale_fragment_ab_owner_wide( + self, + SFA_owner_buf, + SFB_owner_buf, + SFA_buf, + SFB_buf, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + """Load owner scale words for full-tile SM120 with warp_cols=4. + + Each owner word is later selected by the block-scale MMA + `{byte_id, thread_id}` operands. For a 4x4 warp-atom tile this reduces + per-kblock scale loads from 4 A + 4 B + 4 B-rep words to 2 A-owner + + 2 B-owner words. + """ + + warp_rows = self.warp_rows + warp_cols = self.warp_cols + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + if int(warp_rows) != 4 or self.n_dim != 16 or int(warp_cols) != 4: + raise ValueError("ldscale_fragment_ab_owner_wide requires warp_rows=4, warp_cols=4") + if sf_layout not in ("rowmajor", "blockscaled_chunk_kmajor"): + raise ValueError(f"Unsupported SM120 scale layout: {sf_layout}") + + def _cutlass_sf_word(idx, word_k): + return T.call_pure_extern("int32", "tl::detail::sm120_blockscaled_chunk_kmajor_sf_word", idx, word_k) + + @T.macro + def _warp_ldscale_ab_owner_wide(SFA_owner_buf, SFB_owner_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + qlane = tx & 3 + a_owner_in_pair = qlane >> 1 + for g in T.unroll(2): + scale_m = warp_m * warp_row_tiles + g * (2 * micro_size_x) + a_owner_in_pair * micro_size_x + sfa_row + if sf_layout == "blockscaled_chunk_kmajor": + scale_a_word = _cutlass_sf_word(scale_m, scale_a_word_k) + SFA_owner_buf[g] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_a_word // 4, SFA_base_k + scale_a_word % 4)] + else: + SFA_owner_buf[g] = SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)] + for g in T.unroll(2): + scale_n = warp_n * warp_col_tiles + g * (2 * micro_size_y) + sfb_col + qlane * 8 + if sf_layout == "blockscaled_chunk_kmajor": + scale_b_word = _cutlass_sf_word(scale_n, scale_b_word_k) + SFB_owner_buf[g] = SFB_data[tuple(SFB_other) + (SFB_base_n + scale_b_word // 4, SFB_base_k + scale_b_word % 4)] + else: + SFB_owner_buf[g] = SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_k)] + + return _warp_ldscale_ab_owner_wide( + SFA_owner_buf, + SFB_owner_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def ldscale_fragment_kpack( + self, + SFA_pack_buf, + SFB_pack_buf, + SFB_rep_pack_buf, + SFA_buf, + SFB_buf, + num_k_blocks: int, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + ): + """Load a small K-block pack of SM120 scale fragments. + + The retained shared scale layout has contiguous uint32 scale words along + the K-word dimension. This private loader keeps that layout and tries to + expose the four K words as a vectorized copy into registers. + """ + + warp_rows = self.warp_rows + warp_cols = self.warp_cols + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_start = self._scale_word_k(k_start, 0, sf_a_granularity_k) + scale_b_word_start = self._scale_word_k(k_start, 0, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_ldscale_kpack(SFA_pack_buf, SFB_pack_buf, SFB_rep_pack_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + for i in T.unroll(warp_rows): + scale_m = warp_m * warp_row_tiles + i * micro_size_x + sfa_row + for kk in T.vectorized(num_k_blocks): + SFA_pack_buf[i * num_k_blocks + kk] = SFA_data[ + tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_start + kk) + ] + for j in T.unroll(warp_cols): + scale_n = warp_n * warp_col_tiles + j * micro_size_y + sfb_col + for kk in T.vectorized(num_k_blocks): + SFB_pack_buf[j * num_k_blocks + kk] = SFB_data[ + tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_start + kk) + ] + if replicate_b: + SFB_rep_pack_buf[j * num_k_blocks + kk] = SFB_data[ + tuple(SFB_other) + (SFB_base_n + scale_n + 8, SFB_base_k + scale_b_word_start + kk) + ] + + return _warp_ldscale_kpack( + SFA_pack_buf, + SFB_pack_buf, + SFB_rep_pack_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def mma_backend_kblock_fulltile( + self, + A_shared_buf: Buffer | BufferRegion, + B_shared_buf: Buffer | BufferRegion, + C_local_buf: Buffer, + SFA_buf: Buffer | BufferRegion, + SFB_buf: Buffer | BufferRegion, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + ): + """Emit one SM120 NVF4 K=64 full-tile helper call. + + This path deliberately keeps A/B/SFA/SFB fragments out of TIR local + buffers. The CUDA helper owns the ldmatrix loads, scale loads, and the + full set of m16n8k64 block-scaled MMA atoms for one K block. + """ + if int(self.warp_rows) != 4 or int(self.warp_cols) != 4: + raise ValueError("sm120 backend fulltile helper requires warp_rows=4 and warp_cols=4") + if self.n_dim != 16: + raise ValueError("sm120 backend fulltile helper requires replicated B n_dim=16") + if not self.b_transposed: + raise ValueError("sm120 backend fulltile helper currently requires transpose_B=True") + + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_a = self.local_size_a + local_size_b = self.local_size_b + sf_vec_size = self.sf_vec_size + a_transposed = self.a_transposed + + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + if sf_layout not in ("rowmajor", "blockscaled_chunk_kmajor"): + raise ValueError(f"Unsupported SM120 scale layout: {sf_layout}") + + A_region = self._legalize_to_buffer_region(A_shared_buf) + A_buf = A_region.buffer + A_base0 = A_region.region[-2].min + A_base1 = A_region.region[-1].min + A_other = [r.min for r in A_region.region[:-2]] + + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + + @T.macro + def _warp_backend_kblock( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + thread_binding, + ): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + a_row_off, a_col_off = ldmatrix_32x32_to_shared_16x64_layout_a(tx) + b_row_off, b_col_off = ldmatrix_32x32_to_shared_16x64_layout_b(tx) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + wk = ki * micro_size_k + + if a_transposed: + a0 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + a_col_off)], + "r", + extent=local_size_a, + ) + a1 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + a2 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + 2 * micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + a3 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + 3 * micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + else: + a0 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a1 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a2 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + 2 * micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a3 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + 3 * micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + + b0 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b1 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b2 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + 2 * micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b3 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + 3 * micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + + scale_m0 = warp_m * warp_row_tiles + sfa_row + scale_m1 = warp_m * warp_row_tiles + micro_size_x + sfa_row + scale_m2 = warp_m * warp_row_tiles + 2 * micro_size_x + sfa_row + scale_m3 = warp_m * warp_row_tiles + 3 * micro_size_x + sfa_row + scale_n0 = warp_n * warp_col_tiles + sfb_col + scale_n1 = warp_n * warp_col_tiles + micro_size_y + sfb_col + scale_n2 = warp_n * warp_col_tiles + 2 * micro_size_y + sfb_col + scale_n3 = warp_n * warp_col_tiles + 3 * micro_size_y + sfb_col + + if sf_layout == "blockscaled_chunk_kmajor": + sfa0 = T.access_ptr(SFA_data[tuple(SFA_other) + (SFA_base_m, SFA_base_k)], "r") + sfa1 = sfa0 + sfa2 = sfa0 + sfa3 = sfa0 + sfb0 = T.access_ptr(SFB_data[tuple(SFB_other) + (SFB_base_n, SFB_base_k)], "r") + sfb1 = sfb0 + sfb2 = sfb0 + sfb3 = sfb0 + sfb_rep0 = sfb0 + sfb_rep1 = sfb0 + sfb_rep2 = sfb0 + sfb_rep3 = sfb0 + else: + sfa0 = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m0, SFA_base_k + scale_a_word_k)], + "r", + ) + sfa1 = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m1, SFA_base_k + scale_a_word_k)], + "r", + ) + sfa2 = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m2, SFA_base_k + scale_a_word_k)], + "r", + ) + sfa3 = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m3, SFA_base_k + scale_a_word_k)], + "r", + ) + sfb0 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n0, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb1 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n1, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb2 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n2, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb3 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n3, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb_rep0 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n0 + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb_rep1 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n1 + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb_rep2 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n2 + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + sfb_rep3 = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n3 + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + + T.sm120_mma_blockscaled_kblock_fulltile( + C_local_buf.data, + 0, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa0, + sfa1, + sfa2, + sfa3, + sfb0, + sfb1, + sfb2, + sfb3, + sfb_rep0, + sfb_rep1, + sfb_rep2, + sfb_rep3, + ki, + ) + + return _warp_backend_kblock( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def mma_backend_kblock_fulltile_ab_owner_wide( + self, + A_shared_buf: Buffer | BufferRegion, + B_shared_buf: Buffer | BufferRegion, + C_local_buf: Buffer, + SFA_buf: Buffer | BufferRegion, + SFB_buf: Buffer | BufferRegion, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + sf_layout: str = "rowmajor", + backend_op: str = "ab_owner_wide", + ): + """Emit one SM120 K=64 helper call with backend-owned scale owners. + + This is the backend equivalent of the earlier owner-wide TIR probe. It + passes only packed A/B source pointers and SFA/SFB base pointers to C++, + so the owner words are scalar helper locals rather than TIR local arrays. + """ + if backend_op not in ("ab_owner_wide", "afull_bpanel_owner_wide"): + raise ValueError(f"Unsupported SM120 backend owner-wide op: {backend_op}") + if int(self.warp_rows) != 4 or int(self.warp_cols) != 4: + raise ValueError("sm120 backend owner-wide helper requires warp_rows=4 and warp_cols=4") + if self.n_dim != 16: + raise ValueError("sm120 backend owner-wide helper requires replicated B n_dim=16") + if not self.b_transposed: + raise ValueError("sm120 backend owner-wide helper currently requires transpose_B=True") + if sf_layout != "blockscaled_chunk_kmajor": + raise ValueError("sm120 backend owner-wide helper currently requires sf_layout='blockscaled_chunk_kmajor'") + + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_a = self.local_size_a + local_size_b = self.local_size_b + a_transposed = self.a_transposed + thread_binding = self.get_thread_binding() + + A_region = self._legalize_to_buffer_region(A_shared_buf) + A_buf = A_region.buffer + A_base0 = A_region.region[-2].min + A_base1 = A_region.region[-1].min + A_other = [r.min for r in A_region.region[:-2]] + + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + + @T.macro + def _warp_backend_kblock_owner_wide( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + thread_binding, + ): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + a_row_off, a_col_off = ldmatrix_32x32_to_shared_16x64_layout_a(tx) + b_row_off, b_col_off = ldmatrix_32x32_to_shared_16x64_layout_b(tx) + wk = ki * micro_size_k + + if a_transposed: + a0 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + a_col_off)], + "r", + extent=local_size_a, + ) + a1 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + a2 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + 2 * micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + a3 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + wk + a_row_off, A_base1 + warp_m * warp_row_tiles + 3 * micro_size_x + a_col_off)], + "r", + extent=local_size_a, + ) + else: + a0 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a1 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a2 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + 2 * micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + a3 = T.access_ptr( + A_buf[tuple(A_other) + (A_base0 + warp_m * warp_row_tiles + 3 * micro_size_x + a_row_off, A_base1 + wk + a_col_off)], + "r", + extent=local_size_a, + ) + + b0 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b1 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b2 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + 2 * micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + b3 = T.access_ptr( + B_buf[tuple(B_other) + (B_base0 + warp_n * warp_col_tiles + 3 * micro_size_y + b_row_off, B_base1 + wk + b_col_off)], + "r", + extent=local_size_b, + ) + sfa_base = T.access_ptr(SFA_data[tuple(SFA_other) + (SFA_base_m, SFA_base_k)], "r", extent=1) + sfb_base = T.access_ptr(SFB_data[tuple(SFB_other) + (SFB_base_n, SFB_base_k)], "r", extent=1) + + if backend_op == "afull_bpanel_owner_wide": + T.sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide( + C_local_buf.data, + 0, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki, + ) + else: + T.sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide( + C_local_buf.data, + 0, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki, + ) + + return _warp_backend_kblock_owner_wide( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + thread_binding, + ) + + def mma_backend_cute_consumer_bridge( + self, + A_shared_buf: Buffer | BufferRegion, + B_shared_buf: Buffer | BufferRegion, + C_local_buf: Buffer, + SFA_buf: Buffer | BufferRegion, + SFB_buf: Buffer | BufferRegion, + ki: PrimExpr = 0, + ): + """Emit a private SM120 full shared-tile consumer bridge call. + + Unlike mma_backend_kblock_fulltile, this hook passes whole A/B/SFA/SFB + shared tile bases to CUDA. It is the lowering contract needed for a + backend-owned blockscaled operand package. + """ + if int(self.warp_rows) != 4 or int(self.warp_cols) != 4: + raise ValueError("sm120 CuTe consumer bridge requires warp_rows=4 and warp_cols=4") + if self.n_dim != 16: + raise ValueError("sm120 CuTe consumer bridge requires replicated B n_dim=16") + if not self.b_transposed: + raise ValueError("sm120 CuTe consumer bridge currently requires transpose_B=True") + + A_region = self._legalize_to_buffer_region(A_shared_buf) + A_buf = A_region.buffer + A_base0 = A_region.region[-2].min + A_base1 = A_region.region[-1].min + A_other = [r.min for r in A_region.region[:-2]] + + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + + @T.macro + def _warp_backend_cute_consumer_bridge( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + ): + a_base = T.access_ptr( + A_buf[tuple(A_other) + (A_base0, A_base1)], + "r", + extent=1, + ) + b_base = T.access_ptr( + B_buf[tuple(B_other) + (B_base0, B_base1)], + "r", + extent=1, + ) + sfa_base = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m, SFA_base_k)], + "r", + extent=1, + ) + sfb_base = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n, SFB_base_k)], + "r", + extent=1, + ) + T.sm120_mma_blockscaled_cute_consumer_bridge( + C_local_buf.data, + 0, + a_base, + b_base, + sfa_base, + sfb_base, + ki, + ) + + return _warp_backend_cute_consumer_bridge( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + ) + + def mma_backend_kblock_fulltile_package_pingpong( + self, + A_shared_buf: Buffer | BufferRegion, + B_shared_buf: Buffer | BufferRegion, + C_local_buf: Buffer, + SFA_buf: Buffer | BufferRegion, + SFB_buf: Buffer | BufferRegion, + sf_layout: str = "rowmajor", + ): + """Emit one backend-owned SM120 package lifecycle helper call. + + The CUDA helper receives full A/B/SFA/SFB shared K-stage bases and owns + the copy_kblock_package(next) -> gemm_kblock_package(current) schedule. + """ + if int(self.warp_rows) != 4 or int(self.warp_cols) != 4: + raise ValueError("sm120 package pingpong helper requires warp_rows=4 and warp_cols=4") + if self.n_dim != 16: + raise ValueError("sm120 package pingpong helper requires replicated B n_dim=16") + if not self.b_transposed: + raise ValueError("sm120 package pingpong helper currently requires transpose_B=True") + if sf_layout != "blockscaled_chunk_kmajor": + raise ValueError("sm120 package pingpong helper currently requires sf_layout='blockscaled_chunk_kmajor'") + package_contract = SM120BlockScaledFullTilePackageContract.for_package_pingpong( + self, + sf_layout=sf_layout, + kblocks=4, + ) + assert package_contract.issue_count_per_warpgroup == 128 + + A_region = self._legalize_to_buffer_region(A_shared_buf) + A_buf = A_region.buffer + A_base0 = A_region.region[-2].min + A_base1 = A_region.region[-1].min + A_other = [r.min for r in A_region.region[:-2]] + + B_region = self._legalize_to_buffer_region(B_shared_buf) + B_buf = B_region.buffer + B_base0 = B_region.region[-2].min + B_base1 = B_region.region[-1].min + B_other = [r.min for r in B_region.region[:-2]] + + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + + @T.macro + def _warp_backend_package_pingpong( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + ): + a_base = T.access_ptr( + A_buf[tuple(A_other) + (A_base0, A_base1)], + "r", + extent=1, + ) + b_base = T.access_ptr( + B_buf[tuple(B_other) + (B_base0, B_base1)], + "r", + extent=1, + ) + sfa_base = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m, SFA_base_k)], + "r", + extent=1, + ) + sfb_base = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n, SFB_base_k)], + "r", + extent=1, + ) + T.sm120_mma_blockscaled_kblock_fulltile_package_pingpong( + C_local_buf.data, + 0, + a_base, + b_base, + sfa_base, + sfb_base, + ) + + return _warp_backend_package_pingpong( + A_buf, + B_buf, + C_local_buf, + SFA_data, + SFB_data, + ) + + def mma_with_scale_fragments( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + ): + """Issue SM120 block-scaled MMA using prefetched scale fragments.""" + return self.mma_with_prefetched_scales( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + ) + + def mma_atom_with_scale_fragments( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + inst_m_idx: PrimExpr | int, + inst_n_idx: PrimExpr | int, + ): + """Issue one SM120 block-scaled MMA atom with scale fragments.""" + return self.mma_atom_with_prefetched_scales( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + inst_m_idx, + inst_n_idx, + ) + + def mma_full_b_atom_with_scale_fragments( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + inst_m_idx: PrimExpr | int, + inst_n_idx: PrimExpr | int, + ): + """Issue one SM120 block-scaled MMA atom from a full B fragment tile.""" + return self.mma_full_b_atom_with_prefetched_scales( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + inst_m_idx, + inst_n_idx, + ) + + def mma_with_prefetched_scales( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + for i, j in T.grid(warp_rows, warp_cols): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + scale_b_ptr = T.access_ptr(SFB_local_buf[j], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[j], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_block_scale_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_with_prefetched_scales_cutlass_order( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_prefetched_cutlass_order( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + if replicate_b: + for j in T.unroll(warp_cols): + scale_b_ptr = T.access_ptr(SFB_local_buf[j], "r") + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[j], "r") + for rev_i in T.unroll(warp_rows): + i = warp_rows - 1 - rev_i + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + else: + for j in T.unroll(warp_cols): + scale_b_ptr = T.access_ptr(SFB_local_buf[j], "r") + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + + return _warp_mma_block_scale_prefetched_cutlass_order( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_with_prefetched_scales_selector_probe( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_prefetched_selector_probe( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + scale_a_thread_id = i & 1 + for j in T.unroll(warp_cols): + scale_b_ptr = T.access_ptr(SFB_local_buf[j], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_ptr, + 0, + scale_a_thread_id, + 0, + j & 3, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[j], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + 0, + scale_a_thread_id, + 0, + (j + 2) & 3, + ) + + return _warp_mma_block_scale_prefetched_selector_probe( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_with_prefetched_scales_b_owner( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_owner_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + if self.n_dim != 16 or int(warp_cols) != 2: + raise ValueError("mma_with_prefetched_scales_b_owner currently requires replicated B with warp_cols=2") + + @T.macro + def _warp_mma_block_scale_b_owner( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_owner_buf, + ): + scale_b_owner_ptr = T.access_ptr(SFB_owner_buf[0], "r") + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + for j in T.unroll(warp_cols): + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_ptr, + scale_b_owner_ptr, + 0, + 0, + 0, + j * 2, + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_owner_ptr, + 0, + 0, + 0, + j * 2 + 1, + ) + + return _warp_mma_block_scale_b_owner( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_owner_buf, + ) + + def mma_with_prefetched_scales_ab_owner( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + if int(warp_rows) != 4 or self.n_dim != 16 or int(warp_cols) != 2: + raise ValueError("mma_with_prefetched_scales_ab_owner currently requires fulltile warp_rows=4, warp_cols=2") + + @T.macro + def _warp_mma_block_scale_ab_owner( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ): + scale_b_owner_ptr = T.access_ptr(SFB_owner_buf[0], "r") + for g in T.unroll(2): + scale_a_owner_ptr = T.access_ptr(SFA_owner_buf[g], "r") + for ai in T.unroll(2): + for j in T.unroll(warp_cols): + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + (g * 2 + ai) * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + (g * 2 + ai) * warp_cols * local_size_out + j * local_size_out, + scale_a_owner_ptr, + scale_b_owner_ptr, + 0, + ai, + 0, + j * 2, + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + (g * 2 + ai) * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + (g * 2 + ai) * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_owner_ptr, + scale_b_owner_ptr, + 0, + ai, + 0, + j * 2 + 1, + ) + + return _warp_mma_block_scale_ab_owner( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ) + + def mma_with_prefetched_scales_ab_owner_wide( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + if int(warp_rows) != 4 or self.n_dim != 16 or int(warp_cols) != 4: + raise ValueError("mma_with_prefetched_scales_ab_owner_wide requires warp_rows=4, warp_cols=4") + + @T.macro + def _warp_mma_block_scale_ab_owner_wide( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ): + for gm in T.unroll(2): + scale_a_owner_ptr = T.access_ptr(SFA_owner_buf[gm], "r") + for ai in T.unroll(2): + i = gm * 2 + ai + for gn in T.unroll(2): + scale_b_owner_ptr = T.access_ptr(SFB_owner_buf[gn], "r") + for ji in T.unroll(2): + j = gn * 2 + ji + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + scale_a_owner_ptr, + scale_b_owner_ptr, + 0, + ai, + 0, + ji * 2, + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + scale_a_owner_ptr, + scale_b_owner_ptr, + 0, + ai, + 0, + ji * 2 + 1, + ) + + return _warp_mma_block_scale_ab_owner_wide( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_owner_buf, + SFB_owner_buf, + ) + + def mma_atom_with_prefetched_scales( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + inst_m_idx: PrimExpr | int, + inst_n_idx: PrimExpr | int, + ): + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + warp_cols = self.warp_cols + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_atom_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + scale_a_ptr = T.access_ptr(SFA_local_buf[inst_m_idx], "r") + scale_b_ptr = T.access_ptr(SFB_local_buf[inst_n_idx], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + 0, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[inst_n_idx], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + lift(local_size_b) // 2, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_block_scale_atom_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_b_atom_n8_serpentine_with_prefetched_scales( + self, + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + inst_n_idx: PrimExpr | int, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_b_atom_n8_serpentine_prefetched( + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + scale_b_ptr = T.access_ptr(SFB_local_buf[inst_n_idx], "r") + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_atom_buf.data, + 0, + C_local_buf.data, + i * warp_cols * local_size_out + inst_n_idx * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[inst_n_idx], "r") + for rev_i in T.unroll(warp_rows): + i = warp_rows - 1 - rev_i + scale_a_ptr = T.access_ptr(SFA_local_buf[i], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_atom_buf.data, + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_b_atom_n8_serpentine_prefetched( + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_b_atom_n8_serpentine_with_scale_pack( + self, + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_pack_buf, + SFB_pack_buf, + SFB_rep_pack_buf, + k_block: PrimExpr | int, + inst_n_idx: PrimExpr | int, + num_k_blocks: int, + ): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_b_atom_n8_serpentine_scale_pack( + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_pack_buf, + SFB_pack_buf, + SFB_rep_pack_buf, + ): + scale_b_ptr = T.access_ptr(SFB_pack_buf[inst_n_idx * num_k_blocks + k_block], "r") + for i in T.unroll(warp_rows): + scale_a_ptr = T.access_ptr(SFA_pack_buf[i * num_k_blocks + k_block], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_atom_buf.data, + 0, + C_local_buf.data, + i * warp_cols * local_size_out + inst_n_idx * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_pack_buf[inst_n_idx * num_k_blocks + k_block], "r") + for rev_i in T.unroll(warp_rows): + i = warp_rows - 1 - rev_i + scale_a_ptr = T.access_ptr(SFA_pack_buf[i * num_k_blocks + k_block], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + i * local_size_a, + B_atom_buf.data, + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_b_atom_n8_serpentine_scale_pack( + A_local_buf, + B_atom_buf, + C_local_buf, + SFA_pack_buf, + SFB_pack_buf, + SFB_rep_pack_buf, + ) + + def mma_full_b_atom_with_prefetched_scales( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + inst_m_idx: PrimExpr | int, + inst_n_idx: PrimExpr | int, + ): + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + warp_cols = self.warp_cols + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_full_b_atom_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ): + scale_a_ptr = T.access_ptr(SFA_local_buf[inst_m_idx], "r") + scale_b_ptr = T.access_ptr(SFB_local_buf[inst_n_idx], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + inst_n_idx * local_size_b, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr(SFB_rep_local_buf[inst_n_idx], "r") + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + inst_n_idx * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_block_scale_full_b_atom_prefetched( + A_local_buf, + B_local_buf, + C_local_buf, + SFA_local_buf, + SFB_local_buf, + SFB_rep_local_buf, + ) + + def mma_atom( + self, + A_local_buf, + B_local_buf, + C_local_buf, + SFA_buf, + SFB_buf, + inst_m_idx: PrimExpr | int, + inst_n_idx: PrimExpr | int, + ki: PrimExpr = 0, + k_start: PrimExpr = 0, + sf_a_granularity_k: int | None = None, + sf_b_granularity_k: int | None = None, + ): + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + kind = self.kind + scale_vec_size = self.scale_vec_size + stype = self.stype + accum_dtype = self.accum_dtype + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + mma_prefix = self.mma_prefix + warp_cols = self.warp_cols + warp_row_tiles = self.warp_row_tiles + warp_col_tiles = self.warp_col_tiles + micro_size_x = self.micro_size_x + micro_size_y = self.micro_size_y + sf_vec_size = self.sf_vec_size + sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k + sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k + scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k) + scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k) + thread_binding = self.get_thread_binding() + SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf) + SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf) + replicate_b = self.n_dim == 16 + + @T.macro + def _warp_mma_block_scale_atom(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding): + tx, warp_n, warp_m = self.extract_thread_binding(thread_binding) + sfa_row = self._sfa_row_in_atom(tx) + sfb_col = self._sfb_col_in_atom(tx) + scale_m = warp_m * warp_row_tiles + inst_m_idx * micro_size_x + sfa_row + scale_n = warp_n * warp_col_tiles + inst_n_idx * micro_size_y + sfb_col + scale_a_ptr = T.access_ptr( + SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)], + "r", + ) + scale_b_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_k)], + "r", + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + 0, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out, + scale_a_ptr, + scale_b_ptr, + ) + if replicate_b: + scale_b_rep_ptr = T.access_ptr( + SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n + 8, SFB_base_k + scale_b_word_k)], + "r", + ) + T.ptx_mma_block_scale( + accum_dtype, + mma_prefix, + "row", + "col", + kind, + scale_vec_size, + a_dtype_abbrv, + b_dtype_abbrv, + stype, + A_local_buf.data, + inst_m_idx * local_size_a, + B_local_buf.data, + lift(local_size_b) // 2, + C_local_buf.data, + inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2, + scale_a_ptr, + scale_b_rep_ptr, + ) + + return _warp_mma_block_scale_atom(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding) + + +class SM120BlockScaledOperandPackage: + """Internal K-block package for SM120 blockscaled MMA lowering. + + This object intentionally owns no scheduling policy. It only groups the + A/B fragments and matching scale fragments so the GEMM lowering can express + CUTLASS-like copy_kblock/gemm_kblock structure without scattering the + operand lifetime rules across the lowering body. + """ + + def __init__( + self, + mma_emitter: TensorCoreIntrinEmitterWithBlockScale, + A_fragment_buf, + B_fragment_buf, + SFA_fragment_buf, + SFB_fragment_buf, + SFB_rep_fragment_buf, + A_region, + B_region, + C_buf, + SFA_region, + SFB_region, + sf_k_start, + sf_a_granularity_k: int, + sf_b_granularity_k: int, + sf_layout: str = "rowmajor", + ): + self.mma_emitter = mma_emitter + self.A_fragment_buf = A_fragment_buf + self.B_fragment_buf = B_fragment_buf + self.SFA_fragment_buf = SFA_fragment_buf + self.SFB_fragment_buf = SFB_fragment_buf + self.SFB_rep_fragment_buf = SFB_rep_fragment_buf + self.A_region = A_region + self.B_region = B_region + self.C_buf = C_buf + self.SFA_region = SFA_region + self.SFB_region = SFB_region + self.sf_k_start = sf_k_start + self.sf_a_granularity_k = sf_a_granularity_k + self.sf_b_granularity_k = sf_b_granularity_k + self.sf_layout = sf_layout + + def copy_kblock(self, k_block: PrimExpr | int): + self.mma_emitter.ldmatrix_a(self.A_fragment_buf, self.A_region, k_block) + self.mma_emitter.ldmatrix_b(self.B_fragment_buf, self.B_region, k_block) + self.mma_emitter.ldscale_fragment( + self.SFA_fragment_buf, + self.SFB_fragment_buf, + self.SFB_rep_fragment_buf, + self.SFA_region, + self.SFB_region, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=self.sf_a_granularity_k, + sf_b_granularity_k=self.sf_b_granularity_k, + sf_layout=self.sf_layout, + ) + + def gemm_kblock(self): + self.mma_emitter.mma_with_scale_fragments( + self.A_fragment_buf, + self.B_fragment_buf, + self.C_buf, + self.SFA_fragment_buf, + self.SFB_fragment_buf, + self.SFB_rep_fragment_buf, + ) + + def gemm_atom(self, inst_m_idx: PrimExpr | int, inst_n_idx: PrimExpr | int): + self.mma_emitter.mma_full_b_atom_with_scale_fragments( + self.A_fragment_buf, + self.B_fragment_buf, + self.C_buf, + self.SFA_fragment_buf, + self.SFB_fragment_buf, + self.SFB_rep_fragment_buf, + inst_m_idx, + inst_n_idx, + ) diff --git a/tilelang/cuda/op/gemm/gemm_mma.py b/tilelang/cuda/op/gemm/gemm_mma.py index 51b8b5c3d..69ed3a4eb 100644 --- a/tilelang/cuda/op/gemm/gemm_mma.py +++ b/tilelang/cuda/op/gemm/gemm_mma.py @@ -3,7 +3,9 @@ from __future__ import annotations from tilelang.tileop.gemm.gemm_base import GemmBase from tilelang.layout import make_swizzled_layout from tilelang.cuda.intrinsics.macro.mma_macro_generator import ( + SM120BlockScaledOperandPackage, TensorCoreIntrinEmitter, + TensorCoreIntrinEmitterWithBlockScale, ) from tilelang.utils.language import is_shared, is_fragment, is_full_region from tilelang import tvm as tvm @@ -17,14 +19,30 @@ from tilelang.transform.simplify import _Simplify GEMM_INST_MMA = "cuda.mma" +def _is_explicit_non_sm120_cuda(target: Target) -> bool: + if target.kind.name != "cuda": + return False + arch = target.attrs.get("arch", None) + if arch is None: + return False + arch_str = str(arch) + if not arch_str.startswith("sm_"): + return False + arch_int = int(arch_str[3:].rstrip("af")) + return arch_int < 120 or arch_int >= 130 + + class GemmMMA(GemmBase): intrin_emitter_cls = TensorCoreIntrinEmitter def _make_mma_emitter(self, target: Target, thread_nums: int, thread_var: tirx.Var | None = None): + if self.is_blockscaled and _is_explicit_non_sm120_cuda(target): + raise ValueError("T.mma_gemm_blockscaled requires SM120 CUDA target") m_warp, n_warp = self.policy.compute_warp_partition(self.M, self.N, thread_nums, target, GEMM_INST_MMA) warp_row_tiles = int(self.M // m_warp) warp_col_tiles = int(self.N // n_warp) - emitter = self.intrin_emitter_cls( + emitter_cls = TensorCoreIntrinEmitterWithBlockScale if self.is_blockscaled else self.intrin_emitter_cls + emitter = emitter_cls( a_dtype=self.a_dtype, b_dtype=self.b_dtype, accum_dtype=self.accum_dtype, @@ -41,6 +59,8 @@ class GemmMMA(GemmBase): def infer_layout(self, target: Target, thread_nums: int): mma_emitter = self._make_mma_emitter(target, thread_nums) + if self.is_blockscaled and not self.is_gemm_ss(): + raise ValueError("T.mma_gemm_blockscaled supports shared-memory A/B operands only") if self.is_gemm_ss(): return { self.A: make_swizzled_layout(self.A), @@ -100,12 +120,1441 @@ class GemmMMA(GemmBase): C_buf = C_region.buffer clear_accum = self.clear_accum - assert block_K >= micro_size_k, f"block_K ({block_K}) must be >= micro_size_k ({micro_size_k})" assert block_K % micro_size_k == 0, f"block_K ({block_K}) must be a multiple of micro_size_k ({micro_size_k})" assert is_full_region(C_region), "Fragment output C must be a full region" + if self.is_blockscaled: + if not self.is_gemm_ss(): + raise ValueError("T.mma_gemm_blockscaled supports shared-memory A/B operands only") + annotations = getattr(self.gemm_node, "annotations", {}) + sf_a_granularity_k = annotations.get("sf_a_granularity_k") + sf_b_granularity_k = annotations.get("sf_b_granularity_k") + micro_pipeline = annotations.get("micro_pipeline") + sf_layout = annotations.get("sf_layout", "rowmajor") + if sf_layout not in ("rowmajor", "blockscaled_chunk_kmajor"): + raise ValueError(f"Unsupported SM120 scale layout: {sf_layout}") + if sf_a_granularity_k is None or sf_b_granularity_k is None: + raise ValueError("Block-scaled MMA GEMM requires sf_a_granularity_k and sf_b_granularity_k") + + if micro_pipeline == "b_pingpong_scale_prefetch": + + @T.prim_func + def _gemm_ss_blockscaled_b_pingpong_scale_prefetch() -> None: + A_local = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((local_size_b), b_dtype) + B_local_1 = T.alloc_local((local_size_b), b_dtype) + SFA_local = T.alloc_local((warp_rows), "uint32") + SFB_local = T.alloc_local((warp_cols), "uint32") + SFB_rep_local = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + for ki in T.serial(0, (block_K // micro_size_k)): + mma_emitter.ldmatrix_a(A_local, A_region, ki) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=ki, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + ) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, ki, 0) + if warp_cols > 1: + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, ki, 1) + for j in T.unroll(warp_cols): + if j % 2 == 0: + for i in T.unroll(warp_rows): + mma_emitter.mma_atom_with_scale_fragments( + A_local, + B_local_0, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + i, + j, + ) + else: + for i in T.unroll(warp_rows): + mma_emitter.mma_atom_with_scale_fragments( + A_local, + B_local_1, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + i, + j, + ) + next_j = j + 2 + if next_j < warp_cols: + if next_j % 2 == 0: + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, ki, next_j) + else: + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, ki, next_j) + + return _Simplify(_gemm_ss_blockscaled_b_pingpong_scale_prefetch, inline_let=True) + + if micro_pipeline == "b_pingpong": + + @T.prim_func + def _gemm_ss_blockscaled_b_pingpong() -> None: + A_local = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((local_size_b), b_dtype) + B_local_1 = T.alloc_local((local_size_b), b_dtype) + if clear_accum: + T.clear(C_buf) + for ki in T.serial(0, (block_K // micro_size_k)): + mma_emitter.ldmatrix_a(A_local, A_region, ki) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, ki, 0) + if warp_cols > 1: + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, ki, 1) + for j in T.unroll(warp_cols): + if j % 2 == 0: + for i in T.unroll(warp_rows): + mma_emitter.mma_atom( + A_local, + B_local_0, + C_buf, + self.SFARegion, + self.SFBRegion, + i, + j, + ki=ki, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + ) + else: + for i in T.unroll(warp_rows): + mma_emitter.mma_atom( + A_local, + B_local_1, + C_buf, + self.SFARegion, + self.SFBRegion, + i, + j, + ki=ki, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + ) + next_j = j + 2 + if next_j < warp_cols: + if next_j % 2 == 0: + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, ki, next_j) + else: + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, ki, next_j) + + return _Simplify(_gemm_ss_blockscaled_b_pingpong, inline_let=True) + + if micro_pipeline == "k_static_b_pingpong_scale_stream": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static_b_pingpong_scale_stream' currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 1 or int(warp_cols) != 8: + raise ValueError("micro_pipeline='k_static_b_pingpong_scale_stream' currently targets warp_rows=1, warp_cols=8") + + @T.prim_func + def _gemm_ss_blockscaled_k_static_b_pingpong_scale_stream() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((local_size_b), b_dtype) + B_local_1 = T.alloc_local((local_size_b), b_dtype) + SFA_local_0 = T.alloc_local((warp_rows), "uint32") + SFA_local_1 = T.alloc_local((warp_rows), "uint32") + SFB_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_local_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldscale_fragment( + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + mma_emitter.ldscale_fragment( + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 0, 0) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 0, 1) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 0) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 0, 2) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 1) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 0, 3) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 2) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 0, 4) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 3) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 0, 5) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 4) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 0, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 5) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 0, 7) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 7) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + mma_emitter.ldscale_fragment( + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 1, 0) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 1, 1) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 0) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 1, 2) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 1) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 1, 3) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 2) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 1, 4) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 3) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 1, 5) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 4) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 1, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 5) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 1, 7) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 7) + + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + mma_emitter.ldscale_fragment( + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 2, 0) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 2, 1) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 0) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 2, 2) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 1) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 2, 3) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 2) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 2, 4) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 3) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 2, 5) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 4) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 2, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 5) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 2, 7) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_0, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_0, B_local_1, C_buf, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0, 7) + + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 3, 0) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 3, 1) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 0) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 3, 2) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 1) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 3, 3) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 2) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 3, 4) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 3) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 3, 5) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 4) + mma_emitter.ldmatrix_b_atom(B_local_0, B_region, 3, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 5) + mma_emitter.ldmatrix_b_atom(B_local_1, B_region, 3, 7) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_0, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 6) + mma_emitter.mma_atom_with_scale_fragments(A_local_1, B_local_1, C_buf, SFA_local_1, SFB_local_1, SFB_rep_local_1, 0, 7) + + return _Simplify(_gemm_ss_blockscaled_k_static_b_pingpong_scale_stream, inline_let=True) + + if micro_pipeline == "k_static_b_atom_n8_stream": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static_b_atom_n8_stream' currently requires block_K / micro_size_k == 4") + if int(warp_cols) != 8: + raise ValueError("micro_pipeline='k_static_b_atom_n8_stream' currently targets warp_cols=8") + + def copy_a_scale_kblock(A_frag, SFA_frag, SFB_frag, SFB_rep_frag, k_block): + mma_emitter.ldmatrix_a(A_frag, A_region, k_block) + mma_emitter.ldscale_fragment( + SFA_frag, + SFB_frag, + SFB_rep_frag, + self.SFARegion, + self.SFBRegion, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + @T.macro + def gemm_b_atom_stream(A_frag, B_atom_0, B_atom_1, SFA_frag, SFB_frag, SFB_rep_frag, k_block): + mma_emitter.ldmatrix_b_atom(B_atom_0, B_region, k_block, 0) + mma_emitter.ldmatrix_b_atom(B_atom_1, B_region, k_block, 1) + for j in T.unroll(warp_cols): + if j % 2 == 0: + mma_emitter.mma_b_atom_n8_serpentine_with_prefetched_scales( + A_frag, + B_atom_0, + C_buf, + SFA_frag, + SFB_frag, + SFB_rep_frag, + j, + ) + else: + mma_emitter.mma_b_atom_n8_serpentine_with_prefetched_scales( + A_frag, + B_atom_1, + C_buf, + SFA_frag, + SFB_frag, + SFB_rep_frag, + j, + ) + next_j = j + 2 + if next_j < warp_cols: + if next_j % 2 == 0: + mma_emitter.ldmatrix_b_atom(B_atom_0, B_region, k_block, next_j) + else: + mma_emitter.ldmatrix_b_atom(B_atom_1, B_region, k_block, next_j) + + @T.prim_func + def _gemm_ss_blockscaled_k_static_b_atom_n8_stream() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_atom_0 = T.alloc_local((local_size_b), b_dtype) + B_atom_1 = T.alloc_local((local_size_b), b_dtype) + SFA_local_0 = T.alloc_local((warp_rows), "uint32") + SFA_local_1 = T.alloc_local((warp_rows), "uint32") + SFB_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_local_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + copy_a_scale_kblock(A_local_0, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0) + copy_a_scale_kblock(A_local_1, SFA_local_1, SFB_local_1, SFB_rep_local_1, 1) + gemm_b_atom_stream(A_local_0, B_atom_0, B_atom_1, SFA_local_0, SFB_local_0, SFB_rep_local_0, 0) + copy_a_scale_kblock(A_local_0, SFA_local_0, SFB_local_0, SFB_rep_local_0, 2) + gemm_b_atom_stream(A_local_1, B_atom_0, B_atom_1, SFA_local_1, SFB_local_1, SFB_rep_local_1, 1) + copy_a_scale_kblock(A_local_1, SFA_local_1, SFB_local_1, SFB_rep_local_1, 3) + gemm_b_atom_stream(A_local_0, B_atom_0, B_atom_1, SFA_local_0, SFB_local_0, SFB_rep_local_0, 2) + gemm_b_atom_stream(A_local_1, B_atom_0, B_atom_1, SFA_local_1, SFB_local_1, SFB_rep_local_1, 3) + + return _Simplify(_gemm_ss_blockscaled_k_static_b_atom_n8_stream, inline_let=True) + + if micro_pipeline == "k_static_b_atom_n8_stream_vecsf": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static_b_atom_n8_stream_vecsf' currently requires block_K / micro_size_k == 4") + if int(warp_cols) != 8: + raise ValueError("micro_pipeline='k_static_b_atom_n8_stream_vecsf' currently targets warp_cols=8") + + @T.macro + def gemm_b_atom_stream_vecsf(A_frag, B_atom_0, B_atom_1, SFA_pack, SFB_pack, SFB_rep_pack, k_block): + mma_emitter.ldmatrix_b_atom(B_atom_0, B_region, k_block, 0) + mma_emitter.ldmatrix_b_atom(B_atom_1, B_region, k_block, 1) + for j in T.unroll(warp_cols): + if j % 2 == 0: + mma_emitter.mma_b_atom_n8_serpentine_with_scale_pack( + A_frag, + B_atom_0, + C_buf, + SFA_pack, + SFB_pack, + SFB_rep_pack, + k_block, + j, + num_k_blocks, + ) + else: + mma_emitter.mma_b_atom_n8_serpentine_with_scale_pack( + A_frag, + B_atom_1, + C_buf, + SFA_pack, + SFB_pack, + SFB_rep_pack, + k_block, + j, + num_k_blocks, + ) + next_j = j + 2 + if next_j < warp_cols: + if next_j % 2 == 0: + mma_emitter.ldmatrix_b_atom(B_atom_0, B_region, k_block, next_j) + else: + mma_emitter.ldmatrix_b_atom(B_atom_1, B_region, k_block, next_j) + + @T.prim_func + def _gemm_ss_blockscaled_k_static_b_atom_n8_stream_vecsf() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_atom_0 = T.alloc_local((local_size_b), b_dtype) + B_atom_1 = T.alloc_local((local_size_b), b_dtype) + SFA_pack = T.alloc_local((warp_rows * num_k_blocks), "uint32") + SFB_pack = T.alloc_local((warp_cols * num_k_blocks), "uint32") + SFB_rep_pack = T.alloc_local((warp_cols * num_k_blocks), "uint32") + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldscale_fragment_kpack( + SFA_pack, + SFB_pack, + SFB_rep_pack, + self.SFARegion, + self.SFBRegion, + num_k_blocks=num_k_blocks, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + gemm_b_atom_stream_vecsf(A_local_0, B_atom_0, B_atom_1, SFA_pack, SFB_pack, SFB_rep_pack, 0) + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + gemm_b_atom_stream_vecsf(A_local_1, B_atom_0, B_atom_1, SFA_pack, SFB_pack, SFB_rep_pack, 1) + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + gemm_b_atom_stream_vecsf(A_local_0, B_atom_0, B_atom_1, SFA_pack, SFB_pack, SFB_rep_pack, 2) + gemm_b_atom_stream_vecsf(A_local_1, B_atom_0, B_atom_1, SFA_pack, SFB_pack, SFB_rep_pack, 3) + + return _Simplify(_gemm_ss_blockscaled_k_static_b_atom_n8_stream_vecsf, inline_let=True) + + if micro_pipeline == "k_static_full_b_atom_scale_stream": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static_full_b_atom_scale_stream' currently requires block_K / micro_size_k == 4") + + @T.prim_func + def _gemm_ss_blockscaled_k_static_full_b_atom_scale_stream() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_local_0 = T.alloc_local((warp_rows), "uint32") + SFA_local_1 = T.alloc_local((warp_rows), "uint32") + SFB_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_local_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldmatrix_b(B_local_0, B_region, 0) + mma_emitter.ldscale_fragment( + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + mma_emitter.ldmatrix_b(B_local_1, B_region, 1) + mma_emitter.ldscale_fragment( + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + mma_emitter.mma_full_b_atom_with_scale_fragments( + A_local_0, + B_local_0, + C_buf, + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + i, + j, + ) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + mma_emitter.ldmatrix_b(B_local_0, B_region, 2) + mma_emitter.ldscale_fragment( + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + mma_emitter.mma_full_b_atom_with_scale_fragments( + A_local_1, + B_local_1, + C_buf, + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + i, + j, + ) + + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + mma_emitter.ldmatrix_b(B_local_1, B_region, 3) + mma_emitter.ldscale_fragment( + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + mma_emitter.mma_full_b_atom_with_scale_fragments( + A_local_0, + B_local_0, + C_buf, + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + i, + j, + ) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + mma_emitter.mma_full_b_atom_with_scale_fragments( + A_local_1, + B_local_1, + C_buf, + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + i, + j, + ) + + return _Simplify(_gemm_ss_blockscaled_k_static_full_b_atom_scale_stream, inline_let=True) + + if micro_pipeline in ( + "sm120_kblock_fulltile", + "sm120_kblock_fulltile_selector_probe", + "sm120_kblock_fulltile_b_owner_probe", + "sm120_kblock_fulltile_ab_owner_probe", + ): + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 4 or int(warp_cols) != 2: + raise ValueError( + f"micro_pipeline={micro_pipeline!r} currently targets the balanced " + "128x128 full-tile consumer shape with warp_rows=4 and warp_cols=2" + ) + + def copy_kblock(A_frag, B_frag, SFA_frag, SFB_frag, SFB_rep_frag, k_block): + mma_emitter.ldmatrix_a(A_frag, A_region, k_block) + mma_emitter.ldmatrix_b(B_frag, B_region, k_block) + if micro_pipeline == "sm120_kblock_fulltile_ab_owner_probe": + mma_emitter.ldscale_fragment_ab_owner( + SFA_frag, + SFB_frag, + self.SFARegion, + self.SFBRegion, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + elif micro_pipeline == "sm120_kblock_fulltile_b_owner_probe": + mma_emitter.ldscale_fragment_b_owner( + SFA_frag, + SFB_frag, + self.SFARegion, + self.SFBRegion, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + else: + mma_emitter.ldscale_fragment( + SFA_frag, + SFB_frag, + SFB_rep_frag, + self.SFARegion, + self.SFBRegion, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + def gemm_kblock(A_frag, B_frag, SFA_frag, SFB_frag, SFB_rep_frag): + if micro_pipeline == "sm120_kblock_fulltile_ab_owner_probe": + mma_emitter.mma_with_prefetched_scales_ab_owner( + A_frag, + B_frag, + C_buf, + SFA_frag, + SFB_frag, + ) + elif micro_pipeline == "sm120_kblock_fulltile_b_owner_probe": + mma_emitter.mma_with_prefetched_scales_b_owner( + A_frag, + B_frag, + C_buf, + SFA_frag, + SFB_frag, + ) + elif micro_pipeline == "sm120_kblock_fulltile_selector_probe": + mma_emitter.mma_with_prefetched_scales_selector_probe( + A_frag, + B_frag, + C_buf, + SFA_frag, + SFB_frag, + SFB_rep_frag, + ) + else: + mma_emitter.mma_with_scale_fragments( + A_frag, + B_frag, + C_buf, + SFA_frag, + SFB_frag, + SFB_rep_frag, + ) + + @T.prim_func + def _gemm_ss_blockscaled_sm120_kblock_fulltile() -> None: + A_frag_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_frag_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_frag_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_frag_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + if micro_pipeline == "sm120_kblock_fulltile_ab_owner_probe": + SFA_frag_0 = T.alloc_local((2), "uint32") + SFA_frag_1 = T.alloc_local((2), "uint32") + SFB_frag_0 = T.alloc_local((1), "uint32") + SFB_frag_1 = T.alloc_local((1), "uint32") + SFB_rep_frag_0 = T.alloc_local((1), "uint32") + SFB_rep_frag_1 = T.alloc_local((1), "uint32") + elif micro_pipeline == "sm120_kblock_fulltile_b_owner_probe": + SFA_frag_0 = T.alloc_local((warp_rows), "uint32") + SFA_frag_1 = T.alloc_local((warp_rows), "uint32") + SFB_frag_0 = T.alloc_local((1), "uint32") + SFB_frag_1 = T.alloc_local((1), "uint32") + SFB_rep_frag_0 = T.alloc_local((1), "uint32") + SFB_rep_frag_1 = T.alloc_local((1), "uint32") + else: + SFA_frag_0 = T.alloc_local((warp_rows), "uint32") + SFA_frag_1 = T.alloc_local((warp_rows), "uint32") + SFB_frag_0 = T.alloc_local((warp_cols), "uint32") + SFB_frag_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_frag_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_frag_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + copy_kblock(A_frag_0, B_frag_0, SFA_frag_0, SFB_frag_0, SFB_rep_frag_0, 0) + copy_kblock(A_frag_1, B_frag_1, SFA_frag_1, SFB_frag_1, SFB_rep_frag_1, 1) + gemm_kblock(A_frag_0, B_frag_0, SFA_frag_0, SFB_frag_0, SFB_rep_frag_0) + copy_kblock(A_frag_0, B_frag_0, SFA_frag_0, SFB_frag_0, SFB_rep_frag_0, 2) + gemm_kblock(A_frag_1, B_frag_1, SFA_frag_1, SFB_frag_1, SFB_rep_frag_1) + copy_kblock(A_frag_1, B_frag_1, SFA_frag_1, SFB_frag_1, SFB_rep_frag_1, 3) + gemm_kblock(A_frag_0, B_frag_0, SFA_frag_0, SFB_frag_0, SFB_rep_frag_0) + gemm_kblock(A_frag_1, B_frag_1, SFA_frag_1, SFB_frag_1, SFB_rep_frag_1) + + return _Simplify(_gemm_ss_blockscaled_sm120_kblock_fulltile, inline_let=True) + + if micro_pipeline == "sm120_kblock_fulltile_ab_owner_wide_probe": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError( + "micro_pipeline='sm120_kblock_fulltile_ab_owner_wide_probe' currently requires block_K / micro_size_k == 4" + ) + if int(warp_rows) != 4 or int(warp_cols) != 4: + raise ValueError( + "micro_pipeline='sm120_kblock_fulltile_ab_owner_wide_probe' currently targets warp_rows=4, warp_cols=4" + ) + + def copy_kblock(A_frag, B_frag, SFA_owner_frag, SFB_owner_frag, k_block): + mma_emitter.ldmatrix_a(A_frag, A_region, k_block) + mma_emitter.ldmatrix_b(B_frag, B_region, k_block) + mma_emitter.ldscale_fragment_ab_owner_wide( + SFA_owner_frag, + SFB_owner_frag, + self.SFARegion, + self.SFBRegion, + ki=k_block, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + def gemm_kblock(A_frag, B_frag, SFA_owner_frag, SFB_owner_frag): + mma_emitter.mma_with_prefetched_scales_ab_owner_wide( + A_frag, + B_frag, + C_buf, + SFA_owner_frag, + SFB_owner_frag, + ) + + @T.prim_func + def _gemm_ss_blockscaled_sm120_kblock_fulltile_ab_owner_wide() -> None: + A_frag_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_frag_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_frag_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_frag_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_owner_frag_0 = T.alloc_local((2), "uint32") + SFA_owner_frag_1 = T.alloc_local((2), "uint32") + SFB_owner_frag_0 = T.alloc_local((2), "uint32") + SFB_owner_frag_1 = T.alloc_local((2), "uint32") + if clear_accum: + T.clear(C_buf) + + copy_kblock(A_frag_0, B_frag_0, SFA_owner_frag_0, SFB_owner_frag_0, 0) + copy_kblock(A_frag_1, B_frag_1, SFA_owner_frag_1, SFB_owner_frag_1, 1) + gemm_kblock(A_frag_0, B_frag_0, SFA_owner_frag_0, SFB_owner_frag_0) + copy_kblock(A_frag_0, B_frag_0, SFA_owner_frag_0, SFB_owner_frag_0, 2) + gemm_kblock(A_frag_1, B_frag_1, SFA_owner_frag_1, SFB_owner_frag_1) + copy_kblock(A_frag_1, B_frag_1, SFA_owner_frag_1, SFB_owner_frag_1, 3) + gemm_kblock(A_frag_0, B_frag_0, SFA_owner_frag_0, SFB_owner_frag_0) + gemm_kblock(A_frag_1, B_frag_1, SFA_owner_frag_1, SFB_owner_frag_1) + + return _Simplify(_gemm_ss_blockscaled_sm120_kblock_fulltile_ab_owner_wide, inline_let=True) + + if micro_pipeline in ( + "sm120_backend_kblock_fulltile_ab_owner_wide", + "sm120_backend_kblock_fulltile_afull_bpanel_owner_wide", + ): + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 4 or int(warp_cols) != 4: + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently targets warp_rows=4, warp_cols=4") + if sf_layout != "blockscaled_chunk_kmajor": + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires sf_layout='blockscaled_chunk_kmajor'") + backend_op = ( + "afull_bpanel_owner_wide" + if micro_pipeline == "sm120_backend_kblock_fulltile_afull_bpanel_owner_wide" + else "ab_owner_wide" + ) + + @T.prim_func + def _gemm_ss_blockscaled_sm120_backend_kblock_fulltile_ab_owner_wide() -> None: + if clear_accum: + T.clear(C_buf) + + mma_emitter.mma_backend_kblock_fulltile_ab_owner_wide( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + backend_op=backend_op, + ) + mma_emitter.mma_backend_kblock_fulltile_ab_owner_wide( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + backend_op=backend_op, + ) + mma_emitter.mma_backend_kblock_fulltile_ab_owner_wide( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + backend_op=backend_op, + ) + mma_emitter.mma_backend_kblock_fulltile_ab_owner_wide( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + backend_op=backend_op, + ) + + return _Simplify(_gemm_ss_blockscaled_sm120_backend_kblock_fulltile_ab_owner_wide, inline_let=True) + + if micro_pipeline == "sm120_backend_kblock_fulltile": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='sm120_backend_kblock_fulltile' currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 4 or int(warp_cols) != 4: + raise ValueError("micro_pipeline='sm120_backend_kblock_fulltile' currently targets warp_rows=4, warp_cols=4") + + @T.prim_func + def _gemm_ss_blockscaled_sm120_backend_kblock_fulltile() -> None: + if clear_accum: + T.clear(C_buf) + + mma_emitter.mma_backend_kblock_fulltile( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_backend_kblock_fulltile( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_backend_kblock_fulltile( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_backend_kblock_fulltile( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + return _Simplify(_gemm_ss_blockscaled_sm120_backend_kblock_fulltile, inline_let=True) + + if micro_pipeline == "sm120_cute_consumer_bridge_skeleton": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='sm120_cute_consumer_bridge_skeleton' currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 4 or int(warp_cols) != 4: + raise ValueError("micro_pipeline='sm120_cute_consumer_bridge_skeleton' currently targets warp_rows=4, warp_cols=4") + + @T.prim_func + def _gemm_ss_blockscaled_sm120_cute_consumer_bridge_skeleton() -> None: + if clear_accum: + T.clear(C_buf) + + mma_emitter.mma_backend_cute_consumer_bridge( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=0, + ) + mma_emitter.mma_backend_cute_consumer_bridge( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=1, + ) + mma_emitter.mma_backend_cute_consumer_bridge( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=2, + ) + mma_emitter.mma_backend_cute_consumer_bridge( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=3, + ) + + return _Simplify(_gemm_ss_blockscaled_sm120_cute_consumer_bridge_skeleton, inline_let=True) + + if micro_pipeline == "sm120_backend_kblock_fulltile_package_pingpong": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires block_K / micro_size_k == 4") + if int(warp_rows) != 4 or int(warp_cols) != 4: + raise ValueError( + f"micro_pipeline={micro_pipeline!r} currently targets warp_rows=4, warp_cols=4, " + f"got warp_rows={int(warp_rows)}, warp_cols={int(warp_cols)}" + ) + if sf_layout != "blockscaled_chunk_kmajor": + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires sf_layout='blockscaled_chunk_kmajor'") + + @T.prim_func + def _gemm_ss_blockscaled_sm120_backend_package_pingpong() -> None: + if clear_accum: + T.clear(C_buf) + + mma_emitter.mma_backend_kblock_fulltile_package_pingpong( + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + sf_layout=sf_layout, + ) + + return _Simplify(_gemm_ss_blockscaled_sm120_backend_package_pingpong, inline_let=True) + + if micro_pipeline == "sm120_pkg_atom_neutral": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='sm120_pkg_atom_neutral' currently requires block_K / micro_size_k == 4") + + @T.prim_func + def _gemm_ss_blockscaled_sm120_pkg_atom_neutral() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_local_0 = T.alloc_local((warp_rows), "uint32") + SFA_local_1 = T.alloc_local((warp_rows), "uint32") + SFB_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_local_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + pkg0 = SM120BlockScaledOperandPackage( + mma_emitter, + A_local_0, + B_local_0, + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + self.sf_k_start, + int(sf_a_granularity_k), + int(sf_b_granularity_k), + sf_layout, + ) + pkg1 = SM120BlockScaledOperandPackage( + mma_emitter, + A_local_1, + B_local_1, + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + self.sf_k_start, + int(sf_a_granularity_k), + int(sf_b_granularity_k), + sf_layout, + ) + + pkg0.copy_kblock(0) + pkg1.copy_kblock(1) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + pkg0.gemm_atom(i, j) + pkg0.copy_kblock(2) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + pkg1.gemm_atom(i, j) + pkg1.copy_kblock(3) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + pkg0.gemm_atom(i, j) + for i in T.unroll(warp_rows): + for j in T.unroll(warp_cols): + pkg1.gemm_atom(i, j) + + return _Simplify(_gemm_ss_blockscaled_sm120_pkg_atom_neutral, inline_let=True) + + if micro_pipeline == "k_static_scale_prefetch": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError(f"micro_pipeline='{micro_pipeline}' currently requires block_K / micro_size_k == 4") + + @T.prim_func + def _gemm_ss_blockscaled_k_static_scale_prefetch() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_local_0 = T.alloc_local((warp_rows), "uint32") + SFA_local_1 = T.alloc_local((warp_rows), "uint32") + SFB_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_local_1 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_0 = T.alloc_local((warp_cols), "uint32") + SFB_rep_local_1 = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + pkg0 = SM120BlockScaledOperandPackage( + mma_emitter, + A_local_0, + B_local_0, + SFA_local_0, + SFB_local_0, + SFB_rep_local_0, + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + self.sf_k_start, + int(sf_a_granularity_k), + int(sf_b_granularity_k), + sf_layout, + ) + pkg1 = SM120BlockScaledOperandPackage( + mma_emitter, + A_local_1, + B_local_1, + SFA_local_1, + SFB_local_1, + SFB_rep_local_1, + A_region, + B_region, + C_buf, + self.SFARegion, + self.SFBRegion, + self.sf_k_start, + int(sf_a_granularity_k), + int(sf_b_granularity_k), + sf_layout, + ) + + pkg0.copy_kblock(0) + pkg1.copy_kblock(1) + pkg0.gemm_kblock() + pkg0.copy_kblock(2) + pkg1.gemm_kblock() + pkg1.copy_kblock(3) + pkg0.gemm_kblock() + pkg1.gemm_kblock() + + return _Simplify(_gemm_ss_blockscaled_k_static_scale_prefetch, inline_let=True) + + if micro_pipeline in ("k_static_scale_stream", "k_static_scale_stream_cutlass_order"): + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError(f"micro_pipeline={micro_pipeline!r} currently requires block_K / micro_size_k == 4") + issue_mma = ( + mma_emitter.mma_with_prefetched_scales_cutlass_order + if micro_pipeline == "k_static_scale_stream_cutlass_order" + else mma_emitter.mma_with_scale_fragments + ) + + @T.prim_func + def _gemm_ss_blockscaled_k_static_scale_stream() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_local = T.alloc_local((warp_rows), "uint32") + SFB_local = T.alloc_local((warp_cols), "uint32") + SFB_rep_local = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldmatrix_b(B_local_0, B_region, 0) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + mma_emitter.ldmatrix_b(B_local_1, B_region, 1) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + issue_mma( + A_local_0, + B_local_0, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + mma_emitter.ldmatrix_b(B_local_0, B_region, 2) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + issue_mma( + A_local_1, + B_local_1, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + mma_emitter.ldmatrix_b(B_local_1, B_region, 3) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + issue_mma( + A_local_0, + B_local_0, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + issue_mma( + A_local_1, + B_local_1, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + + return _Simplify(_gemm_ss_blockscaled_k_static_scale_stream, inline_let=True) + + if micro_pipeline == "k_static_b_all_prefetch_scale_stream": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static_b_all_prefetch_scale_stream' currently requires block_K / micro_size_k == 4") + + @T.prim_func + def _gemm_ss_blockscaled_k_static_b_all_prefetch_scale_stream() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_2 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_3 = T.alloc_local((warp_cols * local_size_b), b_dtype) + SFA_local = T.alloc_local((warp_rows), "uint32") + SFB_local = T.alloc_local((warp_cols), "uint32") + SFB_rep_local = T.alloc_local((warp_cols), "uint32") + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldmatrix_b(B_local_0, B_region, 0) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + mma_emitter.ldmatrix_b(B_local_1, B_region, 1) + mma_emitter.ldmatrix_b(B_local_2, B_region, 2) + mma_emitter.ldmatrix_b(B_local_3, B_region, 3) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_with_scale_fragments( + A_local_0, + B_local_0, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_with_scale_fragments( + A_local_1, + B_local_1, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_with_scale_fragments( + A_local_0, + B_local_2, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + mma_emitter.ldscale_fragment( + SFA_local, + SFB_local, + SFB_rep_local, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma_with_scale_fragments( + A_local_1, + B_local_3, + C_buf, + SFA_local, + SFB_local, + SFB_rep_local, + ) + + return _Simplify(_gemm_ss_blockscaled_k_static_b_all_prefetch_scale_stream, inline_let=True) + + if micro_pipeline == "k_static": + num_k_blocks = int(block_K // micro_size_k) + if num_k_blocks != 4: + raise ValueError("micro_pipeline='k_static' currently requires block_K / micro_size_k == 4") + + @T.prim_func + def _gemm_ss_blockscaled_k_static() -> None: + A_local_0 = T.alloc_local((warp_rows * local_size_a), a_dtype) + A_local_1 = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local_0 = T.alloc_local((warp_cols * local_size_b), b_dtype) + B_local_1 = T.alloc_local((warp_cols * local_size_b), b_dtype) + if clear_accum: + T.clear(C_buf) + + mma_emitter.ldmatrix_a(A_local_0, A_region, 0) + mma_emitter.ldmatrix_b(B_local_0, B_region, 0) + mma_emitter.ldmatrix_a(A_local_1, A_region, 1) + mma_emitter.ldmatrix_b(B_local_1, B_region, 1) + mma_emitter.mma( + A_local_0, + B_local_0, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=0, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_a(A_local_0, A_region, 2) + mma_emitter.ldmatrix_b(B_local_0, B_region, 2) + mma_emitter.mma( + A_local_1, + B_local_1, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=1, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.ldmatrix_a(A_local_1, A_region, 3) + mma_emitter.ldmatrix_b(B_local_1, B_region, 3) + mma_emitter.mma( + A_local_0, + B_local_0, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=2, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + mma_emitter.mma( + A_local_1, + B_local_1, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=3, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + sf_layout=sf_layout, + ) + + return _Simplify(_gemm_ss_blockscaled_k_static, inline_let=True) + + if micro_pipeline is not None: + raise ValueError(f"Unsupported block-scaled micro_pipeline: {micro_pipeline!r}") + + @T.prim_func + def _gemm_ss_blockscaled() -> None: + A_local = T.alloc_local((warp_rows * local_size_a), a_dtype) + B_local = T.alloc_local((warp_cols * local_size_b), b_dtype) + if clear_accum: + T.clear(C_buf) + for ki in T.serial(0, (block_K // micro_size_k)): + mma_emitter.ldmatrix_a(A_local, A_region, ki) + mma_emitter.ldmatrix_b(B_local, B_region, ki) + mma_emitter.mma( + A_local, + B_local, + C_buf, + self.SFARegion, + self.SFBRegion, + ki=ki, + k_start=self.sf_k_start, + sf_a_granularity_k=int(sf_a_granularity_k), + sf_b_granularity_k=int(sf_b_granularity_k), + ) + + return _Simplify(_gemm_ss_blockscaled, inline_let=True) + if self.is_gemm_ss(): @T.prim_func diff --git a/tilelang/intrinsics/__init__.py b/tilelang/intrinsics/__init__.py index 7d4eee9cc..a782b3a77 100644 --- a/tilelang/intrinsics/__init__.py +++ b/tilelang/intrinsics/__init__.py @@ -6,6 +6,7 @@ from tilelang.cuda.intrinsics.layout.utils import ( from tilelang.cuda.intrinsics.macro.mma_macro_generator import ( TensorCoreIntrinEmitter, # noqa: F401 TensorCoreIntrinEmitterWithLadderTransform, # noqa: F401 + TensorCoreIntrinEmitterWithBlockScale, # noqa: F401 ) from tilelang.cuda.intrinsics.macro.wgmma_macro_generator import ( diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index 6fed6e3c7..991fbfd64 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -74,6 +74,7 @@ from .gemm_op import ( # noqa: F401 wgmma_gemm, tcgen05_gemm, tcgen05_gemm_blockscaled, + mma_gemm_blockscaled, make_blockscaled_gemm_layout, ) from .experimental.gemm_sp_op import ( # noqa: F401 diff --git a/tilelang/language/ast/ir.py b/tilelang/language/ast/ir.py index 04f953c86..7f7cf61b5 100644 --- a/tilelang/language/ast/ir.py +++ b/tilelang/language/ast/ir.py @@ -1882,6 +1882,7 @@ call_llvm_pure_intrin = _dtype_forward(_tir_op.call_llvm_pure_intrin) call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) +ptx_mma_block_scale = _dtype_forward(_tir_op.ptx_mma_block_scale) ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) @@ -2137,6 +2138,7 @@ __all__ = [ "tvm_warp_activemask", "ptx_mma", "ptx_mma_sp", + "ptx_mma_block_scale", "ptx_wgmma_ss", "ptx_wgmma_rs", "ptx_wgmma_sp_ss", diff --git a/tilelang/language/gemm_op.py b/tilelang/language/gemm_op.py index 66173a10f..add4617c4 100644 --- a/tilelang/language/gemm_op.py +++ b/tilelang/language/gemm_op.py @@ -442,6 +442,127 @@ def tcgen05_gemm_blockscaled( ) +def mma_gemm_blockscaled( + A: BufferLikeType, + B: BufferLikeType, + C: BufferLikeType, + SFA: BufferLikeType, + SFB: BufferLikeType, + transpose_A: bool = False, + transpose_B: bool = False, + policy: GemmWarpPolicy = GemmWarpPolicy.Square, + clear_accum: bool = False, + *, + k_start: int | tirx.PrimExpr, + sf_a_granularity_k: int, + sf_b_granularity_k: int, + micro_pipeline: str | None = None, + sf_layout: str | None = None, +) -> tirx.PrimExpr: + """Explicit SM120 warp-level block-scaled MMA GEMM. + + This API follows the same scale-factor model as + ``T.tcgen05_gemm_blockscaled``: users pass the scale tensors, logical + ``k_start``, and K granularity, while the lowering derives the low-level + scale addressing. Unlike TCGEN05, this path is synchronous warp-level + ``mma.sync`` and does not use tensor memory or mbarriers. + + The current supported instruction is SM120 NVF4: + ``m16n8k64.kind::mxf4nvf4.block_scale.scale_vec::4X`` with E2M1 operands, + FP32 accumulation, and UE4M3 scale factors. + """ + + ann = { + "sf_a_granularity_k": int(sf_a_granularity_k), + "sf_b_granularity_k": int(sf_b_granularity_k), + } + if micro_pipeline is not None: + ann["micro_pipeline"] = micro_pipeline + if sf_layout is not None: + ann["sf_layout"] = sf_layout + + def legalize(arg): + if isinstance(arg, tirx.Var) and T.has_let_value(arg): + return T.get_let_value(arg).buffer + return arg + + A = legalize(A) + B = legalize(B) + C = legalize(C) + SFA = legalize(SFA) + SFB = legalize(SFB) + + A_region = to_buffer_region(A) + B_region = to_buffer_region(B) + C_region = to_buffer_region(C) + SFA_region = to_buffer_region(SFA) + SFB_region = to_buffer_region(SFB) + + A_shape = retrieve_shape(A_region) + B_shape = retrieve_shape(B_region) + C_shape = retrieve_shape(C_region) + + assert len(C_shape) == 2, "current only support C as a 2D tensor" + assert len(A_shape) >= 2, "current only support A as a 2D or higher-order tensor" + assert len(B_shape) >= 2, "current only support B as a 2D or higher-order tensor" + + M, N = C_shape + M_A = A_shape[-1] if transpose_A else A_shape[-2] + K = A_shape[-2] if transpose_A else A_shape[-1] + N_B = B_shape[-2] if transpose_B else B_shape[-1] + K_B = B_shape[-1] if transpose_B else B_shape[-2] + assert prim_expr_equal(M_A, M), f"T.mma_gemm_blockscaled M shape check failed: M_A = {M_A}, M_C = {M}" + assert prim_expr_equal(K, K_B), f"T.mma_gemm_blockscaled K shape check failed: K_A = {K}, K_B = {K_B}" + assert prim_expr_equal(N_B, N), f"T.mma_gemm_blockscaled N shape check failed: N_B = {N_B}, N_C = {N}" + + A_stride = retrieve_stride(A_region) + B_stride = retrieve_stride(B_region) + A_offset = retrieve_offset(A_region) + B_offset = retrieve_offset(B_region) + stride_a = A_stride[-2] + stride_b = B_stride[-2] + offset_a = A_offset[-1] + offset_b = B_offset[-1] + C_coords = [r.min for r in C_region.region] + + A_arg = buffer_region_to_tile_region(A_region, "r", [r for r in A_shape]) + B_arg = buffer_region_to_tile_region(B_region, "r", [r for r in B_shape]) + C_arg = buffer_region_to_tile_region(C_region, "rw", [r for r in C_shape]) + SFA_arg = buffer_region_to_tile_region(SFA_region, "r", list(retrieve_shape(SFA_region))) + SFB_arg = buffer_region_to_tile_region(SFB_region, "r", list(retrieve_shape(SFB_region))) + + if not isinstance(k_start, tirx.PrimExpr): + k_start = tirx.const(k_start, dtype="int32") + + return tirx.call_intrin( + "handle", + tirx.op.Op.get("tl.tileop.gemm"), + A_arg, + B_arg, + C_arg, + transpose_A, + transpose_B, + M, + N, + K, + policy, + clear_accum, + stride_a, + stride_b, + offset_a, + offset_b, + 1, # k_pack + 0, # wg_wait + tirx.const(0, dtype="int32"), # no mbarrier for synchronous mma.sync + C_coords[0], + C_coords[1], + SFA_arg, + SFB_arg, + k_start, + annotations=ann, + ) + + def make_blockscaled_gemm_layout( C: BufferLikeType, A: BufferLikeType, diff --git a/tilelang/language/tir/ir.py b/tilelang/language/tir/ir.py index dd77a9a8c..16ae93dee 100644 --- a/tilelang/language/tir/ir.py +++ b/tilelang/language/tir/ir.py @@ -295,6 +295,15 @@ call_llvm_pure_intrin = _dtype_forward(_tir_op.call_llvm_pure_intrin) call_pure_extern = _dtype_forward(_tir_op.call_pure_extern) ptx_mma = _dtype_forward(_tir_op.ptx_mma) ptx_mma_sp = _dtype_forward(_tir_op.ptx_mma_sp) +ptx_mma_block_scale = _dtype_forward(_tir_op.ptx_mma_block_scale) +sm120_mma_blockscaled_kblock_fulltile = _tir_op.sm120_mma_blockscaled_kblock_fulltile +sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide = _tir_op.sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide +sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide = _tir_op.sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide +sm120_mma_blockscaled_kblock_fulltile_package_pingpong = _tir_op.sm120_mma_blockscaled_kblock_fulltile_package_pingpong +sm120_mma_blockscaled_cute_consumer_bridge = _tir_op.sm120_mma_blockscaled_cute_consumer_bridge +sm120_store_full_c_fragment_panel64_bf16 = _tir_op.sm120_store_full_c_fragment_panel64_bf16 +sm120_store_full_c_fragment_panel32_tma_bf16 = _tir_op.sm120_store_full_c_fragment_panel32_tma_bf16 +sm120_store_full_c_fragment_epi64x32_tma_bf16 = _tir_op.sm120_store_full_c_fragment_epi64x32_tma_bf16 ptx_wgmma_ss = _dtype_forward(_tir_op.ptx_wgmma_ss) ptx_wgmma_rs = _dtype_forward(_tir_op.ptx_wgmma_rs) ptx_wgmma_sp_ss = _dtype_forward(_tir_op.ptx_wgmma_sp_ss) diff --git a/tilelang/language/tir/op.py b/tilelang/language/tir/op.py index bd3053dd2..d1ce047fd 100644 --- a/tilelang/language/tir/op.py +++ b/tilelang/language/tir/op.py @@ -1076,6 +1076,327 @@ def ptx_mma_sp( ) +def ptx_mma_block_scale( + accum_dtype, + shape, + A_layout, + B_layout, + kind, + scale_vec_size, + A_dtype, + B_dtype, + scale_type, + multiplicand_a, + a_index, + multiplicand_b, + b_index, + accumulator, + c_index, + scale_a, + scale_b, + scale_a_byte_id=0, + scale_a_thread_id=0, + scale_b_byte_id=0, + scale_b_thread_id=0, +): + """TVM intrinsic for SM120a warp-level NVF4 block-scaled MMA.""" + + def _selector_value(value): + return IntImm("int32", value) if isinstance(value, int) else value + + return call_intrin( + accum_dtype, + _tvm_op.Op.get("tl.ptx_mma_block_scale"), + tvm.tirx.StringImm(str(accum_dtype)), + tvm.tirx.StringImm(str(shape)), + tvm.tirx.StringImm(str(A_layout)), + tvm.tirx.StringImm(str(B_layout)), + tvm.tirx.StringImm(str(kind)), + IntImm("int32", scale_vec_size), + tvm.tirx.StringImm(str(A_dtype)), + tvm.tirx.StringImm(str(B_dtype)), + tvm.tirx.StringImm(str(scale_type)), + multiplicand_a, + a_index, + multiplicand_b, + b_index, + accumulator, + c_index, + scale_a, + scale_b, + _selector_value(scale_a_byte_id), + _selector_value(scale_a_thread_id), + _selector_value(scale_b_byte_id), + _selector_value(scale_b_thread_id), + ) + + +def sm120_mma_blockscaled_kblock_fulltile( + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa0, + sfa1, + sfa2, + sfa3, + sfb0, + sfb1, + sfb2, + sfb3, + sfb_rep0, + sfb_rep1, + sfb_rep2, + sfb_rep3, + ki=0, +): + """Internal SM120 NVF4 full-tile K-block helper. + + This is a private lowering hook for the 128x128x256 diagnostic path. It + keeps A/B ldmatrix loads, SFA/SFB loads, and all m16n8k64 block-scaled MMA + atoms inside one CUDA helper, instead of exposing per-atom local fragments + in TIR. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_mma_blockscaled_kblock_fulltile"), + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa0, + sfa1, + sfa2, + sfa3, + sfb0, + sfb1, + sfb2, + sfb3, + sfb_rep0, + sfb_rep1, + sfb_rep2, + sfb_rep3, + ki, + ) + + +def sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide( + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki=0, +): + """Internal SM120 NVF4 full-tile helper with backend-owned scale words. + + This private hook keeps the verified packed A/B access pointers but lets the + CUDA helper compute the 2x2 scale-owner package from SFA/SFB base pointers. + It avoids TIR local scale arrays used by the earlier owner-wide probe. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_mma_blockscaled_kblock_fulltile_ab_owner_wide"), + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki, + ) + + +def sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide( + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki=0, +): + """Internal SM120 NVF4 helper for A-full/B-panel owner-wide probe.""" + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide"), + accumulator, + c_index, + a0, + a1, + a2, + a3, + b0, + b1, + b2, + b3, + sfa_base, + sfb_base, + ki, + ) + + +def sm120_mma_blockscaled_kblock_fulltile_package_pingpong( + accumulator, + c_index, + a_base, + b_base, + sfa_base, + sfb_base, +): + """Internal SM120 NVF4 K-block package-lifecycle probe. + + The CUDA helper receives full shared K-stage bases, creates two backend + register packages, and runs copy_kblock_package(next) -> + gemm_kblock_package(current) for the four K64 atoms inside block_K=256. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_mma_blockscaled_kblock_fulltile_package_pingpong"), + accumulator, + c_index, + a_base, + b_base, + sfa_base, + sfb_base, + ) + + +def sm120_mma_blockscaled_cute_consumer_bridge( + accumulator, + c_index, + a_base, + b_base, + sfa_base, + sfb_base, + ki=0, +): + """Internal SM120 NVF4 full shared-tile consumer bridge. + + This private hook is intentionally lower-level than the public tile API. It + passes full shared-memory tile bases into the CUDA helper so the backend can + own A/B/SFA/SFB fragment lifetime instead of receiving only per-atom row + pointers. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_mma_blockscaled_cute_consumer_bridge"), + accumulator, + c_index, + a_base, + b_base, + sfa_base, + sfb_base, + ki, + ) + + +def sm120_store_full_c_fragment_panel64_bf16( + accumulator, + c_index, + c_shared, + c_shared_index, + panel, +): + """Internal SM120 NVF4 full-fragment epilogue panel-store helper. + + This private hook lets CUDA codegen consume the full 128x128 accumulator + fragment pointer directly. It avoids forcing the TIR frontend to flatten a + fragment with T.access_ptr(C_local[0, 0]), which conflicts with the fragment + layout used by the full-tile MMA lowering. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_store_full_c_fragment_panel64_bf16"), + accumulator, + c_index, + c_shared, + c_shared_index, + panel, + ) + + +def sm120_store_full_c_fragment_panel32_tma_bf16( + accumulator, + c_index, + c_shared, + c_shared_index, + panel32, +): + """Internal SM120 NVF4 full-fragment epilogue TMA-store helper. + + This private hook writes one 128x32 BF16 output panel from the full + 128x128 accumulator fragment into the swizzled shared-memory layout expected + by TileLang's C TMA-store lowering. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_store_full_c_fragment_panel32_tma_bf16"), + accumulator, + c_index, + c_shared, + c_shared_index, + panel32, + ) + + +def sm120_store_full_c_fragment_epi64x32_tma_bf16( + accumulator, + c_index, + c_shared, + c_shared_index, + row64, + col32, +): + """Internal SM120 NVF4 full-fragment epilogue TMA-store helper. + + This private hook writes one 64x32 BF16 output tile from the full 128x128 + accumulator fragment into the swizzled shared-memory layout expected by a + 64x32 C TMA-store. + """ + return call_intrin( + "handle", + _tvm_op.Op.get("tl.sm120_store_full_c_fragment_epi64x32_tma_bf16"), + accumulator, + c_index, + c_shared, + c_shared_index, + row64, + col32, + ) + + def ptx_wgmma_ss( dtype, wgmma_prefix, diff --git a/tilelang/quantize/__init__.py b/tilelang/quantize/__init__.py index b1bb8daa5..907a532f3 100644 --- a/tilelang/quantize/__init__.py +++ b/tilelang/quantize/__init__.py @@ -16,3 +16,19 @@ from .utils import ( from .lop3 import get_lop3_intrin_group # noqa: F401 from .mxfp import get_mxfp_intrin_group # noqa: F401 +from .nvfp4 import ( # noqa: F401 + blockscaled_chunk_kmajor_word_offset, + decode_packed_fp4_e2m1, + decode_ue4m3_scale_bytes, + encode_fp4_e2m1_values, + encode_ue4m3_scale_bytes, + pack_blockscaled_chunk_kmajor_scale_bytes, + pack_fp4_e2m1_codes, + pack_nvfp4_scale_bytes, + tilelang_quantize_bf16_to_nvfp4_blockscaled, + tilelang_quantize_nvfp4_blockscaled, + quantize_bf16_to_nvfp4_blockscaled, + quantize_nvfp4_blockscaled, + swizzle_blockscaled_chunk_kmajor_scale_words, + unswizzle_blockscaled_chunk_kmajor_scale_words, +) diff --git a/tilelang/quantize/nvfp4.py b/tilelang/quantize/nvfp4.py new file mode 100644 index 000000000..b3fa36bd0 --- /dev/null +++ b/tilelang/quantize/nvfp4.py @@ -0,0 +1,1398 @@ +"""NVFP4 quantization and scale layout utilities.""" + +from __future__ import annotations + +import tilelang +import tilelang.language as T + + +_BLOCKSCALED_CHUNK_ROWS = 128 +_BLOCKSCALED_CHUNK_WORDS = 4 +_SCALE_BYTES_PER_WORD = 4 +_NVFP4_SCALE_BLOCK_K = 16 +_SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR = "blockscaled_chunk_kmajor" +_SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR = 0 +_FP4_E2M1_MAX = 6.0 +_UE4M3_MIN_SUBNORMAL = 2.0**-9 +_UE4M3_MIN_NORMAL = 2.0**-6 +_UE4M3_MAX = 1.75 * (2.0**8) +_FP4_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + +_TILELANG_PASS_CONFIGS = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True, + tilelang.PassConfigKey.TL_ENABLE_LOWER_LDGSTG: True, + tilelang.PassConfigKey.TL_ENABLE_LOWER_LDGSTG_PREDICATED: True, +} + + +def _import_torch(): + try: + import torch + except ImportError as exc: # pragma: no cover - exercised only without torch installed + raise ImportError("NVFP4 scale layout utilities require PyTorch tensors") from exc + return torch + + +def _check_block_shape(block_rows: int, block_words: int) -> None: + if block_rows != _BLOCKSCALED_CHUNK_ROWS or block_words != _BLOCKSCALED_CHUNK_WORDS: + raise ValueError( + "SM120 BlockScaledBasicChunk K-major scale packing currently supports " + f"block_rows={_BLOCKSCALED_CHUNK_ROWS} and block_words={_BLOCKSCALED_CHUNK_WORDS}, " + f"got block_rows={block_rows}, block_words={block_words}" + ) + + +def _scale_layout_id(scale_layout: str) -> int: + if scale_layout == _SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR: + return _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR + raise ValueError(f"Unsupported SM120 NVFP4 scale_layout={scale_layout!r}; expected {_SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR!r}") + + +def decode_ue4m3_scale_bytes(scale_bytes): + """Decode unsigned E4M3 scale bytes to float32 values.""" + + torch = _import_torch() + if not isinstance(scale_bytes, torch.Tensor): + raise TypeError(f"scale_bytes must be a torch.Tensor, got {type(scale_bytes)!r}") + if scale_bytes.dtype != torch.uint8: + raise TypeError(f"scale_bytes must have dtype torch.uint8, got {scale_bytes.dtype}") + + u = scale_bytes.to(torch.int32) & 0x7F + exponent = (u >> 3) & 0x0F + mantissa = u & 0x07 + normal = (1.0 + mantissa.to(torch.float32) / 8.0) * torch.pow(2.0, exponent.to(torch.float32) - 7.0) + subnormal = (mantissa.to(torch.float32) / 8.0) * torch.pow(torch.tensor(2.0, device=scale_bytes.device), -6.0) + value = torch.where(exponent == 0, subnormal, normal) + return torch.where(u == 0x7F, torch.full_like(value, float("nan")), value) + + +def encode_ue4m3_scale_bytes(values, *, rounding: str = "nearest", min_positive: bool = False): + """Encode non-negative float values as unsigned E4M3 scale bytes.""" + + torch = _import_torch() + if rounding not in ("nearest", "ceil"): + raise ValueError(f"rounding must be 'nearest' or 'ceil', got {rounding!r}") + if not isinstance(values, torch.Tensor): + raise TypeError(f"values must be a torch.Tensor, got {type(values)!r}") + + x = torch.nan_to_num(values.to(torch.float32), nan=0.0, posinf=_UE4M3_MAX, neginf=0.0).clamp(min=0.0) + x = torch.minimum(x, torch.tensor(_UE4M3_MAX, device=x.device, dtype=torch.float32)) + positive = x > 0 + + if rounding == "ceil": + sub_mantissa = torch.ceil(x / _UE4M3_MIN_SUBNORMAL).to(torch.int32) + else: + sub_mantissa = torch.floor(x / _UE4M3_MIN_SUBNORMAL + 0.5).to(torch.int32) + sub_code = sub_mantissa.clamp(0, 7).to(torch.int32) + sub_overflow_code = torch.full_like(sub_code, 0x08) + sub_code = torch.where(sub_mantissa > 7, sub_overflow_code, sub_code) + + safe_x = torch.clamp(x, min=_UE4M3_MIN_NORMAL) + exponent_unbiased = torch.floor(torch.log2(safe_x)).clamp(-6, 8) + scale = torch.pow(2.0, exponent_unbiased) + mantissa_real = (safe_x / scale - 1.0) * 8.0 + if rounding == "ceil": + mantissa = torch.ceil(mantissa_real).to(torch.int32) + else: + mantissa = torch.floor(mantissa_real + 0.5).to(torch.int32) + exponent = (exponent_unbiased.to(torch.int32) + 7).clamp(1, 15) + carry = mantissa >= 8 + exponent = torch.where(carry, exponent + 1, exponent).clamp(1, 15) + mantissa = torch.where(carry, torch.zeros_like(mantissa), mantissa).clamp(0, 7) + normal_code = ((exponent << 3) | mantissa).clamp(0, 0x7E).to(torch.int32) + + code = torch.where(x < _UE4M3_MIN_NORMAL, sub_code, normal_code) + if min_positive: + code = torch.where(positive & (code == 0), torch.ones_like(code), code) + return code.to(torch.uint8) + + +def encode_fp4_e2m1_values(values): + """Encode float values to raw FP4 E2M1 nibble codes.""" + + torch = _import_torch() + if not isinstance(values, torch.Tensor): + raise TypeError(f"values must be a torch.Tensor, got {type(values)!r}") + + x = torch.nan_to_num(values.to(torch.float32), nan=0.0, posinf=_FP4_E2M1_MAX, neginf=-_FP4_E2M1_MAX) + abs_x = x.abs() + mag = torch.zeros_like(abs_x, dtype=torch.int32) + mag = torch.where(abs_x >= 0.25, torch.ones_like(mag), mag) + mag = torch.where(abs_x >= 0.75, torch.full_like(mag, 2), mag) + mag = torch.where(abs_x >= 1.25, torch.full_like(mag, 3), mag) + mag = torch.where(abs_x >= 1.75, torch.full_like(mag, 4), mag) + mag = torch.where(abs_x >= 2.5, torch.full_like(mag, 5), mag) + mag = torch.where(abs_x >= 3.5, torch.full_like(mag, 6), mag) + mag = torch.where(abs_x >= 5.0, torch.full_like(mag, 7), mag) + sign = ((x < 0) & (mag != 0)).to(torch.int32) << 3 + return (mag | sign).to(torch.uint8) + + +def pack_fp4_e2m1_codes(codes): + """Pack raw FP4 E2M1 nibble codes into int8 storage accepted by TileLang.""" + + torch = _import_torch() + if not isinstance(codes, torch.Tensor): + raise TypeError(f"codes must be a torch.Tensor, got {type(codes)!r}") + if codes.dtype != torch.uint8: + raise TypeError(f"codes must have dtype torch.uint8, got {codes.dtype}") + if codes.ndim != 2: + raise ValueError(f"codes must be a 2D tensor, got shape {tuple(codes.shape)}") + if codes.shape[1] % 2 != 0: + raise ValueError(f"FP4 packing requires an even K dimension, got {tuple(codes.shape)}") + + low = codes[:, 0::2] & 0x0F + high = (codes[:, 1::2] & 0x0F) << 4 + return (low | high).contiguous().view(torch.int8) + + +def decode_packed_fp4_e2m1(packed, cols: int | None = None): + """Decode packed FP4 E2M1 int8/uint8 storage to float32 values.""" + + torch = _import_torch() + if not isinstance(packed, torch.Tensor): + raise TypeError(f"packed must be a torch.Tensor, got {type(packed)!r}") + if packed.ndim != 2: + raise ValueError(f"packed must be a 2D tensor, got shape {tuple(packed.shape)}") + if cols is None: + cols = packed.shape[1] * 2 + if cols != packed.shape[1] * 2: + raise ValueError(f"cols must equal packed.shape[1] * 2, got cols={cols}, packed shape={tuple(packed.shape)}") + + u = packed.contiguous().view(torch.uint8) + lut = torch.tensor(_FP4_E2M1_VALUES, device=packed.device, dtype=torch.float32) + out = torch.empty((packed.shape[0], cols), device=packed.device, dtype=torch.float32) + out[:, 0::2] = lut[(u & 0x0F).long()] + out[:, 1::2] = lut[((u >> 4) & 0x0F).long()] + return out + + +def blockscaled_chunk_kmajor_word_offset(row: int, k64_word: int, block_rows: int = 128, block_words: int = 4) -> tuple[int, int]: + """Return tile-local uint32 coordinates for CUTLASS BlockScaledBasicChunk K-major. + + ``row`` is a tile-local M/N coordinate in ``[0, 128)``. ``k64_word`` is the + tile-local K64 word in ``[0, 4)``. The returned coordinates index the + compressed uint32 source tile whose bytes hold the four adjacent K/16 scale + groups. Flattening these coordinates with a row-major stride of four words + matches the CUDA helper ``sm120_blockscaled_chunk_kmajor_sf_word``. + """ + + _check_block_shape(block_rows, block_words) + if row < 0 or row >= block_rows: + raise ValueError(f"row must be in [0, {block_rows}), got {row}") + if k64_word < 0 or k64_word >= block_words: + raise ValueError(f"k64_word must be in [0, {block_words}), got {k64_word}") + return k64_word * 32 + (row % 32), row // 32 + + +def _sm120_sfa_row_in_lane(lane: int) -> int: + return 8 * (lane & 1) + (lane >> 2) + + +def _sm120_sfb_col_in_lane(lane: int) -> int: + return lane >> 2 + + +def _sm120_sfa_selector_source_lane(lane: int, scale_a_thread_id: int) -> int: + """Return the source lane selected by OMMA.SF's A-side thread selector.""" + + return (lane & ~3) | ((scale_a_thread_id & 1) << 1) | (lane & 1) + + +def _sm120_sfb_selector_source_lane(lane: int, scale_b_thread_id: int) -> int: + """Return the source lane selected by OMMA.SF's B-side thread selector.""" + + return (lane & ~3) | (scale_b_thread_id & 3) + + +def _sm120_compact_sfa_owner_row(lane: int, warp_m: int, reg_group: int) -> int: + qlane = lane & 3 + return warp_m * 64 + reg_group * 32 + (qlane >> 1) * 16 + _sm120_sfa_row_in_lane(lane) + + +def _sm120_compact_sfb_owner_row(lane: int, warp_n: int, reg_group: int) -> int: + qlane = lane & 3 + return warp_n * 64 + reg_group * 32 + qlane * 8 + _sm120_sfb_col_in_lane(lane) + + +def _sm120_compact_sfa_effective_row(lane: int, warp_m: int, mma_i: int) -> int: + source_lane = _sm120_sfa_selector_source_lane(lane, mma_i & 1) + return _sm120_compact_sfa_owner_row(source_lane, warp_m, mma_i >> 1) + + +def _sm120_compact_sfb_effective_row(lane: int, warp_n: int, mma_j: int, half: int) -> int: + source_lane = _sm120_sfb_selector_source_lane(lane, (mma_j & 1) * 2 + half) + return _sm120_compact_sfb_owner_row(source_lane, warp_n, mma_j >> 1) + + +def _sm120_compact_scale_issue_contract(lane: int, warp_m: int, warp_n: int) -> list[dict[str, int | str]]: + """Return the compact-selector OMMA.SF scale contract for one lane. + + This is a pure specification helper for tests and future lowering work. It + models which lane supplies the SFA/SFB scale register selected by each + ``mma.sync.m16n8k64.kind::mxf4nvf4.block_scale`` issue. + """ + + issues: list[dict[str, int | str]] = [] + for mma_i in range(4): + sa_reg_group = mma_i >> 1 + sa_tid = mma_i & 1 + sa_source_lane = _sm120_sfa_selector_source_lane(lane, sa_tid) + for mma_j in range(4): + sb_reg_group = mma_j >> 1 + for half in range(2): + sb_tid = (mma_j & 1) * 2 + half + sb_source_lane = _sm120_sfb_selector_source_lane(lane, sb_tid) + issues.append( + { + "mma_i": mma_i, + "mma_j": mma_j, + "half": half, + "sa_reg": f"sa{sa_reg_group}", + "sa_tid": sa_tid, + "sa_source_lane": sa_source_lane, + "sfa_row": _sm120_compact_sfa_owner_row(sa_source_lane, warp_m, sa_reg_group), + "sb_reg": f"sb{sb_reg_group}", + "sb_tid": sb_tid, + "sb_source_lane": sb_source_lane, + "sfb_row": _sm120_compact_sfb_owner_row(sb_source_lane, warp_n, sb_reg_group), + } + ) + return issues + + +def _sm120_compact_scale_copy_view_contract(lane: int, warp_m: int, warp_n: int, k64_word: int) -> list[dict[str, object]]: + """Return per-issue scale rows plus BlockScaledBasicChunk word addresses.""" + + issues: list[dict[str, object]] = [] + for issue in _sm120_compact_scale_issue_contract(lane, warp_m, warp_n): + sfa_word_coord = blockscaled_chunk_kmajor_word_offset(int(issue["sfa_row"]), k64_word) + sfb_word_coord = blockscaled_chunk_kmajor_word_offset(int(issue["sfb_row"]), k64_word) + issues.append( + { + **issue, + "k64_word": k64_word, + "sfa_word_coord": sfa_word_coord, + "sfa_word_offset": sfa_word_coord[0] * _BLOCKSCALED_CHUNK_WORDS + sfa_word_coord[1], + "sfb_word_coord": sfb_word_coord, + "sfb_word_offset": sfb_word_coord[0] * _BLOCKSCALED_CHUNK_WORDS + sfb_word_coord[1], + } + ) + return issues + + +def _sm120_compact_scale_producer_contract(warp_m: int, warp_n: int) -> list[dict[str, object]]: + """Return the source-lane register assignment required by current OMMA.SF. + + This is the inverse view of ``_sm120_compact_scale_issue_contract``. It says + which producer lane supplies the scale register consumed by each issue. A + future compact scale-copy lowering must preserve these semantic rows while + changing how the producer lanes load/package the registers. + """ + + rows: list[dict[str, object]] = [] + for mma_i in range(4): + sa_tid = mma_i & 1 + sa_reg_group = mma_i >> 1 + for producer_lane in range(32): + consumers = [lane for lane in range(32) if _sm120_sfa_selector_source_lane(lane, sa_tid) == producer_lane] + if not consumers: + continue + rows.append( + { + "kind": "SFA", + "mma_i": mma_i, + "sa_tid": sa_tid, + "sa_reg": f"sa{sa_reg_group}", + "producer_lane": producer_lane, + "consumers": tuple(consumers), + "row": _sm120_compact_sfa_owner_row(producer_lane, warp_m, sa_reg_group), + } + ) + + for mma_j in range(4): + sb_reg_group = mma_j >> 1 + for half in range(2): + sb_tid = (mma_j & 1) * 2 + half + for producer_lane in range(32): + consumers = [lane for lane in range(32) if _sm120_sfb_selector_source_lane(lane, sb_tid) == producer_lane] + if not consumers: + continue + rows.append( + { + "kind": "SFB", + "mma_j": mma_j, + "half": half, + "sb_tid": sb_tid, + "sb_reg": f"sb{sb_reg_group}", + "producer_lane": producer_lane, + "consumers": tuple(consumers), + "row": _sm120_compact_sfb_owner_row(producer_lane, warp_n, sb_reg_group), + } + ) + return rows + + +def _sm120_compact_package_consumer_byte_offsets(lane: int, warp_m: int, warp_n: int, k64_word: int) -> dict[str, tuple[int, ...]]: + """Return the current package's per-consumer byte offsets for one K64 word.""" + + issues = _sm120_compact_scale_copy_view_contract(lane, warp_m, warp_n, k64_word) + sfa_offsets = tuple(int(issue["sfa_word_offset"]) * 4 for issue in issues[0::8]) + sfb_offsets = tuple(int(issue["sfb_word_offset"]) * 4 for issue in issues[:8]) + return {"SFA": sfa_offsets, "SFB": sfb_offsets} + + +def _sm120_cutlass_issue_site_contract(site: int) -> dict[str, int]: + """Return CUTLASS/CuTe SM120 fulltile issue-site coordinates. + + The traversal covers four K64 atoms. Each K64 atom issues eight N8 atoms, + and each N8 atom visits four M atoms. Odd N8 atoms reverse the M order. + """ + + if site < 0 or site >= 128: + raise ValueError(f"site must be in [0, 128), got {site}") + k_block = site // 32 + site_in_k = site % 32 + n8_atom = site_in_k // 4 + m_serp = site_in_k % 4 + m_atom = m_serp if n8_atom % 2 == 0 else 3 - m_serp + n16_col = n8_atom // 2 + n8_half = n8_atom % 2 + return { + "site": site, + "k_block": k_block, + "m_atom": m_atom, + "n8_atom": n8_atom, + "n16_col": n16_col, + "n8_half": n8_half, + "A_slot0": 128 * m_atom + 32 * k_block, + "B_slot0": 64 * n8_half + 128 * n16_col + 16 * k_block, + "SFA_slot0": 4 * m_atom + 16 * k_block, + "SFB_slot0": 4 * n16_col + 16 * n8_half + 32 * k_block, + "C_slot0": 4 * m_atom + 32 * n16_col + 16 * n8_half, + } + + +def _sm120_tilelang_current_issue_site_contract(site: int) -> dict[str, int]: + """Return the current package helper's issue coordinates for comparison.""" + + if site < 0 or site >= 128: + raise ValueError(f"site must be in [0, 128), got {site}") + k_block = site // 32 + site_in_k = site % 32 + mma_i = site_in_k // 8 + n_site = site_in_k % 8 + mma_j = n_site // 2 + half = n_site & 1 + return { + "site": site, + "k_block": k_block, + "mma_i": mma_i, + "mma_j": mma_j, + "half": half, + "sa_reg_group": mma_i >> 1, + "sa_tid": mma_i & 1, + "sb_reg_group": mma_j >> 1, + "sb_tid": (mma_j & 1) * 2 + half, + } + + +def swizzle_blockscaled_chunk_kmajor_scale_words(words, block_rows: int = 128, block_words: int = 4): + """Convert semantic row-major scale words to SM120 BlockScaledBasicChunk K-major. + + Parameters + ---------- + words: + ``torch.uint32`` tensor with shape ``[rows, K // 64]``. Each word packs + four already-encoded UE4M3 scale bytes for consecutive K/16 groups. + block_rows: + The source-layout tile rows. SM120 dense NVFP4 uses ``128``. + block_words: + The source-layout tile K64 words. SM120 dense NVFP4 uses ``4``. + + Returns + ------- + torch.Tensor + ``torch.uint32`` tensor with the same shape, physically ordered for the + CUTLASS ``BlockScaledBasicChunk`` K-major source layout. + """ + + torch = _import_torch() + _check_block_shape(block_rows, block_words) + if not isinstance(words, torch.Tensor): + raise TypeError(f"words must be a torch.Tensor, got {type(words)!r}") + if words.dtype != torch.uint32: + raise TypeError(f"words must have dtype torch.uint32, got {words.dtype}") + if words.ndim != 2: + raise ValueError(f"words must be a 2D tensor, got shape {tuple(words.shape)}") + + rows, cols = words.shape + if rows % block_rows != 0 or cols % block_words != 0: + raise ValueError( + f"blockscaled_chunk_kmajor scale storage requires rows multiple of {block_rows} " + f"and K/64 words multiple of {block_words}, got {tuple(words.shape)}" + ) + + row_blocks = rows // block_rows + src = words.contiguous().reshape(row_blocks, 4, 32, cols) + return src.permute(0, 3, 2, 1).contiguous().reshape(rows, cols) + + +def unswizzle_blockscaled_chunk_kmajor_scale_words(words, block_rows: int = 128, block_words: int = 4): + """Convert BlockScaledBasicChunk K-major source words back to semantic row-major words. + + This is intended for reference checks and debugging. GEMM kernels should + consume the source-layout tensor produced by the quantizer directly. + """ + + torch = _import_torch() + _check_block_shape(block_rows, block_words) + if not isinstance(words, torch.Tensor): + raise TypeError(f"words must be a torch.Tensor, got {type(words)!r}") + if words.dtype != torch.uint32: + raise TypeError(f"words must have dtype torch.uint32, got {words.dtype}") + if words.ndim != 2: + raise ValueError(f"words must be a 2D tensor, got shape {tuple(words.shape)}") + + rows, cols = words.shape + if rows % block_rows != 0 or cols % block_words != 0: + raise ValueError( + f"blockscaled_chunk_kmajor scale storage requires rows multiple of {block_rows} " + f"and K/64 words multiple of {block_words}, got {tuple(words.shape)}" + ) + + row_blocks = rows // block_rows + src = words.contiguous().reshape(row_blocks, cols, 32, 4) + return src.permute(0, 3, 2, 1).contiguous().reshape(rows, cols) + + +def _pack_scale_bytes_to_words(scale_bytes): + torch = _import_torch() + scale_i64 = scale_bytes.to(torch.int64).reshape( + scale_bytes.shape[0], scale_bytes.shape[1] // _SCALE_BYTES_PER_WORD, _SCALE_BYTES_PER_WORD + ) + words = scale_i64[:, :, 0] + words = words | (scale_i64[:, :, 1] << 8) + words = words | (scale_i64[:, :, 2] << 16) + words = words | (scale_i64[:, :, 3] << 24) + return words.to(torch.uint32) + + +def pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes, block_rows: int = 128, block_words: int = 4): + """Pack UE4M3 scale bytes into SM120 BlockScaledBasicChunk K-major words. + + This is a layout packer, not a numeric quantizer. ``scale_bytes`` must + already contain UE4M3-encoded scale bytes with semantic row-major shape + ``[rows, K / 16]``. The result is the compressed source tensor expected by + the SM120 NVFP4 blockscaled GEMM path: ``torch.uint32[rows, K // 64]``. + + The same function applies to SFA and SFB. For SFA, ``rows`` is logical M; for + SFB, ``rows`` is logical N. + """ + + torch = _import_torch() + _check_block_shape(block_rows, block_words) + if not isinstance(scale_bytes, torch.Tensor): + raise TypeError(f"scale_bytes must be a torch.Tensor, got {type(scale_bytes)!r}") + if scale_bytes.dtype != torch.uint8: + raise TypeError(f"scale_bytes must have dtype torch.uint8, got {scale_bytes.dtype}") + if scale_bytes.ndim != 2: + raise ValueError(f"scale_bytes must be a 2D tensor, got shape {tuple(scale_bytes.shape)}") + + rows, scale_cols = scale_bytes.shape + scale_cols_per_tile = block_words * _SCALE_BYTES_PER_WORD + if rows % block_rows != 0 or scale_cols % scale_cols_per_tile != 0: + raise ValueError( + f"blockscaled_chunk_kmajor scale bytes require rows multiple of {block_rows} " + f"and K/16 columns multiple of {scale_cols_per_tile}, got {tuple(scale_bytes.shape)}" + ) + + words = _pack_scale_bytes_to_words(scale_bytes) + return swizzle_blockscaled_chunk_kmajor_scale_words(words, block_rows, block_words) + + +def pack_nvfp4_scale_bytes( + scale_bytes, + *, + scale_layout: str = _SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR, + block_rows: int = 128, + block_words: int = 4, +): + """Pack semantic UE4M3 scale bytes for a supported SM120 NVFP4 scale layout.""" + + _scale_layout_id(scale_layout) + return pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes, block_rows=block_rows, block_words=block_words) + + +def _tl_encode_ue4m3_scale_byte_ceil(x): + x = T.alloc_var(T.float32, init=T.max(T.min(x, _UE4M3_MAX), 0.0)) + sub_mantissa = T.cast(T.ceil(x[0] * 512.0), T.int32) + sub_code = T.min(T.max(sub_mantissa, 0), 7) + sub_code = T.if_then_else(sub_mantissa > 7, 0x08, sub_code) + + bits = T.reinterpret(x[0], T.uint32) + exponent_unbiased = T.cast((bits >> 23) & 0xFF, T.int32) - 127 + exponent_unbiased = T.min(T.max(exponent_unbiased, -6), 8) + mantissa_bits = T.cast(bits & 0x7FFFFF, T.int32) + mantissa = (mantissa_bits >> 20) + T.if_then_else((mantissa_bits & 0xFFFFF) != 0, 1, 0) + biased_exponent = exponent_unbiased + 7 + carry = mantissa >= 8 + biased_exponent = T.if_then_else(carry, biased_exponent + 1, biased_exponent) + biased_exponent = T.min(T.max(biased_exponent, 1), 15) + mantissa = T.if_then_else(carry, 0, mantissa) + mantissa = T.min(T.max(mantissa, 0), 7) + normal_code = (biased_exponent << 3) | mantissa + + code = T.if_then_else(x[0] < _UE4M3_MIN_NORMAL, sub_code, normal_code) + return T.if_then_else(x[0] <= 0.0, 0, code) + + +def _tl_decode_ue4m3_scale_byte(code): + exponent = (code >> 3) & 0x0F + mantissa = code & 0x07 + mantissa_f32 = T.cast(mantissa, T.float32) + normal_bits = T.cast(((exponent + 120) << 23) | (mantissa << 20), T.uint32) + normal = T.reinterpret(normal_bits, T.float32) + subnormal = mantissa_f32 * 0.125 * _UE4M3_MIN_NORMAL + return T.if_then_else(exponent == 0, subnormal, normal) + + +def _tl_ue4m3_scale_byte_inverse(code): + exponent = (code >> 3) & 0x0F + mantissa = code & 0x07 + + normal_factor = T.if_then_else( + mantissa == 0, + 1.0, + T.if_then_else( + mantissa == 1, + 0.8888888955116272, + T.if_then_else( + mantissa == 2, + 0.800000011920929, + T.if_then_else( + mantissa == 3, + 0.7272727489471436, + T.if_then_else( + mantissa == 4, + 0.6666666865348816, + T.if_then_else( + mantissa == 5, + 0.6153846383094788, + T.if_then_else(mantissa == 6, 0.5714285969734192, 0.5333333611488342), + ), + ), + ), + ), + ), + ) + normal_bits = T.cast((134 - exponent) << 23, T.uint32) + normal_inv = T.reinterpret(normal_bits, T.float32) * normal_factor + + subnormal_inv = T.if_then_else( + mantissa == 0, + 0.0, + T.if_then_else( + mantissa == 1, + 512.0, + T.if_then_else( + mantissa == 2, + 256.0, + T.if_then_else( + mantissa == 3, + 170.6666717529297, + T.if_then_else( + mantissa == 4, + 128.0, + T.if_then_else( + mantissa == 5, + 102.4000015258789, + T.if_then_else(mantissa == 6, 85.33333587646484, 73.14286041259766), + ), + ), + ), + ), + ), + ) + return T.if_then_else(exponent == 0, subnormal_inv, normal_inv) + + +def _tl_encode_fp4_e2m1_code(x): + return T.reinterpret(T.cast(x, T.float4_e2m1fn), T.uint8) + + +def _tl_scale_source_coords(logical_row, k64_word, _scale_layout_id: int, scale_cols): + row_in_chunk = logical_row % _BLOCKSCALED_CHUNK_ROWS + flat_word = ( + (logical_row // _BLOCKSCALED_CHUNK_ROWS) * _BLOCKSCALED_CHUNK_ROWS * scale_cols + + k64_word * _BLOCKSCALED_CHUNK_ROWS + + (row_in_chunk % 32) * 4 + + (row_in_chunk // 32) + ) + physical_row = flat_word // scale_cols + physical_word = flat_word % scale_cols + return physical_row, physical_word + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel( + x, + rows_per_cta: int = 16, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 + BlockScaledBasicChunk K-major scales.""" + + M = T.dynamic("M") + K = T.const("K") + in_dtype = T.bfloat16 + fp4_dtype = T.float4_e2m1fn + compute_dtype = T.float32 + lanes_per_scale = 128 // rows_per_cta + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K), fp4_dtype) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 64, threads=128) as (pid_m, pid_k64): + tx = T.get_thread_binding(0) + scale_lane = tx % lanes_per_scale + x0 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x1 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x2 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x3 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + amax0 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax1 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax2 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax3 = T.alloc_fragment((rows_per_cta,), in_dtype) + scale_inv0 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv1 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv2 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv3 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_code0 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code1 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code2 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code3 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_word = T.alloc_fragment((rows_per_cta,), T.uint32) + + for i, j in T.Parallel(rows_per_cta, 16): + x0[i, j] = x[pid_m * rows_per_cta + i, pid_k64 * 64 + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x1[i, j] = x[pid_m * rows_per_cta + i, pid_k64 * 64 + 16 + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x2[i, j] = x[pid_m * rows_per_cta + i, pid_k64 * 64 + 32 + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x3[i, j] = x[pid_m * rows_per_cta + i, pid_k64 * 64 + 48 + j] + + T.reduce_absmax(x0, amax0, dim=1) + T.reduce_absmax(x1, amax1, dim=1) + T.reduce_absmax(x2, amax2, dim=1) + T.reduce_absmax(x3, amax3, dim=1) + + for i in T.Parallel(rows_per_cta): + scale_inv0[i] = 0.0 + scale_inv1[i] = 0.0 + scale_inv2[i] = 0.0 + scale_inv3[i] = 0.0 + if scale_lane == 0: + scale_code0[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax0[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv0[i] = _tl_ue4m3_scale_byte_inverse(scale_code0[i]) + scale_code1[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax1[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv1[i] = _tl_ue4m3_scale_byte_inverse(scale_code1[i]) + scale_code2[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax2[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv2[i] = _tl_ue4m3_scale_byte_inverse(scale_code2[i]) + scale_code3[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax3[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv3[i] = _tl_ue4m3_scale_byte_inverse(scale_code3[i]) + scale_inv0[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv0[i], 0, lanes_per_scale, 32) + scale_inv1[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv1[i], 0, lanes_per_scale, 32) + scale_inv2[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv2[i], 0, lanes_per_scale, 32) + scale_inv3[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv3[i], 0, lanes_per_scale, 32) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled0 = T.cast(x0[i, j], compute_dtype) * scale_inv0[i] + quant[pid_m * rows_per_cta + i, pid_k64 * 64 + j] = T.cast(scaled0, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled1 = T.cast(x1[i, j], compute_dtype) * scale_inv1[i] + quant[pid_m * rows_per_cta + i, pid_k64 * 64 + 16 + j] = T.cast(scaled1, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled2 = T.cast(x2[i, j], compute_dtype) * scale_inv2[i] + quant[pid_m * rows_per_cta + i, pid_k64 * 64 + 32 + j] = T.cast(scaled2, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled3 = T.cast(x3[i, j], compute_dtype) * scale_inv3[i] + quant[pid_m * rows_per_cta + i, pid_k64 * 64 + 48 + j] = T.cast(scaled3, fp4_dtype) + + for i in T.Parallel(rows_per_cta): + if scale_lane == 0: + scale_word[i] = ( + T.cast(scale_code0[i], T.uint32) + | (T.cast(scale_code1[i], T.uint32) << 8) + | (T.cast(scale_code2[i], T.uint32) << 16) + | (T.cast(scale_code3[i], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + i + physical_row, physical_word = _tl_scale_source_coords(logical_row, pid_k64, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word[i] + + return quant, scale_source + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel_16x256_tma_store( + x, + num_stages: int = 0, + num_threads: int = 128, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 with a 16x256 CTA tile and TMA output store.""" + + M = T.dynamic("M") + K = T.const("K") + rows_per_cta = 16 + in_dtype = T.bfloat16 + fp4_dtype = T.float4_e2m1fn + compute_dtype = T.float32 + lanes_per_scale = num_threads // rows_per_cta + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K), fp4_dtype) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 256, threads=num_threads) as (pid_m, pid_k256): + tx = T.get_thread_binding(0) + scale_lane = tx % lanes_per_scale + quant_shared = T.alloc_shared((rows_per_cta, 256), fp4_dtype) + x0 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x1 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x2 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + x3 = T.alloc_fragment((rows_per_cta, 16), in_dtype) + amax0 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax1 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax2 = T.alloc_fragment((rows_per_cta,), in_dtype) + amax3 = T.alloc_fragment((rows_per_cta,), in_dtype) + scale_inv0 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv1 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv2 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_inv3 = T.alloc_fragment((rows_per_cta,), compute_dtype) + scale_code0 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code1 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code2 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_code3 = T.alloc_fragment((rows_per_cta,), T.int32) + scale_word = T.alloc_fragment((rows_per_cta,), T.uint32) + + for k64_inner in T.Pipelined(4, num_stages=num_stages): + k_base = pid_k256 * 256 + k64_inner * 64 + shared_k_base = k64_inner * 64 + + for i, j in T.Parallel(rows_per_cta, 16): + x0[i, j] = x[pid_m * rows_per_cta + i, k_base + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x1[i, j] = x[pid_m * rows_per_cta + i, k_base + 16 + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x2[i, j] = x[pid_m * rows_per_cta + i, k_base + 32 + j] + + for i, j in T.Parallel(rows_per_cta, 16): + x3[i, j] = x[pid_m * rows_per_cta + i, k_base + 48 + j] + + T.reduce_absmax(x0, amax0, dim=1) + T.reduce_absmax(x1, amax1, dim=1) + T.reduce_absmax(x2, amax2, dim=1) + T.reduce_absmax(x3, amax3, dim=1) + + for i in T.Parallel(rows_per_cta): + scale_inv0[i] = 0.0 + scale_inv1[i] = 0.0 + scale_inv2[i] = 0.0 + scale_inv3[i] = 0.0 + if scale_lane == 0: + scale_code0[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax0[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv0[i] = _tl_ue4m3_scale_byte_inverse(scale_code0[i]) + scale_code1[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax1[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv1[i] = _tl_ue4m3_scale_byte_inverse(scale_code1[i]) + scale_code2[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax2[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv2[i] = _tl_ue4m3_scale_byte_inverse(scale_code2[i]) + scale_code3[i] = _tl_encode_ue4m3_scale_byte_ceil(T.cast(amax3[i], compute_dtype) / _FP4_E2M1_MAX) + scale_inv3[i] = _tl_ue4m3_scale_byte_inverse(scale_code3[i]) + scale_inv0[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv0[i], 0, lanes_per_scale, 32) + scale_inv1[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv1[i], 0, lanes_per_scale, 32) + scale_inv2[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv2[i], 0, lanes_per_scale, 32) + scale_inv3[i] = T.tvm_warp_shuffle(0xFFFFFFFF, scale_inv3[i], 0, lanes_per_scale, 32) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled0 = T.cast(x0[i, j], compute_dtype) * scale_inv0[i] + quant_shared[i, shared_k_base + j] = T.cast(scaled0, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled1 = T.cast(x1[i, j], compute_dtype) * scale_inv1[i] + quant_shared[i, shared_k_base + 16 + j] = T.cast(scaled1, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled2 = T.cast(x2[i, j], compute_dtype) * scale_inv2[i] + quant_shared[i, shared_k_base + 32 + j] = T.cast(scaled2, fp4_dtype) + + for i, j in T.Parallel(rows_per_cta, 16): + scaled3 = T.cast(x3[i, j], compute_dtype) * scale_inv3[i] + quant_shared[i, shared_k_base + 48 + j] = T.cast(scaled3, fp4_dtype) + + for i in T.Parallel(rows_per_cta): + if scale_lane == 0: + scale_word[i] = ( + T.cast(scale_code0[i], T.uint32) + | (T.cast(scale_code1[i], T.uint32) << 8) + | (T.cast(scale_code2[i], T.uint32) << 16) + | (T.cast(scale_code3[i], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + i + k64_word = pid_k256 * _BLOCKSCALED_CHUNK_WORDS + k64_inner + physical_row, physical_word = _tl_scale_source_coords(logical_row, k64_word, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word[i] + + T.sync_threads() + T.tma_copy(quant_shared, quant[pid_m * rows_per_cta, pid_k256 * 256]) + T.tma_store_wait() + + return quant, scale_source + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel_16x256_tma_io( + x, + num_stages: int = 4, + num_threads: int = 512, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 with one warp per row and TMA input/output.""" + + M = T.dynamic("M") + K = T.const("K") + rows_per_cta = 16 + in_dtype = T.bfloat16 + compute_dtype = T.float32 + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K // 2), T.int8) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 256, threads=num_threads) as (pid_m, pid_k256): + lane = T.get_lane_idx() + row = T.get_warp_idx_sync() + x_shared = T.alloc_shared((rows_per_cta, 256), in_dtype) + quant_shared = T.alloc_shared((rows_per_cta, 128), T.int8) + scale_codes = T.alloc_shared((rows_per_cta, 16), T.int32) + + T.copy(x[pid_m * rows_per_cta, pid_k256 * 256], x_shared, prefer_instruction="tma") + T.sync_threads() + + for k64_inner in T.serial(_BLOCKSCALED_CHUNK_WORDS): + elem_col = k64_inner * 64 + lane * 2 + byte_col = k64_inner * 32 + lane + group = k64_inner * 4 + (lane // 8) + group_lane = lane % 8 + value0 = T.cast(x_shared[row, elem_col], compute_dtype) + value1 = T.cast(x_shared[row, elem_col + 1], compute_dtype) + amax = T.alloc_var(compute_dtype, init=T.max(T.abs(value0), T.abs(value1))) + amax = T.max(amax, T.shfl_xor(amax, 4, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 2, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 1, width=8)) + + scale_code = T.alloc_var(T.int32, init=0) + scale_inv = T.alloc_var(compute_dtype, init=0.0) + if group_lane == 0: + scale_code = _tl_encode_ue4m3_scale_byte_ceil(amax / _FP4_E2M1_MAX) + scale_inv = _tl_ue4m3_scale_byte_inverse(scale_code) + scale_codes[row, group] = scale_code + scale_inv = T.shfl_sync(scale_inv, 0, width=8) + fp4_code0 = _tl_encode_fp4_e2m1_code(value0 * scale_inv) + fp4_code1 = _tl_encode_fp4_e2m1_code(value1 * scale_inv) + packed_byte = fp4_code0 | (fp4_code1 << 4) + quant_shared[row, byte_col] = T.cast(packed_byte, T.int8) + + T.sync_threads() + if lane < _BLOCKSCALED_CHUNK_WORDS: + k64_inner = lane + scale_word = ( + T.cast(scale_codes[row, k64_inner * 4], T.uint32) + | (T.cast(scale_codes[row, k64_inner * 4 + 1], T.uint32) << 8) + | (T.cast(scale_codes[row, k64_inner * 4 + 2], T.uint32) << 16) + | (T.cast(scale_codes[row, k64_inner * 4 + 3], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + row + k64_word = pid_k256 * _BLOCKSCALED_CHUNK_WORDS + k64_inner + physical_row, physical_word = _tl_scale_source_coords(logical_row, k64_word, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word + + T.sync_threads() + T.tma_copy(quant_shared, quant[pid_m * rows_per_cta, pid_k256 * 128]) + T.tma_store_wait() + + return quant, scale_source + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel_32x256_tma_io( + x, + num_stages: int = 4, + num_threads: int = 512, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 with a 32x256 CTA tile and TMA input/output.""" + + M = T.dynamic("M") + K = T.const("K") + rows_per_cta = 32 + rows_per_warp = rows_per_cta // (num_threads // 32) + in_dtype = T.bfloat16 + compute_dtype = T.float32 + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K // 2), T.int8) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 256, threads=num_threads) as (pid_m, pid_k256): + lane = T.get_lane_idx() + warp = T.get_warp_idx_sync() + x_shared = T.alloc_shared((rows_per_cta, 256), in_dtype) + quant_shared = T.alloc_shared((rows_per_cta, 128), T.int8) + scale_codes = T.alloc_shared((rows_per_cta, 16), T.int32) + + T.copy(x[pid_m * rows_per_cta, pid_k256 * 256], x_shared, prefer_instruction="tma") + T.sync_threads() + + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + for k64_inner in T.serial(_BLOCKSCALED_CHUNK_WORDS): + elem_col = k64_inner * 64 + lane * 2 + byte_col = k64_inner * 32 + lane + group = k64_inner * 4 + (lane // 8) + group_lane = lane % 8 + value0 = T.cast(x_shared[row, elem_col], compute_dtype) + value1 = T.cast(x_shared[row, elem_col + 1], compute_dtype) + amax = T.alloc_var(compute_dtype, init=T.max(T.abs(value0), T.abs(value1))) + amax = T.max(amax, T.shfl_xor(amax, 4, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 2, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 1, width=8)) + + scale_code = T.alloc_var(T.int32, init=0) + scale_inv = T.alloc_var(compute_dtype, init=0.0) + if group_lane == 0: + scale_code = _tl_encode_ue4m3_scale_byte_ceil(amax / _FP4_E2M1_MAX) + scale_inv = _tl_ue4m3_scale_byte_inverse(scale_code) + scale_codes[row, group] = scale_code + scale_inv = T.shfl_sync(scale_inv, 0, width=8) + fp4_code0 = _tl_encode_fp4_e2m1_code(value0 * scale_inv) + fp4_code1 = _tl_encode_fp4_e2m1_code(value1 * scale_inv) + packed_byte = fp4_code0 | (fp4_code1 << 4) + quant_shared[row, byte_col] = T.cast(packed_byte, T.int8) + + T.sync_threads() + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + if lane < _BLOCKSCALED_CHUNK_WORDS: + k64_inner = lane + scale_word = ( + T.cast(scale_codes[row, k64_inner * 4], T.uint32) + | (T.cast(scale_codes[row, k64_inner * 4 + 1], T.uint32) << 8) + | (T.cast(scale_codes[row, k64_inner * 4 + 2], T.uint32) << 16) + | (T.cast(scale_codes[row, k64_inner * 4 + 3], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + row + k64_word = pid_k256 * _BLOCKSCALED_CHUNK_WORDS + k64_inner + physical_row, physical_word = _tl_scale_source_coords(logical_row, k64_word, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word + + T.sync_threads() + T.tma_copy(quant_shared, quant[pid_m * rows_per_cta, pid_k256 * 128]) + T.tma_store_wait() + + return quant, scale_source + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel_32x256_tma_load_global_store( + x, + num_stages: int = 4, + num_threads: int = 512, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 with TMA input and direct global output stores.""" + + M = T.dynamic("M") + K = T.const("K") + rows_per_cta = 32 + rows_per_warp = rows_per_cta // (num_threads // 32) + in_dtype = T.bfloat16 + compute_dtype = T.float32 + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K // 2), T.int8) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 256, threads=num_threads) as (pid_m, pid_k256): + lane = T.get_lane_idx() + warp = T.get_warp_idx_sync() + x_shared = T.alloc_shared((rows_per_cta, 256), in_dtype) + scale_codes = T.alloc_shared((rows_per_cta, 16), T.int32) + + T.copy(x[pid_m * rows_per_cta, pid_k256 * 256], x_shared, prefer_instruction="tma") + T.sync_threads() + + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + for k64_inner in T.serial(_BLOCKSCALED_CHUNK_WORDS): + elem_col = k64_inner * 64 + lane * 2 + byte_col = k64_inner * 32 + lane + group = k64_inner * 4 + (lane // 8) + group_lane = lane % 8 + value0 = T.cast(x_shared[row, elem_col], compute_dtype) + value1 = T.cast(x_shared[row, elem_col + 1], compute_dtype) + amax = T.alloc_var(compute_dtype, init=T.max(T.abs(value0), T.abs(value1))) + amax = T.max(amax, T.shfl_xor(amax, 4, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 2, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 1, width=8)) + + scale_code = T.alloc_var(T.int32, init=0) + scale_inv = T.alloc_var(compute_dtype, init=0.0) + if group_lane == 0: + scale_code = _tl_encode_ue4m3_scale_byte_ceil(amax / _FP4_E2M1_MAX) + scale_inv = _tl_ue4m3_scale_byte_inverse(scale_code) + scale_codes[row, group] = scale_code + scale_inv = T.shfl_sync(scale_inv, 0, width=8) + fp4_code0 = _tl_encode_fp4_e2m1_code(value0 * scale_inv) + fp4_code1 = _tl_encode_fp4_e2m1_code(value1 * scale_inv) + packed_byte = fp4_code0 | (fp4_code1 << 4) + quant[pid_m * rows_per_cta + row, pid_k256 * 128 + byte_col] = T.cast(packed_byte, T.int8) + + T.sync_threads() + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + if lane < _BLOCKSCALED_CHUNK_WORDS: + k64_inner = lane + scale_word = ( + T.cast(scale_codes[row, k64_inner * 4], T.uint32) + | (T.cast(scale_codes[row, k64_inner * 4 + 1], T.uint32) << 8) + | (T.cast(scale_codes[row, k64_inner * 4 + 2], T.uint32) << 16) + | (T.cast(scale_codes[row, k64_inner * 4 + 3], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + row + k64_word = pid_k256 * _BLOCKSCALED_CHUNK_WORDS + k64_inner + physical_row, physical_word = _tl_scale_source_coords(logical_row, k64_word, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word + + return quant, scale_source + + +@tilelang.jit(pass_configs=_TILELANG_PASS_CONFIGS) +def _tilelang_nvfp4_blockscaled_quantize_kernel_64x256_tma_io( + x, + num_stages: int = 2, + num_threads: int = 512, + scale_layout_id: int = _SCALE_LAYOUT_ID_BLOCKSCALED_CHUNK_KMAJOR, +): + """TileLang BF16 -> packed NVFP4 with a 64x256 CTA tile and TMA input/output.""" + + M = T.dynamic("M") + K = T.const("K") + rows_per_cta = 64 + rows_per_warp = rows_per_cta // (num_threads // 32) + in_dtype = T.bfloat16 + compute_dtype = T.float32 + + x: T.Tensor[(M, K), in_dtype] + + quant = T.empty((M, K // 2), T.int8) + scale_source = T.empty((M, K // 64), T.uint32) + + with T.Kernel(M // rows_per_cta, K // 256, threads=num_threads) as (pid_m, pid_k256): + lane = T.get_lane_idx() + warp = T.get_warp_idx_sync() + x_shared = T.alloc_shared((rows_per_cta, 256), in_dtype) + quant_shared = T.alloc_shared((rows_per_cta, 128), T.int8) + scale_codes = T.alloc_shared((rows_per_cta, 16), T.int32) + + T.copy(x[pid_m * rows_per_cta, pid_k256 * 256], x_shared, prefer_instruction="tma") + T.sync_threads() + + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + for k64_inner in T.serial(_BLOCKSCALED_CHUNK_WORDS): + elem_col = k64_inner * 64 + lane * 2 + byte_col = k64_inner * 32 + lane + group = k64_inner * 4 + (lane // 8) + group_lane = lane % 8 + value0 = T.cast(x_shared[row, elem_col], compute_dtype) + value1 = T.cast(x_shared[row, elem_col + 1], compute_dtype) + amax = T.alloc_var(compute_dtype, init=T.max(T.abs(value0), T.abs(value1))) + amax = T.max(amax, T.shfl_xor(amax, 4, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 2, width=8)) + amax = T.max(amax, T.shfl_xor(amax, 1, width=8)) + + scale_code = T.alloc_var(T.int32, init=0) + scale_inv = T.alloc_var(compute_dtype, init=0.0) + if group_lane == 0: + scale_code = _tl_encode_ue4m3_scale_byte_ceil(amax / _FP4_E2M1_MAX) + scale_inv = _tl_ue4m3_scale_byte_inverse(scale_code) + scale_codes[row, group] = scale_code + scale_inv = T.shfl_sync(scale_inv, 0, width=8) + fp4_code0 = _tl_encode_fp4_e2m1_code(value0 * scale_inv) + fp4_code1 = _tl_encode_fp4_e2m1_code(value1 * scale_inv) + packed_byte = fp4_code0 | (fp4_code1 << 4) + quant_shared[row, byte_col] = T.cast(packed_byte, T.int8) + + T.sync_threads() + for row_iter in T.serial(rows_per_warp): + row = warp * rows_per_warp + row_iter + if lane < _BLOCKSCALED_CHUNK_WORDS: + k64_inner = lane + scale_word = ( + T.cast(scale_codes[row, k64_inner * 4], T.uint32) + | (T.cast(scale_codes[row, k64_inner * 4 + 1], T.uint32) << 8) + | (T.cast(scale_codes[row, k64_inner * 4 + 2], T.uint32) << 16) + | (T.cast(scale_codes[row, k64_inner * 4 + 3], T.uint32) << 24) + ) + logical_row = pid_m * rows_per_cta + row + k64_word = pid_k256 * _BLOCKSCALED_CHUNK_WORDS + k64_inner + physical_row, physical_word = _tl_scale_source_coords(logical_row, k64_word, scale_layout_id, K // 64) + scale_source[physical_row, physical_word] = scale_word + + T.sync_threads() + T.tma_copy(quant_shared, quant[pid_m * rows_per_cta, pid_k256 * 128]) + T.tma_store_wait() + + return quant, scale_source + + +def tilelang_quantize_bf16_to_nvfp4_blockscaled( + x, + *, + rows_per_cta: int = 32, + k64_per_cta: int | None = None, + num_stages: int = 4, + num_threads: int = 512, + use_tma_load: bool = True, + scale_layout: str = _SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR, +): + """Quantize a CUDA BF16 activation tensor with a TileLang kernel. + + Parameters + ---------- + x: + CUDA ``torch.bfloat16`` tensor with shape ``[rows, K]``. ``rows`` must + be a multiple of ``128`` and ``K`` a multiple of ``256``. + k64_per_cta: + Number of adjacent K/64 scale-word groups per CTA. The default ``None`` + selects ``4`` for ``rows_per_cta`` in ``(16, 32, 64)`` and ``1`` otherwise. + The ``4`` path uses a K256 CTA tile. + num_stages: + Pipeline stage count for K256 experimental tiles. + num_threads: + Thread count for K256 experimental tiles. ``16x256`` uses ``512``. + ``32x256`` supports ``512`` and ``1024``. + ``64x256`` supports ``512`` and ``1024``. + use_tma_load: + Use TMA for the BF16 global-to-shared input tile before quantization. + scale_layout: + Scale-source memory layout. ``"blockscaled_chunk_kmajor"`` is the + CUTLASS ``BlockScaledBasicChunk`` K-major source contract used by the + TileLang SM120 NVFP4 GEMM path. + + Returns + ------- + tuple + ``(packed_fp4, scale_source)``. ``packed_fp4`` is ``torch.int8[rows, K/2]``. + ``scale_source`` is ``torch.uint32[rows, K/64]`` in CUTLASS + ``BlockScaledBasicChunk`` K-major source order. + The same contract applies to SFA and SFB; ``rows`` is interpreted as M + for A and N for B. + """ + + torch = _import_torch() + if not isinstance(x, torch.Tensor): + raise TypeError(f"x must be a torch.Tensor, got {type(x)!r}") + if x.ndim != 2: + raise ValueError(f"x must be a 2D tensor, got shape {tuple(x.shape)}") + if not x.is_cuda: + raise ValueError("TileLang NVFP4 quantization requires a CUDA tensor") + if x.dtype != torch.bfloat16: + raise TypeError(f"TileLang NVFP4 quantization requires torch.bfloat16 input, got {x.dtype}") + scale_layout_id = _scale_layout_id(scale_layout) + if rows_per_cta < 4 or _BLOCKSCALED_CHUNK_ROWS % rows_per_cta != 0: + raise ValueError(f"rows_per_cta must be at least 4 and divide {_BLOCKSCALED_CHUNK_ROWS}, got {rows_per_cta}") + if k64_per_cta is None: + k64_per_cta = 4 if rows_per_cta in (16, 32, 64) else 1 + if k64_per_cta not in (1, 4): + raise ValueError(f"k64_per_cta must be 1 or 4, got {k64_per_cta}") + if k64_per_cta == 4 and rows_per_cta not in (16, 32, 64): + raise ValueError(f"k64_per_cta=4 currently requires rows_per_cta=16, 32, or 64, got {rows_per_cta}") + if num_threads % 32 != 0: + raise ValueError(f"num_threads must be a multiple of warp size 32, got {num_threads}") + if k64_per_cta == 4 and use_tma_load: + if rows_per_cta == 16 and num_threads != 512: + raise ValueError("rows_per_cta=16,use_tma_load=True uses the explicit one-warp-per-row schedule and requires num_threads=512") + if rows_per_cta in (32, 64) and num_threads not in (512, 1024): + raise ValueError("rows_per_cta=32/64,use_tma_load=True currently supports num_threads=512 or 1024") + if k64_per_cta == 4 and not use_tma_load: + if rows_per_cta == 16 and num_threads != 128: + raise ValueError( + "use_tma_load=False keeps the legacy T.Parallel 16x256 schedule and currently requires rows_per_cta=16,num_threads=128" + ) + if rows_per_cta in (32, 64): + raise ValueError("rows_per_cta=32/64,use_tma_load=False is not implemented") + + rows, cols = x.shape + k_tile = _NVFP4_SCALE_BLOCK_K * _SCALE_BYTES_PER_WORD * _BLOCKSCALED_CHUNK_WORDS + if rows % _BLOCKSCALED_CHUNK_ROWS != 0 or cols % k_tile != 0: + raise ValueError( + f"SM120 NVFP4 quantization requires rows multiple of {_BLOCKSCALED_CHUNK_ROWS} and K multiple of {k_tile}, got {tuple(x.shape)}" + ) + + if k64_per_cta == 4: + if use_tma_load: + if rows_per_cta == 64: + quant, scale_source = _tilelang_nvfp4_blockscaled_quantize_kernel_64x256_tma_io( + x.contiguous(), + num_stages=num_stages, + num_threads=num_threads, + scale_layout_id=scale_layout_id, + ) + elif rows_per_cta == 32: + quant, scale_source = _tilelang_nvfp4_blockscaled_quantize_kernel_32x256_tma_load_global_store( + x.contiguous(), + num_stages=num_stages, + num_threads=num_threads, + scale_layout_id=scale_layout_id, + ) + else: + quant, scale_source = _tilelang_nvfp4_blockscaled_quantize_kernel_16x256_tma_io( + x.contiguous(), + num_stages=num_stages, + num_threads=num_threads, + scale_layout_id=scale_layout_id, + ) + else: + quant, scale_source = _tilelang_nvfp4_blockscaled_quantize_kernel_16x256_tma_store( + x.contiguous(), + num_stages=num_stages, + num_threads=num_threads, + scale_layout_id=scale_layout_id, + ) + else: + quant, scale_source = _tilelang_nvfp4_blockscaled_quantize_kernel( + x.contiguous(), + rows_per_cta=rows_per_cta, + scale_layout_id=scale_layout_id, + ) + return quant.contiguous().view(torch.int8), scale_source + + +def tilelang_quantize_nvfp4_blockscaled(*args, **kwargs): + """Alias for ``tilelang_quantize_bf16_to_nvfp4_blockscaled``.""" + + return tilelang_quantize_bf16_to_nvfp4_blockscaled(*args, **kwargs) + + +def quantize_bf16_to_nvfp4_blockscaled( + x, + *, + block_rows: int = 128, + block_words: int = 4, + scale_block_k: int = 16, + scale_layout: str = _SCALE_LAYOUT_BLOCKSCALED_CHUNK_KMAJOR, + return_scale_bytes: bool = False, +): + """Quantize a BF16 activation tensor to SM120 NVFP4 blockscaled storage. + + Parameters + ---------- + x: + 2D BF16 activation tensor with shape ``[rows, K]``. FP16/FP32 tensors are + accepted for tests and converted internally, but the intended runtime + contract is BF16 activation input. The SM120 source contract requires + ``rows`` to be a multiple of ``128`` and ``K`` to be a multiple of + ``256`` for the promoted ``blockscaled_chunk_kmajor`` scale layout. + scale_block_k: + Number of consecutive K elements sharing one UE4M3 scale byte. SM120 + dense NVFP4 uses ``16``. + return_scale_bytes: + When true, also return semantic row-major UE4M3 scale bytes with shape + ``[rows, K / 16]`` for reference checks. + scale_layout: + Scale-source memory layout for the returned ``scale_source``. + + Returns + ------- + tuple + ``(packed_fp4, scale_source)`` or + ``(packed_fp4, scale_source, scale_bytes)``. ``packed_fp4`` is + ``torch.int8[rows, K / 2]`` with two FP4 E2M1 values per byte. + ``scale_source`` is ``torch.uint32[rows, K // 64]`` in ``scale_layout`` order. + """ + + torch = _import_torch() + _check_block_shape(block_rows, block_words) + _scale_layout_id(scale_layout) + if scale_block_k != _NVFP4_SCALE_BLOCK_K: + raise ValueError(f"SM120 NVFP4 scale_block_k must be 16, got {scale_block_k}") + if not isinstance(x, torch.Tensor): + raise TypeError(f"x must be a torch.Tensor, got {type(x)!r}") + if x.ndim != 2: + raise ValueError(f"x must be a 2D tensor, got shape {tuple(x.shape)}") + if not x.dtype.is_floating_point: + raise TypeError(f"x must be a floating-point tensor, got {x.dtype}") + + rows, cols = x.shape + scale_cols_per_tile = block_words * _SCALE_BYTES_PER_WORD + if rows % block_rows != 0 or cols % (scale_block_k * scale_cols_per_tile) != 0: + raise ValueError( + f"SM120 NVFP4 quantization requires rows multiple of {block_rows} and K multiple of " + f"{scale_block_k * scale_cols_per_tile}, got {tuple(x.shape)}" + ) + + x_f32 = x.contiguous().to(torch.float32) + blocks = x_f32.reshape(rows, cols // scale_block_k, scale_block_k) + amax = blocks.abs().amax(dim=2) + target_scale = amax / _FP4_E2M1_MAX + scale_bytes = encode_ue4m3_scale_bytes(target_scale, rounding="ceil", min_positive=True) + scale_bytes = torch.where(amax == 0, torch.zeros_like(scale_bytes), scale_bytes) + scale_values = decode_ue4m3_scale_bytes(scale_bytes) + scaled_blocks = torch.where(scale_values[..., None] > 0, blocks / scale_values[..., None], torch.zeros_like(blocks)) + fp4_codes = encode_fp4_e2m1_values(scaled_blocks.reshape(rows, cols)) + packed_fp4 = pack_fp4_e2m1_codes(fp4_codes) + packed_scales = pack_blockscaled_chunk_kmajor_scale_bytes(scale_bytes, block_rows=block_rows, block_words=block_words) + if return_scale_bytes: + return packed_fp4, packed_scales, scale_bytes + return packed_fp4, packed_scales + + +def quantize_nvfp4_blockscaled(*args, **kwargs): + """Backward-compatible alias for ``quantize_bf16_to_nvfp4_blockscaled``.""" + + return quantize_bf16_to_nvfp4_blockscaled(*args, **kwargs)