From 6e12aebefb6b9c120945e7f3b3464a807c541856 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:15:49 +0800 Subject: [PATCH 01/51] Create test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 479 +++++++++++++++++++++++++++++++++++++ 1 file changed, 479 insertions(+) create mode 100644 tests/test_sm90_fp8_fp4.py diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py new file mode 100644 index 0000000000..62e57d9669 --- /dev/null +++ b/tests/test_sm90_fp8_fp4.py @@ -0,0 +1,479 @@ +import time +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.math import cast_back_from_fp4, per_token_cast_to_fp4, per_token_cast_to_fp8 + + +def _align_up(x: int, alignment: int) -> int: + return (x + alignment - 1) // alignment * alignment + + +def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int = 128) -> torch.Tensor: + group_idx = torch.arange(x.size(-1), device=x.device) // gran_k + return x.float() * sf[..., group_idx] + + +def _require_sm90() -> None: + assert torch.cuda.is_available() + major, _ = torch.cuda.get_device_capability() + if major != 9: + raise RuntimeError(f"This fallback test is intended for SM90, got sm_{major}x") + + +def _time_cuda(fn, warmup: int = 5, iters: int = 20) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 + + +def _normal_case(m: int, n: int, k: int, gran_k: int = 128) -> None: + a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) + b = per_token_cast_to_fp4(b_ref_src, use_ue8m0=True, gran_k=gran_k) + c = torch.zeros((m, n), device="cuda", dtype=torch.float) + d = torch.empty((m, n), device="cuda", dtype=torch.float) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) + b_dequant = cast_back_from_fp4(b[0], b[1], gran_k=gran_k) + b_fp8 = per_token_cast_to_fp8(b_dequant, use_ue8m0=False, gran_k=gran_k) + ref = (a_dequant @ b_dequant.t()).to(torch.bfloat16) + + def run(): + deep_gemm.fp8_gemm_nt( + a, + b_fp8, + d, + c=c, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + ) + + d_fp8 = torch.empty_like(d) + d_fused = torch.empty_like(d) + + def run_fp8(): + deep_gemm.fp8_gemm_nt( + a, + b_fp8, + d_fp8, + c=c, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + ) + + d_fused = torch.empty_like(d) + + def run_fused(): + deep_gemm.fp8_fp4_gemm_nt_sm90_cuda_core( + a, + b, + d_fused, + c=c, + gran_k=gran_k, + ) + + d_fused_wgmma = torch.empty_like(d) + + def run_fused_wgmma(): + deep_gemm.fp8_fp4_gemm_nt_sm90_fused_wgmma( + a, + b, + d_fused_wgmma, + c=c, + gran_k=gran_k, + ) + + run() + run_fp8() + if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core"): + run_fused() + if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma"): + run_fused_wgmma() + fallback_diff = calc_diff(d, ref) + fp8_diff = calc_diff(d_fp8, ref) + fallback_vs_fp8_diff = calc_diff(d, d_fp8) + fused_diff = calc_diff(d_fused, ref) if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core") else None + fused_wgmma_diff = ( + calc_diff(d_fused_wgmma, ref) + if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") + else None + ) + fallback_elapsed = _time_cuda(run) + fp8_elapsed = _time_cuda(run_fp8) + fused_elapsed = ( + _time_cuda(run_fused) + if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core") + else None + ) + fused_wgmma_elapsed = ( + _time_cuda(run_fused_wgmma) + if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") + else None + ) + fallback_tflops = 2 * m * n * k / fallback_elapsed / 1e12 + fp8_tflops = 2 * m * n * k / fp8_elapsed / 1e12 + fused_tflops = ( + 2 * m * n * k / fused_elapsed / 1e12 if fused_elapsed is not None else None + ) + fused_wgmma_tflops = ( + 2 * m * n * k / fused_wgmma_elapsed / 1e12 + if fused_wgmma_elapsed is not None + else None + ) + slowdown = fallback_elapsed / fp8_elapsed + message = ( + f"normal m={m} n={n} k={k}:\n" + f" fallback end2end: diff={fallback_diff:.6f}, " + f"time={fallback_elapsed * 1e6:.1f} us, {fallback_tflops:.2f} TFLOPS\n" + f" pure fp8 gemm: diff={fp8_diff:.6f}, " + f"time={fp8_elapsed * 1e6:.1f} us, {fp8_tflops:.2f} TFLOPS\n" + f" fallback vs fp8: diff={fallback_vs_fp8_diff:.6f}\n" + f" slowdown: {slowdown:.2f}x" + ) + if fused_elapsed is not None: + message += ( + f"\n cuda-core fused: diff={fused_diff:.6f}, " + f"time={fused_elapsed * 1e6:.1f} us, {fused_tflops:.2f} TFLOPS, " + f"slowdown_vs_fp8={fused_elapsed / fp8_elapsed:.2f}x" + ) + if fused_wgmma_elapsed is not None: + message += ( + f"\n fused wgmma: diff={fused_wgmma_diff:.6f}, " + f"time={fused_wgmma_elapsed * 1e6:.1f} us, {fused_wgmma_tflops:.2f} TFLOPS, " + f"slowdown_vs_fp8={fused_wgmma_elapsed / fp8_elapsed:.2f}x" + ) + print(message) + assert fallback_diff < 0.015 + assert fp8_diff < 0.015 + assert fallback_vs_fp8_diff < 1e-6 + if fused_diff is not None: + assert fused_diff < 0.015 + if fused_wgmma_diff is not None: + assert fused_wgmma_diff < 0.015 + + +def _m_grouped_contiguous_case( + num_groups: int, + m_per_group: int | None, + n: int, + k: int, + gran_k: int = 128, + group_sizes: list[int] | None = None, + use_psum_layout: bool = False, + autotune_fused: bool = True, +) -> None: + if group_sizes is None: + assert m_per_group is not None + group_sizes = [m_per_group] * num_groups + assert len(group_sizes) == num_groups + print( + f"starting m_grouped groups={num_groups} n={n} k={k} " + f"layout={'psum' if use_psum_layout else 'per-row'} " + f"sizes={group_sizes} autotune={autotune_fused}", + flush=True, + ) + + block_m_alignment = 128 + if use_psum_layout: + group_starts: list[int] = [] + group_ends: list[int] = [] + cursor = 0 + for size in group_sizes: + group_starts.append(cursor) + cursor += size + group_ends.append(cursor) + cursor = _align_up(cursor, block_m_alignment) + m = group_ends[-1] if group_ends else 0 + grouped_layout = torch.tensor(group_ends, device="cuda", dtype=torch.int32) + else: + assert all(size % block_m_alignment == 0 for size in group_sizes) + m = sum(group_sizes) + group_starts = [] + cursor = 0 + layout_chunks = [] + for group_id, size in enumerate(group_sizes): + group_starts.append(cursor) + cursor += size + if size > 0: + layout_chunks.append( + torch.full((size,), group_id, device="cuda", dtype=torch.int32) + ) + grouped_layout = torch.cat(layout_chunks) if layout_chunks else torch.empty((0,), device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) + b_fp4 = torch.empty((num_groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((num_groups, n, k // gran_k), device="cuda", dtype=torch.float) + for group_id in range(num_groups): + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=gran_k + ) + b = (b_fp4, b_sf) + d = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) + b_fp8_data = torch.empty((num_groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((num_groups, n, k // gran_k), device="cuda", dtype=torch.float) + ref = torch.empty_like(d) + for group_id in range(num_groups): + b_dequant = cast_back_from_fp4(b[0][group_id], b[1][group_id], gran_k=gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=gran_k + ) + start = group_starts[group_id] + end = ( + _align_up(group_ends[group_id], block_m_alignment) + if use_psum_layout and group_id + 1 < num_groups + else group_ends[group_id] + ) if use_psum_layout else start + group_sizes[group_id] + end = min(end, m) + if start == end: + continue + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + def run(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d, + grouped_layout, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + use_psum_layout=use_psum_layout, + expected_m_for_psum_layout=m if use_psum_layout else None, + ) + + d_fp8 = torch.empty_like(d) + d_fused = torch.empty_like(d) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + use_psum_layout=use_psum_layout, + expected_m_for_psum_layout=m if use_psum_layout else None, + ) + + def run_fused_wgmma( + block_m_override: int | None = None, + block_n_override: int | None = None, + ): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b, + d_fused, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=use_psum_layout, + expected_m_for_psum_layout=m if use_psum_layout else None, + block_m_override=block_m_override, + block_n_override=block_n_override, + ) + + print(" running fallback...", flush=True) + run() + print(" running pure fp8...", flush=True) + run_fp8() + print(" running fused wgmma...", flush=True) + run_fused_wgmma() + fallback_diff = calc_diff(d, ref) + fp8_diff = calc_diff(d_fp8, ref) + fused_diff = calc_diff(d_fused, ref) + fallback_vs_fp8_diff = calc_diff(d, d_fp8) + fused_vs_fp8_diff = calc_diff(d_fused, d_fp8) + fallback_elapsed = _time_cuda(run) + fp8_elapsed = _time_cuda(run_fp8) + fused_elapsed = _time_cuda(run_fused_wgmma) + best_fused = (fused_elapsed, None, fused_diff) + if autotune_fused: + # Psum grouped_layout is encoded with 128-row alignment, so only BLOCK_M=128 + # is semantically valid. Sweep BLOCK_N only for psum. + block_m_candidates = (128,) if use_psum_layout else (64, 128, 256) + block_n_candidates = (64, 128, 256) + for block_m in block_m_candidates: + if m < block_m: + continue + if any(size > 0 and size < block_m for size in group_sizes) and not use_psum_layout: + continue + for block_n in block_n_candidates: + d_candidate = torch.empty_like(d) + + def run_candidate( + block_m: int = block_m, + block_n: int = block_n, + out: torch.Tensor = d_candidate, + ): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b, + out, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=use_psum_layout, + expected_m_for_psum_layout=m if use_psum_layout else None, + block_m_override=block_m, + block_n_override=block_n, + ) + + try: + run_candidate() + except RuntimeError as exc: + print( + f" fused autotune block_m={block_m} block_n={block_n}: " + f"skipped ({exc})" + ) + continue + candidate_diff = calc_diff(d_candidate, ref) + candidate_elapsed = _time_cuda(run_candidate) + candidate_tflops = 2 * m * n * k / candidate_elapsed / 1e12 + print( + f" fused autotune block_m={block_m} block_n={block_n}: " + f"diff={candidate_diff:.6f}, " + f"time={candidate_elapsed * 1e6:.1f} us, {candidate_tflops:.2f} TFLOPS" + ) + if candidate_diff < 0.015 and candidate_elapsed < best_fused[0]: + best_fused = (candidate_elapsed, (block_m, block_n), candidate_diff) + d_fused.copy_(d_candidate) + fused_elapsed, best_tile, fused_diff = best_fused + fused_vs_fp8_diff = calc_diff(d_fused, d_fp8) + else: + best_tile = None + fallback_tflops = 2 * m * n * k / fallback_elapsed / 1e12 + fp8_tflops = 2 * m * n * k / fp8_elapsed / 1e12 + fused_tflops = 2 * m * n * k / fused_elapsed / 1e12 + slowdown = fallback_elapsed / fp8_elapsed + fused_slowdown = fused_elapsed / fp8_elapsed + print( + f"m_grouped groups={num_groups} m={m} n={n} k={k} " + f"layout={'psum' if use_psum_layout else 'per-row'} sizes={group_sizes}:\n" + f" fallback end2end: diff={fallback_diff:.6f}, " + f"time={fallback_elapsed * 1e6:.1f} us, {fallback_tflops:.2f} TFLOPS\n" + f" pure fp8 gemm: diff={fp8_diff:.6f}, " + f"time={fp8_elapsed * 1e6:.1f} us, {fp8_tflops:.2f} TFLOPS\n" + f" fused wgmma: diff={fused_diff:.6f}, " + f"time={fused_elapsed * 1e6:.1f} us, {fused_tflops:.2f} TFLOPS, " + f"slowdown_vs_fp8={fused_slowdown:.2f}x" + f"{'' if best_tile is None else f', best_tile={best_tile}'}\n" + f" fallback vs fp8: diff={fallback_vs_fp8_diff:.6f}\n" + f" fused vs fp8: diff={fused_vs_fp8_diff:.6f} " + f"(expected: fused keeps original FP4 scales)\n" + f" slowdown: {slowdown:.2f}x" + ) + # Psum edge cases with few effective rows can amplify the FP4->FP8 requant + # baseline error. Keep fused correctness strict; relax only the requantized + # fallback/pure-FP8 baseline sanity threshold. + requant_threshold = 0.25 if use_psum_layout else 0.05 + assert fallback_diff < requant_threshold + assert fp8_diff < requant_threshold + assert fused_diff < 0.015 + assert fallback_vs_fp8_diff < 1e-6 + + +def test_sm90_fp8_fp4_fallback() -> None: + _require_sm90() + torch.manual_seed(0) + + # Small enough for quick CI/local validation, large enough to exercise SM90 FP8 GEMM. + _normal_case(m=128, n=1024, k=1024) + _m_grouped_contiguous_case(num_groups=1, m_per_group=128, n=1024, k=1024) + _m_grouped_contiguous_case(num_groups=4, m_per_group=128, n=1024, k=1024) + _m_grouped_contiguous_case(num_groups=8, m_per_group=128, n=512, k=1024) + _m_grouped_contiguous_case( + num_groups=4, m_per_group=256, n=2048, k=1024, autotune_fused=True + ) + _m_grouped_contiguous_case( + num_groups=4, + m_per_group=None, + n=1024, + k=1024, + group_sizes=[96, 0, 160, 64], + use_psum_layout=True, + ) + _m_grouped_contiguous_case( + num_groups=5, + m_per_group=None, + n=512, + k=2048, + group_sizes=[0, 128, 32, 256, 48], + use_psum_layout=True, + autotune_fused=True, + ) + # Psum stability coverage: + # - aligned groups: sanity baseline where psum should behave like regular contiguous groups + # - leading/trailing zero groups: scheduler fallthrough and empty group handling + # - tiny partial groups: mixed valid/invalid rows inside a BLOCK_M tile + # - larger K/N: exercise the 1d2d decode pipeline and smem SFB cache over more K blocks + _m_grouped_contiguous_case( + num_groups=3, + m_per_group=None, + n=512, + k=1024, + group_sizes=[128, 128, 128], + use_psum_layout=True, + ) + _m_grouped_contiguous_case( + num_groups=4, + m_per_group=None, + n=512, + k=1024, + group_sizes=[0, 64, 128, 0], + use_psum_layout=True, + ) + _m_grouped_contiguous_case( + num_groups=4, + m_per_group=None, + n=768, + k=1024, + group_sizes=[1, 127, 129, 1], + use_psum_layout=True, + ) + _m_grouped_contiguous_case( + num_groups=6, + m_per_group=None, + n=1024, + k=2048, + group_sizes=[64, 0, 64, 0, 192, 32], + use_psum_layout=True, + ) + _m_grouped_contiguous_case( + num_groups=4, + m_per_group=None, + n=1024, + k=1024, + group_sizes=[96, 160, 0, 64], + use_psum_layout=True, + autotune_fused=True, + ) + + +if __name__ == "__main__": + start_time = time.time() + test_sm90_fp8_fp4_fallback() + print(f"done in {time.time() - start_time:.2f}s") From b118d880cc3bff1e2dc41965c86a36e51d872e61 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:17:11 +0800 Subject: [PATCH 02/51] Create sm90_fp8_fp4_gemm_1d1d.cuh --- .../impls/sm90_fp8_fp4_gemm_1d1d.cuh | 396 ++++++++++++++++++ 1 file changed, 396 insertions(+) create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh new file mode 100644 index 0000000000..0aebdf7dca --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d1d.cuh @@ -0,0 +1,396 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + + +__device__ __forceinline__ uint8_t sm90_fp8_fp4_fused_e2m1_to_e4m3_bits(uint8_t code) { + // E2M1 values {0, 0.5, 1, 1.5, 2, 3, 4, 6} map exactly to E4M3. + constexpr uint8_t kE2M1ToE4M3[8] = {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c}; + const uint8_t value_idx = code & 0x07u; + const uint8_t sign = (value_idx != 0u && (code & 0x08u) != 0u) ? 0x80u : 0x00u; + return kE2M1ToE4M3[value_idx] | sign; +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void +sm90_fp8_fp4_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, int8_t* gmem_b_ptr, + cd_dtype_t* gmem_d_ptr, int* grouped_layout, + cute::TmaDescriptor* tensor_map_buffer, + uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a_base, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, + const __grid_constant__ cute::TmaDescriptor tensor_map_cd) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) + // Scaling checks + DG_STATIC_ASSERT(kNumTMAThreads == 128 and kNumMathThreads % 128 == 0, "Invalid Threads"); + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, "Invalid C/D data dtype"); + DG_STATIC_ASSERT(kGemmType == GemmType::Normal or kGemmType == GemmType::MGroupedContiguous or + kGemmType == GemmType::MGroupedContiguousWithPsumLayout, + "SM90 FP8xFP4 fused only supports normal and m-grouped contiguous GEMM"); + + // Types + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size"); + + // Overwrite shape constants if the compiler gives + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + + // Shared memory + static constexpr uint32_t SMEM_TENSOR_MAP_SIZE = (kGemmType == GemmType::KGroupedContiguous ? sizeof(cute::TmaDescriptor) * 2 : 0); + static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(cd_dtype_t); + static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = BLOCK_N * sizeof(float); + static constexpr uint32_t ALIGNED_SMEM_SFB_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFB_SIZE_PER_STAGE, 128u); + DG_STATIC_ASSERT(SMEM_SFA_SIZE_PER_STAGE % 128 == 0, "Invalid TMA alignment"); + + // Configs + const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane_idx = threadIdx.x % 32; + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a_base); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_sfb); + cute::prefetch_tma_descriptor(&tensor_map_cd); + } + __syncwarp(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + + // Tensor maps on shared and global memory + (void)gmem_a_ptr; + (void)tensor_map_buffer; + (void)gmem_d_ptr; + + // Data on shared memory + auto smem_d = reinterpret_cast(smem_buffer + SMEM_TENSOR_MAP_SIZE); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + (SMEM_TENSOR_MAP_SIZE + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE)); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + (SMEM_TENSOR_MAP_SIZE + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE)); + }); + constexpr auto SMEM_SF_OFFSET = SMEM_TENSOR_MAP_SIZE + SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_SF_OFFSET + i * SMEM_SFA_SIZE_PER_STAGE)); + }); + auto smem_sfb = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * ALIGNED_SMEM_SFB_SIZE_PER_STAGE)); + }); + + // Barriers on shared memory + constexpr auto SMEM_BARRIER_OFFSET = SMEM_SF_OFFSET + kNumStages * (SMEM_SFA_SIZE_PER_STAGE + ALIGNED_SMEM_SFB_SIZE_PER_STAGE); + auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_BARRIER_OFFSET + i * static_cast(sizeof(Barrier)))); + }); + auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_BARRIER_OFFSET + (kNumStages + i) * static_cast(sizeof(Barrier)))); + }); + + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // Initialize barriers + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Pipeline unroll control + constexpr uint32_t kNumPipelineUnrolls = (kGemmType == GemmType::KGroupedContiguous ? 0 : kNumStages); + + // Register reconfigurations (more math registers are needed with unrolling) + constexpr uint32_t kNumTMARegisters = (kNumPipelineUnrolls == 0 ? 40 : 24); + constexpr uint32_t kNumMathRegisters = (kNumPipelineUnrolls == 0 ? 232 : 240); + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // TMA and MMA pipeline + const auto get_pipeline = [=](const uint32_t& iter_idx) -> cute::tuple { + return {iter_idx % kNumStages, (iter_idx / kNumStages) & 1}; // Pipeline stage and phase + }; + uint32_t iter_idx = 0; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); + const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); + + const uint32_t num_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + const uint32_t m_idx = m_block_idx * BLOCK_M; + const uint32_t n_idx = n_block_idx * BLOCK_N; + + #pragma unroll kNumPipelineUnrolls + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; ++ k_block_idx) { + // Wait consumer release + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Issue TMA + auto& full_barrier = *full_barriers[stage_idx]; + const uint32_t k_idx = k_block_idx * BLOCK_K; + const uint32_t current_group_idx = get_current_group_idx(); + const uint32_t sf_k_idx_a = k_block_idx; + const uint32_t sf_k_idx_b = current_group_idx * math::ceil_div(shape_k, BLOCK_K) + k_block_idx; + const auto tensor_map_a_ptr = &tensor_map_a_base; + tma::copy(&tensor_map_sfa, &full_barrier, smem_sfa[stage_idx], m_idx, sf_k_idx_a, num_tma_multicast_a); + tma::copy(&tensor_map_sfb, &full_barrier, smem_sfb[stage_idx], n_idx, sf_k_idx_b, num_tma_multicast_b); + tma::copy(tensor_map_a_ptr, &full_barrier, smem_a[stage_idx], k_idx, m_idx, num_tma_multicast_a); + full_barrier.arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + SMEM_SFB_SIZE_PER_STAGE); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + #pragma unroll + for (uint32_t s = 0; s < kNumStages; ++ s) { + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + empty_barriers[stage_idx]->wait(phase ^ 1); + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto r_0 = warp_idx * 16 + row_idx, r_1 = r_0 + 8; + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Accumulation for WGMMA or CUDA promotion + DG_STATIC_ASSERT(BLOCK_M == WGMMA::M * (BLOCK_M <= 64 ? 1 : 2), "Invalid block sizes"); + const uint32_t current_shape_k = shape_k; + const uint32_t current_group_idx = get_current_group_idx(); + const uint32_t num_k_blocks = math::ceil_div(current_shape_k, BLOCK_K); + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0}; + float2 scales_b[WGMMA::kNumAccum / 4]; + + // Empty barrier arrival + auto empty_barrier_arrive = [&](uint32_t s) { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? empty_barriers[s]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta) : void(); + } + }; + + #pragma unroll kNumPipelineUnrolls + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; ++ k_block_idx) { + // Wait TMA arrivals + CUTE_TIE_DECL(get_pipeline(iter_idx ++), stage_idx, phase); + full_barriers[stage_idx]->wait(phase); + + // Decode one packed byte into two FP8 values. SM90 cannot use + // Blackwell's FP4 TMA/UMMA path, so keep B loads coalesced here. + for (uint32_t idx = threadIdx.x; idx < BLOCK_N * (BLOCK_K / 2); idx += kNumMathThreads) { + const uint32_t tile_n = idx / (BLOCK_K / 2); + const uint32_t packed_k = idx % (BLOCK_K / 2); + const uint32_t tile_k = packed_k * 2; + const uint32_t global_n = n_block_idx * BLOCK_N + tile_n; + const uint32_t global_k = k_block_idx * BLOCK_K + tile_k; + constexpr bool kFullStaticNK = SHAPE_N != 0 and SHAPE_K != 0 and + SHAPE_N % BLOCK_N == 0 and SHAPE_K % BLOCK_K == 0; + uint8_t packed; + const uint32_t b_group_offset = current_group_idx * shape_n * (shape_k / 2); + if constexpr (kFullStaticNK) { + packed = static_cast(gmem_b_ptr[b_group_offset + global_n * (shape_k / 2) + global_k / 2]); + } else { + packed = 0; + if (global_n < shape_n and global_k < current_shape_k) + packed = static_cast(gmem_b_ptr[b_group_offset + global_n * (shape_k / 2) + global_k / 2]); + } + + const uint32_t n_group = tile_n / 8; + const uint32_t n_in_group = tile_n % 8; + const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); + const uint32_t smem_idx = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K + swizzled_k; + auto smem_b_bytes = reinterpret_cast(smem_b[stage_idx]); + const uint8_t lo = sm90_fp8_fp4_fused_e2m1_to_e4m3_bits(packed & 0x0fu); + const uint8_t hi = sm90_fp8_fp4_fused_e2m1_to_e4m3_bits((packed >> 4) & 0x0fu); + if constexpr (kFullStaticNK) { + *reinterpret_cast(smem_b_bytes + smem_idx) = static_cast(lo) | (static_cast(hi) << 8); + } else { + if (global_k + 1 < current_shape_k) { + *reinterpret_cast(smem_b_bytes + smem_idx) = static_cast(lo) | (static_cast(hi) << 8); + } else { + smem_b_bytes[smem_idx] = lo; + } + } + } + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + + // Read A scales + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + auto scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + r_0); + auto scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + r_1); + + // Read B scales + #pragma unroll + for (int i = 0; i < WGMMA::kNumAccum / 4; ++i) + scales_b[i] = ptx::ld_shared(reinterpret_cast(smem_sfb[stage_idx] + i * 8 + col_idx * 2)); + + // Commit WGMMA instructions + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc(smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc(smem_b[stage_idx] + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + // Notify barrier arrival + empty_barrier_arrive(stage_idx); + + // Promote with scales + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float &scale_b_0 = scales_b[i].x; + const float &scale_b_1 = scales_b[i].y; + final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + } + } + + if constexpr (kGemmType == GemmType::Normal) { + // Flush previous stores + if (warp_idx % 4 == 0 and cute::elect_one_sync()) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(128, math_wg_idx); + + // Store to D shared memory + const auto smem_d_0 = reinterpret_cast(smem_d + r_0 * BLOCK_N + col_idx * 2); + const auto smem_d_1 = reinterpret_cast(smem_d + r_1 * BLOCK_N + col_idx * 2); + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + ptx::st_shared(smem_d_0 + i * 4, {final_accum[i * 4 + 0], final_accum[i * 4 + 1]}); + ptx::st_shared(smem_d_1 + i * 4, {final_accum[i * 4 + 2], final_accum[i * 4 + 3]}); + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(128, math_wg_idx); + + // Use TMA reduce-add to accumulate C into D for the normal FP32 path. + if (warp_idx % 4 == 0 and cute::elect_one_sync()) { + cute::SM90_TMA_REDUCE_ADD_2D::copy( + &tensor_map_cd, smem_d_0, n_block_idx * BLOCK_N, + m_block_idx * BLOCK_M + r_0); + cute::tma_store_arrive(); + } + } else { + const uint32_t row_0 = m_block_idx * BLOCK_M + r_0; + const uint32_t row_1 = m_block_idx * BLOCK_M + r_1; + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t col = n_block_idx * BLOCK_N + i * 8 + col_idx * 2; + if (row_0 < shape_m and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = cd_dtype_t(final_accum[i * 4 + 0]); + if (row_0 < shape_m and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = cd_dtype_t(final_accum[i * 4 + 1]); + if (row_1 < shape_m and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = cd_dtype_t(final_accum[i * 4 + 2]); + if (row_1 < shape_m and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = cd_dtype_t(final_accum[i * 4 + 3]); + } + } + __syncwarp(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_90a"); +#endif +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop From 561cf215ea3464bfec014b49a0ffbdd47eff0b3b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:17:43 +0800 Subject: [PATCH 03/51] Create sm90_fp8_fp4_gemm_1d2d.cuh --- .../impls/sm90_fp8_fp4_gemm_1d2d.cuh | 692 ++++++++++++++++++ 1 file changed, 692 insertions(+) create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh new file mode 100644 index 0000000000..caf0c1ceab --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh @@ -0,0 +1,692 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +template +CUTLASS_DEVICE void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) { + if (num_former_iters == kNumFormerIters) { + func(cute::Int{}); + return; + } + + if constexpr (kNumFormerIters + kGap <= kEnd) + dispatch_num_former_iters(num_former_iters, func); +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void +sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, + nv_bfloat16* gmem_d_ptr, + uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, + const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_d, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) + // Scaling checks + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT( + math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or + (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); + + // Types + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M, "Invalid block size"); + + // Overwrite shape constants if the compiler gives + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + + // Shared memory + static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0); + static constexpr uint32_t SMEM_D_SIZE = math::constexpr_align(BLOCK_M * BLOCK_N * static_cast(sizeof(__nv_bfloat16)), 1024u); + static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3); + // Packed FP4 B is loaded by TMA into a separate buffer; each row is BLOCK_K / 2 bytes. + static constexpr uint32_t BLOCK_K_PACKED = BLOCK_K / 2; + static constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = BLOCK_N * BLOCK_K_PACKED; + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + const uint32_t shape_k_scales = math::ceil_div(shape_k, BLOCK_K); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); + // SFB cache aliases the smem_d region (BLOCK_M * BLOCK_N * sizeof(bf16) bytes) + // for the entire K loop. smem_d is only used during epilogue (after the K + // loop), and the cooperative SFB prefetch sync + final wgmma wait give us a + // clean handoff. We require sfb_bytes <= smem_d_bytes; with BLOCK_K=128 + // this reduces to BLOCK_M*128 >= shape_k*2 -> BLOCK_M*64 >= shape_k. The + // assert below verifies this at runtime; for typical kernels (BLOCK_M>=64, + // shape_k<=8192) this is always satisfied. + const uint32_t smem_sfb_bytes = math::align(shape_k_scales * BLOCK_N * sizeof(float), 16u); + constexpr uint32_t smem_sfb_size = 0; // SFB aliases smem_d, no separate allocation + + // NOTES: Make sure we have enough shared memory for WGMMA padding + static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); + DG_STATIC_ASSERT(WGMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, "Memory Out of bound for WGMMA"); + + // Configs + const uint32_t num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K); + const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_d); + } + __syncwarp(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + auto fp4_to_e4m3_bits = [](uint32_t code) -> uint32_t { + // E2M1 mag {0..7} -> E4M3 magnitude bytes {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c}. + // Use hardware __byte_perm as a single-cycle 8-entry byte LUT: + // LUT_LO[mag] = {0x00, 0x30, 0x38, 0x3c}, LUT_HI[mag-4] = {0x40, 0x44, 0x48, 0x4c} + // __byte_perm(LUT_LO, LUT_HI, mag) returns 4 copies of the byte at index `mag` + // (mag is in 0..7, fits in the lower nibble, used as the byte selector). + // Sign: bit 3 of code shifted to MSB, masked to 0 if mag==0 to avoid -0. + constexpr uint32_t LUT_LO = 0x3c383000u; + constexpr uint32_t LUT_HI = 0x4c484440u; + const uint32_t mag = code & 0x07u; + const uint32_t mag_nz_m1 = (mag + 0x07u) >> 3; // 1 if mag!=0 else 0 + const uint32_t sign_mask = -mag_nz_m1; // 0xffffffff if mag!=0 else 0 + const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; + const uint32_t sign = ((code & 0x08u) << 4) & sign_mask; + return mag_byte | sign; + }; + auto fp4_pair_to_e4m3_pair = [&](uint32_t packed) { + const uint32_t lo = fp4_to_e4m3_bits(packed & 0x0fu); + const uint32_t hi = fp4_to_e4m3_bits((packed >> 4) & 0x0fu); + return lo | (hi << 8); + }; + + // Data on shared memory + auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_B_PACKED_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + auto smem_b_packed = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_B_PACKED_OFFSET + i * SMEM_B_PACKED_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumPackedStages * SMEM_B_PACKED_SIZE_PER_STAGE; + auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + }); + // SFB cache aliases the smem_d region (smem_buffer base). Safe because + // smem_d is only used in the epilogue, after the K loop completes and + // wgmma writes are drained. + auto smem_sfb = reinterpret_cast(smem_buffer); + DG_TRAP_ONLY_DEVICE_ASSERT(smem_sfb_bytes <= SMEM_D_SIZE); + + // Fill barriers (SFB does not allocate any extra space because it aliases smem_d) + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); + auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); + auto packed_full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + 2 * kNumStages + i; }); + auto packed_empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + 2 * kNumStages + kNumPackedStages + i; }); + + // Initialize barriers + DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast"); + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + #pragma unroll + for (uint32_t i = 0; i < kNumPackedStages; ++ i) { + packed_full_barriers[i]->init(1); + packed_empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Register reconfigurations + constexpr uint32_t kNumTMARegisters = 40; + constexpr uint32_t kNumMathRegisters = kNumMathThreads == 128 ? 248 : 232; + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // `gmem_b_ptr` is no longer used: B is now loaded by TMA into `smem_b_packed`. + (void)gmem_b_ptr; + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // Pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + uint32_t packed_stage_idx = 0, packed_phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + + // Packed-B pipeline runs in lock-step with the K loop but with its own depth + packed_stage_idx = packed_stage_idx == kNumPackedStages - 1 ? 0 : packed_stage_idx + 1; + packed_phase ^= packed_stage_idx == 0; + }; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used. + // We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32`. + if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); + const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); + + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release for A/SFA slot and packed-B slot independently + empty_barriers[stage_idx]->wait(phase ^ 1); + packed_empty_barriers[packed_stage_idx]->wait(packed_phase ^ 1); + + // Issue TMA A + constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); + const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); + + constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; + auto& full_barrier = *full_barriers[stage_idx]; + const uint32_t k_idx = k_block_idx * BLOCK_K; + tma::copy(&tensor_map_a, &full_barrier, + smem_a[stage_idx], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx), + num_tma_multicast_a, batch_idx); + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx), + num_tma_multicast_a); + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); + + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) + auto& packed_full_barrier = *packed_full_barriers[packed_stage_idx]; + const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; + tma::copy(&tensor_map_b, &packed_full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[packed_stage_idx]), + k_idx_packed, + scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b, batch_idx); + packed_full_barrier.arrive_and_expect_tx(SMEM_B_PACKED_SIZE_PER_STAGE); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; ++ i) { + empty_barriers[stage_idx]->wait(phase ^ 1); + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + } + for (uint32_t i = 0; i < kNumPackedStages; ++ i) { + packed_empty_barriers[packed_stage_idx]->wait(packed_phase ^ 1); + packed_stage_idx = packed_stage_idx == kNumPackedStages - 1 ? 0 : packed_stage_idx + 1; + packed_phase ^= packed_stage_idx == 0; + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto r_0 = warp_idx * 16 + row_idx, r_1 = r_0 + 8; + + auto a_desc = mma::sm90::make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1); + auto b_desc = mma::sm90::make_smem_desc(smem_b[0], 1); + const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0); + const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0); + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const uint32_t current_group_idx = get_current_group_idx(); + + // Cooperatively prefetch the SFB tile for this block from gmem to smem + // Layout in smem: [shape_k_scales, BLOCK_N] (k outer, n inner). + // Out-of-bound n is filled with 1.0f to keep `n_idx >= shape_n` neutral. + { + const uint32_t n_block_base = n_block_idx * BLOCK_N; + const uint32_t total = shape_k_scales * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_idx = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float val; + if (n_idx >= shape_n) { + val = 1.0f; + } else if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + val = sfb[current_group_idx * aligned_shape_n_sfb * shape_k_scales + + k_idx * aligned_shape_n_sfb + n_idx]; + } else { + val = sfb[current_group_idx * shape_n * shape_k_scales + + n_idx * shape_k_scales + k_idx]; + } + ptx::st_shared(smem_sfb + k_idx * BLOCK_N + n_off, val); + } + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + + auto load_sfb = [&](uint32_t n_idx, uint32_t k_block_idx) { + // SFB has been staged into smem above; out-of-bound `n_idx` already + // resolves to 1.0f because we wrote 1.0f for those slots. + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + return ptx::ld_shared(smem_sfb + k_block_idx * BLOCK_N + n_off); + }; + + // Decide the number of scales B to load + DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0); + uint32_t num_former_iters = BLOCK_N / 8; + if constexpr (not kMustUseUniformedScaleB) { + num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; + } + + // Accumulation for WGMMA or CUDA promotion + constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); + float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; + + // Pick threads whose WGMMA results are to be stored in shared memory + DG_STATIC_ASSERT(BLOCK_M >= 64 or kNumMathThreads == 128, "Only one math warp group for `BLOCK_M < 64`"); + constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + const bool do_wgmma_store = BLOCK_M >= WGMMA::M or warp_idx < kNumWGMMAStoreThreads / 32; + + // Empty barrier arrival + auto empty_barrier_arrive = [&]() { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void(); + } + }; + auto packed_empty_barrier_arrive_at = [&](uint32_t which_packed_stage) { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? packed_empty_barriers[which_packed_stage]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? packed_empty_barriers[which_packed_stage]->arrive(target_cta) : void(); + } + }; + auto packed_empty_barrier_arrive = [&]() { + packed_empty_barrier_arrive_at(packed_stage_idx); + }; + + // Skip useless computations + const bool is_cta_computation_valid = scheduler.is_computation_valid(m_block_idx, 0); + if (is_cta_computation_valid) { + // The compiler must know the dynamic variable `num_former_iters`'s real value + constexpr bool kShouldOptimize = BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; + constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; + + // Decode helper: wait for full[s] / packed[ps], decode packed FP4 B + // into the swizzled `smem_b[s]`. Caller must subsequently issue a + // NamedBarrier::sync(...,7) and packed_empty_barrier_arrive_at(ps) + // to make the decoded data visible to the wgmma async-proxy and to + // release the packed buffer. + auto wait_and_decode = [&](uint32_t s, uint32_t p, uint32_t ps, uint32_t pp) { + full_barriers[s]->wait(p); + packed_full_barriers[ps]->wait(pp); + auto smem_b_packed_bytes = smem_b_packed[ps]; + auto smem_b_bytes = reinterpret_cast(smem_b[s]); + constexpr uint32_t kVecPackedBytes = 16; // 16 packed bytes -> 32 e4m3 bytes + constexpr uint32_t kVecsPerRow = BLOCK_K_PACKED / kVecPackedBytes; + constexpr uint32_t kNumVecs = BLOCK_N * kVecsPerRow; + DG_STATIC_ASSERT(BLOCK_K_PACKED % kVecPackedBytes == 0, + "Packed K must be multiple of 16-byte vector width"); + DG_STATIC_ASSERT(BLOCK_K == 128, + "Swizzle assumes BLOCK_K == 128 so 32B store stays in-range"); + for (uint32_t idx = threadIdx.x; idx < kNumVecs; idx += kNumMathThreads) { + const uint32_t tile_n = idx / kVecsPerRow; + const uint32_t vec_k = idx % kVecsPerRow; + const uint32_t tile_k = vec_k * (kVecPackedBytes * 2); // 32-byte step in K + const uint4 packed16 = *reinterpret_cast( + smem_b_packed_bytes + tile_n * BLOCK_K_PACKED + vec_k * kVecPackedBytes); + uint64_t decoded[4] = {0, 0, 0, 0}; + auto decode_u32 = [&](uint32_t packed_word, uint64_t& out) { + #pragma unroll + for (uint32_t b = 0; b < 4; ++ b) { + const uint32_t pair = fp4_pair_to_e4m3_pair((packed_word >> (b * 8)) & 0xffu); + out |= static_cast(pair) << (b * 16); + } + }; + decode_u32(packed16.x, decoded[0]); + decode_u32(packed16.y, decoded[1]); + decode_u32(packed16.z, decoded[2]); + decode_u32(packed16.w, decoded[3]); + const uint32_t n_group = tile_n / 8; + const uint32_t n_in_group = tile_n % 8; + const uint32_t row_base = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K; + { + const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); + uint64_t* dst = reinterpret_cast(smem_b_bytes + row_base + swizzled_k); + dst[0] = decoded[0]; + dst[1] = decoded[1]; + } + { + const uint32_t swizzled_k = (tile_k + 16) ^ (n_in_group * 16); + uint64_t* dst = reinterpret_cast(smem_b_bytes + row_base + swizzled_k); + dst[0] = decoded[2]; + dst[1] = decoded[3]; + } + } + }; + + // Dispatch `num_former_iters` and launch MMAs with decode/wgmma + // pipeline overlap: decode runs one stage ahead of wgmma so that + // each iter's wait+decode for stage k+1 hides under iter k's + // wgmma async work. + dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) { + // Lead pointers: track the stage that decode is currently + // working on (one ahead of wgmma's `stage_idx`). + uint32_t lead_stage = stage_idx, lead_phase = phase; + uint32_t lead_packed = packed_stage_idx, lead_packed_phase = packed_phase; + auto advance_lead = [&]() { + lead_stage = lead_stage == kNumStages - 1 ? 0 : lead_stage + 1; + lead_phase ^= lead_stage == 0; + lead_packed = lead_packed == kNumPackedStages - 1 ? 0 : lead_packed + 1; + lead_packed_phase ^= lead_packed == 0; + }; + + // Prologue: decode the first stage so iter 0's wgmma can + // consume it without waiting. + wait_and_decode(stage_idx, phase, packed_stage_idx, packed_phase); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + packed_empty_barrier_arrive_at(packed_stage_idx); + advance_lead(); // lead now points one stage ahead of wgmma + + #pragma unroll 8 + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); + const auto b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16); + + // smem_b[stage_idx] is already decoded (prologue or the + // previous iter's hide-decode) and made visible by the + // matching NamedBarrier::sync below. + // smem_a[stage_idx] / smem_sfa[stage_idx] are guaranteed + // ready by the wait_and_decode that targeted stage_idx. + + float scale_b_0_regs[WGMMA::kNumAccum / 4]; + float scale_b_1_regs[WGMMA::kNumAccum / 4]; + if (do_wgmma_store) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t n_idx = n_block_idx * BLOCK_N + i * 8 + col_idx * 2; + scale_b_0_regs[i] = load_sfb(n_idx, k_block_idx); + scale_b_1_regs[i] = load_sfb(n_idx + 1, k_block_idx); + } + } + + const bool has_next = (k_block_idx + 1) < num_total_k_blocks; + + // TODO: remove some useless computation for unaligned Ms + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto m_offset = local_idx * WAVE_BLOCK_M; + + // Read A scales + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + auto scale_a_0 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_0 + m_offset) : 0; + auto scale_a_1 = do_wgmma_store ? ptx::ld_shared(smem_sfa[stage_idx] + r_1 + m_offset) : 0; + + // Commit WGMMA instructions + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16; + b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_desc, b_desc, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + + // Hide decode of next stage under this wgmma's async + // work. Run only on the last wave so all warps follow + // the same wait/decode/sync sequence per K iter. + const bool is_last_wave = (local_idx == BLOCK_M / WAVE_BLOCK_M - 1); + if (is_last_wave and has_next) { + wait_and_decode(lead_stage, lead_phase, lead_packed, lead_packed_phase); + } + + ptx::warpgroup_wait<0>(); + + // Notify barrier arrival at the last warpgroup wave + if (is_last_wave) + empty_barrier_arrive(); + + // Skip promotion for the unfilled parts + if (not do_wgmma_store) + continue; + + // Promote with scales + // NOTES: making it as predicates is very important for performance, comparing to two loops + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float scale_b_0 = scale_b_0_regs[i]; + const float scale_b_1 = scale_b_1_regs[i]; + shifted_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0]; + shifted_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1]; + shifted_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2]; + shifted_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3]; + } + } + + // Publish next iter's decoded smem_b to the wgmma + // async-proxy and release the packed-B slot. + if (has_next) { + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); + packed_empty_barrier_arrive_at(lead_packed); + advance_lead(); + } + } + }); + } else { + #pragma unroll + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + packed_full_barriers[packed_stage_idx]->wait(packed_phase); + packed_empty_barrier_arrive(); + full_barriers[stage_idx]->wait(phase); + empty_barrier_arrive(); + } + } + + // Psum layout can have a final partial M tile. TMA store writes a + // full BLOCK_M tile and may go out of the tensor-map bounds, so use + // a guarded scalar store for this layout. + if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + const uint32_t psum_store_end = + scheduler.current_group_idx + 1 < kNumGroups ? + math::align(scheduler.current_psum_m, BLOCK_M) : scheduler.current_psum_m; + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + const uint32_t row_0 = m_block_idx * BLOCK_M + m_offset + r_0; + const uint32_t row_1 = m_block_idx * BLOCK_M + m_offset + r_1; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t base_col = epilogue_type_t::template apply_index_n<8>( + n_block_idx * BLOCK_N + i * 8); + const uint32_t col = base_col + col_idx * 2; + const bool row_0_valid = (row_0 >= scheduler.last_psum_m and row_0 < psum_store_end); + const bool row_1_valid = (row_1 >= scheduler.last_psum_m and row_1 < psum_store_end); + if (row_0_valid and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = __float2bfloat16_rn(shifted_accum[i * 4 + 0]); + if (row_0_valid and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = __float2bfloat16_rn(shifted_accum[i * 4 + 1]); + if (row_1_valid and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = __float2bfloat16_rn(shifted_accum[i * 4 + 2]); + if (row_1_valid and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = __float2bfloat16_rn(shifted_accum[i * 4 + 3]); + } + } + __syncwarp(); + continue; + } + + // TMA checks + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes); + constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4; + DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom"); + DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32, + "Unaligned TMA store or too many TMA store instructions"); + DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N"); + + // Skip WGMMA store for the unfilled parts + if (not do_wgmma_store) + continue; + + // Wait last TMA store to be finished + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + // Write back to shared memory using STSM and issue TMA stores + DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); + #pragma unroll + for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx; + #pragma unroll + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + // Swizzle or padding into the correct address + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + // Calculate the swizzling atom offset and in-atom offset + constexpr uint32_t kNumBankGroupBytes = 16; + auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8); + + // Calculate the index of the bank group to be written in the atom + auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes); + + // Reshape the atom in another view and swizzle + // - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)` + // - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)` + constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8; + auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8); + auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8); + col ^= row % (kSwizzleDMode / 16); + + // Add back into the base pointer + // NOTES: think twice before modifying this, as changes may affect the number of instructions + smem_ptr = reinterpret_cast(smem_d) + // Base pointer + warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp offset + m_offset * kSwizzleDMode + // Wave offset + atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom offset (constants) + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset + } else { + // No swizzling, just padding + smem_ptr = reinterpret_cast(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8); + } + + // NOTES: only 16 lanes' addresses are used + ptx::SM90_U32x2_STSM_N::copy( + __float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}), + __float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}), + smem_ptr + ); + } + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + // Use TMA store to write back to global memory + // TODO: compatible with FP32 output + constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked; + DG_STATIC_ASSERT(kNumWGMMAStoreThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks"); + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) { + auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N; + auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M; + auto n_idx = epilogue_type_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); + auto m_idx = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); + if constexpr (kGemmType == GemmType::Batched) { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, + n_idx, m_idx, scheduler.current_group_idx); + } else { + cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr, n_idx, m_idx); + } + cute::tma_store_arrive(); + } + __syncwarp(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_90a"); +#endif +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop From 4122a7d7bd6687be8538871e74c9d562de6ce2d7 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:18:39 +0800 Subject: [PATCH 04/51] Update gemm.hpp --- csrc/apis/gemm.hpp | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42622df7d8..6b618c3ddf 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -5,6 +5,7 @@ #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" +#include "../jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" @@ -140,6 +141,27 @@ static void fp8_fp4_gemm_tt(const std::pair& a, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); } +static void fp8_fp4_gemm_nt_sm90_fused_wgmma(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional& c, + const int& gran_k, + const std::string& compiled_dims) { + const auto [m, k] = get_shape<2>(a.first); + const auto [n, half_k] = get_shape<2>(b.first); + DG_HOST_ASSERT(half_k * 2 == k); + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k); + std::optional> recipe_b = std::make_tuple(1, gran_k); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, false); + DG_HOST_ASSERT(gran_k_a == gran_k and gran_k_b == gran_k); + + sm90_fp8_fp4_gemm_1d1d_fused( + {a.first, sfa}, {b.first, sfb}, d, c, gran_k, compiled_dims); +} + static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -624,6 +646,20 @@ static void register_apis(pybind11::module_& m) { py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "mn", py::arg("disable_ue8m0_cast") = false); + m.def("fp8_fp4_gemm_nt_sm90_fused_wgmma", &fp8_fp4_gemm_nt_sm90_fused_wgmma, + py::arg("a"), py::arg("b"), py::arg("d"), + py::arg("c") = std::nullopt, + py::arg("gran_k") = 128, + py::arg("compiled_dims") = "nk"); + m.def("m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma", + &sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), + py::arg("gran_k") = 128, + py::arg("compiled_dims") = "nk", + py::arg("use_psum_layout") = false, + py::arg("expected_m_for_psum_layout") = std::nullopt, + py::arg("block_m_override") = std::nullopt, + py::arg("block_n_override") = std::nullopt); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, From b3d92e05b599d12d3e564d3220873779e93c4c5b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:19:21 +0800 Subject: [PATCH 05/51] Update __init__.py --- deep_gemm/__init__.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a9542e2f44..37c6def6ee 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -76,6 +76,20 @@ # TODO: remove these later fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked + try: + fp8_fp4_gemm_nt_sm90_cuda_core = _C.fp8_fp4_gemm_nt_sm90_cuda_core + except AttributeError: + pass + try: + fp8_fp4_gemm_nt_sm90_fused_wgmma = _C.fp8_fp4_gemm_nt_sm90_fused_wgmma + except AttributeError: + pass + try: + m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma + ) + except AttributeError: + pass except ImportError: # Expected behavior for CUDA runtime version before 12.1 pass From 825cbdbe2a954764d35eecd3bc0b358577949537 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:25:59 +0800 Subject: [PATCH 06/51] Update __init__.py --- deep_gemm/__init__.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 37c6def6ee..a466b37ef9 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -76,10 +76,6 @@ # TODO: remove these later fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked - try: - fp8_fp4_gemm_nt_sm90_cuda_core = _C.fp8_fp4_gemm_nt_sm90_cuda_core - except AttributeError: - pass try: fp8_fp4_gemm_nt_sm90_fused_wgmma = _C.fp8_fp4_gemm_nt_sm90_fused_wgmma except AttributeError: From 9625f4f8e2dc5470bdc6fe150ff8086418d2af98 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:26:33 +0800 Subject: [PATCH 07/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 62e57d9669..221150bed9 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -67,7 +67,6 @@ def run(): ) d_fp8 = torch.empty_like(d) - d_fused = torch.empty_like(d) def run_fp8(): deep_gemm.fp8_gemm_nt( @@ -79,17 +78,6 @@ def run_fp8(): recipe_b=(1, gran_k), ) - d_fused = torch.empty_like(d) - - def run_fused(): - deep_gemm.fp8_fp4_gemm_nt_sm90_cuda_core( - a, - b, - d_fused, - c=c, - gran_k=gran_k, - ) - d_fused_wgmma = torch.empty_like(d) def run_fused_wgmma(): @@ -103,14 +91,11 @@ def run_fused_wgmma(): run() run_fp8() - if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core"): - run_fused() if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma"): run_fused_wgmma() fallback_diff = calc_diff(d, ref) fp8_diff = calc_diff(d_fp8, ref) fallback_vs_fp8_diff = calc_diff(d, d_fp8) - fused_diff = calc_diff(d_fused, ref) if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core") else None fused_wgmma_diff = ( calc_diff(d_fused_wgmma, ref) if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") @@ -118,11 +103,6 @@ def run_fused_wgmma(): ) fallback_elapsed = _time_cuda(run) fp8_elapsed = _time_cuda(run_fp8) - fused_elapsed = ( - _time_cuda(run_fused) - if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_cuda_core") - else None - ) fused_wgmma_elapsed = ( _time_cuda(run_fused_wgmma) if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") @@ -130,9 +110,6 @@ def run_fused_wgmma(): ) fallback_tflops = 2 * m * n * k / fallback_elapsed / 1e12 fp8_tflops = 2 * m * n * k / fp8_elapsed / 1e12 - fused_tflops = ( - 2 * m * n * k / fused_elapsed / 1e12 if fused_elapsed is not None else None - ) fused_wgmma_tflops = ( 2 * m * n * k / fused_wgmma_elapsed / 1e12 if fused_wgmma_elapsed is not None @@ -148,12 +125,6 @@ def run_fused_wgmma(): f" fallback vs fp8: diff={fallback_vs_fp8_diff:.6f}\n" f" slowdown: {slowdown:.2f}x" ) - if fused_elapsed is not None: - message += ( - f"\n cuda-core fused: diff={fused_diff:.6f}, " - f"time={fused_elapsed * 1e6:.1f} us, {fused_tflops:.2f} TFLOPS, " - f"slowdown_vs_fp8={fused_elapsed / fp8_elapsed:.2f}x" - ) if fused_wgmma_elapsed is not None: message += ( f"\n fused wgmma: diff={fused_wgmma_diff:.6f}, " @@ -164,8 +135,6 @@ def run_fused_wgmma(): assert fallback_diff < 0.015 assert fp8_diff < 0.015 assert fallback_vs_fp8_diff < 1e-6 - if fused_diff is not None: - assert fused_diff < 0.015 if fused_wgmma_diff is not None: assert fused_wgmma_diff < 0.015 From 043f007110d13b2154bfbbca053d7388e8733e4b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 11:32:00 +0800 Subject: [PATCH 08/51] Create sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 501 ++++++++++++++++++ 1 file changed, 501 insertions(+) create mode 100644 csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp new file mode 100644 index 0000000000..2e0c6e2c3d --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -0,0 +1,501 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/layout.hpp" +#include "epilogue.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +class SM90FP8FP4Gemm1D1DRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + void *gmem_b_ptr; + void *gmem_d_ptr; + void *grouped_layout; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_sfa; + CUtensorMap tensor_map_sfb; + CUtensorMap tensor_map_cd; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_fp4_gemm_1d1d_impl< + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, + {}, {} + >); +}}; +)", + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), + args.gemm_desc.num_groups, + args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k, + args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, + args.gemm_config.storage_config.swizzle_cd_mode, + args.gemm_config.pipeline_config.num_stages, + args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, + args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, + args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), + to_string(args.gemm_desc.cd_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + nullptr, args.gmem_b_ptr, + args.gmem_d_ptr, args.grouped_layout, + nullptr, + args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.tensor_map_a, + args.tensor_map_sfa, args.tensor_map_sfb, + args.tensor_map_cd)); + } +}; + +class SM90FP8FP4Gemm1D2DRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + cute::UMMA::Major major_sfb; + int num_packed_stages; + void *gmem_b_ptr; + void *gmem_d_ptr; + void *sfb; + void *grouped_layout; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + CUtensorMap tensor_map_sfa; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_fp4_gemm_1d2d_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, {}, + {}, {}, + {}, {}, + {}, {}, + {} + >); +}}; +)", + to_string(args.major_sfb), + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), + args.gemm_desc.num_groups, + args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k, + args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, + args.gemm_config.storage_config.swizzle_cd_mode, + args.gemm_config.pipeline_config.num_stages, args.num_packed_stages, + args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, + args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, + args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), + get_default_epilogue_type(std::nullopt)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.gmem_b_ptr, args.sfb, args.grouped_layout, args.gmem_d_ptr, + args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.tensor_map_sfa)); + } +}; + +static void sm90_fp8_fp4_gemm_1d1d_fused(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional& c, + const int& gran_k, + const std::string& compiled_dims) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(gran_k == 128); + DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(a.first.is_contiguous()); + DG_HOST_ASSERT(b.first.is_contiguous()); + DG_HOST_ASSERT(d.is_contiguous() and c->is_contiguous()); + + const auto [m, k] = get_shape<2>(a.first); + const auto [n, half_k] = get_shape<2>(b.first); + const auto [m_d, n_d] = get_shape<2>(d); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(a.second.size(0) == m and a.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(b.second.size(0) == n and b.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(c->sizes() == d.sizes()); + + if (m == 0 or n == 0) { + return; + } + if (c->data_ptr() != d.data_ptr()) { + d.copy_(*c); + } + + auto desc = GemmDesc { + .gemm_type = GemmType::Normal, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = k, .num_groups = 1, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = c.has_value(), + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims + }; + auto config = get_best_config(desc); + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, a.second, m, k, + config.layout.block_m, config.layout.block_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, b.second, n, k, + config.layout.block_n, config.layout.block_k, 1, 0); + const auto tensor_map_cd = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), 1, + config.storage_config.swizzle_cd_mode); + + const SM90FP8FP4Gemm1D1DRuntime::Args& args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .grouped_layout = nullptr, + .tensor_map_a = tensor_map_a, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd, + }; + const auto code = SM90FP8FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm90_fp8_fp4_gemm_1d1d_fused", code); + SM90FP8FP4Gemm1D1DRuntime::launch(runtime, args); +} + +static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( + const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& grouped_layout, + const int& gran_k, + const std::string& compiled_dims, + const bool& use_psum_layout, + const std::optional& expected_m_for_psum_layout, + const std::optional& block_m_override, + const std::optional& block_n_override) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(gran_k == 128); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt and grouped_layout.is_contiguous()); + DG_HOST_ASSERT(a.first.is_contiguous() and b.first.is_contiguous() and d.is_contiguous()); + + const auto [m, k] = get_shape<2>(a.first); + const auto [num_groups, n, half_k] = get_shape<3>(b.first); + const auto [m_d, n_d] = get_shape<2>(d); + const auto [layout_size] = get_shape<1>(grouped_layout); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(use_psum_layout ? (layout_size == num_groups) : (layout_size == m)); + if (expected_m_for_psum_layout) { + DG_HOST_ASSERT(use_psum_layout); + } + DG_HOST_ASSERT(a.second.size(0) == m and a.second.size(1) == ceil_div(k, gran_k)); + DG_HOST_ASSERT(b.second.size(0) == num_groups and b.second.size(1) == n and b.second.size(2) == ceil_div(k, gran_k)); + + if (m == 0 or n == 0) { + return; + } + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k); + std::optional> recipe_b = std::make_tuple(1, gran_k); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, + std::nullopt, num_groups, false); + DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128); + + const auto gemm_type = use_psum_layout ? + GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous; + // NOTE: psum layout previously always took the 1d1d fallback path. With the + // 1d2d psum scheduler / SFB indexing aligned, we let psum also flow into the + // common 1d2d code below. This relies on: + // - All groups in `grouped_layout` having the same K (current_shape_k == shape_k); + // true for current call sites where K is shared across groups. + // - SFB physical layout `[num_groups, n, shape_k_scales]` already matches + // 1d2d cooperative ld.global indexing. + // - 0-size groups handled by the scheduler's while-loop fallthrough. + if (false /* use_psum_layout disabled: psum now goes through the 1d2d common path */) { + auto desc = GemmDesc { + .gemm_type = gemm_type, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = expected_m_for_psum_layout.value_or(m), + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups + }; + auto config = get_best_config(desc); + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, config.layout.block_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, + config.layout.block_n, config.layout.block_k, num_groups, 0); + const auto tensor_map_cd = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), 1, + config.storage_config.swizzle_cd_mode); + + const SM90FP8FP4Gemm1D1DRuntime::Args& args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .grouped_layout = grouped_layout.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd, + }; + const auto code = SM90FP8FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused", code); + SM90FP8FP4Gemm1D1DRuntime::launch(runtime, args); + return; + } + + auto desc = GemmDesc { + .gemm_type = gemm_type, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = expected_m_for_psum_layout.value_or(m), + .expected_n = n, + .expected_k = k, + .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 + }; + auto rebuild_config = [&](Layout layout) { + const auto storage_config = SM90ArchSpec::get_storage_config(desc, layout); + const auto pipeline_config = SM90ArchSpec::get_pipeline_config(desc, layout, storage_config); + DG_HOST_ASSERT(pipeline_config.num_stages >= 3); + const auto launch_config = SM90ArchSpec::get_launch_config(desc, layout); + return GemmConfig{layout, storage_config, pipeline_config, launch_config}; + }; + + const auto layout_candidates = SM90ArchSpec::get_layout_candidates(desc); + DG_HOST_ASSERT(not layout_candidates.empty()); + auto layout = layout_candidates[0]; + auto layout_info = SM90ArchSpec::get_layout_info(desc, layout); + bool found_layout = false; + // FP4 B is decoded by each CTA, so only A multicast (cluster_n) is legal here. + for (const auto& candidate: layout_candidates) { + if (candidate.cluster_m != 1) { + continue; + } + const auto candidate_info = SM90ArchSpec::get_layout_info(desc, candidate); + if (not found_layout or SM90ArchSpec::compare(candidate_info, layout_info)) { + layout = candidate; + layout_info = candidate_info; + found_layout = true; + } + } + DG_HOST_ASSERT(found_layout); + // Psum grouped_layout is encoded with 128-row alignment. Keep BLOCK_M fixed + // at 128 so scheduler psum boundaries match the physical layout. + if (use_psum_layout) { + layout.block_m = 128; + layout.cluster_m = 1; + } + if (not use_psum_layout and m >= num_groups * 256 and n >= 1024) { + layout.block_m = 256; + layout.block_n = 64; + layout.cluster_m = 1; + } + auto config = rebuild_config(layout); + + if (block_m_override or block_n_override) { + auto layout = config.layout; + if (block_m_override) { + DG_HOST_ASSERT(not use_psum_layout or *block_m_override == 128); + layout.block_m = *block_m_override; + } + if (block_n_override) { + layout.block_n = *block_n_override; + } + DG_HOST_ASSERT((layout.block_m == 64 or layout.block_m == 128 or layout.block_m == 256) and layout.block_n % 16 == 0); + DG_HOST_ASSERT(layout.block_n <= 256); + layout.cluster_m = 1; + config = rebuild_config(layout); + } + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); + + // Re-derive `num_stages` and `smem_size` to account for: + // (1) the extra packed-FP4 B staging buffer (BLOCK_N * BLOCK_K / 2 bytes per packed stage) + // (2) the enlarged SFB cache: now stores (shape_k_scales, BLOCK_N) floats per block + // so the K loop reads SFB from smem instead of gmem. + // The packed-B pipeline uses an independent (smaller) number of stages so we can keep + // `num_stages` large for the WGMMA consumer pipeline. + int num_packed_stages = 2; + { + const int block_k = config.layout.block_k; + const int block_n = config.layout.block_n; + const int shape_k_scales = ceil_div(static_cast(k), block_k); + const bool uniform_scale_b = (block_k % block_n == 0); + const int sfb_old_bytes = align(shape_k_scales * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + // SFB cache now aliases smem_d (no separate allocation). We only need to + // remove the original SFB bytes from the baseline. + const int sfb_extra = -sfb_old_bytes; + + const int packed_per_stage = block_n * (block_k / 2); + const int smem_a_per_stage = config.storage_config.load_block_m * block_k * + static_cast(c10::elementSize(desc.a_dtype)); + const int smem_b_per_stage = config.storage_config.load_block_n * block_k * + static_cast(c10::elementSize(desc.b_dtype)); + const int smem_sfa_per_stage = align(config.layout.block_m * static_cast(sizeof(float)), 128); + const int original_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage; + const int orig_num_stages = config.pipeline_config.num_stages; + const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage + sfb_extra; + + auto fits = [&](int stages, int p_stages) { + return smem_extra + stages * original_per_stage + p_stages * packed_per_stage + <= SM90ArchSpec::smem_capacity; + }; + + // Try to keep `num_stages = orig_num_stages` with packed_stages=2; otherwise drop num_stages, + // and finally fall back to packed_stages=1 if needed. + int chosen_stages = orig_num_stages; + int chosen_packed = 2; + while (chosen_stages >= 3 and not fits(chosen_stages, chosen_packed)) + -- chosen_stages; + if (chosen_stages < 3) { + chosen_packed = 1; + chosen_stages = orig_num_stages; + while (chosen_stages >= 3 and not fits(chosen_stages, chosen_packed)) + -- chosen_stages; + } + DG_HOST_ASSERT(chosen_stages >= 3); + config.pipeline_config.num_stages = chosen_stages; + num_packed_stages = chosen_packed; + config.pipeline_config.smem_size = smem_extra + chosen_stages * original_per_stage + + chosen_packed * packed_per_stage; + } + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, 1, + config.storage_config.swizzle_a_mode); + // View packed FP4 B as 1-byte FP8 so TMA loads raw packed bytes (no FP4 unpacking). + const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b_bytes, n, half_k, + config.layout.block_n, config.layout.block_k / 2, + half_k, num_groups, 0); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, config.layout.block_k, 1, 0); + const auto tensor_map_d = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), 1, + config.storage_config.swizzle_cd_mode); + + const SM90FP8FP4Gemm1D2DRuntime::Args& args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .major_sfb = get_major_type_ab(sfb), + .num_packed_stages = num_packed_stages, + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .sfb = sfb.data_ptr(), + .grouped_layout = grouped_layout.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .tensor_map_sfa = tensor_map_sfa, + }; + const auto code = SM90FP8FP4Gemm1D2DRuntime::generate(args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_contiguous_1d2d_fused", code); + SM90FP8FP4Gemm1D2DRuntime::launch(runtime, args); +} + +} // namespace deep_gemm From eebc4f473a8b87728fb8f452477b5647efeb261f Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 17:53:36 +0800 Subject: [PATCH 09/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 53 ++++++++----------- 1 file changed, 22 insertions(+), 31 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index 2e0c6e2c3d..bec03e9796 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -83,7 +83,7 @@ class SM90FP8FP4Gemm1D2DRuntime final: public LaunchRuntime); }}; @@ -123,11 +124,12 @@ static void __instantiate_kernel() {{ args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k, args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode, - args.gemm_config.pipeline_config.num_stages, args.num_packed_stages, + args.gemm_config.pipeline_config.num_stages, args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), - get_default_epilogue_type(std::nullopt)); + get_default_epilogue_type(std::nullopt), + args.decode_stub ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -231,7 +233,8 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( const bool& use_psum_layout, const std::optional& expected_m_for_psum_layout, const std::optional& block_m_override, - const std::optional& block_n_override) { + const std::optional& block_n_override, + const bool& decode_stub) { DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); DG_HOST_ASSERT(gran_k == 128); DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); @@ -365,7 +368,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( auto layout = layout_candidates[0]; auto layout_info = SM90ArchSpec::get_layout_info(desc, layout); bool found_layout = false; - // FP4 B is decoded by each CTA, so only A multicast (cluster_n) is legal here. + // FP4 B is decoded by each CTA, so only A multicast (cluster_n) is enabled. for (const auto& candidate: layout_candidates) { if (candidate.cluster_m != 1) { continue; @@ -387,8 +390,8 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( if (not use_psum_layout and m >= num_groups * 256 and n >= 1024) { layout.block_m = 256; layout.block_n = 64; - layout.cluster_m = 1; } + layout.cluster_m = 1; auto config = rebuild_config(layout); if (block_m_override or block_n_override) { @@ -409,20 +412,18 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k); // Re-derive `num_stages` and `smem_size` to account for: - // (1) the extra packed-FP4 B staging buffer (BLOCK_N * BLOCK_K / 2 bytes per packed stage) + // (1) the extra packed-FP4 B staging buffer (BLOCK_N * BLOCK_K / 2 bytes per stage) + // — after the A/packed-B barrier merge it is sized at `num_stages` slots + // (same as A/SFA), so it folds into the per-stage cost. // (2) the enlarged SFB cache: now stores (shape_k_scales, BLOCK_N) floats per block - // so the K loop reads SFB from smem instead of gmem. - // The packed-B pipeline uses an independent (smaller) number of stages so we can keep - // `num_stages` large for the WGMMA consumer pipeline. - int num_packed_stages = 2; + // so the K loop reads SFB from smem instead of gmem. SFB cache aliases + // smem_d (no separate allocation), so we only subtract the original SFB bytes. { const int block_k = config.layout.block_k; const int block_n = config.layout.block_n; const int shape_k_scales = ceil_div(static_cast(k), block_k); const bool uniform_scale_b = (block_k % block_n == 0); const int sfb_old_bytes = align(shape_k_scales * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); - // SFB cache now aliases smem_d (no separate allocation). We only need to - // remove the original SFB bytes from the baseline. const int sfb_extra = -sfb_old_bytes; const int packed_per_stage = block_n * (block_k / 2); @@ -432,31 +433,21 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( static_cast(c10::elementSize(desc.b_dtype)); const int smem_sfa_per_stage = align(config.layout.block_m * static_cast(sizeof(float)), 128); const int original_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage; + const int merged_per_stage = original_per_stage + packed_per_stage; const int orig_num_stages = config.pipeline_config.num_stages; const int smem_extra = config.pipeline_config.smem_size - orig_num_stages * original_per_stage + sfb_extra; - auto fits = [&](int stages, int p_stages) { - return smem_extra + stages * original_per_stage + p_stages * packed_per_stage - <= SM90ArchSpec::smem_capacity; + auto fits = [&](int stages) { + return smem_extra + stages * merged_per_stage <= SM90ArchSpec::smem_capacity; }; - // Try to keep `num_stages = orig_num_stages` with packed_stages=2; otherwise drop num_stages, - // and finally fall back to packed_stages=1 if needed. + // Try to keep `num_stages = orig_num_stages`; otherwise drop until it fits. int chosen_stages = orig_num_stages; - int chosen_packed = 2; - while (chosen_stages >= 3 and not fits(chosen_stages, chosen_packed)) + while (chosen_stages >= 3 and not fits(chosen_stages)) -- chosen_stages; - if (chosen_stages < 3) { - chosen_packed = 1; - chosen_stages = orig_num_stages; - while (chosen_stages >= 3 and not fits(chosen_stages, chosen_packed)) - -- chosen_stages; - } DG_HOST_ASSERT(chosen_stages >= 3); config.pipeline_config.num_stages = chosen_stages; - num_packed_stages = chosen_packed; - config.pipeline_config.smem_size = smem_extra + chosen_stages * original_per_stage + - chosen_packed * packed_per_stage; + config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; } const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, @@ -483,7 +474,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( config.pipeline_config.smem_size, config.layout.get_cluster_size()), .major_sfb = get_major_type_ab(sfb), - .num_packed_stages = num_packed_stages, + .decode_stub = decode_stub, .gmem_b_ptr = b.first.data_ptr(), .gmem_d_ptr = d.data_ptr(), .sfb = sfb.data_ptr(), From 48dec949345927106fa0234027a8d390b79c6976 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 17:54:06 +0800 Subject: [PATCH 10/51] Update ld_st.cuh --- deep_gemm/include/deep_gemm/ptx/ld_st.cuh | 36 +++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index c3e03bec73..398d131a16 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -138,6 +138,42 @@ CUTLASS_DEVICE void st_shared(const __int128_t* ptr, __int128_t val) { asm volatile("st.shared.b128 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "q"(val)); } +CUTLASS_DEVICE void st_shared(const float4* ptr, float4 val) { + asm volatile("st.shared.v4.f32 [%0], {%1, %2, %3, %4};" :: + "l"(__cvta_generic_to_shared(ptr)), + "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)); +} + +// 16-byte vector store as two 64-bit values (st.shared.v2.u64). Useful when +// the data is naturally produced as two 64-bit lanes; saves one instruction +// versus issuing 2 separate `st.shared.u64`. +CUTLASS_DEVICE void st_shared_v2_u64(void* ptr, uint64_t a, uint64_t b) { + asm volatile("st.shared.v2.u64 [%0], {%1, %2};" :: + "l"(__cvta_generic_to_shared(ptr)), "l"(a), "l"(b)); +} + +/// Async copy from global to shared (cp.async) +// 16-byte cache-global async copy +CUTLASS_DEVICE void cp_async_16(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: + "l"(__cvta_generic_to_shared(smem_dst)), "l"(gmem_src)); +} + +// 4-byte cache-all async copy (for scalar / strided fallback) +CUTLASS_DEVICE void cp_async_4(void* smem_dst, const void* gmem_src) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" :: + "l"(__cvta_generic_to_shared(smem_dst)), "l"(gmem_src)); +} + +CUTLASS_DEVICE void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template +CUTLASS_DEVICE void cp_async_wait_group() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); +} + CUTLASS_DEVICE void st_shared_bulk(void* smem_ptr, const uint32_t& num_bytes) { // `size` must be 64-bit before PTX ISA 9.0 asm volatile("st.bulk.weak.shared::cta [%0], %1, 0;" :: From 2520585dce7ab622f5175a7fcf05182df5309854 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 17:54:40 +0800 Subject: [PATCH 11/51] Update sm90_fp8_fp4_gemm_1d2d.cuh --- .../impls/sm90_fp8_fp4_gemm_1d2d.cuh | 243 ++++++++++-------- 1 file changed, 131 insertions(+), 112 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh index caf0c1ceab..0f3777ce50 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh @@ -39,11 +39,12 @@ template + typename epilogue_type_t, + bool kDecodeStub = false> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, nv_bfloat16* gmem_d_ptr, @@ -118,14 +119,15 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, // LUT_LO[mag] = {0x00, 0x30, 0x38, 0x3c}, LUT_HI[mag-4] = {0x40, 0x44, 0x48, 0x4c} // __byte_perm(LUT_LO, LUT_HI, mag) returns 4 copies of the byte at index `mag` // (mag is in 0..7, fits in the lower nibble, used as the byte selector). - // Sign: bit 3 of code shifted to MSB, masked to 0 if mag==0 to avoid -0. + // Sign: bit 3 of code shifted to MSB. We deliberately allow the FP4 "negative zero" + // (mag=0, sign=1) to map to E4M3 0x80 (which is -0, not NaN). WGMMA treats -0 as 0, + // so dropping the explicit mag!=0 mask saves ~4 ops per nibble (~128 per uint4 of + // packed FP4) without affecting numerical results. constexpr uint32_t LUT_LO = 0x3c383000u; constexpr uint32_t LUT_HI = 0x4c484440u; - const uint32_t mag = code & 0x07u; - const uint32_t mag_nz_m1 = (mag + 0x07u) >> 3; // 1 if mag!=0 else 0 - const uint32_t sign_mask = -mag_nz_m1; // 0xffffffff if mag!=0 else 0 - const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; - const uint32_t sign = ((code & 0x08u) << 4) & sign_mask; + const uint32_t mag = code & 0x07u; + const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; + const uint32_t sign = (code & 0x08u) << 4; // 0 or 0x80 return mag_byte | sign; }; auto fp4_pair_to_e4m3_pair = [&](uint32_t packed) { @@ -146,7 +148,7 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, auto smem_b_packed = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_buffer + SMEM_B_PACKED_OFFSET + i * SMEM_B_PACKED_SIZE_PER_STAGE); }); - constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumPackedStages * SMEM_B_PACKED_SIZE_PER_STAGE; + constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumStages * SMEM_B_PACKED_SIZE_PER_STAGE; auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); }); @@ -156,13 +158,12 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, auto smem_sfb = reinterpret_cast(smem_buffer); DG_TRAP_ONLY_DEVICE_ASSERT(smem_sfb_bytes <= SMEM_D_SIZE); - // Fill barriers (SFB does not allocate any extra space because it aliases smem_d) + // Fill barriers (SFB does not allocate any extra space because it aliases smem_d). + // After the A/packed-B barrier merge there is only one set of full/empty barriers. auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); - auto packed_full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + 2 * kNumStages + i; }); - auto packed_empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + 2 * kNumStages + kNumPackedStages + i; }); // Initialize barriers DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast"); @@ -174,11 +175,6 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, full_barriers[i]->init(1); empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); } - #pragma unroll - for (uint32_t i = 0; i < kNumPackedStages; ++ i) { - packed_full_barriers[i]->init(1); - packed_empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32); - } // Make initialized barrier visible in async proxy cutlass::arch::fence_barrier_init(); @@ -210,19 +206,14 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, } }; - // Pipeline and TMA phases + // Pipeline and TMA phases (single shared pipeline for A/SFA/packed-B after the barrier merge) uint32_t stage_idx = 0, phase = 0; - uint32_t packed_stage_idx = 0, packed_phase = 0; auto advance_pipeline = [&](uint32_t& k_block_idx) { ++ k_block_idx; // Flip phases only if reach the next first stage stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; phase ^= stage_idx == 0; - - // Packed-B pipeline runs in lock-step with the K loop but with its own depth - packed_stage_idx = packed_stage_idx == kNumPackedStages - 1 ? 0 : packed_stage_idx + 1; - packed_phase ^= packed_stage_idx == 0; }; if (warp_idx >= kNumMathThreads / 32) { @@ -242,9 +233,8 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { - // Wait consumer release for A/SFA slot and packed-B slot independently + // Wait consumer release for the (now merged) A/SFA/packed-B slot empty_barriers[stage_idx]->wait(phase ^ 1); - packed_empty_barriers[packed_stage_idx]->wait(packed_phase ^ 1); // Issue TMA A constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); @@ -259,17 +249,17 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, tma::copy(&tensor_map_sfa, &full_barrier, smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales, 1, k_block_idx), num_tma_multicast_a); - full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE); - // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) - auto& packed_full_barrier = *packed_full_barriers[packed_stage_idx]; + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; - tma::copy(&tensor_map_b, &packed_full_barrier, - reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[packed_stage_idx]), + tma::copy(&tensor_map_b, &full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), k_idx_packed, scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), num_tma_multicast_b, batch_idx); - packed_full_barrier.arrive_and_expect_tx(SMEM_B_PACKED_SIZE_PER_STAGE); + + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + + SMEM_B_PACKED_SIZE_PER_STAGE); } } @@ -280,11 +270,6 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; phase ^= stage_idx == 0; } - for (uint32_t i = 0; i < kNumPackedStages; ++ i) { - packed_empty_barriers[packed_stage_idx]->wait(packed_phase ^ 1); - packed_stage_idx = packed_stage_idx == kNumPackedStages - 1 ? 0 : packed_stage_idx + 1; - packed_phase ^= packed_stage_idx == 0; - } } } } else { @@ -305,28 +290,74 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, while (scheduler.get_next_block(m_block_idx, n_block_idx)) { const uint32_t current_group_idx = get_current_group_idx(); - // Cooperatively prefetch the SFB tile for this block from gmem to smem + // Cooperatively prefetch the SFB tile for this block from gmem to smem. // Layout in smem: [shape_k_scales, BLOCK_N] (k outer, n inner). // Out-of-bound n is filled with 1.0f to keep `n_idx >= shape_n` neutral. + // + // Optimization: use cp.async to copy gmem->smem directly (no register + // round-trip). For MN-major (n innermost in gmem) we can issue 16-byte + // (float4) cp.async per thread, cutting #instructions by 4x. K-major + // and the OOB tail use scalar 4-byte cp.async / st.shared. { const uint32_t n_block_base = n_block_idx * BLOCK_N; - const uint32_t total = shape_k_scales * BLOCK_N; - for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { - const uint32_t k_idx = i / BLOCK_N; - const uint32_t n_off = i % BLOCK_N; - const uint32_t n_idx = n_block_base + n_off; - float val; - if (n_idx >= shape_n) { - val = 1.0f; - } else if constexpr (kMajorSFB == cute::UMMA::Major::MN) { - val = sfb[current_group_idx * aligned_shape_n_sfb * shape_k_scales + - k_idx * aligned_shape_n_sfb + n_idx]; - } else { - val = sfb[current_group_idx * shape_n * shape_k_scales + - n_idx * shape_k_scales + k_idx]; + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = shape_k_scales * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales; + // Issue cp.async (16B per thread) for fully in-bounds vectors; + // fall back to scalar st.shared with 1.0f-fill for the tail. + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_idx = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + k_idx * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + // Tail: at least one element is OOB; use scalar st with 1.0f fill. + float4 vals; + vals.x = (n_idx + 0 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 0] : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 1] : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 2] : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? + sfb_base[k_idx * aligned_shape_n_sfb + n_idx + 3] : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + // K-major: sfb is strided along n; cannot easily vectorize across n. + // Use scalar 4B cp.async for in-bounds, st.shared with 1.0f for OOB. + const uint32_t total = shape_k_scales * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_idx = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales + + n_idx * shape_k_scales + k_idx; + ptx::cp_async_4(smem_dst, gmem_src); + } } - ptx::st_shared(smem_sfb + k_idx * BLOCK_N + n_off, val); } + // Commit and wait for all cp.async issued above; pair with the + // existing NamedBarrier so the smem cache is visible to all + // math warps before they enter the K loop. + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); } @@ -363,17 +394,6 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void(); } }; - auto packed_empty_barrier_arrive_at = [&](uint32_t which_packed_stage) { - if constexpr (kNumTMAMulticast == 1) { - lane_idx == 0 ? packed_empty_barriers[which_packed_stage]->arrive() : void(); - } else { - auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); - lane_idx < kNumTMAMulticast ? packed_empty_barriers[which_packed_stage]->arrive(target_cta) : void(); - } - }; - auto packed_empty_barrier_arrive = [&]() { - packed_empty_barrier_arrive_at(packed_stage_idx); - }; // Skip useless computations const bool is_cta_computation_valid = scheduler.is_computation_valid(m_block_idx, 0); @@ -383,15 +403,15 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; - // Decode helper: wait for full[s] / packed[ps], decode packed FP4 B - // into the swizzled `smem_b[s]`. Caller must subsequently issue a - // NamedBarrier::sync(...,7) and packed_empty_barrier_arrive_at(ps) - // to make the decoded data visible to the wgmma async-proxy and to - // release the packed buffer. - auto wait_and_decode = [&](uint32_t s, uint32_t p, uint32_t ps, uint32_t pp) { + // Decode helper: wait for the (merged) full[s] barrier, then + // decode packed FP4 B from `smem_b_packed[s]` into the swizzled + // `smem_b[s]`. Caller must subsequently issue + // NamedBarrier::sync(...,7) to publish the decoded data to the + // wgmma async-proxy. The packed slot is released together with + // A/SFA via the single empty_barrier_arrive() after wgmma_wait. + auto wait_and_decode = [&](uint32_t s, uint32_t p) { full_barriers[s]->wait(p); - packed_full_barriers[ps]->wait(pp); - auto smem_b_packed_bytes = smem_b_packed[ps]; + auto smem_b_packed_bytes = smem_b_packed[s]; auto smem_b_bytes = reinterpret_cast(smem_b[s]); constexpr uint32_t kVecPackedBytes = 16; // 16 packed bytes -> 32 e4m3 bytes constexpr uint32_t kVecsPerRow = BLOCK_K_PACKED / kVecPackedBytes; @@ -400,42 +420,47 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, "Packed K must be multiple of 16-byte vector width"); DG_STATIC_ASSERT(BLOCK_K == 128, "Swizzle assumes BLOCK_K == 128 so 32B store stays in-range"); - for (uint32_t idx = threadIdx.x; idx < kNumVecs; idx += kNumMathThreads) { - const uint32_t tile_n = idx / kVecsPerRow; - const uint32_t vec_k = idx % kVecsPerRow; - const uint32_t tile_k = vec_k * (kVecPackedBytes * 2); // 32-byte step in K - const uint4 packed16 = *reinterpret_cast( - smem_b_packed_bytes + tile_n * BLOCK_K_PACKED + vec_k * kVecPackedBytes); - uint64_t decoded[4] = {0, 0, 0, 0}; - auto decode_u32 = [&](uint32_t packed_word, uint64_t& out) { - #pragma unroll - for (uint32_t b = 0; b < 4; ++ b) { - const uint32_t pair = fp4_pair_to_e4m3_pair((packed_word >> (b * 8)) & 0xffu); - out |= static_cast(pair) << (b * 16); + if constexpr (not kDecodeStub) { + for (uint32_t idx = threadIdx.x; idx < kNumVecs; idx += kNumMathThreads) { + const uint32_t tile_n = idx / kVecsPerRow; + const uint32_t vec_k = idx % kVecsPerRow; + const uint32_t tile_k = vec_k * (kVecPackedBytes * 2); // 32-byte step in K + const uint4 packed16 = *reinterpret_cast( + smem_b_packed_bytes + tile_n * BLOCK_K_PACKED + vec_k * kVecPackedBytes); + uint64_t decoded[4] = {0, 0, 0, 0}; + auto decode_u32 = [&](uint32_t packed_word, uint64_t& out) { + #pragma unroll + for (uint32_t b = 0; b < 4; ++ b) { + const uint32_t pair = fp4_pair_to_e4m3_pair((packed_word >> (b * 8)) & 0xffu); + out |= static_cast(pair) << (b * 16); + } + }; + decode_u32(packed16.x, decoded[0]); + decode_u32(packed16.y, decoded[1]); + decode_u32(packed16.z, decoded[2]); + decode_u32(packed16.w, decoded[3]); + const uint32_t n_group = tile_n / 8; + const uint32_t n_in_group = tile_n % 8; + const uint32_t row_base = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K; + // Each segment writes 16 bytes (= 2 u64). Use a + // single st.shared.v2.u64 instead of two + // st.shared.u64, halving the store instruction + // count for the decoded B tile. + { + const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); + ptx::st_shared_v2_u64(smem_b_bytes + row_base + swizzled_k, + decoded[0], decoded[1]); + } + { + const uint32_t swizzled_k = (tile_k + 16) ^ (n_in_group * 16); + ptx::st_shared_v2_u64(smem_b_bytes + row_base + swizzled_k, + decoded[2], decoded[3]); } - }; - decode_u32(packed16.x, decoded[0]); - decode_u32(packed16.y, decoded[1]); - decode_u32(packed16.z, decoded[2]); - decode_u32(packed16.w, decoded[3]); - const uint32_t n_group = tile_n / 8; - const uint32_t n_in_group = tile_n % 8; - const uint32_t row_base = n_group * 8 * BLOCK_K + n_in_group * BLOCK_K; - { - const uint32_t swizzled_k = tile_k ^ (n_in_group * 16); - uint64_t* dst = reinterpret_cast(smem_b_bytes + row_base + swizzled_k); - dst[0] = decoded[0]; - dst[1] = decoded[1]; - } - { - const uint32_t swizzled_k = (tile_k + 16) ^ (n_in_group * 16); - uint64_t* dst = reinterpret_cast(smem_b_bytes + row_base + swizzled_k); - dst[0] = decoded[2]; - dst[1] = decoded[3]; } } }; + // Dispatch `num_former_iters` and launch MMAs with decode/wgmma // pipeline overlap: decode runs one stage ahead of wgmma so that // each iter's wait+decode for stage k+1 hides under iter k's @@ -444,19 +469,15 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, // Lead pointers: track the stage that decode is currently // working on (one ahead of wgmma's `stage_idx`). uint32_t lead_stage = stage_idx, lead_phase = phase; - uint32_t lead_packed = packed_stage_idx, lead_packed_phase = packed_phase; auto advance_lead = [&]() { lead_stage = lead_stage == kNumStages - 1 ? 0 : lead_stage + 1; lead_phase ^= lead_stage == 0; - lead_packed = lead_packed == kNumPackedStages - 1 ? 0 : lead_packed + 1; - lead_packed_phase ^= lead_packed == 0; }; // Prologue: decode the first stage so iter 0's wgmma can // consume it without waiting. - wait_and_decode(stage_idx, phase, packed_stage_idx, packed_phase); + wait_and_decode(stage_idx, phase); cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); - packed_empty_barrier_arrive_at(packed_stage_idx); advance_lead(); // lead now points one stage ahead of wgmma #pragma unroll 8 @@ -514,7 +535,7 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, // the same wait/decode/sync sequence per K iter. const bool is_last_wave = (local_idx == BLOCK_M / WAVE_BLOCK_M - 1); if (is_last_wave and has_next) { - wait_and_decode(lead_stage, lead_phase, lead_packed, lead_packed_phase); + wait_and_decode(lead_stage, lead_phase); } ptx::warpgroup_wait<0>(); @@ -542,10 +563,10 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, } // Publish next iter's decoded smem_b to the wgmma - // async-proxy and release the packed-B slot. + // async-proxy. The packed-B slot is released together + // with A/SFA via the unified empty_barrier_arrive() above. if (has_next) { cutlass::arch::NamedBarrier::sync(kNumMathThreads, 7); - packed_empty_barrier_arrive_at(lead_packed); advance_lead(); } } @@ -553,8 +574,6 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, } else { #pragma unroll for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { - packed_full_barriers[packed_stage_idx]->wait(packed_phase); - packed_empty_barrier_arrive(); full_barriers[stage_idx]->wait(phase); empty_barrier_arrive(); } From 214baafcd97de1bff9b37c2e5d80f8798cfd15af Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 18:02:16 +0800 Subject: [PATCH 12/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 556 +++++++++++++------------------------ 1 file changed, 199 insertions(+), 357 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 221150bed9..7bee97f6cf 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -1,5 +1,5 @@ -import time import sys +import time from pathlib import Path import torch @@ -11,10 +11,6 @@ from deep_gemm.utils.math import cast_back_from_fp4, per_token_cast_to_fp4, per_token_cast_to_fp8 -def _align_up(x: int, alignment: int) -> int: - return (x + alignment - 1) // alignment * alignment - - def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int = 128) -> torch.Tensor: group_idx = torch.arange(x.size(-1), device=x.device) // gran_k return x.float() * sf[..., group_idx] @@ -24,10 +20,10 @@ def _require_sm90() -> None: assert torch.cuda.is_available() major, _ = torch.cuda.get_device_capability() if major != 9: - raise RuntimeError(f"This fallback test is intended for SM90, got sm_{major}x") + raise RuntimeError(f"This benchmark is intended for SM90, got sm_{major}x") -def _time_cuda(fn, warmup: int = 5, iters: int = 20) -> float: +def _time_cuda(fn, warmup: int = 3, iters: int = 10) -> float: for _ in range(warmup): fn() torch.cuda.synchronize() @@ -42,407 +38,253 @@ def _time_cuda(fn, warmup: int = 5, iters: int = 20) -> float: return start.elapsed_time(end) / iters / 1e3 -def _normal_case(m: int, n: int, k: int, gran_k: int = 128) -> None: +def _effective_bytes(groups: int, m_per_group: int, n: int, k: int, gran_k: int, *, fp8_b: bool) -> int: + logical_m = groups * m_per_group + scale_k = (k + gran_k - 1) // gran_k + a_bytes = logical_m * k + logical_m * scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _sfb_alias_fits(block_m: int, k: int, gran_k: int) -> bool: + # The 1d2d kernel aliases the SFB cache onto smem_d. This mirrors the + # kernel-side capacity check and avoids device traps during autotune. + return (k // gran_k) * 4 <= block_m * 2 + + +def _build_grouped_layout(groups: int, m_per_group: int): + m = groups * m_per_group + group_starts = [group_id * m_per_group for group_id in range(groups)] + group_ends = [(group_id + 1) * m_per_group for group_id in range(groups)] + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) + return m, group_starts, group_ends, grouped_layout + + +def _benchmark_case(groups: int, m_per_group: int, n: int, k: int, gran_k: int = 128) -> dict[str, float | int]: + m, group_starts, group_ends, grouped_layout = _build_grouped_layout(groups, m_per_group) a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) - b_ref_src = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) - b = per_token_cast_to_fp4(b_ref_src, use_ue8m0=True, gran_k=gran_k) - c = torch.zeros((m, n), device="cuda", dtype=torch.float) - d = torch.empty((m, n), device="cuda", dtype=torch.float) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for group_id in range(groups): + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=gran_k + ) + b_w4 = (b_fp4, b_sf) a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) - b_dequant = cast_back_from_fp4(b[0], b[1], gran_k=gran_k) - b_fp8 = per_token_cast_to_fp8(b_dequant, use_ue8m0=False, gran_k=gran_k) - ref = (a_dequant @ b_dequant.t()).to(torch.bfloat16) - - def run(): - deep_gemm.fp8_gemm_nt( - a, - b_fp8, - d, - c=c, - recipe_a=(1, gran_k), - recipe_b=(1, gran_k), + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + b_dequant = cast_back_from_fp4(b_w4[0][group_id], b_w4[1][group_id], gran_k=gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=gran_k ) + start = group_starts[group_id] + end = group_ends[group_id] + if start != end: + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) - d_fp8 = torch.empty_like(d) + d_fp8 = torch.empty_like(ref) def run_fp8(): - deep_gemm.fp8_gemm_nt( + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( a, b_fp8, d_fp8, - c=c, + grouped_layout, recipe_a=(1, gran_k), recipe_b=(1, gran_k), + use_psum_layout=False, ) - d_fused_wgmma = torch.empty_like(d) - - def run_fused_wgmma(): - deep_gemm.fp8_fp4_gemm_nt_sm90_fused_wgmma( - a, - b, - d_fused_wgmma, - c=c, - gran_k=gran_k, - ) - - run() run_fp8() - if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma"): - run_fused_wgmma() - fallback_diff = calc_diff(d, ref) fp8_diff = calc_diff(d_fp8, ref) - fallback_vs_fp8_diff = calc_diff(d, d_fp8) - fused_wgmma_diff = ( - calc_diff(d_fused_wgmma, ref) - if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") - else None - ) - fallback_elapsed = _time_cuda(run) fp8_elapsed = _time_cuda(run_fp8) - fused_wgmma_elapsed = ( - _time_cuda(run_fused_wgmma) - if hasattr(deep_gemm, "fp8_fp4_gemm_nt_sm90_fused_wgmma") - else None - ) - fallback_tflops = 2 * m * n * k / fallback_elapsed / 1e12 - fp8_tflops = 2 * m * n * k / fp8_elapsed / 1e12 - fused_wgmma_tflops = ( - 2 * m * n * k / fused_wgmma_elapsed / 1e12 - if fused_wgmma_elapsed is not None - else None - ) - slowdown = fallback_elapsed / fp8_elapsed - message = ( - f"normal m={m} n={n} k={k}:\n" - f" fallback end2end: diff={fallback_diff:.6f}, " - f"time={fallback_elapsed * 1e6:.1f} us, {fallback_tflops:.2f} TFLOPS\n" - f" pure fp8 gemm: diff={fp8_diff:.6f}, " - f"time={fp8_elapsed * 1e6:.1f} us, {fp8_tflops:.2f} TFLOPS\n" - f" fallback vs fp8: diff={fallback_vs_fp8_diff:.6f}\n" - f" slowdown: {slowdown:.2f}x" - ) - if fused_wgmma_elapsed is not None: - message += ( - f"\n fused wgmma: diff={fused_wgmma_diff:.6f}, " - f"time={fused_wgmma_elapsed * 1e6:.1f} us, {fused_wgmma_tflops:.2f} TFLOPS, " - f"slowdown_vs_fp8={fused_wgmma_elapsed / fp8_elapsed:.2f}x" + + best_w4_elapsed = float("inf") + best_w4_diff = float("inf") + best_block_m = None + best_block_n = None + for block_m in (64, 128, 256): + if m_per_group < block_m: + continue + if not _sfb_alias_fits(block_m, k, gran_k): + continue + for block_n in (64, 128, 256): + d_w4 = torch.empty_like(ref) + + def run_w4( + block_m: int = block_m, + block_n: int = block_n, + out: torch.Tensor = d_w4, + ): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b_w4, + out, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=False, + block_m_override=block_m, + block_n_override=block_n, + ) + + try: + run_w4() + except RuntimeError: + continue + w4_diff = calc_diff(d_w4, ref) + w4_elapsed = _time_cuda(run_w4) + if w4_diff < 0.015 and w4_elapsed < best_w4_elapsed: + best_w4_elapsed = w4_elapsed + best_w4_diff = w4_diff + best_block_m = block_m + best_block_n = block_n + + if best_block_m is None or best_block_n is None: + raise RuntimeError(f"No valid W4 candidate for groups={groups}, m/group={m_per_group}, n={n}, k={k}") + + fp8_threshold = 0.05 + assert best_w4_diff < 0.015 + assert fp8_diff < fp8_threshold + + w4_bytes = _effective_bytes(groups, m_per_group, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_bytes(groups, m_per_group, n, k, gran_k, fp8_b=True) + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": best_w4_elapsed * 1e6, + "w4_gbps": w4_bytes / best_w4_elapsed / 1e9, + "w4_diff": best_w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / best_w4_elapsed, + } + + +def _print_markdown_table(rows: list[dict[str, float | int]]) -> None: + print("groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for row in rows: + print( + f"{row['groups']} | {row['m_per_group']} | {row['n']} | {row['k']} | " + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " + f"{row['speedup']:.2f}x" ) - print(message) - assert fallback_diff < 0.015 - assert fp8_diff < 0.015 - assert fallback_vs_fp8_diff < 1e-6 - if fused_wgmma_diff is not None: - assert fused_wgmma_diff < 0.015 -def _m_grouped_contiguous_case( - num_groups: int, - m_per_group: int | None, +def _accuracy_case( + groups: int, + m_per_group: int, n: int, k: int, gran_k: int = 128, - group_sizes: list[int] | None = None, - use_psum_layout: bool = False, - autotune_fused: bool = True, -) -> None: - if group_sizes is None: - assert m_per_group is not None - group_sizes = [m_per_group] * num_groups - assert len(group_sizes) == num_groups - print( - f"starting m_grouped groups={num_groups} n={n} k={k} " - f"layout={'psum' if use_psum_layout else 'per-row'} " - f"sizes={group_sizes} autotune={autotune_fused}", - flush=True, - ) - - block_m_alignment = 128 - if use_psum_layout: - group_starts: list[int] = [] - group_ends: list[int] = [] - cursor = 0 - for size in group_sizes: - group_starts.append(cursor) - cursor += size - group_ends.append(cursor) - cursor = _align_up(cursor, block_m_alignment) - m = group_ends[-1] if group_ends else 0 - grouped_layout = torch.tensor(group_ends, device="cuda", dtype=torch.int32) - else: - assert all(size % block_m_alignment == 0 for size in group_sizes) - m = sum(group_sizes) - group_starts = [] - cursor = 0 - layout_chunks = [] - for group_id, size in enumerate(group_sizes): - group_starts.append(cursor) - cursor += size - if size > 0: - layout_chunks.append( - torch.full((size,), group_id, device="cuda", dtype=torch.int32) - ) - grouped_layout = torch.cat(layout_chunks) if layout_chunks else torch.empty((0,), device="cuda", dtype=torch.int32) - + *, + block_m: int = 128, + block_n: int = 128, +) -> tuple[float, float]: + m, group_starts, group_ends, grouped_layout = _build_grouped_layout(groups, m_per_group) a_ref_src = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) - b_ref_src = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=gran_k) - b_fp4 = torch.empty((num_groups, n, k // 2), device="cuda", dtype=torch.int8) - b_sf = torch.empty((num_groups, n, k // gran_k), device="cuda", dtype=torch.float) - for group_id in range(num_groups): + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for group_id in range(groups): b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( b_ref_src[group_id], use_ue8m0=True, gran_k=gran_k ) - b = (b_fp4, b_sf) - d = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + b_w4 = (b_fp4, b_sf) a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=gran_k) - b_fp8_data = torch.empty((num_groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) - b_fp8_sf = torch.empty((num_groups, n, k // gran_k), device="cuda", dtype=torch.float) - ref = torch.empty_like(d) - for group_id in range(num_groups): - b_dequant = cast_back_from_fp4(b[0][group_id], b[1][group_id], gran_k=gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + ref = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + b_dequant = cast_back_from_fp4(b_w4[0][group_id], b_w4[1][group_id], gran_k=gran_k) b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( b_dequant, use_ue8m0=False, gran_k=gran_k ) start = group_starts[group_id] - end = ( - _align_up(group_ends[group_id], block_m_alignment) - if use_psum_layout and group_id + 1 < num_groups - else group_ends[group_id] - ) if use_psum_layout else start + group_sizes[group_id] - end = min(end, m) - if start == end: - continue - ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + end = group_ends[group_id] + if start != end: + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) b_fp8 = (b_fp8_data, b_fp8_sf) - def run(): - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( - a, - b_fp8, - d, - grouped_layout, - recipe_a=(1, gran_k), - recipe_b=(1, gran_k), - use_psum_layout=use_psum_layout, - expected_m_for_psum_layout=m if use_psum_layout else None, - ) - - d_fp8 = torch.empty_like(d) - d_fused = torch.empty_like(d) - - def run_fp8(): - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( - a, - b_fp8, - d_fp8, - grouped_layout, - recipe_a=(1, gran_k), - recipe_b=(1, gran_k), - use_psum_layout=use_psum_layout, - expected_m_for_psum_layout=m if use_psum_layout else None, - ) - - def run_fused_wgmma( - block_m_override: int | None = None, - block_n_override: int | None = None, - ): - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( - a, - b, - d_fused, - grouped_layout, - gran_k=gran_k, - compiled_dims="nk", - use_psum_layout=use_psum_layout, - expected_m_for_psum_layout=m if use_psum_layout else None, - block_m_override=block_m_override, - block_n_override=block_n_override, - ) + d_w4 = torch.empty_like(ref) + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b_w4, + d_w4, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=False, + block_m_override=block_m, + block_n_override=block_n, + ) - print(" running fallback...", flush=True) - run() - print(" running pure fp8...", flush=True) - run_fp8() - print(" running fused wgmma...", flush=True) - run_fused_wgmma() - fallback_diff = calc_diff(d, ref) - fp8_diff = calc_diff(d_fp8, ref) - fused_diff = calc_diff(d_fused, ref) - fallback_vs_fp8_diff = calc_diff(d, d_fp8) - fused_vs_fp8_diff = calc_diff(d_fused, d_fp8) - fallback_elapsed = _time_cuda(run) - fp8_elapsed = _time_cuda(run_fp8) - fused_elapsed = _time_cuda(run_fused_wgmma) - best_fused = (fused_elapsed, None, fused_diff) - if autotune_fused: - # Psum grouped_layout is encoded with 128-row alignment, so only BLOCK_M=128 - # is semantically valid. Sweep BLOCK_N only for psum. - block_m_candidates = (128,) if use_psum_layout else (64, 128, 256) - block_n_candidates = (64, 128, 256) - for block_m in block_m_candidates: - if m < block_m: - continue - if any(size > 0 and size < block_m for size in group_sizes) and not use_psum_layout: - continue - for block_n in block_n_candidates: - d_candidate = torch.empty_like(d) - - def run_candidate( - block_m: int = block_m, - block_n: int = block_n, - out: torch.Tensor = d_candidate, - ): - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( - a, - b, - out, - grouped_layout, - gran_k=gran_k, - compiled_dims="nk", - use_psum_layout=use_psum_layout, - expected_m_for_psum_layout=m if use_psum_layout else None, - block_m_override=block_m, - block_n_override=block_n, - ) - - try: - run_candidate() - except RuntimeError as exc: - print( - f" fused autotune block_m={block_m} block_n={block_n}: " - f"skipped ({exc})" - ) - continue - candidate_diff = calc_diff(d_candidate, ref) - candidate_elapsed = _time_cuda(run_candidate) - candidate_tflops = 2 * m * n * k / candidate_elapsed / 1e12 - print( - f" fused autotune block_m={block_m} block_n={block_n}: " - f"diff={candidate_diff:.6f}, " - f"time={candidate_elapsed * 1e6:.1f} us, {candidate_tflops:.2f} TFLOPS" - ) - if candidate_diff < 0.015 and candidate_elapsed < best_fused[0]: - best_fused = (candidate_elapsed, (block_m, block_n), candidate_diff) - d_fused.copy_(d_candidate) - fused_elapsed, best_tile, fused_diff = best_fused - fused_vs_fp8_diff = calc_diff(d_fused, d_fp8) - else: - best_tile = None - fallback_tflops = 2 * m * n * k / fallback_elapsed / 1e12 - fp8_tflops = 2 * m * n * k / fp8_elapsed / 1e12 - fused_tflops = 2 * m * n * k / fused_elapsed / 1e12 - slowdown = fallback_elapsed / fp8_elapsed - fused_slowdown = fused_elapsed / fp8_elapsed - print( - f"m_grouped groups={num_groups} m={m} n={n} k={k} " - f"layout={'psum' if use_psum_layout else 'per-row'} sizes={group_sizes}:\n" - f" fallback end2end: diff={fallback_diff:.6f}, " - f"time={fallback_elapsed * 1e6:.1f} us, {fallback_tflops:.2f} TFLOPS\n" - f" pure fp8 gemm: diff={fp8_diff:.6f}, " - f"time={fp8_elapsed * 1e6:.1f} us, {fp8_tflops:.2f} TFLOPS\n" - f" fused wgmma: diff={fused_diff:.6f}, " - f"time={fused_elapsed * 1e6:.1f} us, {fused_tflops:.2f} TFLOPS, " - f"slowdown_vs_fp8={fused_slowdown:.2f}x" - f"{'' if best_tile is None else f', best_tile={best_tile}'}\n" - f" fallback vs fp8: diff={fallback_vs_fp8_diff:.6f}\n" - f" fused vs fp8: diff={fused_vs_fp8_diff:.6f} " - f"(expected: fused keeps original FP4 scales)\n" - f" slowdown: {slowdown:.2f}x" + d_fp8 = torch.empty_like(ref) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, gran_k), + recipe_b=(1, gran_k), + use_psum_layout=False, ) - # Psum edge cases with few effective rows can amplify the FP4->FP8 requant - # baseline error. Keep fused correctness strict; relax only the requantized - # fallback/pure-FP8 baseline sanity threshold. - requant_threshold = 0.25 if use_psum_layout else 0.05 - assert fallback_diff < requant_threshold - assert fp8_diff < requant_threshold - assert fused_diff < 0.015 - assert fallback_vs_fp8_diff < 1e-6 + + return calc_diff(d_w4, ref), calc_diff(d_fp8, ref) def test_sm90_fp8_fp4_fallback() -> None: _require_sm90() torch.manual_seed(0) - # Small enough for quick CI/local validation, large enough to exercise SM90 FP8 GEMM. - _normal_case(m=128, n=1024, k=1024) - _m_grouped_contiguous_case(num_groups=1, m_per_group=128, n=1024, k=1024) - _m_grouped_contiguous_case(num_groups=4, m_per_group=128, n=1024, k=1024) - _m_grouped_contiguous_case(num_groups=8, m_per_group=128, n=512, k=1024) - _m_grouped_contiguous_case( - num_groups=4, m_per_group=256, n=2048, k=1024, autotune_fused=True - ) - _m_grouped_contiguous_case( - num_groups=4, - m_per_group=None, - n=1024, - k=1024, - group_sizes=[96, 0, 160, 64], - use_psum_layout=True, - ) - _m_grouped_contiguous_case( - num_groups=5, - m_per_group=None, - n=512, - k=2048, - group_sizes=[0, 128, 32, 256, 48], - use_psum_layout=True, - autotune_fused=True, - ) - # Psum stability coverage: - # - aligned groups: sanity baseline where psum should behave like regular contiguous groups - # - leading/trailing zero groups: scheduler fallthrough and empty group handling - # - tiny partial groups: mixed valid/invalid rows inside a BLOCK_M tile - # - larger K/N: exercise the 1d2d decode pipeline and smem SFB cache over more K blocks - _m_grouped_contiguous_case( - num_groups=3, - m_per_group=None, - n=512, - k=1024, - group_sizes=[128, 128, 128], - use_psum_layout=True, - ) - _m_grouped_contiguous_case( - num_groups=4, - m_per_group=None, - n=512, - k=1024, - group_sizes=[0, 64, 128, 0], - use_psum_layout=True, - ) - _m_grouped_contiguous_case( - num_groups=4, - m_per_group=None, - n=768, - k=1024, - group_sizes=[1, 127, 129, 1], - use_psum_layout=True, - ) - _m_grouped_contiguous_case( - num_groups=6, - m_per_group=None, - n=1024, - k=2048, - group_sizes=[64, 0, 64, 0, 192, 32], - use_psum_layout=True, - ) - _m_grouped_contiguous_case( - num_groups=4, - m_per_group=None, - n=1024, - k=1024, - group_sizes=[96, 160, 0, 64], - use_psum_layout=True, - autotune_fused=True, + rows = [] + for groups in (8, 16, 24, 32): + for m_per_group in (128, 256, 512, 1024): + rows.append(_benchmark_case(groups, m_per_group, n=4096, k=7168)) + _print_markdown_table(rows) + + +def test_sm90_fp8_fp4_fallback_accuracy() -> None: + _require_sm90() + torch.manual_seed(1) + + cases = ( + (8, 128, 1024, 1024), + (16, 128, 2048, 2048), + (24, 256, 4096, 7168), + (32, 512, 4096, 7168), ) + for groups, m_per_group, n, k in cases: + w4_diff, fp8_diff = _accuracy_case(groups, m_per_group, n, k) + assert w4_diff < 0.015, ( + f"W4 accuracy check failed for groups={groups}, m/group={m_per_group}, " + f"n={n}, k={k}: diff={w4_diff}" + ) + assert fp8_diff < 0.05, ( + f"FP8 accuracy check failed for groups={groups}, m/group={m_per_group}, " + f"n={n}, k={k}: diff={fp8_diff}" + ) if __name__ == "__main__": start_time = time.time() test_sm90_fp8_fp4_fallback() + test_sm90_fp8_fp4_fallback_accuracy() print(f"done in {time.time() - start_time:.2f}s") From 4392c080444e8c57f3593fd96ecdc2eef9fadc99 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 18:04:09 +0800 Subject: [PATCH 13/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 7bee97f6cf..236a81c1a0 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -256,7 +256,7 @@ def test_sm90_fp8_fp4_fallback() -> None: rows = [] for groups in (8, 16, 24, 32): - for m_per_group in (128, 256, 512, 1024): + for m_per_group in (256, 512, 1024, 2048): rows.append(_benchmark_case(groups, m_per_group, n=4096, k=7168)) _print_markdown_table(rows) @@ -267,9 +267,9 @@ def test_sm90_fp8_fp4_fallback_accuracy() -> None: cases = ( (8, 128, 1024, 1024), - (16, 128, 2048, 2048), - (24, 256, 4096, 7168), - (32, 512, 4096, 7168), + (16, 256, 2048, 2048), + (24, 512, 4096, 7168), + (32, 1024, 4096, 7168), ) for groups, m_per_group, n, k in cases: w4_diff, fp8_diff = _accuracy_case(groups, m_per_group, n, k) From 83cd196749481922591dbaac8465a92a57dd7d1b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 11 May 2026 21:58:10 +0800 Subject: [PATCH 14/51] Update gemm.hpp --- csrc/apis/gemm.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 6b618c3ddf..7e823052d0 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -659,7 +659,8 @@ static void register_apis(pybind11::module_& m) { py::arg("use_psum_layout") = false, py::arg("expected_m_for_psum_layout") = std::nullopt, py::arg("block_m_override") = std::nullopt, - py::arg("block_n_override") = std::nullopt); + py::arg("block_n_override") = std::nullopt, + py::arg("decode_stub") = false); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, From 02ff071d1b17609cc09dceb9af1da2317fa07173 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:00:26 +0800 Subject: [PATCH 15/51] Update gemm.hpp --- csrc/apis/gemm.hpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 7e823052d0..576509bb35 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -661,6 +661,19 @@ static void register_apis(pybind11::module_& m) { py::arg("block_m_override") = std::nullopt, py::arg("block_n_override") = std::nullopt, py::arg("decode_stub") = false); + m.def("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", + &sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"), + py::arg("expected_m"), + py::arg("gran_k") = 128, + py::arg("gran_k_a") = std::nullopt, + py::arg("gran_k_b") = std::nullopt, + py::arg("compiled_dims") = "nk", + py::arg("block_m_override") = std::nullopt, + py::arg("block_n_override") = std::nullopt, + py::arg("decode_stub") = false); + m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") = + m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma"); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, From 2f989465f672c1f9e78eaa5d0495b5b5fe33a2bb Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:01:02 +0800 Subject: [PATCH 16/51] Update layout.hpp --- csrc/apis/layout.hpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index b404241a66..b80eacfc61 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -36,8 +36,13 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, // Pre-transform checks check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups); - // (FP32, 1, 128) on SM90: transform to TMA-aligned and MN-major - if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast)) + // (FP32, 1, 32/128) on SM90: transform to TMA-aligned and MN-major + if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and + (arch_major == 9 or disable_ue8m0_cast)) + return get_mn_major_tma_aligned_tensor(sf); + + // (INT packed UE8M0, 1, 32/128) on SM90: transform to TMA-aligned and MN-major. + if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9) return get_mn_major_tma_aligned_tensor(sf); // (FP32, 128, 128) on SM90: no need to transform, check SFB requirements From d2e9012d5e967232678b06d8b4860de80fbd9cb7 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:02:00 +0800 Subject: [PATCH 17/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 355 +++++++++++++++++- 1 file changed, 349 insertions(+), 6 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index bec03e9796..ad1d8d7590 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -10,6 +10,7 @@ #include "../../utils/layout.hpp" #include "epilogue.hpp" #include "runtime_utils.hpp" +#include "sm90_fp8_fp4_gemm_1d2d_rs.hpp" namespace deep_gemm { @@ -152,7 +153,7 @@ static void sm90_fp8_fp4_gemm_1d1d_fused(const std::pairis_contiguous()); @@ -387,13 +388,28 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( layout.block_m = 128; layout.cluster_m = 1; } - if (not use_psum_layout and m >= num_groups * 256 and n >= 1024) { - layout.block_m = 256; - layout.block_n = 64; + if (not use_psum_layout and n >= 1024 and num_groups > 0 and m % num_groups == 0) { + const int expected_m_per_group = m / num_groups; + if (expected_m_per_group >= 256) { + layout.block_m = 256; + layout.block_n = 64; + } else if (expected_m_per_group == 128) { + // Tuned on the SM90 FP8xFP4 contiguous fallback benchmark at + // N=4096,K=7168. Earlier code carved out BN=128 for 16/24 groups, + // but this consistently underperformed BN=64 (16g:0.74x, 24g:0.78x + // vs the BN=64 baseline). Larger BN doubles the per-stage packed-B + // and SFB cache footprint, which forces `num_stages` down without + // recovering wave utilization, so always pick BN=64 here. + layout.block_m = 128; + layout.block_n = 64; + } } layout.cluster_m = 1; auto config = rebuild_config(layout); + // The contiguous fallback benchmark explicitly sweeps BLOCK_M/BLOCK_N via + // these overrides; keep this path independent from masked RS heuristics so + // regressions can be attributed to the selected block shape. if (block_m_override or block_n_override) { auto layout = config.layout; if (block_m_override) { @@ -424,7 +440,9 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( const int shape_k_scales = ceil_div(static_cast(k), block_k); const bool uniform_scale_b = (block_k % block_n == 0); const int sfb_old_bytes = align(shape_k_scales * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); - const int sfb_extra = -sfb_old_bytes; + const int sfb_cache_bytes = align(shape_k_scales * block_n * static_cast(sizeof(float)), 16); + const int smem_d_bytes = align(config.layout.block_m * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; const int packed_per_stage = block_n * (block_k / 2); const int smem_a_per_stage = config.storage_config.load_block_m * block_k * @@ -441,8 +459,15 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( return smem_extra + stages * merged_per_stage <= SM90ArchSpec::smem_capacity; }; - // Try to keep `num_stages = orig_num_stages`; otherwise drop until it fits. + // Packed-FP4 B halves the per-stage B footprint, so the smem freed by + // the FP4 path can usually accommodate more pipeline stages than the + // FP8 baseline. Mirror the masked path's logic: try to push `num_stages` + // up to `kW4DefaultMaxStages` while it still fits, then fall back if + // the chosen value does not fit (e.g. due to large SFB cache at BN=128). + constexpr int kW4DefaultMaxStages = 8; int chosen_stages = orig_num_stages; + while (chosen_stages + 1 <= kW4DefaultMaxStages and fits(chosen_stages + 1)) + ++ chosen_stages; while (chosen_stages >= 3 and not fits(chosen_stages)) -- chosen_stages; DG_HOST_ASSERT(chosen_stages >= 3); @@ -489,4 +514,322 @@ static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( SM90FP8FP4Gemm1D2DRuntime::launch(runtime, args); } +static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( + const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& expected_m, + const int& gran_k, + const std::optional& gran_k_a_override, + const std::optional& gran_k_b_override, + const std::string& compiled_dims, + const std::optional& block_m_override, + const std::optional& block_n_override, + const bool& decode_stub) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + const int gran_k_a_requested = gran_k_a_override.value_or(gran_k); + const int gran_k_b_requested = gran_k_b_override.value_or(gran_k); + DG_HOST_ASSERT(gran_k_a_requested == 128); + DG_HOST_ASSERT(gran_k_b_requested == 32 or gran_k_b_requested == 128); + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat or b.second.scalar_type() == torch::kInt); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt and masked_m.is_contiguous()); + DG_HOST_ASSERT(a.first.is_contiguous() and b.first.is_contiguous() and d.is_contiguous()); + + const auto [num_groups, m, k] = get_shape<3>(a.first); + const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); + const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(num_groups == num_groups_b and num_groups == num_groups_d); + DG_HOST_ASSERT(masked_m.numel() == num_groups); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0); + DG_HOST_ASSERT(a.second.size(0) == num_groups and a.second.size(1) == m and + a.second.size(2) == ceil_div(k, gran_k_a_requested)); + const int gran_k_b_shape = b.second.scalar_type() == torch::kInt ? + gran_k_b_requested * 4 : gran_k_b_requested; + DG_HOST_ASSERT(b.second.size(0) == num_groups and b.second.size(1) == n and + b.second.size(2) == ceil_div(k, gran_k_b_shape)); + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, gran_k_a_requested); + std::optional> recipe_b = std::make_tuple(1, gran_k_b_requested); + const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, + num_groups, num_groups, false); + DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == gran_k_b_requested); + + auto desc = GemmDesc { + .gemm_type = GemmType::MGroupedMasked, + .kernel_type = KernelType::Kernel1D2D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = expected_m, + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups + }; + auto rebuild_config = [&](Layout layout) { + const auto storage_config = SM90ArchSpec::get_storage_config(desc, layout); + const auto pipeline_config = SM90ArchSpec::get_pipeline_config(desc, layout, storage_config); + DG_HOST_ASSERT(pipeline_config.num_stages >= 3); + const auto launch_config = SM90ArchSpec::get_launch_config(desc, layout); + return GemmConfig{layout, storage_config, pipeline_config, launch_config}; + }; + + const auto layout_candidates = SM90ArchSpec::get_layout_candidates(desc); + DG_HOST_ASSERT(not layout_candidates.empty()); + auto layout = layout_candidates[0]; + auto layout_info = SM90ArchSpec::get_layout_info(desc, layout); + bool found_layout = false; + for (const auto& candidate: layout_candidates) { + if (candidate.cluster_m != 1) { + continue; + } + const auto candidate_info = SM90ArchSpec::get_layout_info(desc, candidate); + if (not found_layout or SM90ArchSpec::compare(candidate_info, layout_info)) { + layout = candidate; + layout_info = candidate_info; + found_layout = true; + } + } + DG_HOST_ASSERT(found_layout); + // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages + // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 + // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 + // + // 参数空间：BM ∈ {8, 16, 32, 64, 128}（BM<64 用于 masked small-M）， + // BN ∈ {64, 128, 256}。 + // + // 选择目标：在 (waves <= ceil_div(total_tiles, num_sms)) 的前提下最大化 stages， + // 其次最大化 last_wave 利用率，最后倾向更小的 per-stage（即更小 BN）。 + if (not block_m_override and not block_n_override) { + const int num_sms = desc.num_sms; + const int block_k = layout.block_k; + const int shape_k_scales_b = ceil_div(static_cast(k), gran_k_b); + + auto eval_layout = [&](int bm, int bn) -> std::tuple { + // 返回 (sat_stages, -waves, last_wave_util, -per_stage)，越大越好。 + // stages 在 ~6 之上对 TMA 隐藏几乎饱和，因此用饱和 stage 数比较，避免 + // 小 BN 因 stages=8 击败 wave 利用率更高的候选。 + const int tiles = ceil_div(expected_m, bm) * ceil_div(static_cast(n), bn) * num_groups; + const int waves = ceil_div(tiles, num_sms); + const int last = tiles - (waves - 1) * num_sms; + const int last_util = last <= 0 ? num_sms : last; + const bool uniform_scale_b = (block_k % bn == 0); + const int sfb_old_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + const int sfb_cache_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * bn * static_cast(sizeof(float)), 16); + const int rs_padded_bm = std::max(bm, 64); + const int smem_d_bytes = + align(rs_padded_bm * bn * static_cast(sizeof(nv_bfloat16)), 1024); + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; + const int smem_a_per_stage = rs_padded_bm * block_k * static_cast(c10::elementSize(desc.a_dtype)); + const int smem_sfa_per_stage = + align(rs_padded_bm * static_cast(sizeof(float)), 128); + const int packed_per_stage = bn * (block_k / 2); + const int merged_per_stage = smem_a_per_stage + smem_sfa_per_stage + packed_per_stage; + constexpr int kMaxEvaluatedStages = 10; + constexpr int kBarrierBytes = 16 * kMaxEvaluatedStages * 2; + const int fixed = smem_d_bytes + kBarrierBytes + sfb_extra; + const int max_stages = (SM90ArchSpec::smem_capacity - fixed) / merged_per_stage; + constexpr int kStageSaturation = 6; + const int sat_stages = std::min(std::min(max_stages, kMaxEvaluatedStages), kStageSaturation); + return std::make_tuple(sat_stages, -waves, last_util, -merged_per_stage); + }; + + std::vector> w4_candidates = { + {64, 64}, {64, 128}, {64, 256}, + {128, 64}, {128, 128}, + }; + if (expected_m <= 32) { + w4_candidates.insert(w4_candidates.begin(), {{8, 64}, {16, 64}, {32, 64}}); + } + + std::pair best{layout.block_m, layout.block_n}; + std::tuple best_score{-1, 0, 0, 0}; + bool first = true; + for (const auto& cand : w4_candidates) { + const int bm = cand.first; + const int bn = cand.second; + // 1D2D 内核 unroll 要求 + if (bn > block_k and (bn % (bn - block_k) != 0 and block_k % (bn - block_k) != 0)) + continue; + // masked 路径 multicast 合法性：当前固定 cluster_m=1, cluster_n=1，恒满足 + const auto score = eval_layout(bm, bn); + if (std::get<0>(score) < 3) + continue; + if (first or score > best_score) { + best_score = score; + best = cand; + first = false; + } + } + layout.block_m = best.first; + layout.block_n = best.second; + + // RS masked W4 empirical fallback: + // pick BM close to expected_m to avoid over-computing promotion work, + // and use BN=256 for fewer CTAs while staying within Hopper TMA limits. + if (desc.gemm_type == GemmType::MGroupedMasked) { + if (expected_m <= 8) { + layout.block_m = 8; + } else if (expected_m <= 16) { + layout.block_m = 16; + } else if (expected_m <= 32) { + layout.block_m = 32; + } else if (expected_m <= 64) { + layout.block_m = 64; + } + layout.block_n = 256; + } + } + if (not block_m_override and not block_n_override) { + DG_HOST_ASSERT(layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or + layout.block_m == 64 or layout.block_m == 128); + DG_HOST_ASSERT(layout.block_n == 64 or layout.block_n == 128 or layout.block_n == 256); + } + layout.cluster_m = 1; + auto config = rebuild_config(layout); + + if (block_m_override or block_n_override) { + auto layout = config.layout; + if (block_m_override) { + layout.block_m = *block_m_override; + } + if (block_n_override) { + layout.block_n = *block_n_override; + } + DG_HOST_ASSERT((layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or + layout.block_m == 64 or layout.block_m == 128 or layout.block_m == 256) and layout.block_n % 16 == 0); + DG_HOST_ASSERT(layout.block_n <= 256); + layout.cluster_m = 1; + config = rebuild_config(layout); + } + // Packed FP4 B has half the K bytes of FP8 B. Match PR #287's W4 path: + // TMA writes B with a 64B swizzle and the RS kernel reads it via ldmatrix. + config.storage_config.swizzle_b_mode = config.layout.block_k / 2; + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k / 2); + + { + const int block_k = config.layout.block_k; + const int block_n = config.layout.block_n; + const int shape_k_scales_b = ceil_div(static_cast(k), gran_k_b); + const bool uniform_scale_b = (block_k % block_n == 0); + const int sfb_old_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * (uniform_scale_b ? 1 : 2) * static_cast(sizeof(float)), 16); + const int sfb_cache_bytes = gran_k_b == 32 ? 0 : + align(shape_k_scales_b * block_n * static_cast(sizeof(float)), 16); + const int rs_padded_bm = std::max(config.layout.block_m, 64); + const int base_smem_d_bytes = + align(config.layout.block_m * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int smem_d_bytes = + align(rs_padded_bm * block_n * static_cast(sizeof(nv_bfloat16)), 1024); + const int smem_d_extra = smem_d_bytes - base_smem_d_bytes; + const int sfb_extra = (sfb_cache_bytes > smem_d_bytes ? sfb_cache_bytes : 0) - sfb_old_bytes; + + const int packed_per_stage = block_n * (block_k / 2); + const int base_smem_a_per_stage = config.storage_config.load_block_m * block_k * + static_cast(c10::elementSize(desc.a_dtype)); + const int base_smem_b_per_stage = config.storage_config.load_block_n * block_k * + static_cast(c10::elementSize(desc.b_dtype)); + const int base_smem_sfa_per_stage = + align(config.layout.block_m * static_cast(sizeof(float)), 128); + const int original_per_stage = + base_smem_a_per_stage + base_smem_b_per_stage + base_smem_sfa_per_stage; + const int smem_a_per_stage = rs_padded_bm * block_k * static_cast(c10::elementSize(desc.a_dtype)); + const int smem_sfa_per_stage = + align(rs_padded_bm * static_cast(sizeof(float)), 128); + const int merged_per_stage = smem_a_per_stage + smem_sfa_per_stage + packed_per_stage; + const int orig_num_stages = config.pipeline_config.num_stages; + const int smem_extra = + config.pipeline_config.smem_size - orig_num_stages * original_per_stage + smem_d_extra + sfb_extra; + + auto fits = [&](int stages) { + return smem_extra + stages * merged_per_stage <= SM90ArchSpec::smem_capacity; + }; + + constexpr int kW4DefaultMaxStages = 8; + int max_fitting = orig_num_stages; + while (max_fitting + 1 <= kW4DefaultMaxStages and fits(max_fitting + 1)) + ++ max_fitting; + int chosen_stages = max_fitting; + while (chosen_stages >= 3 and not fits(chosen_stages)) + -- chosen_stages; + if (gran_k_b == 32 and expected_m <= 8) + chosen_stages = std::min(chosen_stages, 6); + DG_HOST_ASSERT(chosen_stages >= 3); + config.pipeline_config.num_stages = chosen_stages; + config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; + } + + // R2b-A swap_ab maps original N onto WGMMA M. Use enough math warpgroups + // to cover the 64-row WGMMA-M strips selected by BLOCK_N. + DG_HOST_ASSERT(config.layout.block_n == 64 or config.layout.block_n == 128 or config.layout.block_n == 256); + const int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; + config.launch_config.num_math_threads = rs_num_math_threads; + config.launch_config.num_threads = config.launch_config.num_tma_threads + rs_num_math_threads; + + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a.first, m, k, + config.storage_config.load_block_m, + config.layout.block_k, k, num_groups, + config.storage_config.swizzle_a_mode); + const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b_bytes, n, half_k, + config.layout.block_n, config.layout.block_k / 2, + half_k, num_groups, + config.storage_config.swizzle_b_mode); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, config.layout.block_k, num_groups, 0); + const auto tensor_map_d = make_tma_cd_desc(d, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), num_groups, + config.storage_config.swizzle_cd_mode); + + const bool scale_b_direct_load = gran_k_b == 32 and expected_m <= 16; + const bool k32_quad_reduce = gran_k_b == 32 and expected_m <= 8; + const bool small_m_simple_sched = + gran_k_b == 32 and expected_m <= 8 and + static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096; + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); + const SM90FP8FP4Gemm1D2DRSRuntime::Args& rs_args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .major_sfb = get_major_type_ab(sfb), + .scale_b_direct_load = scale_b_direct_load, + .k32_quad_reduce = k32_quad_reduce, + .small_m_simple_sched = small_m_simple_sched, + .scale_b_gran_k = static_cast(gran_k_b), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .sfb = sfb.data_ptr(), + .grouped_layout = masked_m.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .tensor_map_sfa = tensor_map_sfa, + }; + const auto code = SM90FP8FP4Gemm1D2DRSRuntime::generate(rs_args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_fused", code); + SM90FP8FP4Gemm1D2DRSRuntime::launch(runtime, rs_args); +} + } // namespace deep_gemm From 307455a9840b0eeec28f3e920b7a529f0e1e43ca Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:02:40 +0800 Subject: [PATCH 18/51] Create sm90_fp8_fp4_gemm_1d2d_rs.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 215 ++++++++++++++++++ 1 file changed, 215 insertions(+) create mode 100644 csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp new file mode 100644 index 0000000000..4cb03b9f4e --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -0,0 +1,215 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/layout.hpp" +#include "epilogue.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +// JIT runtime for the RS-mode SM90 FP8xFP4 1D2D kernel. +// +// S2 stage: identical wire-format to `SM90FP8FP4Gemm1D2DRuntime`, only the +// included header and instantiated kernel name differ. This gives the new +// kernel its own JIT artifact and entry point so subsequent steps (S3+) can +// safely diverge the device code without affecting the SS path or the +// contiguous host route. +class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + cute::UMMA::Major major_sfb; + bool scale_b_direct_load; + bool k32_quad_reduce; + bool small_m_simple_sched; + uint32_t scale_b_gran_k; + void *gmem_b_ptr; + void *gmem_d_ptr; + void *sfb; + void *grouped_layout; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + CUtensorMap tensor_map_sfa; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_fp4_gemm_1d2d_rs_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", + to_string(args.major_sfb), + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), + args.gemm_desc.num_groups, + args.gemm_config.layout.block_m, + args.gemm_config.layout.block_n, + args.gemm_config.layout.block_k, + args.gemm_config.storage_config.swizzle_a_mode, + args.gemm_config.storage_config.swizzle_b_mode, + args.gemm_config.storage_config.swizzle_cd_mode, + args.gemm_config.pipeline_config.num_stages, + args.gemm_config.launch_config.num_tma_threads, + args.gemm_config.launch_config.num_math_threads, + args.gemm_config.layout.get_cluster_size(), + args.gemm_config.layout.cluster_n > 1, + args.gemm_config.launch_config.num_sms, + to_string(args.gemm_desc.gemm_type), + get_default_epilogue_type(std::nullopt), + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + args.scale_b_direct_load ? "true" : "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + args.k32_quad_reduce ? "true" : "false", + "false", + args.k32_quad_reduce ? "true" : "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + args.small_m_simple_sched ? "true" : "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + "false", + 0, + "false", + args.scale_b_gran_k, + 1, + 1, + 0); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.gmem_b_ptr, args.sfb, args.grouped_layout, args.gmem_d_ptr, + args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.tensor_map_sfa)); + } +}; + +} // namespace deep_gemm From d7f3149e21a659696b1313632b421c127d56486c Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:03:30 +0800 Subject: [PATCH 19/51] Update smxx_layout.hpp --- csrc/jit_kernels/impls/smxx_layout.hpp | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 5d1f17b5b8..af207ebc95 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -99,11 +99,17 @@ static void __instantiate_kernel() {{ } }; -static std::tuple preprocess_sf(const torch::Tensor& sf) { +static std::tuple preprocess_sf(const torch::Tensor& sf, + const bool& require_float = true) { // NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA const auto dim = sf.dim(); DG_HOST_ASSERT(dim == 2 or dim == 3); - DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); + if (require_float) { + DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); + } else { + DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat or sf.scalar_type() == torch::kInt); + DG_HOST_ASSERT(sf.element_size() == sizeof(float)); + } const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf; const auto [num_groups, mn, sf_k] = get_shape<3>(batched_sf); @@ -112,7 +118,7 @@ static std::tuple preprocess_sf(const to } static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { - const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf); + const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf, false); // The last kernel already gives a column-major TMA aligned layout if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn) From 79a5c8c030cebb231a00f77635cf1674047bf996 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:04:14 +0800 Subject: [PATCH 20/51] Update layout.hpp --- csrc/utils/layout.hpp | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 928472d35a..780707d426 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,10 +1,12 @@ #pragma once +#include + #include #include -#include "math.hpp" -#include "exception.hpp" +#include "../utils/math.hpp" +#include "../utils/exception.hpp" #include "../jit/device_runtime.hpp" namespace deep_gemm { @@ -95,7 +97,12 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, if (num_groups.has_value()) DG_HOST_ASSERT(sf.size(-3) == num_groups.value()); DG_HOST_ASSERT(sf.size(-2) == ceil_div(mn, gran_mn)); - DG_HOST_ASSERT(sf.size(-1) == ceil_div(k, gran_k * (sf_dtype == torch::kFloat ? 1 : 4))); + const char* scale_b_lut_env = std::getenv("DG_W4_SCALE_B_LUT"); + const bool scale_b_lut = sf_dtype == torch::kInt and + scale_b_lut_env != nullptr and scale_b_lut_env[0] != '\0' and scale_b_lut_env[0] != '0'; + const int expected_sf_k = scale_b_lut ? + ceil_div(k, gran_k) * 2 : ceil_div(k, gran_k * (sf_dtype == torch::kFloat ? 1 : 4)); + DG_HOST_ASSERT(sf.size(-1) == expected_sf_k); // TMA stride checks: TMA aligned and MN-major if (tma_stride_check) { From 08d51f27baac25c5d18cec1ccab3e7343068f435 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:04:51 +0800 Subject: [PATCH 21/51] Update __init__.py --- deep_gemm/__init__.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a466b37ef9..8c4cad5cdc 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -86,6 +86,15 @@ ) except AttributeError: pass + try: + m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma + ) + m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma = ( + _C.m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma + ) + except AttributeError: + pass except ImportError: # Expected behavior for CUDA runtime version before 12.1 pass From 854268094a2f3fa45625387f980b86ca286ecf5c Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:05:52 +0800 Subject: [PATCH 22/51] Update sm90_fp8_fp4_gemm_1d2d.cuh --- .../impls/sm90_fp8_fp4_gemm_1d2d.cuh | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh index 0f3777ce50..dd0a0eaf83 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d.cuh @@ -82,15 +82,14 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); const uint32_t shape_k_scales = math::ceil_div(shape_k, BLOCK_K); const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); - // SFB cache aliases the smem_d region (BLOCK_M * BLOCK_N * sizeof(bf16) bytes) - // for the entire K loop. smem_d is only used during epilogue (after the K - // loop), and the cooperative SFB prefetch sync + final wgmma wait give us a - // clean handoff. We require sfb_bytes <= smem_d_bytes; with BLOCK_K=128 - // this reduces to BLOCK_M*128 >= shape_k*2 -> BLOCK_M*64 >= shape_k. The - // assert below verifies this at runtime; for typical kernels (BLOCK_M>=64, - // shape_k<=8192) this is always satisfied. + // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough + // smem_d capacity, so they fall back to a separate SFB region. const uint32_t smem_sfb_bytes = math::align(shape_k_scales * BLOCK_N * sizeof(float), 16u); - constexpr uint32_t smem_sfb_size = 0; // SFB aliases smem_d, no separate allocation + constexpr uint32_t COMPILED_SHAPE_K_SCALES = SHAPE_K == 0 ? 0 : math::constexpr_ceil_div(SHAPE_K, BLOCK_K); + constexpr uint32_t COMPILED_SMEM_SFB_BYTES = + math::constexpr_align(COMPILED_SHAPE_K_SCALES * BLOCK_N * static_cast(sizeof(float)), 16u); + constexpr bool kUseSeparateSFB = SHAPE_K != 0 and COMPILED_SMEM_SFB_BYTES > SMEM_D_SIZE; + constexpr uint32_t SMEM_SFB_SIZE = kUseSeparateSFB ? COMPILED_SMEM_SFB_BYTES : 0; // NOTES: Make sure we have enough shared memory for WGMMA padding static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); @@ -152,16 +151,17 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); }); - // SFB cache aliases the smem_d region (smem_buffer base). Safe because - // smem_d is only used in the epilogue, after the K loop completes and - // wgmma writes are drained. - auto smem_sfb = reinterpret_cast(smem_buffer); - DG_TRAP_ONLY_DEVICE_ASSERT(smem_sfb_bytes <= SMEM_D_SIZE); + constexpr uint32_t SMEM_SFB_OFFSET = SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE; + // Prefer aliasing SFB onto smem_d. For small-M tiles smem_d is too small, + // so allocate a separate SFB cache to enable BLOCK_M=64 masked kernels. + auto smem_sfb = reinterpret_cast(smem_buffer + (kUseSeparateSFB ? SMEM_SFB_OFFSET : 0)); + if constexpr (not kUseSeparateSFB) { + DG_TRAP_ONLY_DEVICE_ASSERT(smem_sfb_bytes <= SMEM_D_SIZE); + } - // Fill barriers (SFB does not allocate any extra space because it aliases smem_d). + // Fill barriers. // After the A/packed-B barrier merge there is only one set of full/empty barriers. - auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + - kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SFB_OFFSET + SMEM_SFB_SIZE); auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); @@ -199,6 +199,8 @@ sm90_fp8_fp4_gemm_1d2d_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, auto get_current_group_idx = [&]() { if constexpr (kGemmType == GemmType::MGroupedContiguous) { return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedMasked) { + return scheduler.current_group_idx; } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { return scheduler.current_group_idx; } else { From 698155ff09df97d73954b3ee4b019419de964ecc Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:07:21 +0800 Subject: [PATCH 23/51] Create sm90_fp8_fp4_gemm_1d2d_rs.cuh --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh | 3208 +++++++++++++++++ 1 file changed, 3208 insertions(+) create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh new file mode 100644 index 0000000000..24a14a6d8e --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh @@ -0,0 +1,3208 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +// SM90 FP8 x FP4 GEMM, RS-mode variant. +// +// S2 stage: this kernel is a verbatim copy of `sm90_fp8_fp4_gemm_1d2d_impl`, +// renamed to `sm90_fp8_fp4_gemm_1d2d_rs_impl` so it gets its own JIT artifact. +// Functionally identical to the SS variant; later steps will replace the +// inner WGMMA loop with RS-mode (register A) WGMMA + register-resident FP4 +// dequant (using `FP8MMASelectorRS`, ldmatrix, prmt+lop3) without touching +// the SS variant or the contiguous code path. + +// ============================================================================= +// S3 prep (M1): RS-mode utility helpers. +// These helpers are intentionally not yet called by `sm90_fp8_fp4_gemm_1d2d_rs_impl`; +// they are introduced as a self-contained, side-effect-free building block so +// the second-round kernel rewrite can plug them in without further header +// surgery. Numerical equivalence with the SS variant is preserved this round. +// ============================================================================= +namespace fp4_rs_detail { + +template +struct FinalAccumStorage { + using type = float; +}; + +template <> +struct FinalAccumStorage { + using type = nv_bfloat162; +}; + +template +CUTLASS_DEVICE void final_accum_init(dtype_t* base, uint32_t idx) { + if constexpr (kBF16FinalAccum) { + base[idx] = __float22bfloat162_rn({0.0f, 0.0f}); + } else { + base[idx] = 0.0f; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair(dtype_t* base, uint32_t pair_idx, + float scale_0, float value_0, + float scale_1, float value_1) { + if constexpr (kBF16PromoteMath) { + const nv_bfloat162 scale = __float22bfloat162_rn({scale_0, scale_1}); + const nv_bfloat162 value = __float22bfloat162_rn({value_0, value_1}); + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + const nv_bfloat162 out = __hfma2(scale, value, dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = out; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(out); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0, value_0, dst.x); + dst.y = __fmaf_rn(scale_1, value_1, dst.y); + } else { + dst.x += scale_0 * value_0; + dst.y += scale_1 * value_1; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + if constexpr (kFmaPromote) { + base[idx + 0] = __fmaf_rn(scale_0, value_0, base[idx + 0]); + base[idx + 1] = __fmaf_rn(scale_1, value_1, base[idx + 1]); + } else { + base[idx + 0] += scale_0 * value_0; + base[idx + 1] += scale_1 * value_1; + } + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair2(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b) { + if constexpr (kBF16PromoteMath) { + const nv_bfloat162 scale_a = __float22bfloat162_rn({scale_0_a, scale_1_a}); + const nv_bfloat162 value_a = __float22bfloat162_rn({value_0_a, value_1_a}); + const nv_bfloat162 scale_b = __float22bfloat162_rn({scale_0_b, scale_1_b}); + const nv_bfloat162 value_b = __float22bfloat162_rn({value_0_b, value_1_b}); + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(scale_a, value_a, dst); + dst = __hfma2(scale_b, value_b, dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0_a, value_0_a, dst.x); + dst.y = __fmaf_rn(scale_1_a, value_1_a, dst.y); + dst.x = __fmaf_rn(scale_0_b, value_0_b, dst.x); + dst.y = __fmaf_rn(scale_1_b, value_1_b, dst.y); + } else { + dst.x += scale_0_a * value_0_a; + dst.y += scale_1_a * value_1_a; + dst.x += scale_0_b * value_0_b; + dst.y += scale_1_b * value_1_b; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(scale_0_a, value_0_a, dst_0); + dst_1 = __fmaf_rn(scale_1_a, value_1_a, dst_1); + dst_0 = __fmaf_rn(scale_0_b, value_0_b, dst_0); + dst_1 = __fmaf_rn(scale_1_b, value_1_b, dst_1); + } else { + dst_0 += scale_0_a * value_0_a; + dst_1 += scale_1_a * value_1_a; + dst_0 += scale_0_b * value_0_b; + dst_1 += scale_1_b * value_1_b; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b, + float scale_0_c, float value_0_c, + float scale_1_c, float value_1_c, + float scale_0_d, float value_0_d, + float scale_1_d, float value_1_d) { + if constexpr (kBF16PromoteMath) { + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(__float22bfloat162_rn({scale_0_a, scale_1_a}), + __float22bfloat162_rn({value_0_a, value_1_a}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_b, scale_1_b}), + __float22bfloat162_rn({value_0_b, value_1_b}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_c, scale_1_c}), + __float22bfloat162_rn({value_0_c, value_1_c}), dst); + dst = __hfma2(__float22bfloat162_rn({scale_0_d, scale_1_d}), + __float22bfloat162_rn({value_0_d, value_1_d}), dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(scale_0_a, value_0_a, dst.x); + dst.y = __fmaf_rn(scale_1_a, value_1_a, dst.y); + dst.x = __fmaf_rn(scale_0_b, value_0_b, dst.x); + dst.y = __fmaf_rn(scale_1_b, value_1_b, dst.y); + dst.x = __fmaf_rn(scale_0_c, value_0_c, dst.x); + dst.y = __fmaf_rn(scale_1_c, value_1_c, dst.y); + dst.x = __fmaf_rn(scale_0_d, value_0_d, dst.x); + dst.y = __fmaf_rn(scale_1_d, value_1_d, dst.y); + } else { + dst.x += scale_0_a * value_0_a; + dst.y += scale_1_a * value_1_a; + dst.x += scale_0_b * value_0_b; + dst.y += scale_1_b * value_1_b; + dst.x += scale_0_c * value_0_c; + dst.y += scale_1_c * value_1_c; + dst.x += scale_0_d * value_0_d; + dst.y += scale_1_d * value_1_d; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(scale_0_a, value_0_a, dst_0); + dst_1 = __fmaf_rn(scale_1_a, value_1_a, dst_1); + dst_0 = __fmaf_rn(scale_0_b, value_0_b, dst_0); + dst_1 = __fmaf_rn(scale_1_b, value_1_b, dst_1); + dst_0 = __fmaf_rn(scale_0_c, value_0_c, dst_0); + dst_1 = __fmaf_rn(scale_1_c, value_1_c, dst_1); + dst_0 = __fmaf_rn(scale_0_d, value_0_d, dst_0); + dst_1 = __fmaf_rn(scale_1_d, value_1_d, dst_1); + } else { + dst_0 += scale_0_a * value_0_a; + dst_1 += scale_1_a * value_1_a; + dst_0 += scale_0_b * value_0_b; + dst_1 += scale_1_b * value_1_b; + dst_0 += scale_0_c * value_0_c; + dst_1 += scale_1_c * value_1_c; + dst_0 += scale_0_d * value_0_d; + dst_1 += scale_1_d * value_1_d; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4x2(dtype_t* base, uint32_t pair_idx, + float scale_0_a, float value_0_a, + float scale_1_a, float value_1_a, + float scale_0_b, float value_0_b, + float scale_1_b, float value_1_b, + float scale_0_c, float value_0_c, + float scale_1_c, float value_1_c, + float scale_0_d, float value_0_d, + float scale_1_d, float value_1_d, + float scale_2_a, float value_2_a, + float scale_3_a, float value_3_a, + float scale_2_b, float value_2_b, + float scale_3_b, float value_3_b, + float scale_2_c, float value_2_c, + float scale_3_c, float value_3_c, + float scale_2_d, float value_2_d, + float scale_3_d, float value_3_d) { + final_accum_promote_pair4( + base, pair_idx, + scale_0_a, value_0_a, scale_1_a, value_1_a, + scale_0_b, value_0_b, scale_1_b, value_1_b, + scale_0_c, value_0_c, scale_1_c, value_1_c, + scale_0_d, value_0_d, scale_1_d, value_1_d); + final_accum_promote_pair4( + base, pair_idx + 1, + scale_2_a, value_2_a, scale_3_a, value_3_a, + scale_2_b, value_2_b, scale_3_b, value_3_b, + scale_2_c, value_2_c, scale_3_c, value_3_c, + scale_2_d, value_2_d, scale_3_d, value_3_d); +} + +CUTLASS_DEVICE void keep_float_live(float value) { + asm volatile("" :: "f"(value) : "memory"); +} + +struct FP8MMAF16AccumM64N8K32RS { + static constexpr uint32_t M = 64; + static constexpr uint32_t N = 8; + static constexpr uint32_t K = 32; + static constexpr uint32_t kNumAccum = 2; + + template + CUTLASS_DEVICE static void wgmma(uint32_t const* a, GmmaDescriptor const& desc, + uint32_t* d, bool scale_d) { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %7, 0;\n" + " wgmma.mma_async.sync.aligned.m64n8k32.f16.e4m3.e4m3 " + " {%0, %1}, {%2, %3, %4, %5}, %6, p, 1, 1;\n" + "}\n" + : "+r"(d[0]), "+r"(d[1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "l"(desc.desc_), "r"(static_cast(scale_d))); + } +}; + +CUTLASS_DEVICE void unpack_f16_accum_m64n8(uint32_t const* src, float* dst) { + #pragma unroll + for (uint32_t i = 0; i < FP8MMAF16AccumM64N8K32RS::kNumAccum; ++ i) { + const half2 value_h = *reinterpret_cast(src + i); + const float2 value_f = __half22float2(value_h); + dst[i * 2 + 0] = value_f.x; + dst[i * 2 + 1] = value_f.y; + } +} + +CUTLASS_DEVICE float scale_float_by_pow2(float value, float pow2_scale) { + uint32_t value_bits = *reinterpret_cast(&value); + const uint32_t scale_bits = *reinterpret_cast(&pow2_scale); + const int32_t exp_shift = static_cast((scale_bits >> 23) & 0xffu) - 127; + value_bits = static_cast(static_cast(value_bits) + (exp_shift << 23)); + return *reinterpret_cast(&value_bits); +} + +template +CUTLASS_DEVICE void final_accum_promote_pair4_split_scale(dtype_t* base, uint32_t pair_idx, + float scale_a_0, float scale_a_1, + float scale_b_a, float scale_b_b, + float scale_b_c, float scale_b_d, + float value_0_a, float value_1_a, + float value_0_b, float value_1_b, + float value_0_c, float value_1_c, + float value_0_d, float value_1_d) { + auto prod_0 = [&](float scale_b) { + return kScaleBPow2Promote ? scale_float_by_pow2(scale_a_0, scale_b) : scale_a_0 * scale_b; + }; + auto prod_1 = [&](float scale_b) { + return kScaleBPow2Promote ? scale_float_by_pow2(scale_a_1, scale_b) : scale_a_1 * scale_b; + }; + if constexpr (kBF16PromoteMath) { + nv_bfloat162 dst; + if constexpr (kBF16FinalAccum) { + dst = base[pair_idx]; + } else { + const uint32_t idx = pair_idx * 2; + dst = __float22bfloat162_rn({base[idx + 0], base[idx + 1]}); + } + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_a), prod_1(scale_b_a)}), + __float22bfloat162_rn({value_0_a, value_1_a}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_b), prod_1(scale_b_b)}), + __float22bfloat162_rn({value_0_b, value_1_b}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_c), prod_1(scale_b_c)}), + __float22bfloat162_rn({value_0_c, value_1_c}), dst); + dst = __hfma2(__float22bfloat162_rn({prod_0(scale_b_d), prod_1(scale_b_d)}), + __float22bfloat162_rn({value_0_d, value_1_d}), dst); + if constexpr (kBF16FinalAccum) { + base[pair_idx] = dst; + } else { + const uint32_t idx = pair_idx * 2; + const float2 out_f = __bfloat1622float2(dst); + base[idx + 0] = out_f.x; + base[idx + 1] = out_f.y; + } + } else if constexpr (kBF16FinalAccum) { + float2 dst = __bfloat1622float2(base[pair_idx]); + if constexpr (kFmaPromote) { + dst.x = __fmaf_rn(prod_0(scale_b_a), value_0_a, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_a), value_1_a, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_b), value_0_b, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_b), value_1_b, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_c), value_0_c, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_c), value_1_c, dst.y); + dst.x = __fmaf_rn(prod_0(scale_b_d), value_0_d, dst.x); + dst.y = __fmaf_rn(prod_1(scale_b_d), value_1_d, dst.y); + } else { + dst.x += prod_0(scale_b_a) * value_0_a; + dst.y += prod_1(scale_b_a) * value_1_a; + dst.x += prod_0(scale_b_b) * value_0_b; + dst.y += prod_1(scale_b_b) * value_1_b; + dst.x += prod_0(scale_b_c) * value_0_c; + dst.y += prod_1(scale_b_c) * value_1_c; + dst.x += prod_0(scale_b_d) * value_0_d; + dst.y += prod_1(scale_b_d) * value_1_d; + } + base[pair_idx] = __float22bfloat162_rn(dst); + } else { + const uint32_t idx = pair_idx * 2; + float dst_0 = base[idx + 0]; + float dst_1 = base[idx + 1]; + if constexpr (kFmaPromote) { + dst_0 = __fmaf_rn(prod_0(scale_b_a), value_0_a, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_a), value_1_a, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_b), value_0_b, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_b), value_1_b, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_c), value_0_c, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_c), value_1_c, dst_1); + dst_0 = __fmaf_rn(prod_0(scale_b_d), value_0_d, dst_0); + dst_1 = __fmaf_rn(prod_1(scale_b_d), value_1_d, dst_1); + } else { + dst_0 += prod_0(scale_b_a) * value_0_a; + dst_1 += prod_1(scale_b_a) * value_1_a; + dst_0 += prod_0(scale_b_b) * value_0_b; + dst_1 += prod_1(scale_b_b) * value_1_b; + dst_0 += prod_0(scale_b_c) * value_0_c; + dst_1 += prod_1(scale_b_c) * value_1_c; + dst_0 += prod_0(scale_b_d) * value_0_d; + dst_1 += prod_1(scale_b_d) * value_1_d; + } + base[idx + 0] = dst_0; + base[idx + 1] = dst_1; + } +} + +template +CUTLASS_DEVICE float final_accum_load_scalar(const dtype_t* base, uint32_t idx) { + if constexpr (kBF16FinalAccum) { + const float2 value = __bfloat1622float2(base[idx / 2]); + return (idx % 2 == 0) ? value.x : value.y; + } else { + return base[idx]; + } +} + +template +CUTLASS_DEVICE nv_bfloat162 final_accum_load_pair_bf16(const dtype_t* base, uint32_t pair_idx) { + if constexpr (kBF16FinalAccum) { + return base[pair_idx]; + } else { + return __float22bfloat162_rn({base[pair_idx * 2 + 0], base[pair_idx * 2 + 1]}); + } +} + +// `ldmatrix.sync.aligned.x4.m8n8.trans.shared.b16`: load 4 8x8 b16 matrices +// from swizzled smem into 4 b32 registers per lane, transposed. The transpose +// option is the key to feeding RS-mode WGMMA's A operand layout when the +// source tile is laid out K-major in smem (which is the only layout the +// dequant step can produce cheaply from packed FP4 input). +// +// Note: WGMMA RS-mode A still consumes 8-bit (E4M3) pairs, so we phrase the +// granularity as b16 ldmatrix and let the upper layer treat each b32 register +// as 4 packed E4M3 lanes. +struct SM90_U32x4_LDSM_T { + CUTLASS_DEVICE static void + copy(uint32_t& dst_0, uint32_t& dst_1, uint32_t& dst_2, uint32_t& dst_3, void* smem_src) { + asm volatile( + "ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst_0), "=r"(dst_1), "=r"(dst_2), "=r"(dst_3) + : "l"(__cvta_generic_to_shared(smem_src))); + } +}; + +template +CUTLASS_DEVICE uint32_t permute_col(const uint32_t row, const uint32_t col) { + constexpr uint32_t strd = 128 / (COL_BYTES < 128u ? COL_BYTES : 128u); + return ((col / NV) ^ (row % 8 / strd)) * NV; +} + +template +struct SM90_U32x2_STSM_T { + CUTLASS_DEVICE static void + copy(dtype_t src_0, dtype_t src_1, void* smem_dst) { + DG_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); + const uint32_t src[2] = {*reinterpret_cast(&src_0), *reinterpret_cast(&src_1)}; + asm volatile("stmatrix.sync.aligned.x2.m8n8.shared.b16.trans [%0], {%1, %2};\n" + :: "l"(__cvta_generic_to_shared(smem_dst)), "r"(src[0]), "r"(src[1])); + } +}; + +// Decode a single FP4 (E2M1) code (low 4 bits) to its E4M3 byte +// representation. Same byte-LUT formulation as the scalar path inside the SS +// kernel, hoisted to namespace scope so it can also be called from +// register-level dequant routines. +// mag {0..7} -> {0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c} +// sign bit copied from FP4 bit 3 to E4M3 bit 7 (FP4 -0 maps to E4M3 -0; +// WGMMA treats -0 as 0, so this is numerically safe and saves a branch). +CUTLASS_DEVICE uint32_t fp4_to_e4m3_byte(uint32_t code) { + constexpr uint32_t LUT_LO = 0x3c383000u; + constexpr uint32_t LUT_HI = 0x4c484440u; + const uint32_t mag = code & 0x07u; + const uint32_t mag_byte = __byte_perm(LUT_LO, LUT_HI, mag) & 0xffu; + const uint32_t sign = (code & 0x08u) << 4; // 0 or 0x80 + return mag_byte | sign; +} + +// Decode 8 packed FP4 codes (one 32-bit word holding 8 nibbles) into 8 E4M3 +// bytes packed into a single 64-bit value (low byte = code 0). +// +// This is the register-resident analogue of the SS variant's +// `fp4_pair_to_e4m3_pair` chained over 4 bytes. It is the building block the +// second-round kernel will use after `ldmatrix.x4.trans` to materialise A +// operand registers for RS-mode WGMMA without round-tripping through smem. +// +// Implementation kept straightforward (LUT per nibble). PTX-level +// vectorisation via `prmt`/`lop3` is left for a follow-up micro-optimisation +// in round 2 once the rest of the data path is verified correct. +CUTLASS_DEVICE uint64_t fp4x8_to_e4m3x8(uint32_t packed) { + uint64_t out = 0; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t nib = (packed >> (i * 4)) & 0x0fu; + out |= static_cast(fp4_to_e4m3_byte(nib)) << (i * 8); + } + return out; +} + +CUTLASS_DEVICE void fast_fp4_to_e4m3_convert(uint32_t outputs[2], uint32_t input) { + const uint64_t decoded = fp4x8_to_e4m3x8(input); + outputs[0] = static_cast(decoded); + outputs[1] = static_cast(decoded >> 32); +} + +CUTLASS_DEVICE uint32_t fp4x4_to_e4m3x4(uint32_t packed) { + constexpr uint32_t pos0 = 0x3c383000u; + constexpr uint32_t pos1 = 0x4c484440u; + const uint32_t lut_idx = packed & 0x7777u; + const uint32_t sign_shifted = packed << 4; + uint32_t mag_bytes, sign_bytes; + asm volatile( + "{\n" + " prmt .b32 %0, %3, %4, %2;\n" + " prmt .b32 %1, %5, %6, 0xd9c8;\n" + "}\n" + : "=r"(mag_bytes), "=r"(sign_bytes) + : "r"(lut_idx), "r"(pos0), "r"(pos1), "r"(sign_shifted), "r"(packed)); + uint32_t out; + asm volatile( + "{\n" + " lop3.b32 %0, %1, %2, 0x80808080, 0xf8;\n" + "}\n" + : "=r"(out) + : "r"(mag_bytes), "r"(sign_bytes)); + return out; +} + +CUTLASS_DEVICE int32_t pow2_scale_to_exp_shift(float scale) { + const uint32_t scale_bits = *reinterpret_cast(&scale); + return static_cast((scale_bits >> 23) & 0xffu) - 127; +} + +struct ScaledE4M3Lut { + uint32_t lo; + uint32_t hi; +}; + +CUTLASS_DEVICE ScaledE4M3Lut make_scaled_e4m3_lut(uint32_t exp_offset) { + // Random FP4 per-32 scales are overwhelmingly ceil-pow2(max(abs(x))/6) + // in {2^-1, 1}, which maps to exp_offset {5, 6}. Fast-path those to + // avoid the dynamic IMAD chain in the WGMMA issue loop. + if (exp_offset == 5u) + return {0x34302800u, 0x44403c38u}; + if (exp_offset == 6u) + return {0x3c383000u, 0x4c484440u}; + const uint32_t exp_offset_buffer1 = + exp_offset * 0x08080800u + (exp_offset ? 0xfffffc00u : 0u); + const uint32_t exp_offset_buffer2 = exp_offset * 0x08080808u; + constexpr uint32_t mantissa_lo = 0x0c080400u; + constexpr uint32_t mantissa_hi = 0x1c181410u; + return {mantissa_lo + exp_offset_buffer1, mantissa_hi + exp_offset_buffer2}; +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_lut(uint32_t packed, ScaledE4M3Lut lut) { + const uint32_t lut_idx = packed & 0x7777u; + const uint32_t sign_shifted = packed << 4; + uint32_t mantissa_bytes, sign_bytes; + asm volatile( + "{\n" + " prmt .b32 %0, %3, %4, %2;\n" + " prmt .b32 %1, %5, %6, 0xd9c8;\n" + "}\n" + : "=r"(mantissa_bytes), "=r"(sign_bytes) + : "r"(lut_idx), "r"(lut.lo), "r"(lut.hi), "r"(sign_shifted), "r"(packed)); + return (sign_bytes & 0x80808080u) | mantissa_bytes; +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_offset(uint32_t packed, uint32_t exp_offset) { + return fp4x4_to_scaled_e4m3x4_lut(packed, make_scaled_e4m3_lut(exp_offset)); +} + +CUTLASS_DEVICE uint32_t fp4x4_to_scaled_e4m3x4_humming(uint32_t packed_nibbles, uint32_t exp_offset) { + const uint32_t packed = + (packed_nibbles & 0x000fu) | + ((packed_nibbles & 0x00f0u) << 4) | + ((packed_nibbles & 0x0f00u) << 8) | + ((packed_nibbles & 0xf000u) << 12); + + const uint32_t exp_offset_buffer1 = exp_offset * 0x08080800u + (exp_offset ? 0xfffffc00u : 0u); + const uint32_t exp_offset_buffer2 = exp_offset * 0x08080808u; + const uint32_t exp_offsets0 = __byte_perm(exp_offset_buffer1, exp_offset_buffer2, packed); + const uint32_t exp_offsets1 = __byte_perm(exp_offset_buffer1, exp_offset_buffer2, packed >> 16); + + uint32_t scaled_mag; + asm volatile( + "{\n" + " lop3.b32 %0, %1, 0x80808080, %2, 0xf8;\n" + "}\n" + : "=r"(scaled_mag) + : "r"(packed << 4), "r"((packed & 0x07070707u) << 2)); + return scaled_mag + __byte_perm(exp_offsets0, exp_offsets1, 0x6420); +} + +} // namespace fp4_rs_detail + +template +CUTLASS_DEVICE void dispatch_num_former_iters_rs(uint32_t num_former_iters, const func_t& func) { + if (num_former_iters == kNumFormerIters) { + func(cute::Int{}); + return; + } + + if constexpr (kNumFormerIters + kGap <= kEnd) + dispatch_num_former_iters_rs(num_former_iters, func); +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, kLaunchBoundsMinBlocks) void +sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, + nv_bfloat16* gmem_d_ptr, + uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, + const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_d, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__) + // Scaling checks + DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); + DG_STATIC_ASSERT(kScaleBGranK == 32 or kScaleBGranK == 128, + "Only support per-32 or per-128-channel FP4 scaling"); + DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or (kFuseScaleBDecode and kScaleBGranK == 32), + "Packed UE8M0 SFB is only supported by fused per-32 scale_b decode"); + DG_STATIC_ASSERT(kScaleBGranK == 128 or not kOverlapPromote, + "DG_W4_OVERLAP_PROMOTE does not support per-32 FP4 scaling"); + DG_STATIC_ASSERT(kScaleBGranK == 128 or kScaleKGroup == 1, + "DG_W4_SCALE_K_GROUP does not support per-32 FP4 scaling"); + DG_STATIC_ASSERT(kScaleKGroup == 1 or kScaleKGroup == 2 or kScaleKGroup == 4, + "DG_W4_SCALE_K_GROUP only supports 1/2/4"); + DG_STATIC_ASSERT(kFuseScaleBDecodeAssumeExp == 0 or kFuseScaleBDecodeAssumeExp == 5 or + kFuseScaleBDecodeAssumeExp == 6, + "DG_W4_FUSE_SCALE_B_DECODE_ASSUME_EXP only supports 0/5/6"); + DG_STATIC_ASSERT( + math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or + (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); + // Types + using WGMMA = typename mma::sm90::FP8MMASelectorRS::type; + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(BLOCK_M <= 256, "Invalid RS-mode MMA N size"); + DG_STATIC_ASSERT(BLOCK_N % WGMMA::M == 0, "RS-mode swap_ab requires BLOCK_N to be tiled by 64-row WGMMA M"); + + // Overwrite shape constants if the compiler gives + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + + // Shared memory + static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0); + static constexpr uint32_t SMEM_D_ROWS = BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M; + static constexpr uint32_t SMEM_D_SIZE = math::constexpr_align(SMEM_D_ROWS * BLOCK_N * static_cast(sizeof(__nv_bfloat16)), 1024u); + static constexpr uint32_t SMEM_A_TMA_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = + (BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M) * BLOCK_K * sizeof(__nv_fp8_e4m3); + // Packed FP4 B is loaded by TMA into a separate buffer; each row is BLOCK_K / 2 bytes. + static constexpr uint32_t BLOCK_K_PACKED = BLOCK_K / 2; + static constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = BLOCK_N * BLOCK_K_PACKED; + static constexpr uint32_t kNumRSMathWGs = kNumMathThreads / 128; + static constexpr uint32_t SMEM_SFA_TMA_SIZE_PER_STAGE = BLOCK_M * sizeof(float); + static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = + (BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M) * sizeof(float); + static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = + math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + const uint32_t shape_k_scales_a = math::ceil_div(shape_k, BLOCK_K); + const uint32_t shape_k_scales_b = math::ceil_div(shape_k, kScaleBGranK); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); + // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough + // smem_d capacity, so they fall back to a separate SFB region. + const uint32_t smem_sfb_bytes = kScaleBGranK == 32 ? + math::align((BLOCK_K / kScaleBGranK) * + (kFuseScaleBDecode ? math::ceil_div(BLOCK_N, 4u) : BLOCK_N) * + sizeof(float), 16u) : + math::align(shape_k_scales_b * BLOCK_N * sizeof(float), 16u); + // NOTES: Make sure we have enough shared memory for WGMMA padding + static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); + DG_STATIC_ASSERT(WGMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE, "Memory Out of bound for WGMMA"); + + // Configs + const uint32_t num_total_k_blocks = math::ceil_div(shape_k, BLOCK_K); + const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + const uint32_t lane_idx = ptx::get_lane_idx(); + constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); + constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; + constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; + constexpr bool kParallelNWavesEnabled = + kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; + DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, + "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); + constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + constexpr uint32_t kEmptyBarrierMathWarps = + kParallelNWavesEnabled ? kBaseWGMMAStoreThreads / 32 : kNumMathThreads / 32; + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_d); + } + __syncwarp(); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + // Data on shared memory + auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_B_PACKED_OFFSET = SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE; + auto smem_b_packed = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_B_PACKED_OFFSET + i * SMEM_B_PACKED_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SF_OFFSET = SMEM_B_PACKED_OFFSET + kNumStages * SMEM_B_PACKED_SIZE_PER_STAGE; + auto smem_sfa = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_SF_OFFSET + i * ALIGNED_SMEM_SFA_SIZE_PER_STAGE); + }); + constexpr uint32_t SMEM_SFB_OFFSET = SMEM_SF_OFFSET + kNumStages * ALIGNED_SMEM_SFA_SIZE_PER_STAGE; + // Prefer aliasing SFB onto smem_d. Small-M tiles usually have too little + // smem_d, and SHAPE_K can be runtime-only, so choose the separate SFB region + // dynamically to match the host-side smem_size calculation. + const bool use_separate_sfb = smem_sfb_bytes > SMEM_D_SIZE; + auto smem_sfb = reinterpret_cast(smem_buffer + (use_separate_sfb ? SMEM_SFB_OFFSET : 0)); + auto smem_sfb_exp = reinterpret_cast(smem_sfb); + + // Fill barriers. + // After the A/packed-B barrier merge there is only one set of full/empty barriers. + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_SFB_OFFSET + (use_separate_sfb ? smem_sfb_bytes : 0)); + auto full_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; }); + auto empty_barriers = utils::PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; }); + + // Initialize barriers + DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast"); + if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) { + // NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster, + // even with TMA multicast disabled, we want to make the behavior aligned + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(kNumTMAMulticast * kEmptyBarrierMathWarps); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + + // Synchronize all threads to make barrier visible in normal memory model + (kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads(); + + // Register reconfigurations + constexpr uint32_t kNumTMARegisters = 40; + constexpr uint32_t kDefaultMathRegisters = kNumMathThreads == 128 ? 248 : 232; + constexpr uint32_t kNumMathRegisters = kMathRegCap == 0 ? kDefaultMathRegisters : kMathRegCap; + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // `gmem_b_ptr` is no longer used: B is now loaded by TMA into `smem_b_packed`. + (void)gmem_b_ptr; + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + uint32_t simple_sched_linear_idx = blockIdx.x; + constexpr bool kUseSmallMSimpleSched = + kSmallMSimpleSched and kGemmType == GemmType::MGroupedMasked and BLOCK_M <= 8 and kNumTMAMulticast == 1; + auto get_next_block = [&]() { + if constexpr (kUseSmallMSimpleSched) { + const uint32_t n_blocks = math::ceil_div(shape_n, BLOCK_N); + const uint32_t total_blocks = n_blocks * kNumGroups; + while (simple_sched_linear_idx < total_blocks) { + scheduler.current_group_idx = simple_sched_linear_idx / n_blocks; + n_block_idx = simple_sched_linear_idx - scheduler.current_group_idx * n_blocks; + m_block_idx = 0; + simple_sched_linear_idx += gridDim.x; + if (scheduler.is_computation_valid(m_block_idx, 0)) + return true; + } + return false; + } else { + return scheduler.get_next_block(m_block_idx, n_block_idx); + } + }; + auto get_current_group_idx = [&]() { + if constexpr (kGemmType == GemmType::MGroupedContiguous) { + return static_cast(cute::max(0, grouped_layout[m_block_idx * BLOCK_M])); + } else if constexpr (kGemmType == GemmType::MGroupedMasked) { + return scheduler.current_group_idx; + } else if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + return scheduler.current_group_idx; + } else { + return 0u; + } + }; + + // Pipeline and TMA phases (single shared pipeline for A/SFA/packed-B after the barrier merge) + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + if (warp_idx >= kNumMathThreads / 32) { + // TMA warp-group for loading data + cutlass::arch::warpgroup_reg_dealloc(); + + // NOTES: only one thread (or warp) will be used. + // We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32`. + if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { + // Persistently schedule over blocks + while (get_next_block()) { + // Assign TMA multicast number into A and B + // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. + const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); + const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast"); + + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release for the (now merged) A/SFA/packed-B slot + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Issue TMA A + constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); + const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); + + constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked; + auto& full_barrier = *full_barriers[stage_idx]; + const uint32_t k_idx = k_block_idx * BLOCK_K; + tma::copy(&tensor_map_a, &full_barrier, + smem_a[stage_idx], k_idx, scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx), + num_tma_multicast_a, batch_idx); + if constexpr (not kScaleAStub) { + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales_a, 1, k_block_idx), + num_tma_multicast_a); + } + + if constexpr (not kWeightStub) { + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier. + const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; + tma::copy(&tensor_map_b, &full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), + k_idx_packed, + scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b, batch_idx); + } + + constexpr uint32_t kExpectedTxBytes = SMEM_A_TMA_SIZE_PER_STAGE + + (kWeightStub ? 0 : SMEM_B_PACKED_SIZE_PER_STAGE) + + (kScaleAStub ? 0 : SMEM_SFA_TMA_SIZE_PER_STAGE); + full_barrier.arrive_and_expect_tx(kExpectedTxBytes); + } + } + + // To safely deconstruct distributed shared barriers, we need another round of empty waits + if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; ++ i) { + empty_barriers[stage_idx]->wait(phase ^ 1); + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + } + } + } + } else { + // Math warp-groups for WGMMA + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers + const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const auto row_idx = lane_idx / 4, col_idx = lane_idx % 4; + const auto warp_in_wg = warp_idx % 4; + + auto a_desc = mma::sm90::make_smem_desc(smem_a[0], 1); + const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0); + + constexpr uint32_t kLdmatrixVecBytes = 16 / sizeof(__nv_fp8_e4m3); + + // Persistently schedule over blocks + while (get_next_block()) { + const uint32_t current_group_idx = get_current_group_idx(); + + if constexpr (kScaleBGranK == 128) { + // Cooperatively prefetch the SFB tile for this block from gmem to smem. + // Layout in smem: [shape_k_scales_b, BLOCK_N] (k outer, n inner). + // Out-of-bound n is filled with 1.0f to keep `n_idx >= shape_n` neutral. + // + // Optimization: use cp.async to copy gmem->smem directly (no register + // round-trip). For MN-major (n innermost in gmem) we can issue 16-byte + // (float4) cp.async per thread, cutting #instructions by 4x. K-major + // and the OOB tail use scalar 4-byte cp.async / st.shared. + const uint32_t n_block_base = n_block_idx * BLOCK_N; + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = shape_k_scales_b * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + // Issue cp.async (16B per thread) for fully in-bounds vectors; + // fall back to scalar st.shared with 1.0f-fill for the tail. + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_idx = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + k_idx * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + // Tail: at least one element is OOB; use scalar st with 1.0f fill. + float4 vals; + const float* sfb_row = sfb_base + k_idx * aligned_shape_n_sfb + n_idx; + vals.x = (n_idx + 0 < shape_n) ? *(sfb_row + 0) : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? *(sfb_row + 1) : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? *(sfb_row + 2) : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? *(sfb_row + 3) : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + // K-major: sfb is strided along n; cannot easily vectorize across n. + // Use scalar 4B cp.async for in-bounds, st.shared with 1.0f for OOB. + const uint32_t total = shape_k_scales_b * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_idx = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_idx * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_idx; + ptx::cp_async_4(smem_dst, gmem_src); + } + } + } + // Commit and wait for all cp.async issued above; pair with the + // existing NamedBarrier so the smem cache is visible to all + // math warps before they enter the K loop. + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + + auto cache_sfb_k32 = [&](uint32_t k_block_idx) { + if constexpr (kScaleBGranK == 32 and not kScaleBStub and not kFuseScaleBDecodeStub and not kScaleBDirectLoad) { + const uint32_t n_block_base = n_block_idx * BLOCK_N; + const uint32_t scale_k_base = k_block_idx * (BLOCK_K / kScaleBGranK); + constexpr uint32_t kScaleRows = BLOCK_K / kScaleBGranK; + if constexpr (kFuseScaleBDecode) { + DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or kMajorSFB == cute::UMMA::Major::MN, + "Packed UE8M0 SFB path expects MN-major transformed scale layout"); + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for packed K/32 SFB exp cache"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = kScaleRows * kVecsPerRow; + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_off = i / kVecsPerRow; + const uint32_t n_off = (i % kVecsPerRow) * kVec; + const uint32_t n_idx = n_block_base + n_off; + uint32_t packed_offsets = 0; + #pragma unroll + for (uint32_t j = 0; j < kVec; ++ j) { + uint32_t exp_offset = 6; + if constexpr (kScaleBPackedUE8M0) { + if (n_idx + j < shape_n) { + const uint32_t packed_shape_k_scales_b = math::ceil_div(shape_k_scales_b, 4u); + const uint32_t packed_k_idx = (scale_k_base + k_off) / 4u; + const uint32_t byte_idx = (scale_k_base + k_off) % 4u; + const uint32_t* sfb_packed = reinterpret_cast(sfb); + const uint32_t packed_scale = sfb_packed[ + current_group_idx * aligned_shape_n_sfb * packed_shape_k_scales_b + + packed_k_idx * aligned_shape_n_sfb + n_idx + j]; + const uint32_t e8m0_exp = (packed_scale >> (byte_idx * 8)) & 0xffu; + exp_offset = (e8m0_exp > 121u) ? (e8m0_exp - 121u) : 0u; + } + } else { + float val = 1.0f; + if (n_idx + j < shape_n) { + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + const float* ptr = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx + j; + val = *ptr; + } else { + const float* ptr = sfb + + current_group_idx * shape_n * shape_k_scales_b + + (n_idx + j) * shape_k_scales_b + scale_k_base + k_off; + val = *ptr; + } + } + const int32_t exp_shift = fp4_rs_detail::pow2_scale_to_exp_shift(val); + exp_offset = static_cast(exp_shift + 6) & 0xffu; + } + packed_offsets |= exp_offset << (j * 8); + } + ptx::st_shared(smem_sfb_exp + k_off * kVecsPerRow + n_off / kVec, packed_offsets); + } + } else if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + constexpr uint32_t kVec = 4; + DG_STATIC_ASSERT(BLOCK_N % kVec == 0, + "BLOCK_N must be a multiple of 4 for vectorized K/32 SFB load"); + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t total_vecs = kScaleRows * kVecsPerRow; + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { + const uint32_t k_off = i / kVecsPerRow; + const uint32_t vec_n = i % kVecsPerRow; + const uint32_t n_off = vec_n * kVec; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_off * BLOCK_N + n_off; + if (n_idx + kVec <= shape_n) { + const float* gmem_src = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx; + ptx::cp_async_16(smem_dst, gmem_src); + } else { + float4 vals; + const float* sfb_row = sfb_base + + (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx; + vals.x = (n_idx + 0 < shape_n) ? *(sfb_row + 0) : 1.0f; + vals.y = (n_idx + 1 < shape_n) ? *(sfb_row + 1) : 1.0f; + vals.z = (n_idx + 2 < shape_n) ? *(sfb_row + 2) : 1.0f; + vals.w = (n_idx + 3 < shape_n) ? *(sfb_row + 3) : 1.0f; + ptx::st_shared(reinterpret_cast(smem_dst), vals); + } + } + } else { + const uint32_t total = kScaleRows * BLOCK_N; + for (uint32_t i = threadIdx.x; i < total; i += kNumMathThreads) { + const uint32_t k_off = i / BLOCK_N; + const uint32_t n_off = i % BLOCK_N; + const uint32_t n_idx = n_block_base + n_off; + float* smem_dst = smem_sfb + k_off * BLOCK_N + n_off; + if (n_idx >= shape_n) { + ptx::st_shared(smem_dst, 1.0f); + } else { + const float* gmem_src = sfb + + current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + scale_k_base + k_off; + ptx::cp_async_4(smem_dst, gmem_src); + } + } + } + ptx::cp_async_commit_group(); + ptx::cp_async_wait_group<0>(); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0); + } + }; + + auto load_sfb = [&](uint32_t n_idx, uint32_t k_block_idx) { + if (n_idx >= shape_n) + return 1.0f; + if constexpr (kScaleBDirectLoad and kScaleBGranK == 32) { + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; + const float* ptr = sfb_base + k_block_idx * aligned_shape_n_sfb + n_idx; + return *ptr; + } else { + const float* ptr = sfb + current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_block_idx; + return *ptr; + } + } else if constexpr (kScaleBGranK == 32) { + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + const uint32_t k_off = k_block_idx % (BLOCK_K / kScaleBGranK); + return ptx::ld_shared(smem_sfb + k_off * BLOCK_N + n_off); + } else { + // SFB has been staged into smem above; out-of-bound `n_idx` already + // resolves to 1.0f because we wrote 1.0f for those slots. + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + return ptx::ld_shared(smem_sfb + k_block_idx * BLOCK_N + n_off); + } + }; + + auto load_sfb_exp_offset = [&](uint32_t n_idx, uint32_t k_block_idx) { + if (n_idx >= shape_n) + return uint32_t(6); + constexpr uint32_t kVec = 4; + constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; + const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; + const uint32_t k_off = k_block_idx % (BLOCK_K / kScaleBGranK); + const uint32_t packed_offsets = ptx::ld_shared(smem_sfb_exp + k_off * kVecsPerRow + n_off / kVec); + return (packed_offsets >> ((n_off % kVec) * 8)) & 0xffu; + }; + + // Decide the number of scales B to load + DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0); + uint32_t num_former_iters = BLOCK_N / 8; + if constexpr (not kMustUseUniformedScaleB) { + num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8; + } + + // Accumulation for WGMMA or CUDA promotion + constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; + DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); + constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; + constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; + constexpr bool kParallelNWavesEnabled = + kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; + constexpr bool kUseScaleKGroup = (kScaleKGroup > 1 and (WAVE_WGMMA == 1 or WAVE_WGMMA == 2)); + constexpr bool kUseScaleKGroupExact = + (kScaleKGroupExact and kUseScaleKGroup and kScaleKGroup == 2 and not kFusedPromote); + DG_STATIC_ASSERT(kNumMathThreads % 128 == 0, "RS-mode math threads must be whole warpgroups"); + DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, + "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); + const uint32_t wave_group_idx = kParallelNWavesEnabled ? math_wg_idx / kWGsPerNWave : 0; + const uint32_t wave_mwg_idx = kParallelNWavesEnabled ? math_wg_idx % kWGsPerNWave : math_wg_idx; + const uint32_t wave_warp_idx = wave_mwg_idx * 4 + warp_in_wg; + const uint32_t r_0 = wave_warp_idx * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + using final_accum_t = typename fp4_rs_detail::FinalAccumStorage::type; + constexpr uint32_t kNumAccumSets = kUseScaleKGroup ? WAVE_WGMMA : 1; + constexpr uint32_t kFinalAccumStride = + kBF16FinalAccum ? WGMMA::kNumAccum / 2 : WGMMA::kNumAccum; + constexpr uint32_t kNumFinalAccumRegs = kFinalAccumStride * WAVE_WGMMA; + constexpr bool kUseFinalAccumScratch = + kFinalAccumScratch and kDirectStore and kGemmType == GemmType::MGroupedMasked and + BLOCK_M < WGMMA::M and not kUseScaleKGroup; + constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); + constexpr uint32_t kNumWGMMAStoreThreads = + kParallelNWavesEnabled ? kBaseWGMMAStoreThreads * WAVE_WGMMA : kBaseWGMMAStoreThreads; + constexpr uint32_t kFinalAccumScratchBytes = + kNumMathThreads * kNumFinalAccumRegs * sizeof(final_accum_t); + float accum_storage[WGMMA::kNumAccum * kNumAccumSets]; + final_accum_t final_accum_regs[kUseFinalAccumScratch ? 1 : kNumFinalAccumRegs]; + final_accum_t* final_accum = final_accum_regs; + if constexpr (kUseFinalAccumScratch) { + DG_STATIC_ASSERT(kFinalAccumScratchBytes <= SMEM_D_SIZE, + "DG_W4_FINAL_ACCUM_SCRATCH needs more smem_d scratch space"); + final_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * kNumFinalAccumRegs; + } + #pragma unroll + for (uint32_t i = 0; i < kNumFinalAccumRegs; ++ i) + fp4_rs_detail::final_accum_init(final_accum, i); + + // Pick threads whose WGMMA results are to be stored in shared memory + DG_STATIC_ASSERT(BLOCK_N >= 64, "RS-mode swap_ab requires at least one 64-row compute tile"); + const bool do_wgmma_store = warp_idx < kNumWGMMAStoreThreads / 32; + + // Empty barrier arrival + auto empty_barrier_arrive_stage = [&](uint32_t target_stage) { + if constexpr (kParallelNWavesEnabled) + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 2); + if constexpr (kNumTMAMulticast == 1) { + (lane_idx == 0 and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? + empty_barriers[target_stage]->arrive() : void(); + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + (lane_idx < kNumTMAMulticast and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? + empty_barriers[target_stage]->arrive(target_cta) : void(); + } + }; + auto empty_barrier_arrive = [&]() { + empty_barrier_arrive_stage(stage_idx); + }; + + // Skip useless computations + const bool is_cta_computation_valid = scheduler.is_computation_valid(m_block_idx, 0); + if (is_cta_computation_valid) { + // The compiler must know the dynamic variable `num_former_iters`'s real value + constexpr bool kShouldOptimize = BLOCK_K / math::constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB; + constexpr uint32_t kGap = math::constexpr_gcd(BLOCK_K, BLOCK_N) / 8; + constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0; + + auto wait_stage = [&](uint32_t s, uint32_t p) { + full_barriers[s]->wait(p); + }; + + + // Dispatch `num_former_iters` and launch MMAs. + dispatch_num_former_iters_rs<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) { + constexpr uint32_t kAccumScratchBytes = kNumWGMMAStoreThreads * WGMMA::kNumAccum * sizeof(float); + if constexpr (kOverlapPromote and WAVE_WGMMA == 1 and kNumStages >= 2 and + not kWGMMAStub and kAccumScratchBytes <= SMEM_D_SIZE) { + auto accum = accum_storage; + auto smem_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * WGMMA::kNumAccum; + float scale_0_0_regs[WGMMA::kNumAccum / 4]; + float scale_1_0_regs[WGMMA::kNumAccum / 4]; + float scale_0_1_regs[WGMMA::kNumAccum / 4]; + float scale_1_1_regs[WGMMA::kNumAccum / 4]; + bool prev_valid = false; + + auto snapshot_accum = [&]() { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + smem_accum[i] = accum[i]; + asm volatile("" ::: "memory"); + }; + + auto promote_snapshot = [&]() { + if constexpr (not kPromoteStub) { + asm volatile("" ::: "memory"); + auto shifted_accum = final_accum; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_regs[i], smem_accum[i * 4 + 0], + scale_1_0_regs[i], smem_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_regs[i], smem_accum[i * 4 + 2], + scale_1_1_regs[i], smem_accum[i * 4 + 3]); + } + } + }; + + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks;) { + const uint32_t cur_stage = stage_idx; + const uint32_t cur_phase = phase; + const auto a_desc_base_lo = a_desc_lo + cur_stage * (SMEM_A_SIZE_PER_STAGE / 16); + wait_stage(cur_stage, cur_phase); + + if (not do_wgmma_store) { + empty_barrier_arrive_stage(cur_stage); + advance_pipeline(k_block_idx); + continue; + } + + constexpr uint32_t local_idx = 0; + constexpr uint32_t m_offset = 0; + const uint32_t lane_row = lane_idx / 4; + const uint32_t lane_col_pair = lane_idx % 4; + const uint32_t packed_shift = (lane_col_pair & 1u) * 16u; + const uint32_t lane_pair_col = (lane_col_pair & 2u) * 2u; + uint32_t rs_row_offset[4]; + uint32_t rs_raw_col[4]; + uint32_t rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + lane_row; + const uint32_t addr_tid_g = (warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + rs_raw_col[mat] = addr_t_col / 2 + lane_pair_col; + rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + auto load_a_regs = [&](uint32_t k, uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[cur_stage] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[cur_stage] + rs_row_offset[mat] + swizzled_col)); + } + const uint32_t packed_shifted = packed_word >> packed_shift; + if constexpr (kDecodeStub) { + a_regs[mat] = 0x38383838u; + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } + }; + + uint32_t a_regs[4]; + load_a_regs(0, a_regs); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, k); + asm volatile("" ::: "memory"); + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + + if (prev_valid) { + promote_snapshot(); + } + + const uint32_t compute_n_0 = n_block_idx * BLOCK_N + r_0; + const uint32_t compute_n_1 = n_block_idx * BLOCK_N + r_1; + float scale_b_0 = 1.0f; + float scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + scale_b_0 = load_sfb(compute_n_0, k_block_idx); + scale_b_1 = load_sfb(compute_n_1, k_block_idx); + } + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[cur_stage] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[cur_stage] + m_idx + 1); + } + auto scale_product = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else { + return scale_a * scale_b; + } + }; + scale_0_0_regs[i] = scale_product(scale_a_0, scale_b_0); + scale_1_0_regs[i] = scale_product(scale_a_1, scale_b_0); + scale_0_1_regs[i] = scale_product(scale_a_0, scale_b_1); + scale_1_1_regs[i] = scale_product(scale_a_1, scale_b_1); + } + } + + ptx::warpgroup_wait<0>(); + empty_barrier_arrive_stage(cur_stage); + snapshot_accum(); + prev_valid = true; + advance_pipeline(k_block_idx); + } + + if (prev_valid) { + promote_snapshot(); + } + } else { + constexpr uint32_t kScaleSumStride = WGMMA::kNumAccum / 4; + float scale_0_0_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_1_0_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_0_1_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float scale_1_1_sum[(kUseScaleKGroup ? WAVE_WGMMA : 1) * kScaleSumStride]; + float accum_first_storage[kUseScaleKGroupExact ? WAVE_WGMMA * WGMMA::kNumAccum : 1]; + #pragma unroll 8 + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); + wait_stage(stage_idx, phase); + cache_sfb_k32(k_block_idx); + + // Small-N/BM=32 experiments may launch more math WGs than + // the current B tile has 64-row RS slices. Inactive WGs must + // still release the pipeline slot, but must not read B rows + // beyond BLOCK_N. + if (not do_wgmma_store) { + empty_barrier_arrive(); + continue; + } + + // TODO: remove some useless computation for unaligned Ms + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + auto accum = accum_storage + (kUseScaleKGroup ? local_idx * WGMMA::kNumAccum : 0); + auto m_offset = local_idx * WAVE_BLOCK_M; + const uint32_t scale_group_pos = + kUseScaleKGroup ? (k_block_idx % kScaleKGroup) : 0; + const bool should_promote_group = + (not kUseScaleKGroup) or (scale_group_pos == kScaleKGroup - 1) or + (k_block_idx + 1 == num_total_k_blocks); + + // Read scales before `warpgroup_arrive` so the next CTA cannot pollute shared memory. + // NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results + const uint32_t compute_n_0 = n_block_idx * BLOCK_N + m_offset + r_0; + const uint32_t compute_n_1 = n_block_idx * BLOCK_N + m_offset + r_1; + float scale_b_0 = 0.0f; + float scale_b_1 = 0.0f; + if constexpr (kScaleBGranK == 128) { + if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + if constexpr (kScaleBStub) { + scale_b_0 = 1.0f; + scale_b_1 = 1.0f; + } else { + scale_b_0 = load_sfb(compute_n_0, k_block_idx); + scale_b_1 = load_sfb(compute_n_1, k_block_idx); + } + } + } + float scale_a_0_regs[kLateScaleA ? 1 : WGMMA::kNumAccum / 4]; + float scale_a_1_regs[kLateScaleA ? 1 : WGMMA::kNumAccum / 4]; + if constexpr (kLateScaleA) { + // Keep SFA live in shared memory and load it in the promotion loop. + } else if constexpr (kScaleAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + scale_a_0_regs[i] = 1.0f; + scale_a_1_regs[i] = 1.0f; + } + } else if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0_regs[i] = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1_regs[i] = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } + + const uint32_t lane_row = lane_idx / 4; + const uint32_t lane_col_pair = lane_idx % 4; + const uint32_t packed_shift = (lane_col_pair & 1u) * 16u; + const uint32_t lane_pair_col = (lane_col_pair & 2u) * 2u; + uint32_t rs_row_offset[4]; + uint32_t rs_raw_col[4]; + uint32_t rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + lane_row; + const uint32_t addr_tid_g = (warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + rs_raw_col[mat] = addr_t_col / 2 + lane_pair_col; + rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + auto load_a_regs = [&](uint32_t k, uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + } + const uint32_t packed_shifted = packed_word >> packed_shift; + if constexpr (kDecodeStub) { + // Four E4M3 1.0 values. Keeps the LDS/address path but removes FP4 decode. + a_regs[mat] = 0x38383838u; + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } + }; + auto load_a_regs_scaled_b = [&](uint32_t k, + const uint32_t exp_offsets[2], + const fp4_rs_detail::ScaledE4M3Lut scaled_luts[2], + uint32_t a_regs[4]) { + if constexpr (kWeightStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = 0x38383838u; + return; + } + DG_STATIC_ASSERT(BLOCK_K_PACKED == 64 and kLdmatrixVecBytes == 16 and WGMMA::K / 2 == 16, + "RS decode address fast path assumes 64-byte packed K tile"); + auto load_packed_word = [&](uint32_t mat) { + const uint32_t swizzled_col = ((k ^ rs_swizzle_xor[mat]) << 4) + rs_raw_col[mat]; + uint32_t packed_word = 0; + if constexpr (kDecodePairShfl) { + if ((lane_col_pair & 1u) == 0) + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + packed_word = __shfl_sync(0xffffffff, packed_word, lane_idx & ~1u); + } else { + packed_word = ptx::ld_shared(reinterpret_cast( + smem_b_packed[stage_idx] + rs_row_offset[mat] + swizzled_col)); + } + return packed_word; + }; + if constexpr (kDecodeStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + static_cast(load_packed_word(mat)); + a_regs[mat] = 0x38383838u; + } + } else if constexpr (kBLoadStub) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) + a_regs[mat] = load_packed_word(mat); + } else if constexpr (kFuseScaleBDecodeAssumeExp == 6) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } else if constexpr (kFuseScaleBDecodeAssumeExp == 5) { + constexpr fp4_rs_detail::ScaledE4M3Lut kExp5Lut{0x34302800u, 0x44403c38u}; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, kExp5Lut); + } + } else { + if constexpr (kFuseScaleBDecodeFastCommon) { + const bool uniform_exp = exp_offsets[0] == exp_offsets[1]; + if (uniform_exp and exp_offsets[0] == 6u) { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_e4m3x4(packed_shifted); + } + } else if (uniform_exp and exp_offsets[0] == 5u) { + constexpr fp4_rs_detail::ScaledE4M3Lut kExp5Lut{0x34302800u, 0x44403c38u}; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, kExp5Lut); + } + } else { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + if constexpr (kFuseScaleBHummingDecode) { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_humming( + packed_shifted, exp_offsets[mat & 1u]); + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, scaled_luts[mat & 1u]); + } + } + } + } else { + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t packed_shifted = load_packed_word(mat) >> packed_shift; + if constexpr (kFuseScaleBHummingDecode) { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_humming( + packed_shifted, exp_offsets[mat & 1u]); + } else { + a_regs[mat] = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, scaled_luts[mat & 1u]); + } + } + } + } + }; + + // Decode the next RS-A fragment immediately after issuing the + // current WGMMA so decode/address work can overlap the async MMA. + uint32_t a_regs[4]; + if constexpr (not kFuseScaleBDecode) + load_a_regs(0, a_regs); + if constexpr (kScaleBGranK == 32) { + DG_STATIC_ASSERT(WGMMA::K == 32, + "per-32 FP4 scale path assumes each WGMMA K slice is one scale group"); + DG_STATIC_ASSERT(not kUseScaleKGroup, + "per-32 FP4 scale path does not support DG_W4_SCALE_K_GROUP"); + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + if constexpr (kFuseScaleBDecode) { + if constexpr (not kWGMMAStub) { + uint32_t scale_b_exp_offsets[BLOCK_K / WGMMA::K][2]; + fp4_rs_detail::ScaledE4M3Lut scale_b_luts[kFuseScaleBOnDemandLut ? 1 : (BLOCK_K / WGMMA::K)][2]; + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + const uint32_t n_idx = + n_block_idx * BLOCK_N + rs_row_offset[pair] / BLOCK_K_PACKED; + if constexpr (kScaleBStub) + scale_b_exp_offsets[kk][pair] = 6u; + else + scale_b_exp_offsets[kk][pair] = + load_sfb_exp_offset(n_idx, scale_b_k_idx); + } + } + if constexpr (not kFuseScaleBOnDemandLut and not kDecodeStub and not kBLoadStub) { + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + scale_b_luts[kk][pair] = + fp4_rs_detail::make_scaled_e4m3_lut(scale_b_exp_offsets[kk][pair]); + } + } + } + uint32_t staged_a_regs[BLOCK_K / WGMMA::K][4]; + uint32_t staged_pair_a_regs[2][4]; + uint32_t* staged_smem_a_regs = reinterpret_cast(smem_d) + + math::ceil_div(smem_sfb_bytes, 128u) * 32u + + (warp_idx * 32 + lane_idx) * (BLOCK_K / WGMMA::K) * 4u; + uint32_t* ws_decoded_regs = reinterpret_cast(smem_d) + + math::ceil_div(smem_sfb_bytes, 128u) * 32u; + auto load_a_regs_scaled_b_for_k = [&](uint32_t kk, uint32_t regs[4]) { + if constexpr (kFuseScaleBOnDemandLut) { + fp4_rs_detail::ScaledE4M3Lut on_demand_luts[2]; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) + on_demand_luts[pair] = + fp4_rs_detail::make_scaled_e4m3_lut(scale_b_exp_offsets[kk][pair]); + load_a_regs_scaled_b(kk, scale_b_exp_offsets[kk], on_demand_luts, regs); + } else { + load_a_regs_scaled_b(kk, scale_b_exp_offsets[kk], scale_b_luts[kk], regs); + } + }; + if constexpr (kFuseScaleBSlicePromote) { + DG_STATIC_ASSERT(not kFuseScaleBPredecode and not kFuseScaleBPredecodePair and + not kFuseScaleBSharedStage and not kFuseScaleBWSDecode, + "slice-promote is only for the direct fused decode path"); + DG_STATIC_ASSERT(not kPromoteFromSmem, + "slice-promote does not use promote-from-smem"); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + load_a_regs_scaled_b_for_k(k, a_regs); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + } + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (not kScaleBEarlyProduct) { + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, accum[i * 4 + 0], + scale_a_1, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, accum[i * 4 + 2], + scale_a_1, accum[i * 4 + 3]); + } + } + } + } else { + if constexpr (kFuseScaleBWSDecode) { + DG_STATIC_ASSERT(BLOCK_N == 256 and kNumMathThreads == 256, + "DG_W4_FUSE_SCALE_B_WS_DECODE only targets BN256/256-thread experiments"); + DG_STATIC_ASSERT(SMEM_D_SIZE >= 1024 + kNumMathThreads * + (BLOCK_K / WGMMA::K) * 4 * sizeof(uint32_t), + "smem_d is too small for WS decoded FP4 staging"); + if (warp_idx == 0) { + #pragma unroll + for (uint32_t target_tid = lane_idx; target_tid < kNumMathThreads; target_tid += 32) { + const uint32_t target_warp_idx = target_tid / 32; + const uint32_t target_lane_idx = target_tid % 32; + const uint32_t target_math_wg_idx = target_tid / 128; + const uint32_t target_warp_in_wg = target_warp_idx % 4; + const uint32_t target_wave_mwg_idx = + kParallelNWavesEnabled ? target_math_wg_idx % kWGsPerNWave : target_math_wg_idx; + const uint32_t target_lane_row = target_lane_idx / 4; + const uint32_t target_lane_col_pair = target_lane_idx % 4; + const uint32_t target_packed_shift = (target_lane_col_pair & 1u) * 16u; + const uint32_t target_lane_pair_col = (target_lane_col_pair & 2u) * 2u; + uint32_t target_rs_row_offset[4]; + uint32_t target_rs_raw_col[4]; + uint32_t target_rs_swizzle_xor[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++mat) { + const uint32_t addr_lane = mat * 8 + target_lane_row; + const uint32_t addr_tid_g = (target_warp_idx % 4) * 32 + addr_lane; + const uint32_t addr_t_row = (addr_tid_g & 15) | ((addr_tid_g >> 5) << 4); + const uint32_t addr_t_col = ((addr_tid_g >> 4) & 1) * kLdmatrixVecBytes; + const uint32_t src_row = + target_wave_mwg_idx * WGMMA::M + addr_t_row + m_offset; + target_rs_row_offset[mat] = src_row * BLOCK_K_PACKED; + target_rs_raw_col[mat] = addr_t_col / 2 + target_lane_pair_col; + target_rs_swizzle_xor[mat] = (addr_t_row & 7u) >> 1; + } + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + uint32_t target_exp_offsets[2]; + fp4_rs_detail::ScaledE4M3Lut target_luts[2]; + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + #pragma unroll + for (uint32_t pair = 0; pair < 2; ++ pair) { + const uint32_t n_idx = + n_block_idx * BLOCK_N + + target_rs_row_offset[pair] / BLOCK_K_PACKED; + if constexpr (kScaleBStub) + target_exp_offsets[pair] = 6u; + else + target_exp_offsets[pair] = + load_sfb_exp_offset(n_idx, scale_b_k_idx); + target_luts[pair] = + fp4_rs_detail::make_scaled_e4m3_lut(target_exp_offsets[pair]); + } + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + const uint32_t swizzled_col = + ((kk ^ target_rs_swizzle_xor[mat]) << 4) + + target_rs_raw_col[mat]; + const uint32_t packed_word = ptx::ld_shared( + reinterpret_cast( + smem_b_packed[stage_idx] + + target_rs_row_offset[mat] + swizzled_col)); + const uint32_t packed_shifted = packed_word >> target_packed_shift; + uint32_t decoded = 0x38383838u; + if constexpr (not kDecodeStub) { + decoded = fp4_rs_detail::fp4x4_to_scaled_e4m3x4_lut( + packed_shifted, target_luts[mat & 1u]); + } + ptx::st_shared(ws_decoded_regs + + target_tid * (BLOCK_K / WGMMA::K) * 4u + + kk * 4u + mat, + decoded); + } + } + } + } + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 1); + } else if constexpr (kFuseScaleBSharedStage) { + DG_STATIC_ASSERT(BLOCK_N == 256, + "DG_W4_FUSE_SCALE_B_SHARED_STAGE only targets BN256 experiments"); + DG_STATIC_ASSERT(SMEM_D_SIZE >= 1024 + kNumMathThreads * + (BLOCK_K / WGMMA::K) * 4 * sizeof(uint32_t), + "smem_d is too small for shared decoded FP4 staging"); + uint32_t tmp_a_regs[4]; + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + load_a_regs_scaled_b_for_k(kk, tmp_a_regs); + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + ptx::st_shared(staged_smem_a_regs + kk * 4 + mat, tmp_a_regs[mat]); + } + } + } else if constexpr (kFuseScaleBPredecode) { + #pragma unroll + for (uint32_t kk = 0; kk < BLOCK_K / WGMMA::K; ++ kk) { + load_a_regs_scaled_b_for_k(kk, staged_a_regs[kk]); + } + } else if constexpr (kFuseScaleBPredecodePair) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_FUSE_SCALE_B_PREDECODE_PAIR assumes four K=32 slices"); + load_a_regs_scaled_b_for_k(0, staged_pair_a_regs[0]); + load_a_regs_scaled_b_for_k(1, staged_pair_a_regs[1]); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + if constexpr (not kFuseScaleBPredecode and not kFuseScaleBPredecodePair and + not kFuseScaleBSharedStage and not kFuseScaleBWSDecode) { + load_a_regs_scaled_b_for_k(0, a_regs); + } + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + if constexpr (kFuseScaleBWSDecode) { + uint32_t ws_a_regs[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) { + ws_a_regs[mat] = ptx::ld_shared( + ws_decoded_regs + threadIdx.x * (BLOCK_K / WGMMA::K) * 4u + + k * 4u + mat); + } + WGMMA::wgmma(ws_a_regs, a_desc, accum, static_cast(k)); + } else if constexpr (kFuseScaleBSharedStage) { + uint32_t smem_a_regs[4]; + #pragma unroll + for (uint32_t mat = 0; mat < 4; ++ mat) + smem_a_regs[mat] = ptx::ld_shared(staged_smem_a_regs + k * 4 + mat); + WGMMA::wgmma(smem_a_regs, a_desc, accum, static_cast(k)); + } else if constexpr (kFuseScaleBPredecode) + WGMMA::wgmma(staged_a_regs[k], a_desc, accum, static_cast(k)); + else if constexpr (kFuseScaleBPredecodePair) + WGMMA::wgmma(staged_pair_a_regs[k & 1u], a_desc, accum, + static_cast(k)); + else + WGMMA::wgmma(a_regs, a_desc, accum, static_cast(k)); + asm volatile("" ::: "memory"); + if constexpr (kFuseScaleBPredecodePair) { + if (k + 2 < BLOCK_K / WGMMA::K) { + load_a_regs_scaled_b_for_k(k + 2, staged_pair_a_regs[k & 1u]); + } + } else if constexpr (not kFuseScaleBPredecode) { + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) { + load_a_regs_scaled_b_for_k(k + 1, a_regs); + } + } + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + } + } + if constexpr (kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + + if constexpr (not kPromoteStub) { + float* smem_promote_accum = reinterpret_cast(smem_d) + + (warp_idx * 32 + lane_idx) * WGMMA::kNumAccum; + if constexpr (kPromoteFromSmem) { + DG_STATIC_ASSERT(BLOCK_N == 256 and kNumMathThreads == 256, + "DG_W4_PROMOTE_FROM_SMEM only targets BN256/256-thread experiments"); + DG_STATIC_ASSERT(kNumMathThreads * WGMMA::kNumAccum * sizeof(float) <= SMEM_D_SIZE, + "smem_d is too small for promote-from-smem accum scratch"); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::st_shared(smem_promote_accum + i, accum[i]); + asm volatile("" ::: "memory"); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 0) : accum[i * 4 + 0], + scale_a_1, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 1) : accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 2) : accum[i * 4 + 2], + scale_a_1, kPromoteFromSmem ? ptx::ld_shared(smem_promote_accum + i * 4 + 3) : accum[i * 4 + 3]); + } + } + } else if constexpr (kFuseScaleBDecodeStub) { + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, static_cast(k)); + asm volatile("" ::: "memory"); + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + } else { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0, accum[i * 4 + 0], + scale_a_1, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0, accum[i * 4 + 2], + scale_a_1, accum[i * 4 + 3]); + } + } + } else { + if constexpr (kK32Pingpong) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_K32_PINGPONG currently assumes four K=32 slices"); + float accum_alt[WGMMA::kNumAccum]; + auto promote_k32_slice = [&](float* slice_accum, uint32_t k_slice) { + if constexpr (not kPromoteStub) { + float slice_scale_b_0 = 1.0f; + float slice_scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + const float prod_0_0 = scale_prod(scale_a_0, slice_scale_b_0); + const float prod_1_0 = scale_prod(scale_a_1, slice_scale_b_0); + const float prod_0_1 = scale_prod(scale_a_0, slice_scale_b_1); + const float prod_1_1 = scale_prod(scale_a_1, slice_scale_b_1); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod_0_0 + slice_accum[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod_1_0 + slice_accum[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod_0_1 + slice_accum[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod_1_1 + slice_accum[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + 1.0f, slice_accum[i * 4 + 0], + 1.0f, slice_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + 1.0f, slice_accum[i * 4 + 2], + 1.0f, slice_accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + prod_0_0, slice_accum[i * 4 + 0], + prod_1_0, slice_accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + prod_0_1, slice_accum[i * 4 + 2], + prod_1_1, slice_accum[i * 4 + 3]); + } + } + } + }; + auto promote_k32_pair = [&](float* slice_accum_0, float* slice_accum_1, + uint32_t k_slice_0, uint32_t k_slice_1) { + if constexpr (not kPromoteStub) { + float slice0_scale_b_0 = 1.0f; + float slice0_scale_b_1 = 1.0f; + float slice1_scale_b_0 = 1.0f; + float slice1_scale_b_1 = 1.0f; + if constexpr (not kScaleBStub) { + const uint32_t scale_b_k_idx_0 = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice_0; + const uint32_t scale_b_k_idx_1 = + k_block_idx * (BLOCK_K / kScaleBGranK) + k_slice_1; + slice0_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx_0); + slice0_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx_0); + slice1_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx_1); + slice1_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx_1); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice0_scale_b_0); + fp4_rs_detail::keep_float_live(slice0_scale_b_1); + fp4_rs_detail::keep_float_live(slice1_scale_b_0); + fp4_rs_detail::keep_float_live(slice1_scale_b_1); + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + const float prod0_0_0 = scale_prod(scale_a_0, slice0_scale_b_0); + const float prod0_1_0 = scale_prod(scale_a_1, slice0_scale_b_0); + const float prod0_0_1 = scale_prod(scale_a_0, slice0_scale_b_1); + const float prod0_1_1 = scale_prod(scale_a_1, slice0_scale_b_1); + const float prod1_0_0 = scale_prod(scale_a_0, slice1_scale_b_0); + const float prod1_1_0 = scale_prod(scale_a_1, slice1_scale_b_0); + const float prod1_0_1 = scale_prod(scale_a_0, slice1_scale_b_1); + const float prod1_1_1 = scale_prod(scale_a_1, slice1_scale_b_1); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod0_0_0 + slice_accum_0[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod0_1_0 + slice_accum_0[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod0_0_1 + slice_accum_0[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod0_1_1 + slice_accum_0[i * 4 + 3]); + fp4_rs_detail::keep_float_live(prod1_0_0 + slice_accum_1[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod1_1_0 + slice_accum_1[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod1_0_1 + slice_accum_1[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod1_1_1 + slice_accum_1[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 0, + 1.0f, slice_accum_0[i * 4 + 0], + 1.0f, slice_accum_0[i * 4 + 1], + 1.0f, slice_accum_1[i * 4 + 0], + 1.0f, slice_accum_1[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 1, + 1.0f, slice_accum_0[i * 4 + 2], + 1.0f, slice_accum_0[i * 4 + 3], + 1.0f, slice_accum_1[i * 4 + 2], + 1.0f, slice_accum_1[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 0, + prod0_0_0, slice_accum_0[i * 4 + 0], + prod0_1_0, slice_accum_0[i * 4 + 1], + prod1_0_0, slice_accum_1[i * 4 + 0], + prod1_1_0, slice_accum_1[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair2( + shifted_accum, i * 2 + 1, + prod0_0_1, slice_accum_0[i * 4 + 2], + prod0_1_1, slice_accum_0[i * 4 + 3], + prod1_0_1, slice_accum_1[i * 4 + 2], + prod1_1_1, slice_accum_1[i * 4 + 3]); + } + } + } + }; + auto load_k32_quad_scale_b = [&](float* scale_b_0, float* scale_b_1) { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = 1.0f; + scale_b_1[kk] = 1.0f; + } + if constexpr (not kScaleBStub) { + if constexpr ((kK32QuadScaleBInline or kK32QuadScaleBVec4) and + kScaleBDirectLoad and kScaleBGranK == 32) { + const uint32_t scale_b_k_base = + k_block_idx * (BLOCK_K / kScaleBGranK); + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + const float* sfb_base = sfb + + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b + + scale_b_k_base * aligned_shape_n_sfb; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = compute_n_0 < shape_n ? + *(sfb_base + kk * aligned_shape_n_sfb + compute_n_0) : 1.0f; + scale_b_1[kk] = compute_n_1 < shape_n ? + *(sfb_base + kk * aligned_shape_n_sfb + compute_n_1) : 1.0f; + } + } else { + const float* sfb_base = sfb + + current_group_idx * shape_n * shape_k_scales_b + scale_b_k_base; + if constexpr (kK32QuadScaleBVec4) { + if (compute_n_0 < shape_n) { + const float4 v0 = *reinterpret_cast( + sfb_base + compute_n_0 * shape_k_scales_b); + scale_b_0[0] = v0.x; + scale_b_0[1] = v0.y; + scale_b_0[2] = v0.z; + scale_b_0[3] = v0.w; + } + if (compute_n_1 < shape_n) { + const float4 v1 = *reinterpret_cast( + sfb_base + compute_n_1 * shape_k_scales_b); + scale_b_1[0] = v1.x; + scale_b_1[1] = v1.y; + scale_b_1[2] = v1.z; + scale_b_1[3] = v1.w; + } + } else { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + scale_b_0[kk] = compute_n_0 < shape_n ? + *(sfb_base + compute_n_0 * shape_k_scales_b + kk) : 1.0f; + scale_b_1[kk] = compute_n_1 < shape_n ? + *(sfb_base + compute_n_1 * shape_k_scales_b + kk) : 1.0f; + } + } + } + } else { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + kk; + scale_b_0[kk] = load_sfb(compute_n_0, scale_b_k_idx); + scale_b_1[kk] = load_sfb(compute_n_1, scale_b_k_idx); + } + } + } + }; + auto promote_k32_quad = [&](float* accum_0, float* accum_1, + float* accum_2, float* accum_3, + float* scale_b_0, float* scale_b_1) { + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + fp4_rs_detail::keep_float_live(scale_b_0[kk]); + fp4_rs_detail::keep_float_live(scale_b_1[kk]); + } + } + auto scale_prod = [&](float scale_a, float scale_b) { + if constexpr (kScaleProductStub) { + fp4_rs_detail::keep_float_live(scale_a + scale_b); + return 1.0f; + } else if constexpr (kScaleBMulStub or kPromoteMulStub) { + return scale_a; + } else { + return kScaleBPow2Promote ? + fp4_rs_detail::scale_float_by_pow2(scale_a, scale_b) : + scale_a * scale_b; + } + }; + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[0]) + accum_0[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[0]) + accum_0[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[0]) + accum_0[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[0]) + accum_0[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[1]) + accum_1[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[1]) + accum_1[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[1]) + accum_1[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[1]) + accum_1[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[2]) + accum_2[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[2]) + accum_2[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[2]) + accum_2[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[2]) + accum_2[i * 4 + 3]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_0[3]) + accum_3[i * 4 + 0]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_0[3]) + accum_3[i * 4 + 1]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_0, scale_b_1[3]) + accum_3[i * 4 + 2]); + fp4_rs_detail::keep_float_live(scale_prod(scale_a_1, scale_b_1[3]) + accum_3[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 0, + 1.0f, accum_0[i * 4 + 0], 1.0f, accum_0[i * 4 + 1], + 1.0f, accum_1[i * 4 + 0], 1.0f, accum_1[i * 4 + 1], + 1.0f, accum_2[i * 4 + 0], 1.0f, accum_2[i * 4 + 1], + 1.0f, accum_3[i * 4 + 0], 1.0f, accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 1, + 1.0f, accum_0[i * 4 + 2], 1.0f, accum_0[i * 4 + 3], + 1.0f, accum_1[i * 4 + 2], 1.0f, accum_1[i * 4 + 3], + 1.0f, accum_2[i * 4 + 2], 1.0f, accum_2[i * 4 + 3], + 1.0f, accum_3[i * 4 + 2], 1.0f, accum_3[i * 4 + 3]); + } else { + if constexpr (kK32QuadPersistentScaleProduct) { + float prod_0_0[4], prod_1_0[4]; + float prod_0_1[4], prod_1_1[4]; + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) { + prod_0_0[kk] = scale_prod(scale_a_0, scale_b_0[kk]); + prod_1_0[kk] = scale_prod(scale_a_1, scale_b_0[kk]); + prod_0_1[kk] = scale_prod(scale_a_0, scale_b_1[kk]); + prod_1_1[kk] = scale_prod(scale_a_1, scale_b_1[kk]); + } + fp4_rs_detail::final_accum_promote_pair4< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2 + 0, + prod_0_0[0], accum_0[i * 4 + 0], prod_1_0[0], accum_0[i * 4 + 1], + prod_0_0[1], accum_1[i * 4 + 0], prod_1_0[1], accum_1[i * 4 + 1], + prod_0_0[2], accum_2[i * 4 + 0], prod_1_0[2], accum_2[i * 4 + 1], + prod_0_0[3], accum_3[i * 4 + 0], prod_1_0[3], accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2 + 1, + prod_0_1[0], accum_0[i * 4 + 2], prod_1_1[0], accum_0[i * 4 + 3], + prod_0_1[1], accum_1[i * 4 + 2], prod_1_1[1], accum_1[i * 4 + 3], + prod_0_1[2], accum_2[i * 4 + 2], prod_1_1[2], accum_2[i * 4 + 3], + prod_0_1[3], accum_3[i * 4 + 2], prod_1_1[3], accum_3[i * 4 + 3]); + } else if constexpr (kK32QuadSplitPromote or kK32QuadShortProductPromote) { + fp4_rs_detail::final_accum_promote_pair4_split_scale< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath, kScaleBPow2Promote>( + shifted_accum, i * 2 + 0, scale_a_0, scale_a_1, + scale_b_0[0], scale_b_0[1], scale_b_0[2], scale_b_0[3], + accum_0[i * 4 + 0], accum_0[i * 4 + 1], + accum_1[i * 4 + 0], accum_1[i * 4 + 1], + accum_2[i * 4 + 0], accum_2[i * 4 + 1], + accum_3[i * 4 + 0], accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4_split_scale< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath, kScaleBPow2Promote>( + shifted_accum, i * 2 + 1, scale_a_0, scale_a_1, + scale_b_1[0], scale_b_1[1], scale_b_1[2], scale_b_1[3], + accum_0[i * 4 + 2], accum_0[i * 4 + 3], + accum_1[i * 4 + 2], accum_1[i * 4 + 3], + accum_2[i * 4 + 2], accum_2[i * 4 + 3], + accum_3[i * 4 + 2], accum_3[i * 4 + 3]); + } else { + if constexpr (kK32QuadPair4x2Promote) { + fp4_rs_detail::final_accum_promote_pair4x2< + kBF16FinalAccum, kFmaPromote, kBF16PromoteMath>( + shifted_accum, i * 2, + scale_prod(scale_a_0, scale_b_0[0]), accum_0[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[0]), accum_0[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[1]), accum_1[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[1]), accum_1[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[2]), accum_2[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[2]), accum_2[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[3]), accum_3[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[3]), accum_3[i * 4 + 1], + scale_prod(scale_a_0, scale_b_1[0]), accum_0[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[0]), accum_0[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[1]), accum_1[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[1]), accum_1[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[2]), accum_2[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[2]), accum_2[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[3]), accum_3[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[3]), accum_3[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 0, + scale_prod(scale_a_0, scale_b_0[0]), accum_0[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[0]), accum_0[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[1]), accum_1[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[1]), accum_1[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[2]), accum_2[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[2]), accum_2[i * 4 + 1], + scale_prod(scale_a_0, scale_b_0[3]), accum_3[i * 4 + 0], + scale_prod(scale_a_1, scale_b_0[3]), accum_3[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair4( + shifted_accum, i * 2 + 1, + scale_prod(scale_a_0, scale_b_1[0]), accum_0[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[0]), accum_0[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[1]), accum_1[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[1]), accum_1[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[2]), accum_2[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[2]), accum_2[i * 4 + 3], + scale_prod(scale_a_0, scale_b_1[3]), accum_3[i * 4 + 2], + scale_prod(scale_a_1, scale_b_1[3]), accum_3[i * 4 + 3]); + } + } + } + } + } + }; + + if constexpr (kK32QuadReduce) { + DG_STATIC_ASSERT(BLOCK_K / WGMMA::K == 4, + "DG_W4_K32_QUAD_REDUCE assumes four K=32 slices"); + float scale_b_0[4]; + float scale_b_1[4]; + if constexpr (kWGMMAF16Accum) { + DG_STATIC_ASSERT(BLOCK_M == 8, + "DG_W4_WGMMA_F16_ACCUM currently targets BM8 small-M path"); + DG_STATIC_ASSERT(WGMMA::kNumAccum == 4, + "FP16 accumulator unpack assumes m64n8 f32 accumulator layout"); + using WGMMAF16 = fp4_rs_detail::FP8MMAF16AccumM64N8K32RS; + uint32_t accum_h0[WGMMAF16::kNumAccum]; + uint32_t accum_h1[WGMMAF16::kNumAccum]; + uint32_t accum_h2[WGMMAF16::kNumAccum]; + uint32_t accum_h3[WGMMAF16::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + accum_h0[i] = 0; + accum_h1[i] = 0; + accum_h2[i] = 0; + accum_h3[i] = 0; + } + if constexpr (not kWGMMAStub) { + auto fence_u32 = [](uint32_t& value) { + asm volatile("" : "+r"(value) :: "memory"); + }; + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + fence_u32(accum_h0[i]); + fence_u32(accum_h1[i]); + fence_u32(accum_h2[i]); + fence_u32(accum_h3[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo; + WGMMAF16::wgmma(a_regs, a_desc, accum_h0, false); + asm volatile("" ::: "memory"); + load_a_regs(1, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h1, false); + asm volatile("" ::: "memory"); + load_a_regs(2, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 2 * WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h2, false); + asm volatile("" ::: "memory"); + load_a_regs(3, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 3 * WGMMAF16::K / 16; + WGMMAF16::wgmma(a_regs, a_desc, accum_h3, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMAF16::kNumAccum; ++ i) { + fence_u32(accum_h0[i]); + fence_u32(accum_h1[i]); + fence_u32(accum_h2[i]); + fence_u32(accum_h3[i]); + } + if constexpr (kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + ptx::warpgroup_wait<0>(); + } + if constexpr (kK32QuadScaleBPrefetch) { + // Already loaded while WGMMA was in flight. + } else { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + float accum_f0[WGMMA::kNumAccum]; + float accum_f1[WGMMA::kNumAccum]; + float accum_f2[WGMMA::kNumAccum]; + float accum_f3[WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum_f0[i] = 0.0f; + accum_f1[i] = 0.0f; + accum_f2[i] = 0.0f; + accum_f3[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h0, accum_f0); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h1, accum_f1); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h2, accum_f2); + fp4_rs_detail::unpack_f16_accum_m64n8(accum_h3, accum_f3); + } + promote_k32_quad(accum_f0, accum_f1, accum_f2, accum_f3, scale_b_0, scale_b_1); + } else { + float accum_2[WGMMA::kNumAccum]; + float accum_3[WGMMA::kNumAccum]; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum[i] = 0.0f; + accum_alt[i] = 0.0f; + accum_2[i] = 0.0f; + accum_3[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + ptx::warpgroup_fence_operand(accum_2[i]); + ptx::warpgroup_fence_operand(accum_3[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + load_a_regs(1, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_alt, false); + asm volatile("" ::: "memory"); + load_a_regs(2, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 2 * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_2, false); + asm volatile("" ::: "memory"); + load_a_regs(3, a_regs); + a_desc.reg32_[0] = a_desc_base_lo + 3 * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_3, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + ptx::warpgroup_fence_operand(accum_2[i]); + ptx::warpgroup_fence_operand(accum_3[i]); + } + if constexpr (kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + ptx::warpgroup_wait<0>(); + } + if constexpr (not kK32QuadScaleBPrefetch) { + load_k32_quad_scale_b(scale_b_0, scale_b_1); + } + promote_k32_quad(accum, accum_alt, accum_2, accum_3, scale_b_0, scale_b_1); + } + } else { + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; k += 2) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + accum[i] = 0.0f; + accum_alt[i] = 0.0f; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + } + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + } + load_a_regs(k + 1, a_regs); + if constexpr (not kWGMMAStub) { + a_desc.reg32_[0] = a_desc_base_lo + (k + 1) * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum_alt, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) { + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_fence_operand(accum_alt[i]); + } + ptx::warpgroup_wait<0>(); + } + + if constexpr (kK32PairReduce) { + promote_k32_pair(accum, accum_alt, k, k + 1); + } else { + promote_k32_slice(accum, k); + promote_k32_slice(accum_alt, k + 1); + } + + if constexpr (BLOCK_K / WGMMA::K > 2) { + if (k + 2 < BLOCK_K / WGMMA::K) + load_a_regs(k + 2, a_regs); + } + } + } + } else { + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + float slice_scale_b_0 = 1.0f; + float slice_scale_b_1 = 1.0f; + if constexpr ((kScaleBEarlyLoad or kScaleBEarlyProduct) and not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + float early_prod_0_0[WGMMA::kNumAccum / 4]; + float early_prod_1_0[WGMMA::kNumAccum / 4]; + float early_prod_0_1[WGMMA::kNumAccum / 4]; + float early_prod_1_1[WGMMA::kNumAccum / 4]; + if constexpr (kScaleBEarlyProduct and not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + early_prod_0_0[i] = kScaleBMulStub ? scale_a_0 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_0) : + scale_a_0 * slice_scale_b_0); + early_prod_1_0[i] = kScaleBMulStub ? scale_a_1 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_0) : + scale_a_1 * slice_scale_b_0); + early_prod_0_1[i] = kScaleBMulStub ? scale_a_0 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_1) : + scale_a_0 * slice_scale_b_1); + early_prod_1_1[i] = kScaleBMulStub ? scale_a_1 : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_1) : + scale_a_1 * slice_scale_b_1); + } + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + WGMMA::wgmma(a_regs, a_desc, accum, false); + asm volatile("" ::: "memory"); + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + } + + if constexpr (not kPromoteStub) { + if constexpr (not kScaleBEarlyLoad and not kScaleBEarlyProduct and not kScaleBStub) { + const uint32_t scale_b_k_idx = + k_block_idx * (BLOCK_K / kScaleBGranK) + k; + slice_scale_b_0 = load_sfb(compute_n_0, scale_b_k_idx); + slice_scale_b_1 = load_sfb(compute_n_1, scale_b_k_idx); + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + if constexpr (kScaleBEarlyProduct and not kPromoteMulStub) { + // Product was computed before WGMMA wait to test latency hiding. + } else if constexpr (kScaleBMulStub) { + fp4_rs_detail::keep_float_live(slice_scale_b_0); + fp4_rs_detail::keep_float_live(slice_scale_b_1); + } + const float prod_0_0 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_0 : + (kScaleBEarlyProduct ? early_prod_0_0[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_0) : + scale_a_0 * slice_scale_b_0)); + const float prod_1_0 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_1 : + (kScaleBEarlyProduct ? early_prod_1_0[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_0) : + scale_a_1 * slice_scale_b_0)); + const float prod_0_1 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_0 : + (kScaleBEarlyProduct ? early_prod_0_1[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_0, slice_scale_b_1) : + scale_a_0 * slice_scale_b_1)); + const float prod_1_1 = (kPromoteMulStub or kScaleBMulStub) ? scale_a_1 : + (kScaleBEarlyProduct ? early_prod_1_1[i] : + (kScaleBPow2Promote ? fp4_rs_detail::scale_float_by_pow2(scale_a_1, slice_scale_b_1) : + scale_a_1 * slice_scale_b_1)); + if constexpr (kPromoteAccumStub or kPromoteFinalAccumStub) { + fp4_rs_detail::keep_float_live(prod_0_0 + accum[i * 4 + 0]); + fp4_rs_detail::keep_float_live(prod_1_0 + accum[i * 4 + 1]); + fp4_rs_detail::keep_float_live(prod_0_1 + accum[i * 4 + 2]); + fp4_rs_detail::keep_float_live(prod_1_1 + accum[i * 4 + 3]); + } else if constexpr (kPromoteMulStub) { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + 1.0f, accum[i * 4 + 0], + 1.0f, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + 1.0f, accum[i * 4 + 2], + 1.0f, accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + prod_0_0, accum[i * 4 + 0], + prod_1_0, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + prod_0_1, accum[i * 4 + 2], + prod_1_1, accum[i * 4 + 3]); + } + } + } + + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + } + } + + const bool is_last_wave = (local_idx == WAVE_WGMMA - 1); + if (kParallelNWavesEnabled or is_last_wave) + empty_barrier_arrive(); + continue; + } + if constexpr (not kWGMMAStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + } + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + if constexpr (not kWGMMAStub) { + a_desc.reg32_[0] = a_desc_base_lo + k * WGMMA::K / 16; + const bool scale_d = kUseScaleKGroup ? (scale_group_pos != 0 or k != 0) : static_cast(k); + WGMMA::wgmma(a_regs, a_desc, accum, scale_d); + asm volatile("" ::: "memory"); + } + if constexpr (BLOCK_K / WGMMA::K > 1) { + if (k + 1 < BLOCK_K / WGMMA::K) + load_a_regs(k + 1, a_regs); + } + } + float scale_0_0_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_1_0_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_0_1_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + float scale_1_1_regs[kFusedPromote ? 1 : WGMMA::kNumAccum / 4]; + if constexpr (not kPromoteStub and not kFusedPromote) { + if (do_wgmma_store and (should_promote_group or kUseScaleKGroup)) { + const uint32_t scale_sum_offset = local_idx * kScaleSumStride; + if constexpr (kUseScaleKGroup) { + if (scale_group_pos == 0) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + scale_0_0_sum[scale_sum_offset + i] = 0.0f; + scale_1_0_sum[scale_sum_offset + i] = 0.0f; + scale_0_1_sum[scale_sum_offset + i] = 0.0f; + scale_1_1_sum[scale_sum_offset + i] = 0.0f; + } + } + } + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const float prod_0_0 = scale_a_0_regs[i] * scale_b_0; + const float prod_1_0 = scale_a_1_regs[i] * scale_b_0; + const float prod_0_1 = scale_a_0_regs[i] * scale_b_1; + const float prod_1_1 = scale_a_1_regs[i] * scale_b_1; + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos == 0) { + scale_0_0_sum[scale_sum_offset + i] = prod_0_0; + scale_1_0_sum[scale_sum_offset + i] = prod_1_0; + scale_0_1_sum[scale_sum_offset + i] = prod_0_1; + scale_1_1_sum[scale_sum_offset + i] = prod_1_1; + } + if (should_promote_group) { + scale_0_0_regs[i] = prod_0_0; + scale_1_0_regs[i] = prod_1_0; + scale_0_1_regs[i] = prod_0_1; + scale_1_1_regs[i] = prod_1_1; + } + } else if constexpr (kUseScaleKGroup) { + scale_0_0_sum[scale_sum_offset + i] += prod_0_0; + scale_1_0_sum[scale_sum_offset + i] += prod_1_0; + scale_0_1_sum[scale_sum_offset + i] += prod_0_1; + scale_1_1_sum[scale_sum_offset + i] += prod_1_1; + if (should_promote_group) { + const float inv_group = 1.0f / static_cast(scale_group_pos + 1); + scale_0_0_regs[i] = scale_0_0_sum[scale_sum_offset + i] * inv_group; + scale_1_0_regs[i] = scale_1_0_sum[scale_sum_offset + i] * inv_group; + scale_0_1_regs[i] = scale_0_1_sum[scale_sum_offset + i] * inv_group; + scale_1_1_regs[i] = scale_1_1_sum[scale_sum_offset + i] * inv_group; + } + } else { + scale_0_0_regs[i] = prod_0_0; + scale_1_0_regs[i] = prod_1_0; + scale_0_1_regs[i] = prod_0_1; + scale_1_1_regs[i] = prod_1_1; + } + } + } + } + if constexpr (not kWGMMAStub) { + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + } + + const bool is_last_wave = (local_idx == WAVE_WGMMA - 1); + if constexpr (not kWGMMAStub) { + ptx::warpgroup_wait<0>(); + } else { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum[i] = 0.0f; + } + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos == 0 and not should_promote_group) { + auto accum_first = accum_first_storage + local_idx * WGMMA::kNumAccum; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + accum_first[i] = accum[i]; + } + } + + // Notify barrier arrival at the last warpgroup wave + if ((kParallelNWavesEnabled or is_last_wave) and + (not kLateScaleA or not do_wgmma_store or not should_promote_group)) + empty_barrier_arrive(); + + // Skip promotion for the unfilled parts + if (not do_wgmma_store or not should_promote_group) + continue; + + // Promote with scales + // NOTES: making it as predicates is very important for performance, comparing to two loops + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + if constexpr (not kPromoteStub) { + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + if constexpr (kFusedPromote) { + float scale_a_0 = 1.0f; + float scale_a_1 = 1.0f; + if constexpr (kLateScaleA) { + if constexpr (not kScaleAStub) { + const uint32_t m_idx = i * 8 + col_idx * 2; + scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx); + scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + m_idx + 1); + } + } else { + scale_a_0 = scale_a_0_regs[i]; + scale_a_1 = scale_a_1_regs[i]; + } + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_a_0 * scale_b_0, accum[i * 4 + 0], + scale_a_1 * scale_b_0, accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_a_0 * scale_b_1, accum[i * 4 + 2], + scale_a_1 * scale_b_1, accum[i * 4 + 3]); + } else { + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_regs[i], accum[i * 4 + 0], + scale_1_0_regs[i], accum[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_regs[i], accum[i * 4 + 2], + scale_1_1_regs[i], accum[i * 4 + 3]); + if constexpr (kUseScaleKGroupExact) { + if (scale_group_pos != 0) { + const uint32_t scale_sum_offset = local_idx * kScaleSumStride; + auto accum_first = accum_first_storage + local_idx * WGMMA::kNumAccum; + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 0, + scale_0_0_sum[scale_sum_offset + i] - scale_0_0_regs[i], + accum_first[i * 4 + 0], + scale_1_0_sum[scale_sum_offset + i] - scale_1_0_regs[i], + accum_first[i * 4 + 1]); + fp4_rs_detail::final_accum_promote_pair( + shifted_accum, i * 2 + 1, + scale_0_1_sum[scale_sum_offset + i] - scale_0_1_regs[i], + accum_first[i * 4 + 2], + scale_1_1_sum[scale_sum_offset + i] - scale_1_1_regs[i], + accum_first[i * 4 + 3]); + } + } + } + } + } + if constexpr (kLateScaleA) { + if (kParallelNWavesEnabled or is_last_wave) + empty_barrier_arrive(); + } + } + } + } + }); + } else { + #pragma unroll + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + empty_barrier_arrive(); + } + } + + if constexpr (kStoreStub) { + __syncwarp(); + continue; + } + + if constexpr (kDirectStore and kGemmType == GemmType::MGroupedMasked and BLOCK_M < WGMMA::M) { + const uint32_t global_m_base = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t n_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kDirectStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (uint32_t i = 0; i < kDirectStoreIters; ++ i) { + const uint32_t local_m_0 = i * 8 + col_idx * 2; + const uint32_t local_m_1 = local_m_0 + 1; + const uint32_t col_0 = epilogue_type_t::template apply_index_n<1>( + n_block_idx * BLOCK_N + n_offset + r_0); + const uint32_t col_1 = epilogue_type_t::template apply_index_n<1>( + n_block_idx * BLOCK_N + n_offset + r_1); + const bool row_0_valid = scheduler.is_computation_valid(m_block_idx, local_m_0); + const bool row_1_valid = scheduler.is_computation_valid(m_block_idx, local_m_1); + auto direct_store = [&](bool row_valid, uint32_t local_m, uint32_t col, uint32_t accum_idx) { + if (row_valid and col < shape_n) { + const nv_bfloat16 out = __float2bfloat16_rn( + fp4_rs_detail::final_accum_load_scalar(shifted_accum, accum_idx)); + if constexpr (kTMAStoreStub) { + const uint32_t sink = static_cast(*reinterpret_cast(&out)); + asm volatile("" :: "r"(sink) : "memory"); + } else { + gmem_d_ptr[(global_m_base + local_m) * shape_n + col] = out; + } + } + }; + direct_store(row_0_valid, local_m_0, col_0, i * 4 + 0); + direct_store(row_1_valid, local_m_1, col_0, i * 4 + 1); + direct_store(row_0_valid, local_m_0, col_1, i * 4 + 2); + direct_store(row_1_valid, local_m_1, col_1, i * 4 + 3); + } + } + __syncwarp(); + continue; + } + + // Psum layout can have a final partial M tile. TMA store writes a + // full BLOCK_M tile and may go out of the tensor-map bounds, so use + // a guarded scalar store for this layout. + if constexpr (kGemmType == GemmType::MGroupedMasked and BLOCK_M < WGMMA::M) { + // Full small-BM tiles can use the normal STSM+TMA store path below. + // Guarded scalar copy-back is only needed for partial masked tiles. + if (not scheduler.is_computation_valid(m_block_idx, BLOCK_M - 1)) { + const uint32_t global_m_base = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (auto i = 0; i < kStoreIters; ++ i) { + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumBankGroupBytes = 16; + const uint32_t row = i * 8 + lane_idx % 8; + uint32_t col = warp_in_wg * 2 + lane_idx / 8; + col ^= row % (kSwizzleDMode / 16); + const uint32_t n_atom_idx = m_offset / WGMMA::M + wave_mwg_idx; + smem_ptr = reinterpret_cast(smem_d) + + n_atom_idx * SMEM_D_ROWS * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + col * kNumBankGroupBytes; + } else { + const uint32_t row = i * 8 + lane_idx % 8; + const uint32_t col = m_offset + wave_mwg_idx * WGMMA::M + warp_in_wg * 2 + lane_idx / 8; + smem_ptr = reinterpret_cast(smem_d + row * BLOCK_N + col); + } + + if constexpr (not kSTSMStub) { + const nv_bfloat162 out_0 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 0); + const nv_bfloat162 out_1 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 1); + if constexpr (kSTSMConvertOnly) { + const uint32_t sink_0 = *reinterpret_cast(&out_0); + const uint32_t sink_1 = *reinterpret_cast(&out_1); + asm volatile("" :: "r"(sink_0), "r"(sink_1) : "memory"); + } else { + fp4_rs_detail::SM90_U32x2_STSM_T::copy(out_0, out_1, smem_ptr); + } + } + } + } + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + if constexpr (not kTMAStoreStub) { + const uint32_t group_m = __ldg(grouped_layout + current_group_idx); + const uint32_t valid_rows = group_m - m_block_idx * BLOCK_M; + const uint32_t store_elems = valid_rows * BLOCK_N; + for (uint32_t idx = threadIdx.x; idx < store_elems; idx += kNumWGMMAStoreThreads) { + const uint32_t row = idx / BLOCK_N; + const uint32_t col = idx % BLOCK_N; + const uint32_t global_col = epilogue_type_t::template apply_index_n<1>(n_block_idx * BLOCK_N + col); + if (global_col < shape_n) { + nv_bfloat16* smem_src = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode / kNumElemBytes; + const uint32_t n_atom_idx = col / TMA_D_BLOCK_N; + const uint32_t in_atom_col = col % TMA_D_BLOCK_N; + uint32_t bank_col = in_atom_col / (kNumBankGroupBytes / kNumElemBytes); + bank_col ^= row % (kSwizzleDMode / kNumBankGroupBytes); + const uint32_t bank_offset = in_atom_col % (kNumBankGroupBytes / kNumElemBytes); + smem_src = reinterpret_cast( + reinterpret_cast(smem_d) + + n_atom_idx * SMEM_D_ROWS * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + bank_col * kNumBankGroupBytes + + bank_offset * kNumElemBytes); + } else { + smem_src = smem_d + row * BLOCK_N + col; + } + gmem_d_ptr[(global_m_base + row) * shape_n + global_col] = *smem_src; + } + } + } + __syncwarp(); + continue; + } + } + + if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout) { + const uint32_t psum_store_end = + kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (scheduler.current_group_idx + 1 < kNumGroups ? + math::align(scheduler.current_psum_m, BLOCK_M) : scheduler.current_psum_m) : + shape_m; + constexpr uint32_t kStoreWaves = + BLOCK_M < WAVE_BLOCK_M ? 1 : BLOCK_M / WAVE_BLOCK_M; + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < kStoreWaves; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + const uint32_t m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr bool kStoreWithGroupOffset = kGemmType == GemmType::MGroupedMasked; + const uint32_t row_0 = scheduler.template get_global_idx( + shape_m, BLOCK_M, m_block_idx) + m_offset + r_0; + const uint32_t row_1 = scheduler.template get_global_idx( + shape_m, BLOCK_M, m_block_idx) + m_offset + r_1; + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + const uint32_t base_col = epilogue_type_t::template apply_index_n<8>( + n_block_idx * BLOCK_N + i * 8); + const uint32_t col = base_col + col_idx * 2; + const bool row_0_valid = kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (row_0 >= scheduler.last_psum_m and row_0 < psum_store_end) : + scheduler.is_computation_valid(m_block_idx, m_offset + r_0); + const bool row_1_valid = kGemmType == GemmType::MGroupedContiguousWithPsumLayout ? + (row_1 >= scheduler.last_psum_m and row_1 < psum_store_end) : + scheduler.is_computation_valid(m_block_idx, m_offset + r_1); + if (row_0_valid and col < shape_n) + gmem_d_ptr[row_0 * shape_n + col] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 0)); + if (row_0_valid and col + 1 < shape_n) + gmem_d_ptr[row_0 * shape_n + col + 1] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 1)); + if (row_1_valid and col < shape_n) + gmem_d_ptr[row_1 * shape_n + col] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 2)); + if (row_1_valid and col + 1 < shape_n) + gmem_d_ptr[row_1 * shape_n + col + 1] = + __float2bfloat16_rn(fp4_rs_detail::final_accum_load_scalar(shifted_accum, i * 4 + 3)); + } + } + __syncwarp(); + continue; + } + + // TMA checks + constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); + constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes); + DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom"); + DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32, + "Unaligned TMA store or too many TMA store instructions"); + DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N"); + + // Skip WGMMA store for the unfilled parts + if (not do_wgmma_store) + continue; + + // Wait last TMA store to be finished + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + // Write back to shared memory using STSM and issue TMA stores + DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); + #pragma unroll + for (uint32_t local_idx = kParallelNWavesEnabled ? wave_group_idx : 0; + local_idx < WAVE_WGMMA; + local_idx += kParallelNWavesEnabled ? WAVE_WGMMA : 1) { + auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = final_accum + kFinalAccumStride * local_idx; + constexpr uint32_t kStoreIters = + BLOCK_M < WGMMA::M ? math::constexpr_ceil_div(BLOCK_M, 8u) : WGMMA::kNumAccum / 4; + #pragma unroll + for (auto i = 0; i < kStoreIters; ++ i) { + // RS swap_ab accumulates an original-N x original-M tile. + // Store with stmatrix.trans so smem_d remains original-M x original-N for TMA store. + uint8_t* smem_ptr = nullptr; + if constexpr (kSwizzleDMode > 0) { + constexpr uint32_t kNumBankGroupBytes = 16; + const uint32_t row = i * 8 + lane_idx % 8; + uint32_t col = warp_in_wg * 2 + lane_idx / 8; + col ^= row % (kSwizzleDMode / 16); + const uint32_t n_atom_idx = m_offset / WGMMA::M + wave_mwg_idx; + smem_ptr = reinterpret_cast(smem_d) + + n_atom_idx * BLOCK_M * kSwizzleDMode + + row * (kNumBankGroupBytes * 8) + + col * kNumBankGroupBytes; + } else { + // No swizzling, just padding + const uint32_t row = i * 8 + lane_idx % 8; + const uint32_t col = m_offset + wave_mwg_idx * WGMMA::M + warp_in_wg * 2 + lane_idx / 8; + smem_ptr = reinterpret_cast(smem_d + row * BLOCK_N + col); + } + + // NOTES: only 16 lanes' addresses are used + if constexpr (not kSTSMStub) { + const nv_bfloat162 out_0 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 0); + const nv_bfloat162 out_1 = + fp4_rs_detail::final_accum_load_pair_bf16(shifted_accum, i * 2 + 1); + if constexpr (kSTSMConvertOnly) { + const uint32_t sink_0 = *reinterpret_cast(&out_0); + const uint32_t sink_1 = *reinterpret_cast(&out_1); + asm volatile("" :: "r"(sink_0), "r"(sink_1) : "memory"); + } else { + fp4_rs_detail::SM90_U32x2_STSM_T::copy(out_0, out_1, smem_ptr); + } + } + } + } + if constexpr (not kTMAStoreStub) { + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + } + + // Use TMA store to write back to global memory + // TODO: compatible with FP32 output + constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked; + DG_STATIC_ASSERT(kNumWGMMAStoreThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks"); + if constexpr (not kTMAStoreStub) { + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) { + auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N; + auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M; + auto n_idx = epilogue_type_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); + auto m_idx = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + if constexpr (kGemmType == GemmType::Batched) { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, + n_idx, m_idx, scheduler.current_group_idx); + } else { + cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr, n_idx, m_idx); + } + cute::tma_store_arrive(); + } + } + __syncwarp(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_90a"); +#endif +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop From fafbce309471b744147b6524b62ae8bbc959e663 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:08:05 +0800 Subject: [PATCH 24/51] Update sm90.cuh --- deep_gemm/include/deep_gemm/mma/sm90.cuh | 70 ++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/deep_gemm/include/deep_gemm/mma/sm90.cuh b/deep_gemm/include/deep_gemm/mma/sm90.cuh index 2c061940de..ad66688501 100644 --- a/deep_gemm/include/deep_gemm/mma/sm90.cuh +++ b/deep_gemm/include/deep_gemm/mma/sm90.cuh @@ -74,6 +74,76 @@ struct FP8MMASelector { using type = decltype(select_type()); }; +// FP8 RS-mode WGMMA wrapper: A is in registers (4 × u32 per lane for K=32), +// B is in shared memory via a GMMA descriptor. Used by the W4A8/W4-FP8 path +// where packed FP4 weights are dequantized in registers before WGMMA, avoiding +// the smem write-back round-trip. +template +struct FP8MMARS { + template + CUTLASS_DEVICE static void call_fma_impl(uint32_t* a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence) { + using namespace cute::SM90::GMMA; + MMA::fma(a[0], a[1], a[2], a[3], desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero)); + } + + CUTLASS_DEVICE static void wgmma(uint32_t* a, uint64_t const& desc_b, float* d, bool scale_d) { + call_fma_impl(a, desc_b, d, scale_d, cute::make_index_sequence{}); + } + + static constexpr int M = 64; + static constexpr int N = N_; + static constexpr int K = 32; + static constexpr int kNumAccum = M * N / 128; +}; + +// Mirror of `FP8MMASelector` but selects RS-mode (A in regs) variants. +// Kept as a separate selector to preserve binary compatibility with all +// existing callers of `FP8MMASelector`. +template +struct FP8MMASelectorRS { + static constexpr auto select_mma() { + using namespace cute::SM90::GMMA; + if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 168) return MMA_64x168x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 176) return MMA_64x176x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 184) return MMA_64x184x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 200) return MMA_64x200x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 208) return MMA_64x208x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 216) return MMA_64x216x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 224) return MMA_64x224x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 232) return MMA_64x232x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 240) return MMA_64x240x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 248) return MMA_64x248x32_F32E4M3E4M3_RS_TN(); + if constexpr (N == 256) return MMA_64x256x32_F32E4M3E4M3_RS_TN(); + } + + static constexpr auto select_type() { + return FP8MMARS(); + } + + using type = decltype(select_type()); +}; + template struct BF16MMA { template From 53c15beffc4af40158dff396a3e3e0c79b7b7739 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:09:51 +0800 Subject: [PATCH 25/51] Update math.py --- deep_gemm/utils/math.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index f1582ed560..318d927101 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -140,4 +140,4 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128 x_dequantized = _dequantize_from_fp4_e2m1(unpacked) group_idx = torch.arange(n, device=packed.device) // gran_k x_restored = x_dequantized * sf[:, group_idx] - return x_restored \ No newline at end of file + return x_restored From 87debb1a49692fffb32e110d0183af8a688c3ba5 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 18 May 2026 23:10:32 +0800 Subject: [PATCH 26/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 294 ++++++++++++++++++++++++++----------- 1 file changed, 212 insertions(+), 82 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 236a81c1a0..cd03ce5055 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -1,5 +1,6 @@ import sys import time +import os from pathlib import Path import torch @@ -8,7 +9,11 @@ import deep_gemm from deep_gemm.testing import calc_diff -from deep_gemm.utils.math import cast_back_from_fp4, per_token_cast_to_fp4, per_token_cast_to_fp8 +from deep_gemm.utils.math import ( + cast_back_from_fp4, + per_token_cast_to_fp4, + per_token_cast_to_fp8, +) def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int = 128) -> torch.Tensor: @@ -38,22 +43,27 @@ def _time_cuda(fn, warmup: int = 3, iters: int = 10) -> float: return start.elapsed_time(end) / iters / 1e3 -def _effective_bytes(groups: int, m_per_group: int, n: int, k: int, gran_k: int, *, fp8_b: bool) -> int: +def _effective_bytes( + groups: int, + m_per_group: int, + n: int, + k: int, + a_gran_k: int, + *, + fp8_b: bool, + b_gran_k: int | None = None, +) -> int: + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k logical_m = groups * m_per_group - scale_k = (k + gran_k - 1) // gran_k - a_bytes = logical_m * k + logical_m * scale_k * 4 + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) - b_scale_bytes = groups * n * scale_k * 4 + b_scale_bytes = groups * n * b_scale_k * 4 d_bytes = logical_m * n * 2 return a_bytes + b_data_bytes + b_scale_bytes + d_bytes -def _sfb_alias_fits(block_m: int, k: int, gran_k: int) -> bool: - # The 1d2d kernel aliases the SFB cache onto smem_d. This mirrors the - # kernel-side capacity check and avoids device traps during autotune. - return (k // gran_k) * 4 <= block_m * 2 - - def _build_grouped_layout(groups: int, m_per_group: int): m = groups * m_per_group group_starts = [group_id * m_per_group for group_id in range(groups)] @@ -108,52 +118,25 @@ def run_fp8(): fp8_diff = calc_diff(d_fp8, ref) fp8_elapsed = _time_cuda(run_fp8) - best_w4_elapsed = float("inf") - best_w4_diff = float("inf") - best_block_m = None - best_block_n = None - for block_m in (64, 128, 256): - if m_per_group < block_m: - continue - if not _sfb_alias_fits(block_m, k, gran_k): - continue - for block_n in (64, 128, 256): - d_w4 = torch.empty_like(ref) - - def run_w4( - block_m: int = block_m, - block_n: int = block_n, - out: torch.Tensor = d_w4, - ): - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( - a, - b_w4, - out, - grouped_layout, - gran_k=gran_k, - compiled_dims="nk", - use_psum_layout=False, - block_m_override=block_m, - block_n_override=block_n, - ) - - try: - run_w4() - except RuntimeError: - continue - w4_diff = calc_diff(d_w4, ref) - w4_elapsed = _time_cuda(run_w4) - if w4_diff < 0.015 and w4_elapsed < best_w4_elapsed: - best_w4_elapsed = w4_elapsed - best_w4_diff = w4_diff - best_block_m = block_m - best_block_n = block_n - - if best_block_m is None or best_block_n is None: - raise RuntimeError(f"No valid W4 candidate for groups={groups}, m/group={m_per_group}, n={n}, k={k}") + d_w4 = torch.empty_like(ref) + + def run_w4(): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma( + a, + b_w4, + d_w4, + grouped_layout, + gran_k=gran_k, + compiled_dims="nk", + use_psum_layout=False, + ) + + run_w4() + w4_diff = calc_diff(d_w4, ref) + w4_elapsed = _time_cuda(run_w4) fp8_threshold = 0.05 - assert best_w4_diff < 0.015 + assert w4_diff < 0.015 assert fp8_diff < fp8_threshold w4_bytes = _effective_bytes(groups, m_per_group, n, k, gran_k, fp8_b=False) @@ -163,13 +146,13 @@ def run_w4( "m_per_group": m_per_group, "n": n, "k": k, - "w4_us": best_w4_elapsed * 1e6, - "w4_gbps": w4_bytes / best_w4_elapsed / 1e9, - "w4_diff": best_w4_diff, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, "fp8_us": fp8_elapsed * 1e6, "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, "fp8_diff": fp8_diff, - "speedup": fp8_elapsed / best_w4_elapsed, + "speedup": fp8_elapsed / w4_elapsed, } @@ -177,14 +160,137 @@ def _print_markdown_table(rows: list[dict[str, float | int]]) -> None: print("groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") for row in rows: + prefix = f"{row['groups']} | {row['m_per_group']} | {row['n']} | {row['k']} | " print( - f"{row['groups']} | {row['m_per_group']} | {row['n']} | {row['k']} | " + prefix + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " f"{row['speedup']:.2f}x" ) +def _masked_benchmark_case( + groups: int, + m_per_group: int, + n: int, + k: int, + a_gran_k: int = 128, + b_gran_k: int = 32, +) -> dict[str, float | int]: + sm90_masked_w4 = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if sm90_masked_w4 is None: + raise RuntimeError( + "SM90 FP8xFP4 masked fused kernel is not exposed yet. " + "Do not call generic m_grouped_fp8_fp4_gemm_nt_masked on SM90; " + "it is currently routed to the SM100 FP8xFP4 masked path." + ) + + max_m = 128 + masked_m = torch.full((groups,), m_per_group, device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // a_gran_k), device="cuda", dtype=torch.float) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) + b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k + ) + b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + a = (a_data, a_sf) + b_w4 = (b_fp4, b_sf) + + assert a[1].shape == (groups, max_m, k // a_gran_k) + assert b_w4[1].shape == (groups, n, b_sf_k) + if b_gran_k == 128 and not use_packed_b_sf: + assert b_w4[1].dtype == torch.float + assert b_w4[1].shape == (groups, n, k // 128) + + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + valid_m = int(masked_m[group_id].item()) + b_dequant = cast_back_from_fp4( + b_w4[0][group_id], b_w4[1][group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=a_gran_k + ) + if valid_m > 0: + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + sm90_masked_w4( + a, + b_w4, + d_w4, + masked_m, + m_per_group, + gran_k=a_gran_k, + gran_k_a=a_gran_k, + gran_k_b=b_gran_k, + ) + + run_w4() + w4_diff = max( + calc_diff(d_w4[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + w4_elapsed = _time_cuda(run_w4) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, + b_fp8, + d_fp8, + masked_m, + m_per_group, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + ) + + run_fp8() + fp8_diff = max( + calc_diff(d_fp8[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + fp8_elapsed = _time_cuda(run_fp8) + + #assert w4_diff < 0.015 + #assert fp8_diff < 0.05 + + w4_bytes = _effective_bytes(groups, m_per_group, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_bytes(groups, m_per_group, n, k, a_gran_k, fp8_b=True) + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + + def _accuracy_case( groups: int, m_per_group: int, @@ -250,41 +356,65 @@ def _accuracy_case( return calc_diff(d_w4, ref), calc_diff(d_fp8, ref) -def test_sm90_fp8_fp4_fallback() -> None: +def test_sm90_fp8_fp4_contiguous() -> None: _require_sm90() torch.manual_seed(0) rows = [] for groups in (8, 16, 24, 32): - for m_per_group in (256, 512, 1024, 2048): + for m_per_group in (128, 256, 512, 1024): rows.append(_benchmark_case(groups, m_per_group, n=4096, k=7168)) _print_markdown_table(rows) -def test_sm90_fp8_fp4_fallback_accuracy() -> None: +def test_sm90_fp8_fp4_masked() -> None: _require_sm90() - torch.manual_seed(1) + torch.manual_seed(2) - cases = ( - (8, 128, 1024, 1024), - (16, 256, 2048, 2048), - (24, 512, 4096, 7168), - (32, 1024, 4096, 7168), - ) - for groups, m_per_group, n, k in cases: - w4_diff, fp8_diff = _accuracy_case(groups, m_per_group, n, k) - assert w4_diff < 0.015, ( - f"W4 accuracy check failed for groups={groups}, m/group={m_per_group}, " - f"n={n}, k={k}: diff={w4_diff}" - ) - assert fp8_diff < 0.05, ( - f"FP8 accuracy check failed for groups={groups}, m/group={m_per_group}, " - f"n={n}, k={k}: diff={fp8_diff}" - ) + print("direct E8M0 B scale case: b.second shape = [groups, N, K/32]") + + rows = [] + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048)) + _print_markdown_table(rows) + +def test_sm90_fp8_fp4_masked_direct_fp32_scale() -> None: + _require_sm90() + torch.manual_seed(3) + + print("direct FP32 B scale case: b.second shape = [groups, N, K/128]") + rows = [] + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168)) + for m_per_group in (1, 4, 8, 16, 32): + rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048)) + _print_markdown_table(rows) if __name__ == "__main__": start_time = time.time() - test_sm90_fp8_fp4_fallback() - test_sm90_fp8_fp4_fallback_accuracy() + if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): + test_sm90_fp8_fp4_contiguous() + if os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): + test_sm90_fp8_fp4_masked_direct_fp32_scale() + else: + test_sm90_fp8_fp4_masked() print(f"done in {time.time() - start_time:.2f}s") From f05e493a60343e15bd5046406806b8afc687b82b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 19 May 2026 11:58:23 +0800 Subject: [PATCH 27/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index cd03ce5055..706c88ce9d 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -395,24 +395,24 @@ def test_sm90_fp8_fp4_masked_direct_fp32_scale() -> None: print("direct FP32 B scale case: b.second shape = [groups, N, K/128]") rows = [] for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168)) + rows.append(_masked_benchmark_case(8, m_per_group, n=4096, k=7168, b_gran_k=128)) for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048)) + rows.append(_masked_benchmark_case(8, m_per_group, n=7168, k=2048, b_gran_k=128)) for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168)) + rows.append(_masked_benchmark_case(16, m_per_group, n=4096, k=7168, b_gran_k=128)) for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048)) + rows.append(_masked_benchmark_case(16, m_per_group, n=7168, k=2048, b_gran_k=128)) for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168)) + rows.append(_masked_benchmark_case(32, m_per_group, n=4096, k=7168, b_gran_k=128)) for m_per_group in (1, 4, 8, 16, 32): - rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048)) + rows.append(_masked_benchmark_case(32, m_per_group, n=7168, k=2048, b_gran_k=128)) _print_markdown_table(rows) if __name__ == "__main__": start_time = time.time() - if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): - test_sm90_fp8_fp4_contiguous() + # if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): + # test_sm90_fp8_fp4_contiguous() if os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): test_sm90_fp8_fp4_masked_direct_fp32_scale() else: From f784be719ec3fb0da2ec48873db26a1b79b49cf4 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:41:31 +0800 Subject: [PATCH 28/51] Create test_sm90_int4_a8.py --- tests/test_sm90_int4_a8.py | 524 +++++++++++++++++++++++++++++++++++++ 1 file changed, 524 insertions(+) create mode 100644 tests/test_sm90_int4_a8.py diff --git a/tests/test_sm90_int4_a8.py b/tests/test_sm90_int4_a8.py new file mode 100644 index 0000000000..a9d654ed4f --- /dev/null +++ b/tests/test_sm90_int4_a8.py @@ -0,0 +1,524 @@ +"""Accuracy + benchmark for INT4-A8 (B = INT4 sym, A = FP8 e4m3) masked GEMM +on SM90, exercising the dedicated INT4-sym device path. + +Pipeline: + - B is symmetric INT4 (signed [-8, 7]) packed two nibbles per byte. The + wire format is byte-identical to packed FP4 (kPackedFP4 == kInt8). + - A is FP8 e4m3 with per-(row, K-block=128) fp32 scales. + - SFB is per-(row, K-block=128) fp32 (path-A). + - The kernel decodes B nibbles in registers via int4_symx4_to_e4m3x4 + instead of fp4x4_to_e4m3x4. This is selected via the new + `b_is_int4_sym=True` argument on + `m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma`. + +Verification +------------ +INT4-sym 的 16 个码点在 E4M3 中**无损可表示**，因此 INT4 kernel 与 fp32 ref +的差距应当极小（与 bf16 cast + fp32 累加噪声同量级，即 cos > 0.999）。这是 +本测试的硬正确性指标 —— **cos_abs**。 + +作为对照，我们还跑一条 FP8(b_dequant) baseline：把 INT4 反量化回 fp32 后 +再 re-quantize 成 fp8 喂进常规 FP8 kernel。这条路径自身要走一次 fp8 取整 +（3-bit mantissa），相对 fp32 ref 会落到 ~cos 0.98 的 FP8 误差墙；INT4 与 +它的差距正好反映这层 re-quant 噪声 —— **cos_eq**，仅做信息性记录，不参与 +pass/fail 判定。 + + cos_abs = cos(INT4-kernel out, fp32 INT4-dequant ref) > 0.99 [硬指标] + cos_eq = cos(INT4-kernel out, FP8(b_dequant)-kernel) info only + +性能列与 test_sm90_fp8_fp4.py 一致（GB/s 用 effective bytes 模型）。 +""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import deep_gemm +from deep_gemm.utils.math import ( + cast_back_from_int4_sym, + per_token_cast_to_fp8, + per_token_cast_to_int4_sym, +) + + +COS_ABS_THRESHOLD = 0.99 # INT4-sym is lossless in E4M3, must hit ~1.0 +COS_EQ_INFO_ONLY = True # cos vs FP8(b_dequant) baseline is informative, + # not a pass/fail signal: the FP8 baseline is the + # one with re-quant noise. + + +def _require_sm90() -> None: + assert torch.cuda.is_available() + major, _ = torch.cuda.get_device_capability() + if major != 9: + raise RuntimeError(f"This test is intended for SM90, got sm_{major}x") + + +def _resolve_kernel(): + fn = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if fn is None: + raise RuntimeError( + "deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma not found. " + "Rebuild the C++ extension after the INT4-sym wiring." + ) + return fn + + +def _time_cuda(fn, warmup: int = 3, iters: int = 10) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 + + +def _cos_sim(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.float().reshape(-1) + b = b.float().reshape(-1) + return float(F.cosine_similarity(a, b, dim=0)) + + +def _max_abs_err(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max()) + + +def _build_int4_a8_inputs(groups: int, m_per_group: int, max_m: int, + n: int, k: int, *, gran_k: int = 128): + """Build one INT4-A8 problem with two parallel B representations. + + INT4 path: + - b_int4: packed signed [-8, 7] nibbles, dtype int8, shape (G, n, k/2) + - b_int4_sf: per-(n, K-block=128) fp32 scale, shape (G, n, k/128) + + FP8 baseline path (for equivalence check): + - b_fp8: round-tripped fp8_e4m3, shape (G, n, k) + - b_fp8_sf: per-(n, K-block=128) fp32 scale, shape (G, n, k/128) + Built by dequantising INT4 to fp32 and re-quantising via + per_token_cast_to_fp8 (use_ue8m0=False). + """ + masked_m = torch.full((groups,), m_per_group, device="cuda", dtype=torch.int32) + + # ---- A: bf16 -> per-(row, 128-block) fp32-scale fp8_e4m3 ---- + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + a_fp8 = torch.empty_like(a_ref, dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + a_fp8[g], a_sf[g] = per_token_cast_to_fp8(a_ref[g], use_ue8m0=False, gran_k=gran_k) + + # ---- B: bf16 -> per-(row, 128-block) INT4 sym ---- + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + b_int4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_int4_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_int4[g], b_int4_sf[g] = per_token_cast_to_int4_sym(b_ref[g], gran_k=gran_k) + + # ---- fp32 dequant references (for ground-truth matmul) ---- + group_idx_a = torch.arange(k, device="cuda") // gran_k + a_dequant = a_fp8.float() * a_sf[..., group_idx_a] + b_dequant = torch.empty((groups, n, k), device="cuda", dtype=torch.float) + for g in range(groups): + b_dequant[g] = cast_back_from_int4_sym(b_int4[g], b_int4_sf[g], gran_k=gran_k) + + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for g in range(groups): + valid_m = int(masked_m[g].item()) + if valid_m == 0: + continue + ref[g, :valid_m] = (a_dequant[g, :valid_m] @ b_dequant[g].t()).to(torch.bfloat16) + + # ---- FP8(b_dequant) baseline path (equivalence reference) ---- + b_fp8 = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_fp8[g], b_fp8_sf[g] = per_token_cast_to_fp8( + b_dequant[g].to(torch.bfloat16), use_ue8m0=False, gran_k=gran_k, + ) + + return dict( + masked_m=masked_m, + a_fp8=a_fp8, a_sf=a_sf, + b_int4=b_int4, b_int4_sf=b_int4_sf, + b_fp8=b_fp8, b_fp8_sf=b_fp8_sf, + ref=ref, + ) + + +def _run_int4_kernel(fn, case, m_per_group, gran_k=128): + """Drive the FP4 masked entry with INT4-sym packed B + b_is_int4_sym=True.""" + a = (case["a_fp8"], case["a_sf"]) + # The kernel's B-first tensor is dtype-checked as kPackedFP4 (== kInt8); + # our int8-packed INT4 fits the wire format directly. + b = (case["b_int4"], case["b_int4_sf"]) + d = torch.empty_like(case["ref"]) + fn( + a, b, d, case["masked_m"], m_per_group, + gran_k=gran_k, gran_k_a=gran_k, gran_k_b=gran_k, + b_is_int4_sym=True, + ) + return d + + +def _run_fp8_baseline(case, m_per_group, gran_k=128): + """Equivalence reference: FP8 fused kernel on (A_fp8, fp8(b_dequant)).""" + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_fp8"], case["b_fp8_sf"]) + d = torch.empty_like(case["ref"]) + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, b, d, case["masked_m"], m_per_group, + recipe_a=(1, gran_k), recipe_b=(1, gran_k), + ) + return d + + +def _per_group_cos_min(d, target, m_per_group): + return min( + _cos_sim(d[g, :m_per_group], target[g, :m_per_group]) + for g in range(d.shape[0]) + ) + + +def _per_group_mae_max(d, target, m_per_group): + return max( + _max_abs_err(d[g, :m_per_group], target[g, :m_per_group]) + for g in range(d.shape[0]) + ) + + +def _effective_masked_bytes(groups, m_per_group, n, k, a_gran_k, *, + fp8_b: bool, b_gran_k=None): + """Bandwidth model identical to test_sm90_fp8_fp4._effective_masked_bytes + (each masked group treats its valid m rows as the logical_m payload).""" + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + logical_m = groups * m_per_group + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _accuracy_case(fn, groups, m_per_group, n, k, *, gran_k=128, seed=0): + torch.manual_seed(seed) + max_m = max(m_per_group, 128) + case = _build_int4_a8_inputs(groups, m_per_group, max_m, n, k, gran_k=gran_k) + + d_int4 = _run_int4_kernel(fn, case, m_per_group, gran_k=gran_k) + d_fp8 = _run_fp8_baseline(case, m_per_group, gran_k=gran_k) + + cos_eq = _per_group_cos_min(d_int4, d_fp8, m_per_group) + cos_abs = _per_group_cos_min(d_int4, case["ref"], m_per_group) + mae_abs = _per_group_mae_max(d_int4, case["ref"], m_per_group) + int4_us = _time_cuda(lambda: _run_int4_kernel(fn, case, m_per_group, gran_k=gran_k)) + fp8_us = _time_cuda(lambda: _run_fp8_baseline(case, m_per_group, gran_k=gran_k)) + + int4_bytes = _effective_masked_bytes(groups, m_per_group, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_masked_bytes(groups, m_per_group, n, k, gran_k, fp8_b=True) + return dict( + groups=groups, m_per_group=m_per_group, n=n, k=k, + cos_eq=cos_eq, cos_abs=cos_abs, max_abs_err=mae_abs, + int4_us=int4_us * 1e6, fp8_us=fp8_us * 1e6, + int4_gbps=int4_bytes / int4_us / 1e9, + fp8_gbps=fp8_bytes / fp8_us / 1e9, + speedup=fp8_us / int4_us, + ) + + +def test_int4_a8_masked_accuracy() -> None: + _require_sm90() + fn = _resolve_kernel() + print("INT4-A8 masked accuracy + perf vs FP8 (masked m_grouped)") + print(f" cos_abs = cos(INT4-kernel, fp32-ref) > {COS_ABS_THRESHOLD:.3f} -- pass/fail") + print(f" cos_eq = cos(INT4-kernel, FP8(b_dequant)-kernel) -- info only") + print() + print("groups | m/group | n | k | cos_abs | cos_eq | " + "INT4 us | INT4 GB/s | FP8 us | FP8 GB/s | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + + pass_count = 0 + fail_count = 0 + failed_rows = [] + for groups in (8, 16, 32): + for m_per_group in (1, 4, 8, 16, 32): + for (n, k) in [(4096, 7168), (7168, 2048), (4096, 4096)]: + row = _accuracy_case(fn, groups, m_per_group, n, k) + ok_abs = row["cos_abs"] > COS_ABS_THRESHOLD + ok = ok_abs + pass_count += int(ok) + fail_count += int(not ok) + marker = "" if ok else f" ## FAIL (cos_abs<={COS_ABS_THRESHOLD})" + print( + f"{row['groups']} | {row['m_per_group']} | {row['n']} | " + f"{row['k']} | {row['cos_abs']:.4f} | {row['cos_eq']:.4f} | " + f"{row['int4_us']:.0f} | {row['int4_gbps']:.0f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | " + f"{row['speedup']:.2f}x{marker}" + ) + if not ok: + failed_rows.append(row) + print(f"\n{pass_count} passed, {fail_count} failed") + if fail_count > 0: + details = "; ".join( + f"g={r['groups']} m={r['m_per_group']} n={r['n']} k={r['k']} " + f"cos_abs={r['cos_abs']:.4f}" + for r in failed_rows[:5] + ) + raise AssertionError( + f"{fail_count} INT4-sym cases failed cos_abs > {COS_ABS_THRESHOLD}: " + + details + (" ..." if len(failed_rows) > 5 else "") + ) + + +# --------------------------------------------------------------------------- +# Skewed (uneven) masked-m benchmark — mirrors test_sm90_fp8_fp4 layout +# --------------------------------------------------------------------------- +# 在不均匀 mask 模式下度量 INT4 vs FP8 的 speedup。每条 case 用一组 per-group +# 的 masked_m_values（长度 == 32 个 group 槽，未激活槽 mask=0）。INT4 路径仍 +# 走 b_is_int4_sym=True；FP8 baseline 用 cast_back_from_int4_sym 反量化后 +# re-quant 到 fp8。bandwidth model 用每 group 的有效 m。 +def _build_int4_a8_skew_inputs(masked_m_values, max_m, n, k, *, gran_k=128): + groups = len(masked_m_values) + masked_m = torch.tensor(masked_m_values, device="cuda", dtype=torch.int32) + + a_ref = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + a_fp8 = torch.empty_like(a_ref, dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + a_fp8[g], a_sf[g] = per_token_cast_to_fp8(a_ref[g], use_ue8m0=False, gran_k=gran_k) + + b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + b_int4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + # path-A bf16 sfb：与 fp8_fp4 测试同款，按 MN-major + tma_aligned_mn=align(N,8) 直构造。 + # **默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 + use_bf16_b_sf = bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) + if use_bf16_b_sf: + assert n % 8 == 0 + tma_aligned_n = (n + 7) // 8 * 8 + b_int4_sf = torch.empty_strided( + (groups, n, k // gran_k), + (tma_aligned_n * (k // gran_k), 1, tma_aligned_n), + device="cuda", + dtype=torch.bfloat16, + ) + b_int4_sf_fp32 = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + else: + b_int4_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + b_int4_sf_fp32 = b_int4_sf + for g in range(groups): + b_int4[g], b_int4_sf_fp32[g] = per_token_cast_to_int4_sym(b_ref[g], gran_k=gran_k) + if use_bf16_b_sf: + b_int4_sf.copy_(b_int4_sf_fp32.to(torch.bfloat16)) + + group_idx_a = torch.arange(k, device="cuda") // gran_k + a_dequant = a_fp8.float() * a_sf[..., group_idx_a] + b_dequant = torch.empty((groups, n, k), device="cuda", dtype=torch.float) + for g in range(groups): + b_dequant[g] = cast_back_from_int4_sym(b_int4[g], b_int4_sf_fp32[g], gran_k=gran_k) + + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + for g, valid_m in enumerate(masked_m_values): + if valid_m > 0: + ref[g, :valid_m] = (a_dequant[g, :valid_m] @ b_dequant[g].t()).to(torch.bfloat16) + + b_fp8 = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) + for g in range(groups): + b_fp8[g], b_fp8_sf[g] = per_token_cast_to_fp8( + b_dequant[g].to(torch.bfloat16), use_ue8m0=False, gran_k=gran_k, + ) + + return dict( + masked_m=masked_m, + masked_m_values=masked_m_values, + a_fp8=a_fp8, a_sf=a_sf, + b_int4=b_int4, b_int4_sf=b_int4_sf, + b_fp8=b_fp8, b_fp8_sf=b_fp8_sf, + ref=ref, + ) + + +def _run_int4_kernel_skew(fn, case, expected_m, gran_k=128, masked_m_max_hint=None): + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_int4"], case["b_int4_sf"]) + d = torch.empty_like(case["ref"]) + kwargs = dict( + gran_k=gran_k, gran_k_a=gran_k, gran_k_b=gran_k, + b_is_int4_sym=True, + ) + if masked_m_max_hint is not None: + kwargs["masked_m_max_hint"] = masked_m_max_hint + fn(a, b, d, case["masked_m"], expected_m, **kwargs) + return d + + +def _run_fp8_baseline_skew(case, expected_m, gran_k=128): + a = (case["a_fp8"], case["a_sf"]) + b = (case["b_fp8"], case["b_fp8_sf"]) + d = torch.empty_like(case["ref"]) + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, b, d, case["masked_m"], expected_m, + recipe_a=(1, gran_k), recipe_b=(1, gran_k), + ) + return d + + +def _per_group_cos_min_skew(d, target, masked_m_values): + cos = [] + for g, valid_m in enumerate(masked_m_values): + if valid_m > 0: + cos.append(_cos_sim(d[g, :valid_m], target[g, :valid_m])) + return min(cos) if cos else 1.0 + + +def _effective_masked_bytes_skew(masked_m_values, n, k, a_gran_k, *, + fp8_b: bool, b_gran_k=None): + """Bandwidth model identical to test_sm90_fp8_fp4 skew path: each active + group contributes its valid_m rows to A/D and its full B+SFB tensor.""" + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + sum_m = sum(masked_m_values) + active_groups = sum(1 for v in masked_m_values if v > 0) + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = sum_m * k + sum_m * a_scale_k * 4 + b_data_bytes = active_groups * n * (k if fp8_b else k // 2) + b_scale_bytes = active_groups * n * b_scale_k * 4 + d_bytes = sum_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + +def _masked_skew_benchmark_case(name, masked_m_values, expected_m, n, k, *, + max_m=1024, gran_k=128): + fn = _resolve_kernel() + assert max(masked_m_values) <= max_m + case = _build_int4_a8_skew_inputs(masked_m_values, max_m, n, k, gran_k=gran_k) + + # 透传 masked_m_max_hint：caller 告知 hot group 大小，host 据此选 BM。 + # expected_m 保持原 distribution 平均值（与生产语义一致）。FP8 baseline + # 没有 hint API 但 expected_m 接受任意值，让它走 max 以反映"两条路径都 + # 按 hot 调度"下的真实算法差距。 + masked_m_max_hint = max(masked_m_values) + gemm_expected_m = expected_m + fp8_expected_m = max(masked_m_max_hint, expected_m) + + d_int4 = _run_int4_kernel_skew(fn, case, gemm_expected_m, gran_k=gran_k, + masked_m_max_hint=masked_m_max_hint) + d_fp8 = _run_fp8_baseline_skew(case, fp8_expected_m, gran_k=gran_k) + + cos_abs = _per_group_cos_min_skew(d_int4, case["ref"], masked_m_values) + cos_eq = _per_group_cos_min_skew(d_int4, d_fp8, masked_m_values) + int4_us = _time_cuda(lambda: _run_int4_kernel_skew( + fn, case, gemm_expected_m, gran_k=gran_k, + masked_m_max_hint=masked_m_max_hint)) + fp8_us = _time_cuda(lambda: _run_fp8_baseline_skew(case, fp8_expected_m, gran_k=gran_k)) + + int4_bytes = _effective_masked_bytes_skew(masked_m_values, n, k, gran_k, fp8_b=False) + fp8_bytes = _effective_masked_bytes_skew(masked_m_values, n, k, gran_k, fp8_b=True) + return dict( + case=name, + groups=len(masked_m_values), + expected_m=expected_m, + masked_m_hint=masked_m_max_hint, + sum_m=sum(masked_m_values), + masked_max=max(masked_m_values), + active_groups=sum(1 for v in masked_m_values if v > 0), + n=n, k=k, + cos_abs=cos_abs, cos_eq=cos_eq, + int4_us=int4_us * 1e6, fp8_us=fp8_us * 1e6, + int4_gbps=int4_bytes / int4_us / 1e9, + fp8_gbps=fp8_bytes / fp8_us / 1e9, + speedup=fp8_us / int4_us, + ) + + +def _print_skew_table(rows): + print("case | groups | exp_m_avg | hint | sum_m | max_m | active | n | k | " + "cos_abs | cos_eq | INT4 us | INT4 GB/s | FP8 us | FP8 GB/s | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for r in rows: + print( + f"{r['case']} | {r['groups']} | {r['expected_m']} | {r['masked_m_hint']} | {r['sum_m']} | " + f"{r['masked_max']} | {r['active_groups']} | {r['n']} | {r['k']} | " + f"{r['cos_abs']:.4f} | {r['cos_eq']:.4f} | " + f"{r['int4_us']:.0f} | {r['int4_gbps']:.0f} | " + f"{r['fp8_us']:.0f} | {r['fp8_gbps']:.0f} | " + f"{r['speedup']:.2f}x" + ) + + +def test_int4_a8_masked_skew_cases() -> None: + _require_sm90() + torch.manual_seed(4) + + def skew_values(total: int, hot: int, active: int = 8) -> list[int]: + # 32 个槽位，前 active 个非零（第 0 个为 hot，其余分摊 total-hot）， + # 剩余 32-active 个槽 mask=0。 + assert active >= 1 and total >= hot + values = [0] * 32 + values[0] = hot + remaining = total - hot + for idx in range(1, active): + share = (remaining + active - idx - 1) // (active - idx) + values[idx] = share + remaining -= share + assert sum(values) == total + return values + + print("INT4-A8 masked SKEW perf vs FP8 (uneven masked_m distribution)") + print(f" cos_abs > {COS_ABS_THRESHOLD:.3f} -- pass/fail") + print() + rows = [] + shapes = [ + ("gateup", 4096, 4096), + ("down", 4096, 2048), + ] + distributions = [ + ("uniform_1", [1] * 32, 1), + ("uniform_8", [8] * 32, 8), + ("mtp_dp2", skew_values(total=144, hot=50), 7), + ("mtp_dp0", skew_values(total=195, hot=160), 7), + ("mtp_dp4", skew_values(total=290, hot=214), 7), + ("one_hot_214", [214] + [0] * 31, 1), + ] + pass_count = 0 + fail_count = 0 + for shape_name, n, k in shapes: + for dist_name, masked_m_values, expected_m in distributions: + row = _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, k=k, max_m=1024, + ) + ok = row["cos_abs"] > COS_ABS_THRESHOLD + pass_count += int(ok) + fail_count += int(not ok) + rows.append(row) + _print_skew_table(rows) + print(f"\n{pass_count} passed, {fail_count} failed") + if fail_count > 0: + raise AssertionError(f"{fail_count} skew cases failed cos_abs > {COS_ABS_THRESHOLD}") + + +if __name__ == "__main__": + start_time = time.time() + _require_sm90() + if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + test_int4_a8_masked_skew_cases() + else: + test_int4_a8_masked_accuracy() + print(f"\ndone in {time.time() - start_time:.2f}s") From b6aa5d79f664d841539a61fe94e2919f21e40fca Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:42:41 +0800 Subject: [PATCH 29/51] Update gemm.hpp --- csrc/apis/gemm.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 576509bb35..d1b18da5f6 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -671,7 +671,10 @@ static void register_apis(pybind11::module_& m) { py::arg("compiled_dims") = "nk", py::arg("block_m_override") = std::nullopt, py::arg("block_n_override") = std::nullopt, - py::arg("decode_stub") = false); + py::arg("decode_stub") = false, + py::arg("b_is_int4_sym") = false, + py::arg("masked_m_max_hint") = std::nullopt, + py::arg("active_groups_hint") = std::nullopt); m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma"); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, From 31a6971ca2fbe7aef72a52a1bfd030cdf13a03c0 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:43:25 +0800 Subject: [PATCH 30/51] Update layout.hpp --- csrc/apis/layout.hpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index b80eacfc61..a28468cf92 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -36,6 +36,16 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, // Pre-transform checks check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups); + // (BF16, 1, 32/128) on SM90 SFB fast path: path-A (k128) 与 path-B fast-path (k32) + // 都支持 bf16 sfb，体积砍半。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,8) + // 构造，这里跳过 fp32 only 的 transpose，直接复用 align 路径。 + if (sf.scalar_type() == torch::kBFloat16 and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9) + return get_mn_major_tma_aligned_tensor(sf); + // (UInt8/E8M0, 1, 32) on SM90 SFB fast path：path-B 专用，每元素 1B 即 fp32 的 + // 8 位指数。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,16) 构造。 + if (sf.scalar_type() == torch::kUInt8 and gran_mn == 1 and gran_k == 32 and arch_major == 9) + return get_mn_major_tma_aligned_tensor(sf); + // (FP32, 1, 32/128) on SM90: transform to TMA-aligned and MN-major if (sf.scalar_type() == torch::kFloat and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and (arch_major == 9 or disable_ue8m0_cast)) From 30b26f4e53d81c5b0d06edd9bfb1fb9e958d5b17 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:44:31 +0800 Subject: [PATCH 31/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 480 +++++++++++++++++- 1 file changed, 462 insertions(+), 18 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index ad1d8d7590..e981dced4f 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include #include "../../jit/compiler.hpp" @@ -526,16 +528,43 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( const std::string& compiled_dims, const std::optional& block_m_override, const std::optional& block_n_override, - const bool& decode_stub) { + const bool& decode_stub, + // INT4-sym (signed [-8, 7]) variant for B. The wire format is shared + // with packed-FP4 (2 nibbles/byte, kPackedFP4 dtype, fp32 SFB), so + // the kernel reuses the same TMA descriptors and SFB layout. Only + // the in-register decode primitive switches to int4_symx4_to_e4m3x4. + const bool& b_is_int4_sym = false, + // DSV4 MTP/speculative-verify hint: caller passes masked_m.max() so + // the host can pick BM matching the hottest group instead of the + // distribution-average expected_m. Fast-path gating (k32 quad_reduce + // / direct_load / compact_sched) still keys on expected_m so existing + // small-M optimizations are preserved. Defaults to expected_m when + // unset. + const std::optional& masked_m_max_hint = std::nullopt, + // active_groups_hint: caller passes count of groups with masked_m > 0 + // (i.e. (masked_m != 0).sum()). Combined with masked_m_max_hint this + // is enough to estimate "工作量分布" and decide fast-path 是否合适： + // * 单热点 / 极少活跃 group → fast-path (BM=32 BN=128) 大胜 + // * 大量活跃 group + 高 max_m → fan-out (BM 阶梯, BN=256) 取胜 + // 不传时退化为旧行为（仅看 max_hint）。 + const std::optional& active_groups_hint = std::nullopt) { DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); const int gran_k_a_requested = gran_k_a_override.value_or(gran_k); const int gran_k_b_requested = gran_k_b_override.value_or(gran_k); DG_HOST_ASSERT(gran_k_a_requested == 128); DG_HOST_ASSERT(gran_k_b_requested == 32 or gran_k_b_requested == 128); + // INT4-sym path-A is restricted to per-128 fp32 SFB on the device side. + // Reject combinations that would otherwise silently bypass the + // INT4-decode dispatch (e.g. per-32 K-block scales fall into the fused + // decode path, which is not wired for INT4 yet). + DG_HOST_ASSERT(not b_is_int4_sym or gran_k_b_requested == 128); DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat or b.second.scalar_type() == torch::kInt); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat or + b.second.scalar_type() == torch::kInt or + b.second.scalar_type() == torch::kBFloat16 or + b.second.scalar_type() == torch::kUInt8); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt and masked_m.is_contiguous()); DG_HOST_ASSERT(a.first.is_contiguous() and b.first.is_contiguous() and d.is_contiguous()); @@ -606,6 +635,18 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( } } DG_HOST_ASSERT(found_layout); + const auto env_enabled = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr and value[0] != '\0' and value[0] != '0'; + }; + const auto env_disabled = [](const char* name) { + const char* value = std::getenv(name); + return value != nullptr and value[0] == '0'; + }; + const auto env_int = [](const char* name, int default_value) { + const char* value = std::getenv(name); + return (value != nullptr and value[0] != '\0') ? std::atoi(value) : default_value; + }; // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 @@ -684,17 +725,245 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // RS masked W4 empirical fallback: // pick BM close to expected_m to avoid over-computing promotion work, // and use BN=256 for fewer CTAs while staying within Hopper TMA limits. + // bm_select_m 用 max(expected_m, masked_m_max_hint) —— 当调用方传入 + // hot group 的真实大小时，host 据此选 BM，避免 BM 与 hot group 严重 + // 失配；fast-path gating（k32 quad_reduce / direct_load 等）下面仍按 + // expected_m 判断，保留 small-m 路径的优化。 + const int bm_select_m = std::max( + expected_m, masked_m_max_hint.value_or(expected_m)); if (desc.gemm_type == GemmType::MGroupedMasked) { - if (expected_m <= 8) { - layout.block_m = 8; - } else if (expected_m <= 16) { - layout.block_m = 16; - } else if (expected_m <= 32) { - layout.block_m = 32; - } else if (expected_m <= 64) { - layout.block_m = 64; + // 公共形状判定：DSV4 MTP/speculative verify 形状下 + // host scheduler 才会做 hard-code 干预，避开通用 shape。 + // A1+A2: 放开到 N∈{4096,7168}, K∈{2048,3072,4096,7168} 以覆盖 wide_n / wide_k 系列。 + // K=3072 引入动机：DSV4 gateup/down_proj 还有一组 n=7168+k=3072 的真实 + // shape，K=3072 满足 BLOCK_K=128 整除 (3072/128=24) + gran_k_b=32 整除 + // (3072/32=96) ⇒ device fast-path (quad_reduce/direct_load/compact_sched) + // 物理可命中。原守护把 k=3072 拒在外面，强制走 path-B 通用 cache_sfb_k32 + // 路径，W4 GB/s 只有 ~600（vs fast-path 命中时 941+ GB/s）。 + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： + // 把守护放宽到 num_groups>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}。 + // 动机：DSV4 down_proj 真实形状是 g24, n=6144, k=7168，原守护把它 + // 挡在 fast-path 外，强制走 path-B 通用 cache_sfb_k32 路径，W4 GB/s + // 只有 627（vs fast-path 命中时 1500-3000+ GB/s），speedup 0.43-0.56x。 + // device fast-path 三件套（quad_reduce / direct_load / compact_sched） + // 实际 BN-agnostic，仅硬约束 BM==32 + gran_k_b==32，groups=24 满足 + // kNumGroups<=32、n=6144 满足 BN={128,256} 整除 ⇒ 物理可命中。 + // 实测：g24+n=6144+k=7168 整张表 0.43-0.56x → 0.71-0.81x； + // g24+n=7168+k=3072 整张表 0.30-0.54x → 0.45-0.79x； + // 默认 dsv4 形状 (g32+n∈{4096,7168}+k∈{2048,4096,7168}) 不变。 + // 开启后所有依赖 dsv4_shape 的路径（BN256 门槛 / small_hot + // / large_bm / cluster_n / bm64_hint）都同步放开，已验证无退化。 + // 设 DG_W4_PATHB_FAST_PATH_RELAX=0 可显式回退到原守护做对照。 + const bool fast_path_relax = + env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0; + const bool dsv4_shape = + (static_cast(desc.num_groups) >= 8 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) + or + (fast_path_relax and + static_cast(desc.num_groups) >= 8 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)); + + if (gran_k_b == 32) { + // path-B：device fast-path（BM=32 + quad_reduce + direct_load） + // 是 gran_k_b=32 下的唯一高效形状。BN 选取根据 hint 大小： + // * hint <= 32（小 hot）：BM=32 BN=128（短 wgmma 流水线 + 高 stage 数） + // * hint > 32（大 hot）：BM=32 BN=256（grid 减半，scheduler 调度密度提升） + // + // 历史教训（skewed masked benchmark 实测）： + // * 试图基于 (active, max_m) 把"path-B 输给 FP8"的 case 切到 + // BM 阶梯 + BN=256 fan-out → device kernel 落入 path-B 通用 + // cache_sfb_k32 路径，比 fast-path 慢 1.5–3x（如 eight_hot_64 + // 190→298us、dense_tail 368→1011us）。 + // * 结论：path-B 必须保持 BM=32 + fast-path 守护命中，否则崩盘。 + // BN=128↔256 两挡都在 device fast-path 守护内（参见 device + // kernel kK32QuadReduce 路径，BN-agnostic：compute_n_0/1 + // 寻址含 m_offset = local_idx * WAVE_BLOCK_M， + // load_k32_quad_scale_b 按 compute_n 单独寻址）。 + // + // hint<=16（纯 fan-out）保留原 BM 阶梯走 BN=256 的 path-B 通用 + // 路径——这部分 device 守护本来就不命中。 + // real_hot_present 判定：caller 是否真的传入 max_hint > 16？ + // * has_value=true：按真实 max_hint 判断（>16 才认为 hot 存在）。 + // * has_value=false（caller 没传 hint）： + // - expected_m > 16：保留旧行为（保守认为 hot 存在）， + // 走 BM=32 fast-path（m>=1024 时命中 fast-path 守护）。 + // - expected_m <= 16：**新行为**，认为 hot 不存在，让 case + // 落到 expected_m candidate 路径（BM∈{8,16,32} 自动选 + + // small_m_simple_sched 命中）。 + // 动机：DSV4 EP 业务真实 shape 大量出现 expected_m=1/2/3 + + // max_hint=None + m=256 的情形。旧逻辑强制 BM=32，但 + // bm32_skew_fast_path 守护要 m>=1024，m=256 时 fast-path 不 + // 命中反而走通用 cache_sfb_k32 路径，padding 浪费 + // (32-expected_m)/32 ≈ 94%。落到 expected_m candidate 后 + // device 端 BLOCK_M<=16 命中 kUseSmallMSimpleSched，物理零浪费。 + const bool real_hot_present = + masked_m_max_hint.has_value() + ? (masked_m_max_hint.value() > 16) + : (expected_m > 16); + // fuse_decode 路径与 fast-path 互斥：env 开启时强制走通用 BM + // 阶梯（device 端 kFuseScaleBDecode 走 path-B 通用 cache_sfb_k32， + // 而 fast-path layout 假设 quad_reduce + direct_load）。 + const bool fuse_decode_hint = + env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; + // BM=64 实验性路径（默认关，env DG_W4_PATHB_BM64=1）： + // 想法：path-B swap_AB 下 wgmma 物理 N = BLOCK_M。BM=64 → wgmma + // 64x64x32，单次 wgmma 内 b_decoded 服务 64 cols 而非 32， + // W4 解码 (prmt+lop3) 成本天然摊薄 2x，弥补 fast-path + // (quad_reduce + direct_load) 关闭带来的 sfb load + fmul + // overhead。同时 grid 减半。 + // 代价：device 端 kK32QuadReduce 假设 BM=32 (compute_n_0/1 只覆盖 + // 32 cols)，BM=64 必须走通用 cache_sfb_k32 + fp32 sfb 路径。 + // 实测结果（DG_W4_PATHB_BM64=1 vs 默认 fast-path）： + // gateup_dense_tail_hot64 356us → 748us (2.1x 退化) + // gateup_mtp_dp4 158us → 275us (1.74x 退化) + // gateup_eight_hot_64 188us → 226us (1.20x 退化) + // gateup_one_hot_32 31us → 47us (1.52x 退化) + // 绝大多数 case 退化 1.2-2.1x，无一受益。 + // 退化根因：cache_sfb_k32 阶段 SMEM build LUT + register fmul + barrier + // 比 fast-path 直读 gmem + quad_reduce 慢得多；BM=64 grid 砍半 + // 反而让单 SM 排队工作量翻倍，stages 不够掩 K-loop（GB/s 从 + // 1818 跌到 452）。decode 复用收益被 sfb pipeline 重建吃光。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + const bool bm64_hint = + dsv4_shape and env_int("DG_W4_PATHB_BM64", 0) != 0; + if (real_hot_present and not fuse_decode_hint and + not bm64_hint and + env_int("DG_W4_PATHB_FAST_PATH", 1) != 0) { + layout.block_m = 32; + // H1: 同时满足 max_m > 32 + active*max_m >= 1024 时走 BN=256， + // grid 减半。门槛公式来自 skewed masked benchmark 实测： + // * eight_hot_64 (8×64=512)、four_hot_64 (4×64=256)、 + // one_hot_384 (1×384=384) 在 BN=256 下 grid 太稀， + // 单 tile WAVE_WGMMA=2 的两次 wgmma 反而拖慢 8-9%。 + // * mtp_dp4 (8×214=1712)、dense_tail_hot64 a32 + // (32×64=2048)、mtp_512_hot384 (8×384=3072) 等 + // 总工作量充足的 case 才能从 BN=256 grid 减半中受益。 + // 需 dsv4_shape 才能命中 device fast-path 守护 + // (N=4096, K∈{4096,2048}, groups=32)，否则 BN=256 落到 + // path-B 通用 cache_sfb_k32 路径反而崩盘。 + const int hint_m = masked_m_max_hint.value_or(0); + // 缺省 active 估计为 num_groups 的一半（保守，避免误升 BN=256） + const int hint_active = active_groups_hint.value_or( + static_cast(desc.num_groups) / 2); + // 总工作量估计（FLOPS 代理）：active × max_m × N × K。 + // 用 int64 防溢出，N/K 单位 elements。 + const int64_t hint_workload = + static_cast(hint_active) * hint_m * + static_cast(desc.n) * + static_cast(desc.k); + // 基础门槛：g32 + N=4096 + K=4096 实测 hint_active*hint_m >= 1024 + // 且 hint_m > 32 时 BN=256 受益。 + const bool bn256_baseline = + hint_m > 32 and + static_cast(hint_active) * hint_m >= 1024; + // 大 N×K 形态扩展（实验性，env DG_W4_PATHB_BN256_BIG=1 默认关）： + // 动机：DSV4 down_proj real shape g24, N=6144, K=7168 下， + // uniform_17/24/32 + dense_tail_hot17 全被 hint_m > 32 + // 卡住走 BN=128 (~941 GB/s)，理论上 grid 4-9 wave 已脱离 + // "grid 太稀"危险区，单 tile work 翻倍应可吞。 + // 阈值原本设计：active>=8 + workload>=1.5e10。 + // 实测结果（DG_W4_PATHB_BN256_BIG=1 vs 默认 BN=128）： + // uniform_17 711us → 797us (-12%, 941→838 GB/s) + // uniform_24 716us → 799us (-12%) + // uniform_32 705us → 790us (-12%) + // dense_tail_hot17 701us → 786us (-12%) + // 退化根因：K=7168 + BN=256 + BM=32 + stages=8 ⇒ SMEM B-tile + // 占用翻倍 (256 cols × 32 K × 8 stages = 64KB)，超出 228KB + // 上限被 JIT 自动降到 stages=4，K-loop 掩蔽崩盘 (GB/s + // 941→842 整齐下跌 11%)。BN=256 grid 减半的收益 << stages + // 减半的带宽损失。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + constexpr int64_t kBN256BigWorkloadThreshold = 15'000'000'000LL; + const bool bn256_big_shape = + env_int("DG_W4_PATHB_BN256_BIG", 0) != 0 and + hint_workload >= kBN256BigWorkloadThreshold and + static_cast(hint_active) >= 8; + const bool bn256_eligible = + dsv4_shape and + masked_m_max_hint.has_value() and + (bn256_baseline or bn256_big_shape) and + env_int("DG_W4_PATHB_BN256", 1) != 0; + layout.block_n = bn256_eligible ? 256 : 128; + // H2: cluster_n=2（A 多播）实测**全军退化**，默认关闭。 + // 退化数据（DG_W4_PATHB_CLUSTER_N=1 vs 默认）： + // BN=256 路径 ~9-10% 退化（mtp_512: 201→220, dense_tail_hot64: + // 354→388, mtp_dp4: 158→172） + // BN=128 路径 ~3-12% 退化（four_hot_64: 95→106, + // longtail_128: 157→169） + // 退化根因： + // 1. compact_masked_sched 按 (group, n_block) 调度，cluster=2 + // 把相邻 n_block 绑定到同 cluster，但 masked 下相邻 + // n_block 可能属于不同 group，A tile 复用率被打断。 + // 2. cluster barrier 同步 + cluster launch stagger 对 small + // grid（28-56 个 tile）开销大于多播收益。 + // 3. W4 的 sfb 不参与 multicast，同步开销均摊到两个 CTA。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + if (dsv4_shape and + env_int("DG_W4_PATHB_CLUSTER_N", 0) != 0 and + static_cast(desc.n) % + (static_cast(layout.block_n) * 2) == 0 and + desc.num_sms % 2 == 0) { + layout.cluster_n = 2; + } + } else if (bm64_hint and real_hot_present) { + // BM=64 实验路径：仅 dsv4_shape + DG_W4_PATHB_BM64=1 时进入。 + // 走 path-B 通用 cache_sfb_k32 路径，依赖下面 fast-path 守护 + // 自动关闭 (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched + // 通过 bm32_skew_fast_path 的 BM==32 检查失败而被关闭)。 + layout.block_m = 64; + // BN=128 默认：BM=64 + BN=128 → wgmma 64x64x32 × 2 wgmma per + // K_block，b_decoded 在每次 wgmma 内部服务 64 cols。 + // env DG_W4_PATHB_BM64_BN=256 时尝试更大 grid 减半。 + const int bm64_bn_override = env_int("DG_W4_PATHB_BM64_BN", 0); + layout.block_n = (bm64_bn_override == 256) ? 256 : 128; + } else { + // 纯 fan-out（hint<=16）：base 阶梯，BN=256 + if (bm_select_m <= 8) layout.block_m = 8; + else if (bm_select_m <= 16) layout.block_m = 16; + else if (bm_select_m <= 32) layout.block_m = 32; + else if (bm_select_m <= 64) layout.block_m = 64; + layout.block_n = 256; + } + } else { + // path-A：cooperative prefetch + sfb→smem，BM 阶梯有效。 + if (bm_select_m <= 8) layout.block_m = 8; + else if (bm_select_m <= 16) layout.block_m = 16; + else if (bm_select_m <= 32) layout.block_m = 32; + else if (bm_select_m <= 64) layout.block_m = 64; + layout.block_n = 256; + + // DSV4 + 大 BM：BN=256→128 减小单 tile promote/store 链长， + // 让 grid 上 n-tile 翻倍提升 SM 间负载均衡。 + if (layout.block_m >= 64 and dsv4_shape and + env_int("DG_W4_LARGE_BM_BN128", 1) != 0) { + layout.block_n = 128; + } + // DSV4 + small hot (bm_select∈(32,64])：BM=64 padding 浪费太大， + // 改 BM=32 BN=128 兼顾 hot/small group。 + if (bm_select_m > 32 and bm_select_m <= 64 and dsv4_shape and + env_int("DG_W4_SMALL_HOT_BM32", 1) != 0) { + layout.block_m = 32; + layout.block_n = 128; + } } - layout.block_n = 256; + // 历史经验注记： + // * 曾尝试 path-B BM=32→64 升级（hot hint），device kernel 的 + // kK32QuadReduce 在 BM>=64 + swap_AB 下 sfb 寻址 / 4-wgmma fence + // 跟 BM=32 紧耦合，dp4 退化到 0.37x。 + // * 已知天花板（DSV4 skew, groups=32, N=4096）：进一步提升须从 + // device 侧扩展 BM=64 quad_reduce，工作量大且收益不确定。 } } if (not block_m_override and not block_n_override) { @@ -770,8 +1039,40 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( int chosen_stages = max_fitting; while (chosen_stages >= 3 and not fits(chosen_stages)) -- chosen_stages; - if (gran_k_b == 32 and expected_m <= 8) + if (gran_k_b == 32 and expected_m <= 16 and + static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) chosen_stages = std::min(chosen_stages, 6); + const bool bm32_skew_layout = gran_k_b == 32 and + config.layout.block_m == 32 and + (config.layout.block_n == 128 or config.layout.block_n == 256) and + static_cast(desc.m) >= 1024 and + ((desc.num_groups == 32 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) or + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： + // 放开到 g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168} + (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 32 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + const int bm32_skew_stage_override = env_int("DG_W4_BM32_SKEW_STAGES", 0); + if (bm32_skew_layout and bm32_skew_stage_override > 0) { + DG_HOST_ASSERT(bm32_skew_stage_override >= 3 and bm32_skew_stage_override <= kW4DefaultMaxStages); + DG_HOST_ASSERT(fits(bm32_skew_stage_override)); + chosen_stages = bm32_skew_stage_override; + } else if (bm32_skew_layout and fits(kW4DefaultMaxStages)) { + chosen_stages = kW4DefaultMaxStages; + } DG_HOST_ASSERT(chosen_stages >= 3); config.pipeline_config.num_stages = chosen_stages; config.pipeline_config.smem_size = smem_extra + chosen_stages * merged_per_stage; @@ -780,7 +1081,12 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // R2b-A swap_ab maps original N onto WGMMA M. Use enough math warpgroups // to cover the 64-row WGMMA-M strips selected by BLOCK_N. DG_HOST_ASSERT(config.layout.block_n == 64 or config.layout.block_n == 128 or config.layout.block_n == 256); - const int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; + int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; + const int rs_num_math_threads_override = env_int("DG_W4_RS_MATH_THREADS", 0); + if (rs_num_math_threads_override > 0) { + DG_HOST_ASSERT(rs_num_math_threads_override == 128 or rs_num_math_threads_override == 256); + rs_num_math_threads = rs_num_math_threads_override; + } config.launch_config.num_math_threads = rs_num_math_threads; config.launch_config.num_threads = config.launch_config.num_tma_threads + rs_num_math_threads; @@ -801,12 +1107,137 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - const bool scale_b_direct_load = gran_k_b == 32 and expected_m <= 16; - const bool k32_quad_reduce = gran_k_b == 32 and expected_m <= 8; + const bool bm32_skew_fast_path = gran_k_b == 32 and + config.layout.block_m == 32 and + (config.layout.block_n == 128 or config.layout.block_n == 256) and + static_cast(desc.m) >= 1024 and + ((desc.num_groups == 32 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168)) or + // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）， + // 与 bm32_skew_layout / dsv4_shape 同步生效。device fast-path 三件套 + // (quad_reduce / direct_load / compact_sched) 仅硬约束 BM==32 + g<=32， + // 与 N/K 数值无关，物理可命中。 + (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 32 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + // Fused-decode 路径（path-B 通用 cache_sfb_k32 + LUT decode）： + // * scale 在 cache 阶段编进 LUT，wgmma 后省一次 fmul（按 quad 算 = 4 次） + // * sfb 4x packed UE8M0：每 4 个 e8m0 打成 1 个 int32，体积 = fp32 的 1/4 + // 与 fast-path（kK32QuadReduce + kScaleBDirectLoad）**互斥**（cuh 强制 + // not kScaleBDirectLoad），开启后强制走通用路径。需要 sfb 是 packed int 布局。 + // 默认关，由 DG_W4_PATHB_FUSE_DECODE=1 显式开启。 + const bool fuse_scale_b_decode = + gran_k_b == 32 and + sfb.scalar_type() == torch::kInt and + get_major_type_ab(sfb) == cute::UMMA::Major::MN and + env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; + const bool scale_b_packed_ue8m0 = fuse_scale_b_decode; + // fuse_scale_b_decode 一旦开启，所有 fast-path / direct-load / e8m0 / bf16 + // 互斥关闭——device 那边 static_assert 会拒绝同时开启。 + const bool scale_b_direct_load = + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and + not fuse_scale_b_decode; + const bool k32_quad_reduce = + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and + not fuse_scale_b_decode; + // RS baseline for direct E8M0 B scales: keep scale products in split form by default. + // The exponent-adjust path is experimental and can regress some small-M shapes. + const bool scale_b_pow2_promote = + k32_quad_reduce and env_enabled("DG_W4_SCALE_B_POW2_PROMOTE"); + const bool k32_quad_split_promote = + k32_quad_reduce and not env_disabled("DG_W4_K32_QUAD_SPLIT_PROMOTE"); + const bool k32_quad_scale_b_inline = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_INLINE"); + const bool k32_quad_scale_b_prefetch = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_PREFETCH"); + const bool k32_quad_scale_b_vec4 = + k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_VEC4"); + const bool k32_quad_pair4x2_promote = + k32_quad_reduce and not k32_quad_split_promote and + env_enabled("DG_W4_K32_QUAD_PAIR4X2_PROMOTE"); + // small_m_simple_sched：device 端 kUseSmallMSimpleSched 仅编译期检查 + // BLOCK_M<=16 + GroupedMasked + multicast=1，与 N/K 数值无关。 + // 默认守护 (k>=4096 + n<=4096) 来自历史 g32 + n=4096 dsv4 形状的保守覆盖。 + // 放开到 RELAX 形状集 (g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}) + // 让 DSV4 EP 业务真实 shape (g24 + n∈{6144,7168} + k∈{3072,7168} + expected_m=1/2/3) + // 也能命中 simple_sched，避开通用 masked scheduler 的额外开销。 const bool small_m_simple_sched = - gran_k_b == 32 and expected_m <= 8 and - static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096; - DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); + gran_k_b == 32 and expected_m <= 16 and + ((static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) or + (static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 32 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168))); + const bool compact_masked_sched = + bm32_skew_fast_path and not env_disabled("DG_W4_COMPACT_MASKED_SCHED"); + // compact_masked_sched 按 m_max 降序遍历 active group（实验性扩展，默认关）： + // 想法：compact_masked_sched 默认按 group_idx 升序遍历 active group。 + // 当各 group masked_m 不均时，long-tail group 决定 last-wave 收敛节奏。 + // 开启后 inner loop 每步在剩余 active mask 内 O(active) 扫一次找 group_m + // 最大者，让 wave 0 优先吃重 group，wave 末尾留小 group 收尾。 + // 实测结果（DG_W4_PATHB_REORDER_BY_MM=1 vs 默认 fast-path）： + // gateup_dense_tail_hot64 356us → 356us (0%) + // gateup_dense_tail_hot128 399us → 396us (+0.8%) + // gateup_dense_tail_hot214 401us → 400us (-0.2%) + // gateup_mtp_dp4 157us → 157us (0%) + // down_dense_tail_hot64 186us → 185us (+0.5%) + // down_mtp_512_hot384_a8 98us → 100us (-2%) + // down_dense_tail_hot214 206us → 206us (0%) + // gateup_one_hot_32 31us → 31us (+1%, 边缘 case) + // 全军 ±2% 噪声振荡，无任何 case 出现 5-10% 实质收益。 + // 失败根因：compact_masked_sched 已经做了 (group, m_block, n_block) 三维紧凑 + // 展开，dense_tail_hot64 总 tile = 32 group × 2 m_blocks × 32 n_blocks = + // 2048 ≈ 15.5 wave × 132 SM。当 total_tiles >> num_sms 时，wave 边界由 + // ceil(total/132) 决定，与 group 顺序无关；重排只改变"哪些 SM 在 wave 0 + // 拿到 hot group"，last-wave 余数 = (total mod 132) 不变。 + // 边缘 case (one_hot_32 total=32 tile = 1 wave × 32 SM) 的 +1-2% 收益 + // 是 SM 提前进入 idle 的小幅噪声，不构成可推广收益。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + // 仅在 compact_masked_sched 已开启时生效（即 bm32_skew_fast_path）， + // 否则即便 env 开启也不会传到 device。 + const bool reorder_masked_by_max_m = + compact_masked_sched and env_int("DG_W4_PATHB_REORDER_BY_MM", 0) != 0; + // bf16 SFB：体积砍半，**默认开启**。两条 fast-path： + // - path-B (gran_k_b==32 + direct-load)：load_sfb / load_k32_quad_scale_b 直读 gmem。 + // - path-A (gran_k_b==128)：cooperative prefetch 经 smem，smem 也存 bf16。 + // 仍要避开 packed-UE8M0 / fused-decode 等假设 fp32 位级布局的分支。 + // 显式 `DG_W4_SCALE_B_BF16=0` 时回退 fp32；用户传的 sfb 实际是 fp32 也自然回退。 + const bool scale_b_bf16 = + ((scale_b_direct_load and gran_k_b == 32 and not k32_quad_scale_b_prefetch) or + gran_k_b == 128) and + not env_disabled("DG_W4_SCALE_B_BF16") and + sfb.scalar_type() == torch::kBFloat16; + // E8M0 SFB（仅 path-B fast-path）：每元素 1B = fp32 的 8 位指数，体积再砍 2x。 + // 解码 `__uint_as_float(uint32(e) << 23)` 零误差。**默认开启**：当用户传入 + // uint8 sfb 时自动启用；显式 `DG_W4_SCALE_B_E8M0=0` 时回退。 + // 与 bf16 互斥（因 sfb 实际 dtype 已不同），不会同时命中。 + const bool scale_b_e8m0 = + scale_b_direct_load and gran_k_b == 32 and + not k32_quad_scale_b_prefetch and + (not env_disabled("DG_W4_SCALE_B_E8M0") or + env_enabled("DG_W4_SCALE_B_E8M0_ONLY")) and + sfb.scalar_type() == torch::kUInt8; + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat or + (scale_b_bf16 and sfb.scalar_type() == torch::kBFloat16) or + (scale_b_e8m0 and sfb.scalar_type() == torch::kUInt8) or + (fuse_scale_b_decode and sfb.scalar_type() == torch::kInt)); const SM90FP8FP4Gemm1D2DRSRuntime::Args& rs_args = { .gemm_desc = desc, .gemm_config = config, @@ -815,9 +1246,22 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( config.layout.get_cluster_size()), .major_sfb = get_major_type_ab(sfb), .scale_b_direct_load = scale_b_direct_load, + .scale_b_pow2_promote = scale_b_pow2_promote, .k32_quad_reduce = k32_quad_reduce, + .k32_quad_split_promote = k32_quad_split_promote, + .k32_quad_scale_b_inline = k32_quad_scale_b_inline, + .k32_quad_scale_b_prefetch = k32_quad_scale_b_prefetch, + .k32_quad_scale_b_vec4 = k32_quad_scale_b_vec4, + .k32_quad_pair4x2_promote = k32_quad_pair4x2_promote, .small_m_simple_sched = small_m_simple_sched, + .compact_masked_sched = compact_masked_sched, + .fuse_scale_b_decode = fuse_scale_b_decode, + .scale_b_packed_ue8m0 = scale_b_packed_ue8m0, .scale_b_gran_k = static_cast(gran_k_b), + .b_is_int4_sym = b_is_int4_sym, + .scale_b_bf16 = scale_b_bf16, + .scale_b_e8m0 = scale_b_e8m0, + .reorder_masked_by_max_m = reorder_masked_by_max_m, .gmem_b_ptr = b.first.data_ptr(), .gmem_d_ptr = d.data_ptr(), .sfb = sfb.data_ptr(), From 8d4894f8c91249488a1654bcc41a01f8662eebc0 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:44:57 +0800 Subject: [PATCH 32/51] Update sm90_fp8_fp4_gemm_1d2d_rs.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 54 +++++++++++++++---- 1 file changed, 45 insertions(+), 9 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 4cb03b9f4e..6851476d11 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -29,9 +29,35 @@ class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime); }}; @@ -147,7 +178,7 @@ static void __instantiate_kernel() {{ "false", "false", "false", - "false", + args.scale_b_pow2_promote ? "true" : "false", "false", "false", args.scale_b_direct_load ? "true" : "false", @@ -176,17 +207,18 @@ static void __instantiate_kernel() {{ args.k32_quad_reduce ? "true" : "false", "false", args.k32_quad_reduce ? "true" : "false", - "false", - "false", - "false", - "false", - "false", + args.k32_quad_split_promote ? "true" : "false", + args.k32_quad_scale_b_inline ? "true" : "false", + args.k32_quad_scale_b_prefetch ? "true" : "false", + args.k32_quad_scale_b_vec4 ? "true" : "false", + args.k32_quad_pair4x2_promote ? "true" : "false", "false", "false", "false", "false", args.small_m_simple_sched ? "true" : "false", - "false", + args.compact_masked_sched ? "true" : "false", + args.fuse_scale_b_decode ? "true" : "false", "false", "false", "false", @@ -197,11 +229,15 @@ static void __instantiate_kernel() {{ "false", "false", 0, - "false", + args.scale_b_packed_ue8m0 ? "true" : "false", args.scale_b_gran_k, 1, 1, - 0); + 0, + args.b_is_int4_sym ? "true" : "false", + args.scale_b_bf16 ? "true" : "false", + args.scale_b_e8m0 ? "true" : "false", + args.reorder_masked_by_max_m ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { From 29cfbf6e71261134965a97891b42d691aff29c8e Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:45:31 +0800 Subject: [PATCH 33/51] Update smxx_layout.hpp --- csrc/jit_kernels/impls/smxx_layout.hpp | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index af207ebc95..ef022e6c07 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -107,8 +107,14 @@ static std::tuple preprocess_sf(const to if (require_float) { DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); } else { - DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat or sf.scalar_type() == torch::kInt); - DG_HOST_ASSERT(sf.element_size() == sizeof(float)); + // 放过 bf16 / E8M0 (UInt8) sfb（path-B fast-path，体积砍半 / 砍 4 倍）。 + DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat or + sf.scalar_type() == torch::kInt or + sf.scalar_type() == torch::kBFloat16 or + sf.scalar_type() == torch::kUInt8); + DG_HOST_ASSERT(sf.element_size() == sizeof(float) or + sf.scalar_type() == torch::kBFloat16 or + sf.scalar_type() == torch::kUInt8); } const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf; @@ -124,6 +130,10 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn) return (dim == 2) ? batched_sf.squeeze(0) : batched_sf; + // bf16 / E8M0 (UInt8) sfb 必须由调用方按 MN-major + tma_aligned_mn 直接构造， + // 不能落入下面的 fp32-only TransposeFP32Runtime。 + DG_HOST_ASSERT(sf.scalar_type() != torch::kBFloat16 and sf.scalar_type() != torch::kUInt8); + const auto out = torch::empty_strided({num_groups, mn, sf_k}, {tma_aligned_mn * sf_k, 1, tma_aligned_mn}, batched_sf.options()); From 8d8aaf61e570ec4d2c5d69b651cb0ca390fa73e6 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:48:18 +0800 Subject: [PATCH 34/51] Update layout.hpp --- csrc/utils/layout.hpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 780707d426..17b08e4831 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -92,7 +92,9 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, // Always do shape checks const auto sf_dtype = sf.scalar_type(); - DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt); + // BF16 / E8M0 (UInt8) SFB fast path 也允许；其它情况维持 fp32/int 限制。 + DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt or + sf_dtype == torch::kBFloat16 or sf_dtype == torch::kUInt8); DG_HOST_ASSERT(sf.dim() == static_cast(num_groups.has_value()) + 2); if (num_groups.has_value()) DG_HOST_ASSERT(sf.size(-3) == num_groups.value()); @@ -101,7 +103,8 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, const bool scale_b_lut = sf_dtype == torch::kInt and scale_b_lut_env != nullptr and scale_b_lut_env[0] != '\0' and scale_b_lut_env[0] != '0'; const int expected_sf_k = scale_b_lut ? - ceil_div(k, gran_k) * 2 : ceil_div(k, gran_k * (sf_dtype == torch::kFloat ? 1 : 4)); + ceil_div(k, gran_k) * 2 : + ceil_div(k, gran_k * (sf_dtype == torch::kInt ? 4 : 1)); DG_HOST_ASSERT(sf.size(-1) == expected_sf_k); // TMA stride checks: TMA aligned and MN-major From 06037919fa2b62a7d4b7e0b717e1c40c7c7eb5b2 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:50:33 +0800 Subject: [PATCH 35/51] Update math.py --- deep_gemm/utils/math.py | 81 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index 318d927101..3153d41cd0 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -141,3 +141,84 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128 group_idx = torch.arange(n, device=packed.device) // gran_k x_restored = x_dequantized * sf[:, group_idx] return x_restored + + +# --------------------------------------------------------------------------- +# INT4 symmetric (signed [-8, 7]) helpers used by sm90 INT4-A8 path. +# 4-bit two's-complement codes: 0..7 -> 0..7, 8..15 -> -8..-1. +# Two nibbles are packed into one int8 byte, low nibble first. +# +# All 16 INT4 sym values are exactly representable in FP8 E4M3, so we provide +# a lossless INT4 -> packed E4M3 byte converter that lets us drive the existing +# FP8xFP8 fused kernel with INT4-quantised B for end-to-end accuracy testing +# while a dedicated INT4-A8 fused kernel is being landed. +# --------------------------------------------------------------------------- + +def _int4_sym_decode(nibble: torch.Tensor) -> torch.Tensor: + """Map low-4-bit two's-complement nibble to signed int in [-8, 7].""" + nib = (nibble & 0x0F).to(torch.int32) + return torch.where(nib >= 8, nib - 16, nib) + + +def per_token_cast_to_int4_sym(x: torch.Tensor, gran_k: int = 128 + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-(row, K-block) symmetric INT4 quantisation. + + Returns (packed, sf): + packed: int8 tensor of shape (m, n // 2), two nibbles per byte (low first). + sf: fp32 tensor of shape (m, n // gran_k), absmax / 7. + """ + assert x.dim() == 2 + m, n = x.shape + assert n % 2 == 0 + padded_n = align(n, gran_k) + x_padded = torch.zeros((m, padded_n), dtype=x.dtype, device=x.device) + x_padded[:, :n] = x + x_view = x_padded.view(m, padded_n // gran_k, gran_k) + x_amax = x_view.abs().float().amax(dim=2).clamp_min(1e-4) + sf = x_amax / 7.0 + x_scaled = x_view.float() * (1.0 / sf.unsqueeze(2)) + # Round-half-to-even and clamp to int4 sym range. + q = torch.round(x_scaled).clamp_(-8, 7).to(torch.int8).view(m, padded_n) + # Pack two nibbles per byte; mask with 0x0F first to drop the sign extension. + q2 = q.view(m, padded_n // 2, 2) + packed = ((q2[:, :, 0] & 0x0F) | ((q2[:, :, 1] & 0x0F) << 4)).to(torch.int8) + return packed[:, : n // 2].contiguous(), sf + + +def cast_back_from_int4_sym(packed: torch.Tensor, sf: torch.Tensor, + gran_k: int = 128) -> torch.Tensor: + """Dequantise INT4-sym packed weights back to fp32.""" + m, n2 = packed.shape + n = n2 * 2 + lo = _int4_sym_decode(packed) + hi = _int4_sym_decode(packed >> 4) + unpacked = torch.empty((m, n), dtype=torch.float, device=packed.device) + unpacked[:, 0::2] = lo.float() + unpacked[:, 1::2] = hi.float() + group_idx = torch.arange(n, device=packed.device) // gran_k + return unpacked * sf[:, group_idx] + + +# Lossless 16-entry LUT: INT4 sym (4-bit two's-complement) -> FP8 E4M3 byte. +# Verified: every signed int in [-8, 7] is exactly representable in E4M3, so +# casting q.float().to(torch.float8_e4m3fn) round-trips identically. +def int4_sym_to_e4m3_bytes(packed_int4: torch.Tensor) -> torch.Tensor: + """Convert packed INT4-sym (2 nibbles / byte) to unpacked E4M3 bytes. + + Args: + packed_int4: int8 tensor (..., n // 2). Low nibble is element 2*i, + high nibble is element 2*i+1. + Returns: + torch.float8_e4m3fn tensor (..., n) with byte-identical encoding to + a fresh fp32 -> fp8_e4m3 cast of the dequantised integer values. + """ + assert packed_int4.dtype == torch.int8 + *prefix, n2 = packed_int4.shape + n = n2 * 2 + lo = _int4_sym_decode(packed_int4) + hi = _int4_sym_decode(packed_int4 >> 4) + unpacked = torch.empty((*prefix, n), dtype=torch.int32, device=packed_int4.device) + unpacked[..., 0::2] = lo + unpacked[..., 1::2] = hi + return unpacked.float().to(torch.float8_e4m3fn) From 48689a1a13187df1ceb8d68519b0583f52f37f88 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Wed, 27 May 2026 13:52:07 +0800 Subject: [PATCH 36/51] Update test_sm90_fp8_fp4.py --- tests/test_sm90_fp8_fp4.py | 503 ++++++++++++++++++++++++++++++++++++- 1 file changed, 502 insertions(+), 1 deletion(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 706c88ce9d..9f2f687d89 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -64,6 +64,27 @@ def _effective_bytes( return a_bytes + b_data_bytes + b_scale_bytes + d_bytes +def _effective_masked_bytes( + groups: int, + masked_m_values: list[int], + n: int, + k: int, + a_gran_k: int, + *, + fp8_b: bool, + b_gran_k: int | None = None, +) -> int: + b_gran_k = a_gran_k if b_gran_k is None else b_gran_k + logical_m = sum(masked_m_values) + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + + def _build_grouped_layout(groups: int, m_per_group: int): m = groups * m_per_group group_starts = [group_id * m_per_group for group_id in range(groups)] @@ -291,6 +312,238 @@ def run_fp8(): } +def _masked_skew_benchmark_case( + name: str, + masked_m_values: list[int], + expected_m: int, + n: int, + k: int, + max_m: int = 1024, + a_gran_k: int = 128, + b_gran_k: int = 32, + pass_hints: bool = True, +) -> dict[str, float | int | str]: + sm90_masked_w4 = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if sm90_masked_w4 is None: + raise RuntimeError("SM90 FP8xFP4 masked fused kernel is not exposed yet.") + + groups = len(masked_m_values) + assert groups > 0 + assert max(masked_m_values) <= max_m + masked_m_max_hint = max(masked_m_values) if pass_hints else None + active_groups_hint = ( + sum(1 for v in masked_m_values if v > 0) if pass_hints else None + ) + masked_m = torch.tensor(masked_m_values, device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // a_gran_k), device="cuda", dtype=torch.float) + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) + b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) + block_m_override = int(os.getenv("DG_W4_BLOCK_M_OVERRIDE", "0")) or None + block_n_override = int(os.getenv("DG_W4_BLOCK_N_OVERRIDE", "0")) or None + bm32_skew_fast_path = ( + b_gran_k == 32 + and masked_m_max_hint is not None + and masked_m_max_hint > 16 + and os.getenv("DG_W4_PATHB_FUSE_DECODE", "0") == "0" + and os.getenv("DG_W4_PATHB_FAST_PATH", "1") != "0" + and os.getenv("DG_W4_PATHB_BM64", "0") == "0" + and (block_m_override is None or block_m_override == 32) + and (block_n_override is None or block_n_override in (128, 256)) + and max_m >= 1024 + and groups == 32 + and n in (4096, 7168) + and k in (2048, 3072, 4096, 7168) + ) + scale_b_direct_load = b_gran_k == 32 and ( + expected_m <= 16 or bm32_skew_fast_path + ) + scale_b_dtype_fast_path = ( + scale_b_direct_load + and os.getenv("DG_W4_K32_QUAD_SCALE_B_PREFETCH", "0") == "0" + and os.getenv("DG_W4_SCALE_B_POW2_PROMOTE", "0") == "0" + ) + # bf16 sfb：path-A (gran_k_b=128) 与 path-B fast-path (gran_k_b=32) 都支持， + # 体积砍半。按 MN-major + tma_aligned_mn=align(N, 8) 直接构造，避开 host 端 + # fp32-only 的 transpose 路径。**默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 + use_bf16_b_sf = ( + ((b_gran_k == 32 and scale_b_dtype_fast_path) or b_gran_k == 128) + and not use_packed_b_sf + and bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) + ) + # E8M0 sfb（uint8）：仅 path-B (gran_k_b=32)。每元素 = fp32 pow2 scale 的指数位。 + # per_token_cast_to_fp4(..., use_ue8m0=True) 已保证 sf 是严格 pow2，因此抽指数无损。 + # **默认开启**，DG_W4_SCALE_B_E8M0=0 时回退。优先级高于 bf16（互斥）。 + use_e8m0_b_sf = ( + b_gran_k == 32 + and scale_b_dtype_fast_path + and not use_packed_b_sf + and bool(int(os.getenv("DG_W4_SCALE_B_E8M0", "1"))) + ) + if use_e8m0_b_sf: + use_bf16_b_sf = False # e8m0 优先 + if use_e8m0_b_sf: + # uint8: tma_aligned_mn = ceil(N, 16)，要求 N % 16 == 0。 + assert n % 16 == 0 + tma_aligned_n = (n + 15) // 16 * 16 + b_sf = torch.empty_strided( + (groups, n, b_sf_k), + (tma_aligned_n * b_sf_k, 1, tma_aligned_n), + device="cuda", + dtype=torch.uint8, + ) + b_sf_fp32 = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) + elif use_bf16_b_sf: + # bf16: tma_aligned_mn = ceil(N, 8)，要求 N % 8 == 0。 + assert n % 8 == 0 + tma_aligned_n = (n + 7) // 8 * 8 + b_sf = torch.empty_strided( + (groups, n, b_sf_k), + (tma_aligned_n * b_sf_k, 1, tma_aligned_n), + device="cuda", + dtype=torch.bfloat16, + ) + b_sf_fp32 = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) + else: + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + b_sf_fp32 = b_sf + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k + ) + b_fp4[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp4( + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + if use_e8m0_b_sf: + # b_sf_fp32 已是严格 pow2（per_token_cast_to_fp4 use_ue8m0=True）。 + # 抽 fp32 指数位 = ((bits >> 23) & 0xff)：sign=0、mantissa=0 时 fp32 == 2^(e-127)。 + b_sf_bits = b_sf_fp32.view(torch.int32) + b_sf_e8m0 = ((b_sf_bits >> 23) & 0xff).to(torch.uint8) + b_sf.copy_(b_sf_e8m0) + elif use_bf16_b_sf: + # 写入 MN-major bf16 buffer：fp32 → bf16 round-down，再按目标 stride copy。 + b_sf.copy_(b_sf_fp32.to(torch.bfloat16)) + a = (a_data, a_sf) + b_w4 = (b_fp4, b_sf) + + # caller hint：把 hot group 的真实大小传给 host，host 据此选 BM。 + # path-A 通过 masked_m_max_hint 接收 hint；FP8 baseline 没有 hint API， + # 但它的 expected_m 可以直接传 max（FP8 路径无 expected_m<=8 fast-path 副作用）， + # 这样两条路径都按 hot 调度，speedup 反映的是算法差距而非 caller-side gap。 + gemm_expected_m = expected_m + # pass_hints=False 时 max_hint=None：两边都用 expected_m（mirror 业务真实 + # caller 不传 hint 的状态，host 走没有 hint 的 small-m candidate 路径）。 + fp8_expected_m = ( + max(masked_m_max_hint, expected_m) + if masked_m_max_hint is not None + else expected_m + ) + + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + # ref 反量化必须用 fp32 scale（cast_back_from_fp4 直接做 sf 乘法，不会解码 e8m0/bf16）。 + # use_e8m0_b_sf / use_bf16_b_sf 时 b_w4[1] 是 uint8/bf16 编码，需走 b_sf_fp32 兜底。 + b_sf_for_ref = b_sf_fp32 if (use_e8m0_b_sf or use_bf16_b_sf) else b_w4[1] + for group_id, valid_m in enumerate(masked_m_values): + b_dequant = cast_back_from_fp4( + b_w4[0][group_id], b_sf_for_ref[group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + ) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant, use_ue8m0=False, gran_k=a_gran_k + ) + if valid_m > 0: + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + sm90_masked_w4( + a, + b_w4, + d_w4, + masked_m, + gemm_expected_m, + gran_k=a_gran_k, + gran_k_a=a_gran_k, + gran_k_b=b_gran_k, + masked_m_max_hint=masked_m_max_hint, + active_groups_hint=active_groups_hint, + ) + + run_w4() + w4_diff = max( + calc_diff(d_w4[group_id, :valid_m], ref[group_id, :valid_m]) if valid_m > 0 else 0.0 + for group_id, valid_m in enumerate(masked_m_values) + ) + w4_elapsed = _time_cuda(run_w4) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, + b_fp8, + d_fp8, + masked_m, + fp8_expected_m, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + ) + + run_fp8() + fp8_diff = max( + calc_diff(d_fp8[group_id, :valid_m], ref[group_id, :valid_m]) if valid_m > 0 else 0.0 + for group_id, valid_m in enumerate(masked_m_values) + ) + fp8_elapsed = _time_cuda(run_fp8) + + w4_bytes = _effective_masked_bytes(groups, masked_m_values, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_masked_bytes(groups, masked_m_values, n, k, a_gran_k, fp8_b=True) + return { + "case": name, + "groups": groups, + "expected_m": expected_m, + "masked_m_hint": masked_m_max_hint, + "b_gran_k": b_gran_k, + "max_m": max_m, + "sum_m": sum(masked_m_values), + "masked_max": max(masked_m_values), + "active_groups": sum(1 for value in masked_m_values if value > 0), + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + + +def _print_skew_table(rows: list[dict[str, float | int | str]]) -> None: + print("case | groups | b_gran_k | exp_m_avg | hint | sum_m | max_m | active | n | k | " + "W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + for row in rows: + print( + f"{row['case']} | {row['groups']} | {row['b_gran_k']} | {row['expected_m']} | " + f"{row['masked_m_hint']} | {row['sum_m']} | " + f"{row['masked_max']} | {row['active_groups']} | {row['n']} | {row['k']} | " + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " + f"{row['speedup']:.2f}x" + ) + + def _accuracy_case( groups: int, m_per_group: int, @@ -409,11 +662,259 @@ def test_sm90_fp8_fp4_masked_direct_fp32_scale() -> None: _print_markdown_table(rows) +def test_sm90_fp8_fp4_masked_skew_cases() -> None: + _require_sm90() + torch.manual_seed(4) + + def skew_values( + total: int, hot: int, active: int = 8, groups: int = 32 + ) -> list[int]: + assert active >= 1 + assert active <= groups + assert total >= hot + values = [0] * groups + values[0] = hot + remaining = total - hot + for idx in range(1, active): + share = (remaining + active - idx - 1) // (active - idx) + values[idx] = share + remaining -= share + assert sum(values) == total + return values + + def values_from_active(active_values: list[int], groups: int = 32) -> list[int]: + assert len(active_values) <= groups + assert all(value >= 0 for value in active_values) + return active_values + [0] * (groups - len(active_values)) + + def values_from_repeated(value: int, active: int, groups: int = 32) -> list[int]: + assert active <= groups + return values_from_active([value] * active, groups=groups) + + def long_tail_values( + hot: int, tail_start: int, active: int, groups: int = 32 + ) -> list[int]: + values = [hot] + next_value = tail_start + for _ in range(active - 1): + values.append(max(1, next_value)) + next_value = max(1, next_value // 2) + return values_from_active(values, groups=groups) + + print("skewed masked case: shape mirrors DSV4 MTP verify; " + "groups=32, b_gran_k=32 walks path-B (k32 fast path) " + "with b.second shape [groups, N, K/32]") + rows = [] + shapes = [ + ("gateup", 4096, 4096), + ("down", 4096, 2048), + # Non-DSV4 dimensions keep the same group/masked pattern but stress + # different N/K ratios that can expose scheduler or scale-load cliffs. + ("wide_n", 7168, 2048), + ("wide_k", 4096, 7168), + ] + distributions = [ + # Uniform small-M cases validate that hint-aware BM32 selection does not + # regress the common non-skew path. + ("uniform_1", [1] * 32, 1), + ("uniform_2", [2] * 32, 2), + ("uniform_4", [4] * 32, 4), + ("uniform_8", [8] * 32, 8), + ("uniform_16", [16] * 32, 16), + ("uniform_32", [32] * 32, 32), + # Around the BM16 -> BM32 hot threshold. + ("one_hot_15", values_from_active([15]), 1), + ("one_hot_16", values_from_active([16]), 1), + ("one_hot_17", values_from_active([17]), 1), + ("one_hot_24", values_from_active([24]), 1), + ("one_hot_32", values_from_active([32]), 1), + ("one_hot_48", values_from_active([48]), 2), + ("one_hot_64", values_from_active([64]), 2), + ("one_hot_128", values_from_active([128]), 4), + ("one_hot_214", values_from_active([214]), 1), + ("one_hot_384", values_from_active([384]), 12), + # Original MTP verify distributions copied from observed DP logs. + ("mtp_dp2", skew_values(total=144, hot=50), 7), + ("mtp_dp0", skew_values(total=195, hot=160), 7), + ("mtp_dp4", skew_values(total=290, hot=214), 7), + # Same total token count but different active/hotness patterns. + ("mtp_144_hot96_a4", skew_values(total=144, hot=96, active=4), 7), + ("mtp_144_hot96_a8", skew_values(total=144, hot=96, active=8), 7), + ("mtp_195_hot96_a16", skew_values(total=195, hot=96, active=16), 7), + ("mtp_290_hot160_a16", skew_values(total=290, hot=160, active=16), 7), + ("mtp_384_hot256_a8", skew_values(total=384, hot=256, active=8), 12), + ("mtp_512_hot384_a8", skew_values(total=512, hot=384, active=8), 16), + # Multi-hot cases mimic router concentration on a few experts rather + # than a single dominant expert. + ("two_hot_64_64", values_from_active([64, 64]), 4), + ("two_hot_128_64", values_from_active([128, 64]), 6), + ("two_hot_160_96", values_from_active([160, 96]), 8), + ("four_hot_32", values_from_repeated(32, active=4), 4), + ("four_hot_64", values_from_repeated(64, active=4), 8), + ("eight_hot_32", values_from_repeated(32, active=8), 8), + ("eight_hot_64", values_from_repeated(64, active=8), 16), + # Long tails stress compact masked scheduling and active-group scanning. + ("longtail_128_a8", long_tail_values(hot=128, tail_start=64, active=8), 8), + ("longtail_214_a8", long_tail_values(hot=214, tail_start=48, active=8), 7), + ("longtail_256_a16", long_tail_values(hot=256, tail_start=64, active=16), 12), + # Dense active but skewed cases can happen when all experts receive a + # few tokens and one or two experts still become hot. + ( + "dense_tail_hot64", + values_from_active([64, 32, 16, 8] + [4] * 28), + 8, + ), + ( + "dense_tail_hot128", + values_from_active([128, 64, 32, 16] + [4] * 28), + 8, + ), + ( + "dense_tail_hot214", + values_from_active([214, 64, 32, 16] + [4] * 28), + 8, + ), + ] + # for shape_name, n, k in shapes: + # for dist_name, masked_m_values, expected_m in distributions: + # rows.append( + # _masked_skew_benchmark_case( + # f"{shape_name}_{dist_name}", + # masked_m_values, + # expected_m=expected_m, + # n=n, + # k=k, + # max_m=1024, + # b_gran_k=32, + # ) + # ) + + group24_shapes = [ + ("g24_m4096_n6144_k7168", 6144, 7168, 4096), + ("g24_m4096_n7168_k3072", 7168, 3072, 4096), + ] + group24_distributions = [ + # Reproduce the warmup cliff that showed up as: + # m=17, max_m=4096, n=6144, k=7168, num_groups=24. + ("uniform_1", values_from_repeated(1, active=24, groups=24), 1), + ("uniform_8", values_from_repeated(8, active=24, groups=24), 8), + ("uniform_16", values_from_repeated(16, active=24, groups=24), 16), + ("uniform_17", values_from_repeated(17, active=24, groups=24), 17), + ("uniform_24", values_from_repeated(24, active=24, groups=24), 24), + ("uniform_32", values_from_repeated(32, active=24, groups=24), 32), + ("uniform_64", values_from_repeated(64, active=24, groups=24), 64), + # Around the E8M0 direct-load threshold and BM16/BM32 transition. + ("one_hot_16", values_from_active([16], groups=24), 16), + ("one_hot_17", values_from_active([17], groups=24), 17), + ("one_hot_32", values_from_active([32], groups=24), 32), + ("one_hot_64", values_from_active([64], groups=24), 64), + ("one_hot_128", values_from_active([128], groups=24), 128), + ("one_hot_256", values_from_active([256], groups=24), 256), + ("one_hot_512", values_from_active([512], groups=24), 512), + ("two_hot_17", values_from_active([17, 17], groups=24), 17), + ("two_hot_64_64", values_from_active([64, 64], groups=24), 64), + ("two_hot_128_64", values_from_active([128, 64], groups=24), 128), + ("four_hot_17", values_from_repeated(17, active=4, groups=24), 17), + ("four_hot_32", values_from_repeated(32, active=4, groups=24), 32), + ("four_hot_64", values_from_repeated(64, active=4, groups=24), 64), + ("eight_hot_17", values_from_repeated(17, active=8, groups=24), 17), + ("eight_hot_32", values_from_repeated(32, active=8, groups=24), 32), + ("eight_hot_64", values_from_repeated(64, active=8, groups=24), 64), + # MTP-like skew for 24 local groups, including the m=17 average. + ("mtp_408_hot160_a8", skew_values(408, 160, active=8, groups=24), 17), + ("mtp_408_hot256_a8", skew_values(408, 256, active=8, groups=24), 17), + ("mtp_576_hot384_a8", skew_values(576, 384, active=8, groups=24), 24), + ("mtp_768_hot512_a8", skew_values(768, 512, active=8, groups=24), 32), + # Long tails and dense active tails stress compact scheduling when + # active_groups is neither tiny nor fully uniform. + ( + "longtail_128_a8", + long_tail_values(128, tail_start=64, active=8, groups=24), + 17, + ), + ( + "longtail_256_a12", + long_tail_values(256, tail_start=96, active=12, groups=24), + 24, + ), + ( + "dense_tail_hot17", + values_from_active([17, 16, 8, 4] + [1] * 20, groups=24), + 17, + ), + ( + "dense_tail_hot64", + values_from_active([64, 32, 16, 8] + [4] * 20, groups=24), + 17, + ), + ( + "dense_tail_hot128", + values_from_active([128, 64, 32, 16] + [4] * 20, groups=24), + 24, + ), + ( + "dense_tail_hot256", + values_from_active([256, 128, 64, 32] + [8] * 20, groups=24), + 32, + ), + ] + for shape_name, n, k, max_m in group24_shapes: + for dist_name, masked_m_values, expected_m in group24_distributions: + rows.append( + _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, + k=k, + max_m=max_m, + b_gran_k=32, + ) + ) + + # DSV4 EP 真实业务 shape mirror：m=256 + expected_m∈{1,2,3} + max_hint=None + # （caller 没传 hint，business dispatch_output 不携带 max_hint/active_hint）。 + # 三组 shape：gateup(g24, n=6144, k=7168) / down(g24, n=7168, k=3072) + # 业务比例：expected_m=2 占 71%, expected_m=3 占 18%, expected_m=1 占 11%。 + # 物理 m=256 但 caller 不知道每个 group 的真实 masked_m，masked_m_values 设 + # expected_m * groups 模拟 uniform 分布，用 pass_hints=False 让 host 走没有 + # hint 的 small-m candidate + simple_sched 路径。 + business_shapes = [ + ("biz_gateup_g24_m256_n6144_k7168", 6144, 7168, 256), + ("biz_down_g24_m256_n7168_k3072", 7168, 3072, 256), + ] + business_distributions = [ + ("expected_m_1", values_from_repeated(1, active=24, groups=24), 1), + ("expected_m_2", values_from_repeated(2, active=24, groups=24), 2), + ("expected_m_3", values_from_repeated(3, active=24, groups=24), 3), + ] + for shape_name, n, k, max_m in business_shapes: + for dist_name, masked_m_values, expected_m in business_distributions: + rows.append( + _masked_skew_benchmark_case( + f"{shape_name}_{dist_name}", + masked_m_values, + expected_m=expected_m, + n=n, + k=k, + max_m=max_m, + b_gran_k=32, + pass_hints=False, + ) + ) + _print_skew_table(rows) + + # print("\nworst W4/FP8 speedup cases") + # _print_skew_table(sorted(rows, key=lambda row: float(row["speedup"]))[:12]) + + if __name__ == "__main__": start_time = time.time() # if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): # test_sm90_fp8_fp4_contiguous() - if os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): + if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + test_sm90_fp8_fp4_masked_skew_cases() + elif os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): test_sm90_fp8_fp4_masked_direct_fp32_scale() else: test_sm90_fp8_fp4_masked() From 720da9f9b3b1cd3092691a4e59077aa13eeb148b Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 2 Jun 2026 15:35:07 +0800 Subject: [PATCH 37/51] Update sm90_fp8_fp4_gemm_1d2d_rs.hpp --- csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 6851476d11..3b3310d394 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -149,7 +149,6 @@ static void __instantiate_kernel() {{ {}, {}, {}, - {}, {} >); }}; @@ -217,7 +216,6 @@ static void __instantiate_kernel() {{ "false", "false", args.small_m_simple_sched ? "true" : "false", - args.compact_masked_sched ? "true" : "false", args.fuse_scale_b_decode ? "true" : "false", "false", "false", From 180e71cd6b07fff4a80760d2d830d0aafce85315 Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Tue, 2 Jun 2026 15:35:45 +0800 Subject: [PATCH 38/51] Update sm90_fp8_fp4_gemm_1d2d_rs.cuh --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh | 46 +++++++++++++++---- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh index 24a14a6d8e..ab68ef4013 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh @@ -686,7 +686,11 @@ template + uint32_t kMathRegCap = 0, + bool kBIsInt4Sym = false, + bool kScaleBBF16 = false, + bool kScaleBE8M0 = false, + bool kReorderMaskedByMaxM = false> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, kLaunchBoundsMinBlocks) void sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, nv_bfloat16* gmem_d_ptr, @@ -708,6 +712,10 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo "DG_W4_SCALE_K_GROUP does not support per-32 FP4 scaling"); DG_STATIC_ASSERT(kScaleKGroup == 1 or kScaleKGroup == 2 or kScaleKGroup == 4, "DG_W4_SCALE_K_GROUP only supports 1/2/4"); + DG_STATIC_ASSERT(not kBIsInt4Sym, "RS-mode FP8xFP4 kernel does not support INT4-sym B"); + DG_STATIC_ASSERT(not (kScaleBBF16 and kScaleBE8M0), "Scale-B cannot be both BF16 and E8M0"); + DG_STATIC_ASSERT((not kScaleBBF16 and not kScaleBE8M0) or (kScaleBDirectLoad and kScaleBGranK == 32), + "Compressed Scale-B dtypes are only supported by direct-load per-32 path"); DG_STATIC_ASSERT(kFuseScaleBDecodeAssumeExp == 0 or kFuseScaleBDecodeAssumeExp == 5 or kFuseScaleBDecodeAssumeExp == 6, "DG_W4_FUSE_SCALE_B_DECODE_ASSUME_EXP only supports 0/5/6"); @@ -741,9 +749,13 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo (BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M) * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + static constexpr uint32_t SCALE_B_ELEMENT_SIZE = + kScaleBE8M0 ? static_cast(sizeof(uint8_t)) : + (kScaleBBF16 ? static_cast(sizeof(nv_bfloat16)) : + static_cast(sizeof(float))); const uint32_t shape_k_scales_a = math::ceil_div(shape_k, BLOCK_K); const uint32_t shape_k_scales_b = math::ceil_div(shape_k, kScaleBGranK); - const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / SCALE_B_ELEMENT_SIZE); // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough // smem_d capacity, so they fall back to a separate SFB region. const uint32_t smem_sfb_bytes = kScaleBGranK == 32 ? @@ -1145,14 +1157,30 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo return 1.0f; if constexpr (kScaleBDirectLoad and kScaleBGranK == 32) { if constexpr (kMajorSFB == cute::UMMA::Major::MN) { - const float* sfb_base = sfb + - current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; - const float* ptr = sfb_base + k_block_idx * aligned_shape_n_sfb + n_idx; - return *ptr; + const uint32_t offset = + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b + + k_block_idx * aligned_shape_n_sfb + n_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } } else { - const float* ptr = sfb + current_group_idx * shape_n * shape_k_scales_b + - n_idx * shape_k_scales_b + k_block_idx; - return *ptr; + const uint32_t offset = current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_block_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } } } else if constexpr (kScaleBGranK == 32) { const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; From 40c4fb2dd5252fe5dbe8d423b34db46e44cf3ddd Mon Sep 17 00:00:00 2001 From: zhangxiaolei Date: Mon, 8 Jun 2026 17:53:28 +0800 Subject: [PATCH 39/51] Support redundant expert groups in FP4 fast path --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 155 ++++++++++++++++-- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 44 ++++- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh | 46 +++++- 3 files changed, 215 insertions(+), 30 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index e981dced4f..febe4bbb26 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -647,6 +647,24 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( const char* value = std::getenv(name); return (value != nullptr and value[0] != '\0') ? std::atoi(value) : default_value; }; + // BM=64 fast-path 是否启用（函数作用域统一判据，供三处共用：layout 选择 / + // bm32_skew_layout(stages) / bm32_skew_fast_path(device 三件套总闸)）。 + // 两个触发源： + // (1) DG_W4_PATHB_BM64_FASTPATH=1：手动强制（对照实验用）。 + // (2) DG_W4_PATHB_BM64_AUTO=1（默认开）+ 高置信子集 max_hint>=128 且 + // active<=8：自动触发，抓“少数活跃组 + 大 max_m”（MTP verify spec-len + // 命中形态）。详见 layout 选择处 bm64_fastpath 的长注释。 + // (3) DG_W4_PATHB_BM64_FORCE_ALL=1：无条件全部走 BM=64 fast-path + // （用于在线业务峰值吞吐 A/B 对照，绕开 hint 判据；hint 在生产 cuda graph + // 下恒为 None 时，是验证“假设拿到 hint 全命中”的速度上界）。 + // 注意：必须同时让 stages / fast-path 总闸认账，否则 layout.block_m=64 却 + // 走通用慢路径会退化 ~2x（见 DG_W4_PATHB_BM64 非 fastpath 的坑）。 + const bool bm64_fastpath_enabled = + env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 or + env_int("DG_W4_PATHB_BM64_FASTPATH", 0) != 0 or + (env_int("DG_W4_PATHB_BM64_AUTO", 1) != 0 and + masked_m_max_hint.value_or(0) >= 128 and + active_groups_hint.value_or(static_cast(desc.num_groups)) <= 8); // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 @@ -746,8 +764,8 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // 挡在 fast-path 外，强制走 path-B 通用 cache_sfb_k32 路径，W4 GB/s // 只有 627（vs fast-path 命中时 1500-3000+ GB/s），speedup 0.43-0.56x。 // device fast-path 三件套（quad_reduce / direct_load / compact_sched） - // 实际 BN-agnostic，仅硬约束 BM==32 + gran_k_b==32，groups=24 满足 - // kNumGroups<=32、n=6144 满足 BN={128,256} 整除 ⇒ 物理可命中。 + // 实际 BN-agnostic，仅硬约束 BM==32 + gran_k_b==32，冗余专家场景 + // kNumGroups<=36 也可按同一路径处理。 // 实测：g24+n=6144+k=7168 整张表 0.43-0.56x → 0.71-0.81x； // g24+n=7168+k=3072 整张表 0.30-0.54x → 0.45-0.79x； // 默认 dsv4 形状 (g32+n∈{4096,7168}+k∈{2048,4096,7168}) 不变。 @@ -838,10 +856,61 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // 保留 env 入口仅做对照实验，生产环境永不开启。 const bool bm64_hint = dsv4_shape and env_int("DG_W4_PATHB_BM64", 0) != 0; + // BM=64 fast-path 实验路径（默认关，env DG_W4_PATHB_BM64_FASTPATH=1）： + // 与 DG_W4_PATHB_BM64 不同——后者刻意让 device 落回通用 + // cache_sfb_k32 路径；本开关**保持** fast-path 三件套 + // (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched) 开启， + // 仅把 host BM 从 32 抬到 64，验证“weight 解码次数减半 + grid 减半” + // 能否吃掉 max_m>=64 差场景。 + // 依赖：device 端 swap_AB 下 WGMMA N=BLOCK_M=64 指令已存在 + // (FP8MMASelectorRS<64> = MMA_64x64x32_RS_TN)，scale 寻址 BM-agnostic。 + // BM=64 强制 BN=128，避免 A-smem 叠加 BN=256 把 stages 压崩。 + // + // 实测结论（DG_W4_PATHB_BM64_FASTPATH=1 vs 默认 BM=32，ncu 对照， + // g24 N=6144 K=7168 / N=7168 K=3072 skew benchmark）：**全面退化，已证伪**。 + // uniform_17 784us/0.75x → 1391us/0.43x (max_m=17) + // uniform_32 775us/0.77x → 1387us/0.42x (max_m=32) + // dense_tail_hot17 767/0.77x → 1381/0.43x + // uniform_64 1541us/0.38x → 1390us/0.43x (唯一微升) + // one_hot_512 542us/0.58x → 542us/0.58x (完全无变化) + // ncu 四项硬指标 BM32 vs BM64 **逐字节相同**： + // Registers/Thread 168=168, Dyn SMEM 149984=149984, + // Grid (78,1,1)=(78,1,1), Achieved Occupancy 14.0%≈13.9%。 + // 证伪根因（结构性，非实现 bug；正确性 OK diff=0、无 spill）： + // 1. persistent kernel grid 恒 = num_sms，BLOCK_M 不影响 grid。 + // 2. swap_AB 下 WGMMA 形状 / kNumAccum / WAVE_BLOCK_M 全由 + // BLOCK_N 决定，BLOCK_M 加倍不动寄存器、不动 smem。 + // 3. M-tile 数 = ceil(max_m/BM)。DSV4 真实 max_m 多 <=32， + // BM 32→64 时 tile 数不减，但每 tile WGMMA 覆盖 64 行、 + // ~47 行是 padding 空算 → 工作量翻倍、收益为零 → 耗时近翻倍。 + // 4. 仅 max_m=64 时 tile 2→1 抵消 padding 才微升；max_m 更大 + // (one_hot_512 active=1) 瓶颈不在 tile 数，故毫无变化。 + // 结论：当前 DSV4 负载 (max_m 集中 17-64) 下 BM=64 是负优化。 + // 保留 env 入口仅做对照实验，生产环境永不开启。 + // 真正受益方向见下方思路 2（active=1 大 max_m 时跨 m_block 复用 + // weight decode），不受 padding 问题影响。 + // BM=64 fast-path **自动触发规则**（默认开，env 可关）： + // 前述 DG_W4_PATHB_BM64_FASTPATH 手动实验证明，BM 32→64 让 + // M-tile 数 ceil(m/32)→ceil(m/64) 减半，仅当“每活跃组平均 + // token 数”够大（每组 m>=64，padding 占比可忽略）才净赚；组小则 + // tile 数不减、padding 翻倍纯亏。判别量本应是 sum_m/active，但 + // host API 只有 max_hint / active（拿不到 sum_m，且不能同步读 + // device 上的 masked_m）。 + // 折中：用高置信子集 `max_hint>=128 且 active<=8`（少数活跃组 + + // 大 max_m，正是 MTP verify 的 spec-len 命中形态）触发，刻意避开 + // max_m=64 / 多活跃组的模糊带（uniform_64 赢、dense_tail_hot64 + // 亏，二者 host 信号 (max_hint=64,active=24) 撞车无法区分）。 + // 实测（g24 N6144/7168 与 g32 N4096，BM64+pair vs 默认 quad）两族 + // shape 一致：命中子集 one_hot_128/384、mtp_384/512_hot、 + // one_hot_512、mtp_768_hot512 等普降 5~26%，被避开的 + // dense_tail/uniform/eight_hot 无一受损。 + // BM=64 下 quad 4 累加器会 spill 爆炸，故命中时**必须**配 + // pair_reduce（见下方 k32_pair_reduce 联动）。 + const bool bm64_fastpath = dsv4_shape and bm64_fastpath_enabled; if (real_hot_present and not fuse_decode_hint and not bm64_hint and env_int("DG_W4_PATHB_FAST_PATH", 1) != 0) { - layout.block_m = 32; + layout.block_m = bm64_fastpath ? 64 : 32; // H1: 同时满足 max_m > 32 + active*max_m >= 1024 时走 BN=256， // grid 减半。门槛公式来自 skewed masked benchmark 实测： // * eight_hot_64 (8×64=512)、four_hot_64 (4×64=256)、 @@ -895,7 +964,9 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( masked_m_max_hint.has_value() and (bn256_baseline or bn256_big_shape) and env_int("DG_W4_PATHB_BN256", 1) != 0; - layout.block_n = bn256_eligible ? 256 : 128; + // BM=64 fast-path 强制 BN=128：A-smem 已随 BM 翻倍，再叠加 + // BN=256 会把 B-tile + A-tile 占用一起推高，stages 被压崩。 + layout.block_n = (bm64_fastpath or not bn256_eligible) ? 128 : 256; // H2: cluster_n=2（A 多播）实测**全军退化**，默认关闭。 // 退化数据（DG_W4_PATHB_CLUSTER_N=1 vs 默认）： // BN=256 路径 ~9-10% 退化（mtp_512: 201→220, dense_tail_hot64: @@ -967,6 +1038,16 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( } } if (not block_m_override and not block_n_override) { + // DG_W4_PATHB_BM64_FORCE_ALL=1 时，把 BLOCK_N 强制锁到 128（与 BM=64 + // fast-path 联动）。动机：BM=64 fast-path 下 final_accum 寄存器量 + // ∝ BLOCK_N/64，BN=256 → 4 份 tile 累加器逼出 168reg/thread、occupancy + // 14%；BN=128 砍半 final_accum，腾出寄存器预算给更多 warp，对大 batch / + // 高 active 场景 occupancy 上限更高。属于在线吞吐对照实验的联动开关， + // 不影响 hint 判据自身路径。 + if (env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 and + bm64_fastpath_enabled and layout.block_m == 64) { + layout.block_n = 128; + } DG_HOST_ASSERT(layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or layout.block_m == 64 or layout.block_m == 128); DG_HOST_ASSERT(layout.block_n == 64 or layout.block_n == 128 or layout.block_n == 256); @@ -1043,10 +1124,12 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) chosen_stages = std::min(chosen_stages, 6); const bool bm32_skew_layout = gran_k_b == 32 and - config.layout.block_m == 32 and + (config.layout.block_m == 32 or + (config.layout.block_m == 64 and bm64_fastpath_enabled)) and (config.layout.block_n == 128 or config.layout.block_n == 256) and static_cast(desc.m) >= 1024 and - ((desc.num_groups == 32 and + ((static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 7168) and (static_cast(desc.k) == 2048 or @@ -1057,7 +1140,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // 放开到 g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168} (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 32 and + static_cast(desc.num_groups) <= 36 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 6144 or static_cast(desc.n) == 7168) and @@ -1108,10 +1191,12 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( config.storage_config.swizzle_cd_mode); const bool bm32_skew_fast_path = gran_k_b == 32 and - config.layout.block_m == 32 and + (config.layout.block_m == 32 or + (config.layout.block_m == 64 and bm64_fastpath_enabled)) and (config.layout.block_n == 128 or config.layout.block_n == 256) and static_cast(desc.m) >= 1024 and - ((desc.num_groups == 32 and + ((static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 7168) and (static_cast(desc.k) == 2048 or @@ -1120,11 +1205,12 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( static_cast(desc.k) == 7168)) or // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）， // 与 bm32_skew_layout / dsv4_shape 同步生效。device fast-path 三件套 - // (quad_reduce / direct_load / compact_sched) 仅硬约束 BM==32 + g<=32， + // (quad_reduce / direct_load / compact_sched) 仅硬约束 BM==32，g<=36 + // 冗余专家场景按同一路径处理， // 与 N/K 数值无关，物理可命中。 (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 32 and + static_cast(desc.num_groups) <= 36 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 6144 or static_cast(desc.n) == 7168) and @@ -1144,6 +1230,15 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( get_major_type_ab(sfb) == cute::UMMA::Major::MN and env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; const bool scale_b_packed_ue8m0 = fuse_scale_b_decode; + // 思路2 性能探针：把 scale_a 从 promote 内移除（device kScaleAStub 置 scale_a=1 + // 并跳过 SFA 的 TMA load），用于测 fuse_scale_b_decode 单累加器 + scale_a 后移到 + // silu_and_mul 的速度上界。开启后结果数值错误，仅做性能对照。默认关。 + const bool scale_a_stub = env_int("DG_W4_SCALE_A_STUB", 0) != 0; + // 思路2 串行依赖实验：双缓冲预解码（staged_pair_a_regs，issue k 时预取 k+2）， + // 让 decode 与 WGMMA 多重叠一级，代价 +8 reg/线程。仅在 fuse_scale_b_decode 开启 + // 时有意义，用于测"解 issue 串行依赖"能否补回 fuse 路径在大 M 的退化。默认关。 + const bool fuse_predecode_pair = + fuse_scale_b_decode and env_int("DG_W4_FUSE_PREDECODE_PAIR", 0) != 0; // fuse_scale_b_decode 一旦开启，所有 fast-path / direct-load / e8m0 / bf16 // 互斥关闭——device 那边 static_assert 会拒绝同时开启。 const bool scale_b_direct_load = @@ -1152,6 +1247,37 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( const bool k32_quad_reduce = gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and not fuse_scale_b_decode; + // 杠杆3 / BM=64 联动：把 4 累加器 quad-reduce（峰值 64 float，在 + // __launch_bounds__(384,1) 的 168reg 硬上限下触发 local spill）退化为已存在的 + // 2 累加器串行 pair-reduce（峰值 32 float）降低寄存器峰值活跃量。仅改单线程内 + // 累加器生命周期、不动线程模型/barrier，数值 bit 级等价。 + // 两种启用： + // (A) BM=32 下**单独**开 pair（DG_W4_K32_PAIR_REDUCE=1）：已证伪——ncu spill + // 223680→134016(-40%) 但 wall-clock 反而 +10~25%（串行累加拉长依赖链， + // spill 不在 BM=32 关键路径上）。默认关。 + // (B) BM=64 fast-path 命中时**必须**配 pair（自动，下面 block_m==64 分支）： + // BM=64 让 kNumAccum 32→quad 4 份 = 128 fp32，spill 爆炸退化 ~2x；pair + // 2 份 = 64 fp32 受控。BM=64 靠 M-tile 数减半换吞吐，pair 是其必要前提。 + // 实测高置信子集（max_hint>=128 且 active<=8）净降 5~26%。 + const bool k32_pair_reduce = + k32_quad_reduce and + ((config.layout.block_m == 64 and bm64_fastpath_enabled) or + env_enabled("DG_W4_K32_PAIR_REDUCE")); + // 杠杆3 续（已证伪，默认关）：pair-reduce 之上把长驻 final_accum 从 fp32 改 bf16 + // 存储（nv_bfloat162）再减半长驻寄存器、继续压 spill（ncu 134016→6480，-97% vs + // quad）。精度无损（W4 diff 0.0001/0.0000）。但 wall-clock 仍慢于 quad baseline + // （随 pair-reduce 一起退化，叠加 bf16 promote 反而更长依赖链）。默认关，需 + // DG_W4_K32_BF16_FINAL_ACCUM=1 显式开启复现。 + const bool k32_bf16_final_accum = + k32_pair_reduce and env_enabled("DG_W4_K32_BF16_FINAL_ACCUM"); + // 杠杆3 续2（已证伪，默认关）：把长驻 final_accum 整体搬进 smem_d（kDirectStore + // + kFinalAccumScratch），final_accum_regs 缩成 1 元素。实测 spill 虽归零，但 + // (1) 寄存器仍顶 168（launch_bounds 硬顶，省下的 reg 没还给 occupancy，Block + // Limit Registers 仍为 1）；(2) DirectStore 逐元素非合并写 gmem 严重拖累大 + // max_m：one_hot_256 0.77x→0.50x、one_hot_512 0.49x→0.29x。净负收益，仅保留 + // 开关供复现，默认关。需 DG_W4_K32_FINAL_ACCUM_SMEM=1 显式开启。 + const bool k32_final_accum_smem = + k32_pair_reduce and env_enabled("DG_W4_K32_FINAL_ACCUM_SMEM"); // RS baseline for direct E8M0 B scales: keep scale products in split form by default. // The exponent-adjust path is experimental and can regress some small-M shapes. const bool scale_b_pow2_promote = @@ -1177,7 +1303,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( gran_k_b == 32 and expected_m <= 16 and ((static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) or (static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 32 and + static_cast(desc.num_groups) <= 36 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 6144 or static_cast(desc.n) == 7168) and @@ -1248,6 +1374,9 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( .scale_b_direct_load = scale_b_direct_load, .scale_b_pow2_promote = scale_b_pow2_promote, .k32_quad_reduce = k32_quad_reduce, + .k32_pair_reduce = k32_pair_reduce, + .k32_bf16_final_accum = k32_bf16_final_accum, + .k32_final_accum_smem = k32_final_accum_smem, .k32_quad_split_promote = k32_quad_split_promote, .k32_quad_scale_b_inline = k32_quad_scale_b_inline, .k32_quad_scale_b_prefetch = k32_quad_scale_b_prefetch, @@ -1256,6 +1385,8 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( .small_m_simple_sched = small_m_simple_sched, .compact_masked_sched = compact_masked_sched, .fuse_scale_b_decode = fuse_scale_b_decode, + .scale_a_stub = scale_a_stub, + .fuse_predecode_pair = fuse_predecode_pair, .scale_b_packed_ue8m0 = scale_b_packed_ue8m0, .scale_b_gran_k = static_cast(gran_k_b), .b_is_int4_sym = b_is_int4_sym, diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 6851476d11..65653763be 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -31,6 +31,22 @@ class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime); }}; @@ -174,7 +201,7 @@ static void __instantiate_kernel() {{ to_string(args.gemm_desc.gemm_type), get_default_epilogue_type(std::nullopt), "false", - "false", + args.scale_a_stub ? "true" : "false", "false", "false", "false", @@ -185,6 +212,7 @@ static void __instantiate_kernel() {{ "false", "false", "false", + args.k32_final_accum_smem ? "true" : "false", "false", "false", "false", @@ -197,33 +225,31 @@ static void __instantiate_kernel() {{ "false", "false", "false", - "false", - "false", + args.k32_bf16_final_accum ? "true" : "false", "false", "false", "false", "false", "false", args.k32_quad_reduce ? "true" : "false", - "false", - args.k32_quad_reduce ? "true" : "false", + args.k32_pair_reduce ? "true" : "false", + (args.k32_quad_reduce and not args.k32_pair_reduce) ? "true" : "false", args.k32_quad_split_promote ? "true" : "false", args.k32_quad_scale_b_inline ? "true" : "false", args.k32_quad_scale_b_prefetch ? "true" : "false", args.k32_quad_scale_b_vec4 ? "true" : "false", args.k32_quad_pair4x2_promote ? "true" : "false", "false", - "false", + args.k32_final_accum_smem ? "true" : "false", "false", "false", args.small_m_simple_sched ? "true" : "false", - args.compact_masked_sched ? "true" : "false", args.fuse_scale_b_decode ? "true" : "false", "false", "false", "false", "false", - "false", + args.fuse_predecode_pair ? "true" : "false", "false", "false", "false", diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh index 24a14a6d8e..ab68ef4013 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh @@ -686,7 +686,11 @@ template + uint32_t kMathRegCap = 0, + bool kBIsInt4Sym = false, + bool kScaleBBF16 = false, + bool kScaleBE8M0 = false, + bool kReorderMaskedByMaxM = false> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, kLaunchBoundsMinBlocks) void sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, nv_bfloat16* gmem_d_ptr, @@ -708,6 +712,10 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo "DG_W4_SCALE_K_GROUP does not support per-32 FP4 scaling"); DG_STATIC_ASSERT(kScaleKGroup == 1 or kScaleKGroup == 2 or kScaleKGroup == 4, "DG_W4_SCALE_K_GROUP only supports 1/2/4"); + DG_STATIC_ASSERT(not kBIsInt4Sym, "RS-mode FP8xFP4 kernel does not support INT4-sym B"); + DG_STATIC_ASSERT(not (kScaleBBF16 and kScaleBE8M0), "Scale-B cannot be both BF16 and E8M0"); + DG_STATIC_ASSERT((not kScaleBBF16 and not kScaleBE8M0) or (kScaleBDirectLoad and kScaleBGranK == 32), + "Compressed Scale-B dtypes are only supported by direct-load per-32 path"); DG_STATIC_ASSERT(kFuseScaleBDecodeAssumeExp == 0 or kFuseScaleBDecodeAssumeExp == 5 or kFuseScaleBDecodeAssumeExp == 6, "DG_W4_FUSE_SCALE_B_DECODE_ASSUME_EXP only supports 0/5/6"); @@ -741,9 +749,13 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo (BLOCK_M < WGMMA::M ? WGMMA::M : BLOCK_M) * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_SFA_SIZE_PER_STAGE = math::constexpr_align(SMEM_SFA_SIZE_PER_STAGE, 128u); + static constexpr uint32_t SCALE_B_ELEMENT_SIZE = + kScaleBE8M0 ? static_cast(sizeof(uint8_t)) : + (kScaleBBF16 ? static_cast(sizeof(nv_bfloat16)) : + static_cast(sizeof(float))); const uint32_t shape_k_scales_a = math::ceil_div(shape_k, BLOCK_K); const uint32_t shape_k_scales_b = math::ceil_div(shape_k, kScaleBGranK); - const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / sizeof(float)); + const uint32_t aligned_shape_n_sfb = math::align(shape_n, 16u / SCALE_B_ELEMENT_SIZE); // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough // smem_d capacity, so they fall back to a separate SFB region. const uint32_t smem_sfb_bytes = kScaleBGranK == 32 ? @@ -1145,14 +1157,30 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo return 1.0f; if constexpr (kScaleBDirectLoad and kScaleBGranK == 32) { if constexpr (kMajorSFB == cute::UMMA::Major::MN) { - const float* sfb_base = sfb + - current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; - const float* ptr = sfb_base + k_block_idx * aligned_shape_n_sfb + n_idx; - return *ptr; + const uint32_t offset = + current_group_idx * aligned_shape_n_sfb * shape_k_scales_b + + k_block_idx * aligned_shape_n_sfb + n_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } } else { - const float* ptr = sfb + current_group_idx * shape_n * shape_k_scales_b + - n_idx * shape_k_scales_b + k_block_idx; - return *ptr; + const uint32_t offset = current_group_idx * shape_n * shape_k_scales_b + + n_idx * shape_k_scales_b + k_block_idx; + if constexpr (kScaleBE8M0) { + const auto* sfb_u8 = reinterpret_cast(sfb); + return __uint_as_float(static_cast(sfb_u8[offset]) << 23); + } else if constexpr (kScaleBBF16) { + const auto* sfb_bf16 = reinterpret_cast(sfb); + return __bfloat162float(sfb_bf16[offset]); + } else { + return sfb[offset]; + } } } else if constexpr (kScaleBGranK == 32) { const uint32_t n_off = n_idx - n_block_idx * BLOCK_N; From d5b0ae6b5859819b879dbcfabbcc1cddddf2e0f3 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Tue, 14 Jul 2026 12:22:52 +0000 Subject: [PATCH 40/51] create test_sm90_fp8_fp4_g128.py --- tests/test_sm90_fp8_fp4_g128.py | 607 ++++++++++++++++++++++++++++++++ 1 file changed, 607 insertions(+) create mode 100644 tests/test_sm90_fp8_fp4_g128.py diff --git a/tests/test_sm90_fp8_fp4_g128.py b/tests/test_sm90_fp8_fp4_g128.py new file mode 100644 index 0000000000..fcfca50c59 --- /dev/null +++ b/tests/test_sm90_fp8_fp4_g128.py @@ -0,0 +1,607 @@ +"""SM90 FP8xFP4(g128, fp32-scale) accuracy + benchmark script. + +这份脚本用于你们新的权重编码语义： + +1. A 仍然是 FP8(e4m3) + per-token/per-128 scale +2. B 仍然是 packed FP4(两个 4-bit nibble 打包成一个 int8) +3. B 的 scale 改成 `group_size=128`，且 dtype 明确为 `torch.float32` + +它和原始 `test_sm90_fp8_fp4.py` 的关键区别是： + +- 不再走偏向 `gran_k_b=32` 的 fast path 测试逻辑 +- 直接构造与你们转换脚本一致的 FP4(g128, fp32-scale) 权重 +- 同时覆盖 contiguous / masked 两条常用调用路径 + +用法： +python /sgl-workspace/sglang/test_sm90_fp8_fp4_g128.py --repo-root /root/work/DeepGEMM + +""" +from __future__ import annotations + +import argparse +import importlib +import sys +import time +from pathlib import Path +from typing import Callable + +import torch + +FP4_TABLE = torch.tensor( + [ + 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, + ], + dtype=torch.float32, +) + +DEFAULT_SHAPES = [(4096, 7168), (7168, 2048), (4096, 4096)] +DEFAULT_GROUPS = (8, 16, 32) +DEFAULT_M_PER_GROUP = (1, 4, 8, 16, 32, 64) + +W4_DIFF_THRESHOLD = 0.015 +FP8_DIFF_THRESHOLD = 0.05 + +deep_gemm = None +calc_diff = None +per_token_cast_to_fp8 = None + +def _parse_csv_ints(value: str) -> tuple[int, ...]: + return tuple(int(item.strip()) for item in value.split(",") if item.strip()) + +def _parse_shapes(value: str) -> list[tuple[int, int]]: + shapes: list[tuple[int, int]] = [] + for item in value.split(","): + item = item.strip().lower() + if not item: + continue + if "x" not in item: + raise ValueError(f"invalid shape item: {item!r}, expected like 4096x7168") + n_str, k_str = item.split("x", 1) + shapes.append((int(n_str), int(k_str))) + return shapes + +def _discover_repo_root(explicit_repo_root: str | None) -> Path: + candidates: list[Path] = [] + if explicit_repo_root: + candidates.append(Path(explicit_repo_root).resolve()) + + script_dir = Path(__file__).resolve().parent + cwd = Path.cwd().resolve() + + candidates.extend( + [ + cwd, + script_dir, + script_dir.parent, + cwd.parent, + ] + ) + candidates.extend(script_dir.parents) + candidates.extend(cwd.parents) + + seen: set[Path] = set() + for candidate in candidates: + if candidate in seen: + continue + seen.add(candidate) + if (candidate / "deep_gemm").exists(): + return candidate + + raise FileNotFoundError( + "cannot locate DeepGEMM repo root. Please pass --repo-root /path/to/DeepGEMM" + ) + +def _load_deep_gemm(repo_root: Path) -> None: + global deep_gemm, calc_diff, per_token_cast_to_fp8 + + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + deep_gemm = importlib.import_module("deep_gemm") + calc_diff = importlib.import_module("deep_gemm.testing").calc_diff + per_token_cast_to_fp8 = importlib.import_module("deep_gemm.utils.math").per_token_cast_to_fp8 + +def _require_sm90() -> None: + assert torch.cuda.is_available() + major, _ = torch.cuda.get_device_capability() + if major != 9: + raise RuntimeError(f"This benchmark is intended for SM90, got sm_{major}x") + +def _time_cuda(fn: Callable[[], None], warmup: int = 3, iters: int = 10) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 + +def _effective_bytes( + groups: int, + logical_m: int, + n: int, + k: int, + a_gran_k: int, + *, + fp8_b: bool, + b_gran_k: int, +) -> int: + a_scale_k = (k + a_gran_k - 1) // a_gran_k + b_scale_k = (k + b_gran_k - 1) // b_gran_k + a_bytes = logical_m * k + logical_m * a_scale_k * 4 + b_data_bytes = groups * n * k if fp8_b else groups * n * (k // 2) + b_scale_bytes = groups * n * b_scale_k * 4 + d_bytes = logical_m * n * 2 + return a_bytes + b_data_bytes + b_scale_bytes + d_bytes + +def _build_grouped_layout(groups: int, m_per_group: int): + m = groups * m_per_group + group_starts = [group_id * m_per_group for group_id in range(groups)] + group_ends = [(group_id + 1) * m_per_group for group_id in range(groups)] + grouped_layout = torch.arange(groups, device="cuda", dtype=torch.int32).repeat_interleave(m_per_group) + return m, group_starts, group_ends, grouped_layout + +def _get_m_alignment_for_contiguous_layout() -> int: + """DeepGEMM contiguous grouped layout 要求每个 group 的 M 段按 block 对齐。 + + `deep_gemm.get_mk_alignment_for_contiguous_layout()` 的返回值在不同版本里可能是： + - 一个 int(只代表 M 对齐） + - 一个 tuple/list(例如 (align_m, align_k) + """ + align = deep_gemm.get_mk_alignment_for_contiguous_layout() + if isinstance(align, (tuple, list)): + return int(align[0]) + return int(align) + +def _align_up(x: int, a: int) -> int: + return ((x + a - 1) // a) * a + +def _unpack_fp4_values(packed: torch.Tensor) -> torch.Tensor: + if packed.dtype != torch.int8 or packed.ndim != 2: + raise ValueError(f"expected int8 2D packed tensor, got {packed.dtype}, {tuple(packed.shape)}") + + out_dim, packed_in_dim = packed.shape + packed_u8 = packed.view(torch.uint8) + low = packed_u8 & 0x0F + high = (packed_u8 >> 4) & 0x0F + table = FP4_TABLE.to(device=packed.device) + values = torch.stack((table[low.long()], table[high.long()]), dim=-1).reshape(out_dim, packed_in_dim * 2) + return values + +def _dequant_fp4_grouped(packed: torch.Tensor, scale: torch.Tensor, group_size: int) -> torch.Tensor: + fp4 = _unpack_fp4_values(packed) + out_dim, in_dim = fp4.shape + expected_scale_shape = (out_dim, in_dim // group_size) + if tuple(scale.shape) != expected_scale_shape: + raise ValueError( + f"scale shape mismatch: got {tuple(scale.shape)}, expected {expected_scale_shape}" + ) + full_scale = scale.float().repeat_interleave(group_size, dim=1) + return fp4 * full_scale + +def _quantize_real_to_fp4_grouped( + real: torch.Tensor, + group_size: int, + eps: float = 1e-12, +) -> tuple[torch.Tensor, torch.Tensor]: + if real.ndim != 2: + raise ValueError(f"expected 2D tensor, got shape {tuple(real.shape)}") + + out_dim, in_dim = real.shape + if in_dim % group_size != 0: + raise ValueError(f"in_dim={in_dim} must be divisible by group_size={group_size}") + if in_dim % 2 != 0: + raise ValueError(f"in_dim={in_dim} must be even for nibble packing") + + num_groups = in_dim // group_size + grouped = real.float().view(out_dim, num_groups, group_size) + max_abs = grouped.abs().amax(dim=-1) + scale = torch.clamp(max_abs / 6.0, min=eps).float() + + codebook = FP4_TABLE.to(device=real.device).view(1, 1, 1, 16) + normalized = grouped / scale.unsqueeze(-1) + dist = (normalized.unsqueeze(-1) - codebook).abs() + nibble = dist.argmin(dim=-1).to(torch.uint8).view(out_dim, in_dim) + + low = nibble[:, 0::2] + high = nibble[:, 1::2] + packed_u8 = low | (high << 4) + packed_i8 = packed_u8.view(torch.int8) + return packed_i8, scale + +def _cast_back_from_fp8_1d(x: torch.Tensor, sf: torch.Tensor, gran_k: int = 128) -> torch.Tensor: + group_idx = torch.arange(x.size(-1), device=x.device) // gran_k + return x.float() * sf[..., group_idx] + +def _make_b_fp4_g128(real_b: torch.Tensor, group_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: + return _quantize_real_to_fp4_grouped(real_b, group_size=group_size) + +def _resolve_contiguous_kernel(): + fn = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma", None) + if fn is None: + raise RuntimeError( + "cannot find `m_grouped_fp8_fp4_gemm_nt_contiguous_sm90_fused_wgmma` in deep_gemm. " + "Please use the DeepGEMM branch/build that contains the SM90 FP8xFP4 kernels." + ) + return fn + +def _resolve_masked_kernel(): + fn = getattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma", None) + if fn is None: + raise RuntimeError( + "cannot find `m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma` in deep_gemm. " + "Please use the DeepGEMM branch/build that contains the SM90 FP8xFP4 kernels." + ) + return fn + +def _benchmark_case_contiguous( + groups: int, + m_per_group: int, + n: int, + k: int, + *, + a_gran_k: int = 128, + b_gran_k: int = 128, +) -> dict[str, float | int]: + if b_gran_k != 128: + raise ValueError("this g128 script expects b_gran_k == 128") + + kernel = _resolve_contiguous_kernel() + # NOTE: contiguous grouped layout 要求每个 group 的 token 段在 M 维按 block 对齐； + # 否则在某些配置下可能出现错误结果（而不是直接报错）。 + align_m = _get_m_alignment_for_contiguous_layout() + m_per_group_aligned = _align_up(m_per_group, align_m) + m, group_starts, group_ends, grouped_layout = _build_grouped_layout(groups, m_per_group_aligned) + + # 每个 group 只有前 m_per_group 行是有效 token，后面 padding 行填 0，保证参考值与 kernel 一致。 + a_ref_src = torch.zeros((m, k), device="cuda", dtype=torch.bfloat16) + for group_id in range(groups): + start = group_id * m_per_group_aligned + a_ref_src[start:start + m_per_group] = torch.randn( + (m_per_group, k), device="cuda", dtype=torch.bfloat16 + ) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a = per_token_cast_to_fp8(a_ref_src, use_ue8m0=False, gran_k=a_gran_k) + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // b_gran_k), device="cuda", dtype=torch.float) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((m, n), device="cuda", dtype=torch.bfloat16) + + for group_id in range(groups): + packed, scale = _make_b_fp4_g128(b_ref_src[group_id], group_size=b_gran_k) + b_fp4[group_id] = packed + b_sf[group_id] = scale + + b_dequant = _dequant_fp4_grouped(packed, scale, group_size=b_gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant.to(torch.bfloat16), use_ue8m0=False, gran_k=a_gran_k + ) + # 仅对每个 group 的有效 token 行计算参考值；padding 行保持 0。 + start = group_id * m_per_group_aligned + end = start + m_per_group + if m_per_group > 0: + ref[start:end] = (a_dequant[start:end] @ b_dequant.t()).to(torch.bfloat16) + + b_w4 = (b_fp4, b_sf) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b_fp8, + d_fp8, + grouped_layout, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + use_psum_layout=False, + ) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + kernel( + a, + b_w4, + d_w4, + grouped_layout, + gran_k=b_gran_k, + compiled_dims="nk", + use_psum_layout=False, + ) + + run_fp8() + run_w4() + + fp8_diff = calc_diff(d_fp8, ref) + w4_diff = calc_diff(d_w4, ref) + fp8_elapsed = _time_cuda(run_fp8) + w4_elapsed = _time_cuda(run_w4) + + # 带宽模型按“实际参与计算/写回的 M 行数”计：contiguous 场景下 padding 行也会发生 A/D 的读写。 + logical_m = groups * m_per_group_aligned + w4_bytes = _effective_bytes(groups, logical_m, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_bytes(groups, logical_m, n, k, a_gran_k, fp8_b=True, b_gran_k=a_gran_k) + + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + +def _benchmark_case_masked( + groups: int, + m_per_group: int, + n: int, + k: int, + *, + max_m: int = 128, + a_gran_k: int = 128, + b_gran_k: int = 128, +) -> dict[str, float | int]: + if b_gran_k != 128: + raise ValueError("this g128 script expects b_gran_k == 128") + + kernel = _resolve_masked_kernel() + masked_m = torch.full((groups,), m_per_group, device="cuda", dtype=torch.int32) + + a_ref_src = torch.randn((groups, max_m, k), device="cuda", dtype=torch.bfloat16) + b_ref_src = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) + + a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) + a_sf = torch.empty((groups, max_m, k // a_gran_k), device="cuda", dtype=torch.float) + for group_id in range(groups): + a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( + a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k + ) + a = (a_data, a_sf) + a_dequant = _cast_back_from_fp8_1d(a[0], a[1], gran_k=a_gran_k) + + b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) + b_sf = torch.empty((groups, n, k // b_gran_k), device="cuda", dtype=torch.float) + b_fp8_data = torch.empty((groups, n, k), device="cuda", dtype=torch.float8_e4m3fn) + b_fp8_sf = torch.empty((groups, n, k // a_gran_k), device="cuda", dtype=torch.float) + ref = torch.zeros((groups, max_m, n), device="cuda", dtype=torch.bfloat16) + + for group_id in range(groups): + packed, scale = _make_b_fp4_g128(b_ref_src[group_id], group_size=b_gran_k) + b_fp4[group_id] = packed + b_sf[group_id] = scale + + b_dequant = _dequant_fp4_grouped(packed, scale, group_size=b_gran_k) + b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( + b_dequant.to(torch.bfloat16), use_ue8m0=False, gran_k=a_gran_k + ) + + valid_m = int(masked_m[group_id].item()) + if valid_m > 0: + ref[group_id, :valid_m] = (a_dequant[group_id, :valid_m] @ b_dequant.t()).to(torch.bfloat16) + + b_w4 = (b_fp4, b_sf) + b_fp8 = (b_fp8_data, b_fp8_sf) + + d_w4 = torch.empty_like(ref) + + def run_w4(): + kernel( + a, + b_w4, + d_w4, + masked_m, + m_per_group, + gran_k=b_gran_k, + gran_k_a=a_gran_k, + gran_k_b=b_gran_k, + ) + + d_fp8 = torch.empty_like(ref) + + def run_fp8(): + deep_gemm.m_grouped_fp8_gemm_nt_masked( + a, + b_fp8, + d_fp8, + masked_m, + m_per_group, + recipe_a=(1, a_gran_k), + recipe_b=(1, a_gran_k), + ) + + run_w4() + run_fp8() + + w4_diff = max( + calc_diff(d_w4[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + fp8_diff = max( + calc_diff(d_fp8[group_id, :m_per_group], ref[group_id, :m_per_group]) + for group_id in range(groups) + ) + w4_elapsed = _time_cuda(run_w4) + fp8_elapsed = _time_cuda(run_fp8) + + logical_m = groups * m_per_group + w4_bytes = _effective_bytes(groups, logical_m, n, k, a_gran_k, fp8_b=False, b_gran_k=b_gran_k) + fp8_bytes = _effective_bytes(groups, logical_m, n, k, a_gran_k, fp8_b=True, b_gran_k=a_gran_k) + + return { + "groups": groups, + "m_per_group": m_per_group, + "n": n, + "k": k, + "w4_us": w4_elapsed * 1e6, + "w4_gbps": w4_bytes / w4_elapsed / 1e9, + "w4_diff": w4_diff, + "fp8_us": fp8_elapsed * 1e6, + "fp8_gbps": fp8_bytes / fp8_elapsed / 1e9, + "fp8_diff": fp8_diff, + "speedup": fp8_elapsed / w4_elapsed, + } + +def _print_header(title: str) -> None: + print() + print(title) + print("=" * len(title)) + print("B scale format: group_size=128, dtype=torch.float32, non-UE8M0") + print("groups | m/group | n | k | W4 us | W4 GB/s | W4 diff | FP8 us | FP8 GB/s | FP8 diff | Speedup") + print("-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --") + +def _print_row(row: dict[str, float | int]) -> None: + print( + f"{row['groups']} | {row['m_per_group']} | {row['n']} | {row['k']} | " + f"{row['w4_us']:.0f} | {row['w4_gbps']:.0f} | {row['w4_diff']:.4f} | " + f"{row['fp8_us']:.0f} | {row['fp8_gbps']:.0f} | {row['fp8_diff']:.4f} | " + f"{row['speedup']:.2f}x" + ) + +def _run_suite( + mode: str, + groups_list: tuple[int, ...], + m_per_group_list: tuple[int, ...], + shapes: list[tuple[int, int]], + *, + a_gran_k: int, + b_gran_k: int, + max_m: int, + fail_on_threshold: bool, +) -> None: + if mode in ("contiguous", "both"): + _print_header("Contiguous FP8xFP4(g128) benchmark") + for groups in groups_list: + for m_per_group in m_per_group_list: + for n, k in shapes: + row = _benchmark_case_contiguous( + groups, + m_per_group, + n, + k, + a_gran_k=a_gran_k, + b_gran_k=b_gran_k, + ) + _print_row(row) + if fail_on_threshold: + assert row["w4_diff"] < W4_DIFF_THRESHOLD, row + assert row["fp8_diff"] < FP8_DIFF_THRESHOLD, row + + if mode in ("masked", "both"): + _print_header("Masked FP8xFP4(g128) benchmark") + for groups in groups_list: + for m_per_group in m_per_group_list: + for n, k in shapes: + row = _benchmark_case_masked( + groups, + m_per_group, + n, + k, + max_m=max_m, + a_gran_k=a_gran_k, + b_gran_k=b_gran_k, + ) + _print_row(row) + if fail_on_threshold: + assert row["w4_diff"] < W4_DIFF_THRESHOLD, row + assert row["fp8_diff"] < FP8_DIFF_THRESHOLD, row + +def main() -> None: + parser = argparse.ArgumentParser(description="SM90 FP8xFP4 g128 benchmark for fp32-scale weights") + parser.add_argument("--repo-root", type=str, default=None, help="DeepGEMM 仓库根目录") + parser.add_argument( + "--mode", + choices=("contiguous", "masked", "both"), + default="both", + help="运行 contiguous、masked 或两者都跑", + ) + parser.add_argument( + "--groups", + type=str, + default="8,16,32", + help="要测试的 group 数，逗号分隔，如 8,16,32", + ) + parser.add_argument( + "--m-per-group", + type=str, + default="1,4,8,16,32, 64", + ) + parser.add_argument( + "--shapes", + type=str, + default="4096x7168,7168x2048,4096x4096", + help="测试形状，格式 NxK, 逗号分隔", + ) + parser.add_argument("--a-gran-k", type=int, default=128) + parser.add_argument("--b-gran-k", type=int, default=128) + parser.add_argument("--max-m", type=int, default=128, help="masked 路径下每组分配的最大 m") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--no-fail-on-threshold", + action="store_true", + help="仅打印结果，不对 diff 阈值做 assert", + ) + args = parser.parse_args() + + if args.b_gran_k != 128: + raise ValueError("this script is specifically for FP4 + group_size=128") + + repo_root = _discover_repo_root(args.repo_root) + _load_deep_gemm(repo_root) + + torch.manual_seed(args.seed) + _require_sm90() + + groups_list = _parse_csv_ints(args.groups) + m_per_group_list = _parse_csv_ints(args.m_per_group) + shapes = _parse_shapes(args.shapes) + + print(f"repo_root = {repo_root}") + print("running SM90 FP8xFP4 g128 benchmark ...") + started = time.time() + + _run_suite( + mode=args.mode, + groups_list=groups_list, + m_per_group_list=m_per_group_list, + shapes=shapes, + a_gran_k=args.a_gran_k, + b_gran_k=args.b_gran_k, + max_m=args.max_m, + fail_on_threshold=not args.no_fail_on_threshold, + ) + + print() + print(f"done in {time.time() - started:.1f}s") + +if __name__ == "__main__": + main() From d88dc7f9c22e7fa6ef123bd6101de95828297c7c Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Tue, 14 Jul 2026 14:02:21 +0000 Subject: [PATCH 41/51] update sm90_fp8_fp4_g128 kernel --- csrc/apis/gemm.hpp | 53 +++++++ .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 141 ++++++++++++++++++ .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 82 ++++++++++ 3 files changed, 276 insertions(+) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index d1b18da5f6..816996e655 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -241,6 +241,51 @@ static void m_grouped_fp8_fp4_gemm_nn_contiguous(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& expected_m, + const std::string& compiled_dims) { + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 9); + + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); + DG_HOST_ASSERT(masked_m.is_contiguous()); + + const auto [num_groups, m, k] = get_shape<3>(a.first); + const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); + const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); + DG_HOST_ASSERT(num_groups == num_groups_b && num_groups == num_groups_d); + DG_HOST_ASSERT(m == m_d && n == n_d); + DG_HOST_ASSERT(half_k * 2 == k); + DG_HOST_ASSERT(static_cast(masked_m.numel()) == num_groups); + DG_HOST_ASSERT(expected_m > 0); + + std::optional> recipe = std::nullopt; + std::optional> recipe_a = std::make_tuple(1, 128); + std::optional> recipe_b = std::make_tuple(1, 128); + + const auto [sfa, sfb, gran_k_a, gran_k_b] = + layout::transform_sf_pair_into_required_layout( + a.second, b.second, m, n, k, + recipe, recipe_a, recipe_b, + num_groups, num_groups, false); + + DG_HOST_ASSERT(gran_k_a == 128 && gran_k_b == 128); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); + + sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_g128( + {a.first, sfa}, {b.first, sfb}, d, masked_m, expected_m, compiled_dims); +} + static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -675,6 +720,14 @@ static void register_apis(pybind11::module_& m) { py::arg("b_is_int4_sym") = false, py::arg("masked_m_max_hint") = std::nullopt, py::arg("active_groups_hint") = std::nullopt); + m.def("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma_g128", + &m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma_g128, + py::arg("a"), + py::arg("b"), + py::arg("d"), + py::arg("masked_m"), + py::arg("expected_m"), + py::arg("compiled_dims") = "nk"); m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma"); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index febe4bbb26..4d9a706dca 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -226,6 +226,147 @@ static void sm90_fp8_fp4_gemm_1d1d_fused(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& expected_m, + const std::string& compiled_dims) { + DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + + DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); + DG_HOST_ASSERT(masked_m.is_contiguous()); + + DG_HOST_ASSERT(a.first.is_contiguous()); + DG_HOST_ASSERT(b.first.is_contiguous()); + DG_HOST_ASSERT(d.is_contiguous()); + + const auto [num_groups, m, k] = get_shape<3>(a.first); + const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); + const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); + const auto num_groups_mask = static_cast(masked_m.numel()); + + DG_HOST_ASSERT(num_groups == num_groups_b and num_groups == num_groups_d and num_groups == num_groups_mask); + DG_HOST_ASSERT(m == m_d and n == n_d); + DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); + DG_HOST_ASSERT(expected_m > 0); + + DG_HOST_ASSERT(a.second.size(0) == num_groups); + DG_HOST_ASSERT(a.second.size(1) == m); + DG_HOST_ASSERT(a.second.size(2) == ceil_div(k, 128)); + + DG_HOST_ASSERT(b.second.size(0) == num_groups); + DG_HOST_ASSERT(b.second.size(1) == n); + DG_HOST_ASSERT(b.second.size(2) == ceil_div(k, 128)); + + check_major_type_cd(d); + + if (m == 0 or n == 0) { + return; + } + + auto desc = GemmDesc { + .gemm_type = GemmType::MGroupedMasked, + .kernel_type = KernelType::Kernel1D2D, + .m = m, + .n = n, + .k = k, + .num_groups = num_groups, + .a_dtype = a.first.scalar_type(), + .b_dtype = torch::kFloat8_e4m3fn, + .cd_dtype = d.scalar_type(), + .major_a = cute::UMMA::Major::K, + .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = expected_m, + .expected_n = n, + .expected_k = k, + .expected_num_groups = num_groups + }; + + auto config = get_best_config(desc); + + // Packed FP4 B has half the K bytes of FP8 B. Match the W4 path: + config.storage_config.swizzle_b_mode = config.layout.block_k / 2; + DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); + DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k / 2); + + const auto tensor_map_a = make_tma_a_desc( + cute::UMMA::Major::K, + a.first, + m, + k, + config.storage_config.load_block_m, + config.layout.block_k, + k, + num_groups, + config.storage_config.swizzle_a_mode); + + const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); + const auto tensor_map_b = make_tma_b_desc( + cute::UMMA::Major::K, + b_bytes, + n, + half_k, + config.layout.block_n, + config.layout.block_k / 2, + half_k, + num_groups, + config.storage_config.swizzle_b_mode); + + const auto tensor_map_sfa = make_tma_sf_desc( + cute::UMMA::Major::MN, + a.second, + m, + k, + config.layout.block_m, + config.layout.block_k, + num_groups, + 0); + + const auto tensor_map_d = make_tma_cd_desc( + d, + m, + n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(d.stride(-2)), + num_groups, + config.storage_config.swizzle_cd_mode); + + const SM90FP8FP4Gemm1D2DRSG128Runtime::Args& rs_args = { + .gemm_desc = desc, + .gemm_config = config, + .launch_args = LaunchArgs( + config.launch_config.num_sms, + config.launch_config.num_threads, + config.pipeline_config.smem_size, + config.layout.get_cluster_size()), + .major_sfb = get_major_type_ab(b.second), + .gmem_b_ptr = b.first.data_ptr(), + .gmem_d_ptr = d.data_ptr(), + .sfb = b.second.data_ptr(), + .grouped_layout = masked_m.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .tensor_map_sfa = tensor_map_sfa, + }; + + const auto code = SM90FP8FP4Gemm1D2DRSG128Runtime::generate(rs_args); + const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_g128", code); + SM90FP8FP4Gemm1D2DRSG128Runtime::launch(runtime, rs_args); +} + static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( const std::pair& a, const std::pair& b, diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 65653763be..8789f67ca3 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -274,4 +274,86 @@ static void __instantiate_kernel() {{ } }; +class SM90FP8FP4Gemm1D2DRSG128Runtime final + : public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + + cute::UMMA::Major major_sfb; + + void *gmem_b_ptr; + void *gmem_d_ptr; + void *sfb; + void *grouped_layout; + + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + CUtensorMap tensor_map_sfa; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_fp4_gemm_1d2d_rs_g128_impl< + {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, + {}, + {} + >); +}}; +)", + to_string(args.major_sfb), + get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), + get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), + args.gemm_desc.num_groups, + args.gemm_config.layout.block_m, + args.gemm_config.layout.block_n, + args.gemm_config.layout.block_k, + args.gemm_config.storage_config.swizzle_a_mode, + args.gemm_config.storage_config.swizzle_b_mode, + args.gemm_config.storage_config.swizzle_cd_mode, + args.gemm_config.pipeline_config.num_stages, + args.gemm_config.launch_config.num_tma_threads, + args.gemm_config.launch_config.num_math_threads, + args.gemm_config.layout.get_cluster_size(), + args.gemm_config.layout.cluster_n > 1, + args.gemm_config.launch_config.num_sms, + to_string(args.gemm_desc.gemm_type)); + } + + static void launch_impl(const KernelHandle& kernel, + const LaunchConfigHandle& config, + Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel( + kernel, + config, + args.gmem_b_ptr, + args.sfb, + args.grouped_layout, + args.gmem_d_ptr, + args.gemm_desc.m, + args.gemm_desc.n, + args.gemm_desc.k, + args.tensor_map_a, + args.tensor_map_b, + args.tensor_map_d, + args.tensor_map_sfa)); + } +}; + } // namespace deep_gemm From 24aa91a962f8aa225e3d2ff4c1d8dad0310403c8 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Wed, 15 Jul 2026 08:34:43 +0000 Subject: [PATCH 42/51] update sm90_fp8_fp4_g128 kernel --- csrc/apis/gemm.hpp | 53 ------- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 141 ------------------ .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 83 ----------- tests/test_sm90_fp8_fp4_g128.py | 2 +- 4 files changed, 1 insertion(+), 278 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 816996e655..d1b18da5f6 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -241,51 +241,6 @@ static void m_grouped_fp8_fp4_gemm_nn_contiguous(const std::pair& a, - const std::pair& b, - const torch::Tensor& d, - const torch::Tensor& masked_m, - const int& expected_m, - const std::string& compiled_dims) { - const auto arch_major = device_runtime->get_arch_major(); - DG_HOST_ASSERT(arch_major == 9); - - DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); - DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); - DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); - DG_HOST_ASSERT(masked_m.is_contiguous()); - - const auto [num_groups, m, k] = get_shape<3>(a.first); - const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); - const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); - DG_HOST_ASSERT(num_groups == num_groups_b && num_groups == num_groups_d); - DG_HOST_ASSERT(m == m_d && n == n_d); - DG_HOST_ASSERT(half_k * 2 == k); - DG_HOST_ASSERT(static_cast(masked_m.numel()) == num_groups); - DG_HOST_ASSERT(expected_m > 0); - - std::optional> recipe = std::nullopt; - std::optional> recipe_a = std::make_tuple(1, 128); - std::optional> recipe_b = std::make_tuple(1, 128); - - const auto [sfa, sfb, gran_k_a, gran_k_b] = - layout::transform_sf_pair_into_required_layout( - a.second, b.second, m, n, k, - recipe, recipe_a, recipe_b, - num_groups, num_groups, false); - - DG_HOST_ASSERT(gran_k_a == 128 && gran_k_b == 128); - DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); - DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); - - sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_g128( - {a.first, sfa}, {b.first, sfb}, d, masked_m, expected_m, compiled_dims); -} - static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -720,14 +675,6 @@ static void register_apis(pybind11::module_& m) { py::arg("b_is_int4_sym") = false, py::arg("masked_m_max_hint") = std::nullopt, py::arg("active_groups_hint") = std::nullopt); - m.def("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma_g128", - &m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma_g128, - py::arg("a"), - py::arg("b"), - py::arg("d"), - py::arg("masked_m"), - py::arg("expected_m"), - py::arg("compiled_dims") = "nk"); m.attr("m_grouped_fp8_fp4_gemm_nt_mask_sm90_fused_wgmma") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked_sm90_fused_wgmma"); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index 4d9a706dca..febe4bbb26 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -226,147 +226,6 @@ static void sm90_fp8_fp4_gemm_1d1d_fused(const std::pair& a, - const std::pair& b, - const torch::Tensor& d, - const torch::Tensor& masked_m, - const int& expected_m, - const std::string& compiled_dims) { - DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); - - DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(a.second.scalar_type() == torch::kFloat); - DG_HOST_ASSERT(b.first.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(b.second.scalar_type() == torch::kFloat); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); - DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); - DG_HOST_ASSERT(masked_m.is_contiguous()); - - DG_HOST_ASSERT(a.first.is_contiguous()); - DG_HOST_ASSERT(b.first.is_contiguous()); - DG_HOST_ASSERT(d.is_contiguous()); - - const auto [num_groups, m, k] = get_shape<3>(a.first); - const auto [num_groups_b, n, half_k] = get_shape<3>(b.first); - const auto [num_groups_d, m_d, n_d] = get_shape<3>(d); - const auto num_groups_mask = static_cast(masked_m.numel()); - - DG_HOST_ASSERT(num_groups == num_groups_b and num_groups == num_groups_d and num_groups == num_groups_mask); - DG_HOST_ASSERT(m == m_d and n == n_d); - DG_HOST_ASSERT(k % 2 == 0 and half_k * 2 == k); - DG_HOST_ASSERT(expected_m > 0); - - DG_HOST_ASSERT(a.second.size(0) == num_groups); - DG_HOST_ASSERT(a.second.size(1) == m); - DG_HOST_ASSERT(a.second.size(2) == ceil_div(k, 128)); - - DG_HOST_ASSERT(b.second.size(0) == num_groups); - DG_HOST_ASSERT(b.second.size(1) == n); - DG_HOST_ASSERT(b.second.size(2) == ceil_div(k, 128)); - - check_major_type_cd(d); - - if (m == 0 or n == 0) { - return; - } - - auto desc = GemmDesc { - .gemm_type = GemmType::MGroupedMasked, - .kernel_type = KernelType::Kernel1D2D, - .m = m, - .n = n, - .k = k, - .num_groups = num_groups, - .a_dtype = a.first.scalar_type(), - .b_dtype = torch::kFloat8_e4m3fn, - .cd_dtype = d.scalar_type(), - .major_a = cute::UMMA::Major::K, - .major_b = cute::UMMA::Major::K, - .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), - .compiled_dims = compiled_dims, - .expected_m = expected_m, - .expected_n = n, - .expected_k = k, - .expected_num_groups = num_groups - }; - - auto config = get_best_config(desc); - - // Packed FP4 B has half the K bytes of FP8 B. Match the W4 path: - config.storage_config.swizzle_b_mode = config.layout.block_k / 2; - DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k); - DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k / 2); - - const auto tensor_map_a = make_tma_a_desc( - cute::UMMA::Major::K, - a.first, - m, - k, - config.storage_config.load_block_m, - config.layout.block_k, - k, - num_groups, - config.storage_config.swizzle_a_mode); - - const auto b_bytes = b.first.view(torch::kFloat8_e4m3fn); - const auto tensor_map_b = make_tma_b_desc( - cute::UMMA::Major::K, - b_bytes, - n, - half_k, - config.layout.block_n, - config.layout.block_k / 2, - half_k, - num_groups, - config.storage_config.swizzle_b_mode); - - const auto tensor_map_sfa = make_tma_sf_desc( - cute::UMMA::Major::MN, - a.second, - m, - k, - config.layout.block_m, - config.layout.block_k, - num_groups, - 0); - - const auto tensor_map_d = make_tma_cd_desc( - d, - m, - n, - config.storage_config.store_block_m, - config.storage_config.store_block_n, - static_cast(d.stride(-2)), - num_groups, - config.storage_config.swizzle_cd_mode); - - const SM90FP8FP4Gemm1D2DRSG128Runtime::Args& rs_args = { - .gemm_desc = desc, - .gemm_config = config, - .launch_args = LaunchArgs( - config.launch_config.num_sms, - config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .major_sfb = get_major_type_ab(b.second), - .gmem_b_ptr = b.first.data_ptr(), - .gmem_d_ptr = d.data_ptr(), - .sfb = b.second.data_ptr(), - .grouped_layout = masked_m.data_ptr(), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .tensor_map_d = tensor_map_d, - .tensor_map_sfa = tensor_map_sfa, - }; - - const auto code = SM90FP8FP4Gemm1D2DRSG128Runtime::generate(rs_args); - const auto runtime = compiler->build("sm90_m_grouped_fp8_fp4_gemm_masked_1d2d_rs_g128", code); - SM90FP8FP4Gemm1D2DRSG128Runtime::launch(runtime, rs_args); -} - static void sm90_m_grouped_fp8_fp4_gemm_contiguous_1d1d_fused( const std::pair& a, const std::pair& b, diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 8789f67ca3..dad8c18d91 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -273,87 +273,4 @@ static void __instantiate_kernel() {{ args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.tensor_map_sfa)); } }; - -class SM90FP8FP4Gemm1D2DRSG128Runtime final - : public LaunchRuntime { -public: - struct Args { - GemmDesc gemm_desc; - GemmConfig gemm_config; - LaunchArgs launch_args; - - cute::UMMA::Major major_sfb; - - void *gmem_b_ptr; - void *gmem_d_ptr; - void *sfb; - void *grouped_layout; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - CUtensorMap tensor_map_d; - CUtensorMap tensor_map_sfa; - }; - - static std::string generate_impl(const Args& args) { - return fmt::format(R"( -#include - -using namespace deep_gemm; - -static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm90_fp8_fp4_gemm_1d2d_rs_g128_impl< - {}, - {}, {}, {}, - {}, - {}, {}, {}, - {}, {}, {}, - {}, - {}, {}, - {}, {}, - {}, - {} - >); -}}; -)", - to_string(args.major_sfb), - get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims), - get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims), - get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims), - args.gemm_desc.num_groups, - args.gemm_config.layout.block_m, - args.gemm_config.layout.block_n, - args.gemm_config.layout.block_k, - args.gemm_config.storage_config.swizzle_a_mode, - args.gemm_config.storage_config.swizzle_b_mode, - args.gemm_config.storage_config.swizzle_cd_mode, - args.gemm_config.pipeline_config.num_stages, - args.gemm_config.launch_config.num_tma_threads, - args.gemm_config.launch_config.num_math_threads, - args.gemm_config.layout.get_cluster_size(), - args.gemm_config.layout.cluster_n > 1, - args.gemm_config.launch_config.num_sms, - to_string(args.gemm_desc.gemm_type)); - } - - static void launch_impl(const KernelHandle& kernel, - const LaunchConfigHandle& config, - Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel( - kernel, - config, - args.gmem_b_ptr, - args.sfb, - args.grouped_layout, - args.gmem_d_ptr, - args.gemm_desc.m, - args.gemm_desc.n, - args.gemm_desc.k, - args.tensor_map_a, - args.tensor_map_b, - args.tensor_map_d, - args.tensor_map_sfa)); - } -}; - } // namespace deep_gemm diff --git a/tests/test_sm90_fp8_fp4_g128.py b/tests/test_sm90_fp8_fp4_g128.py index fcfca50c59..270c57c0b7 100644 --- a/tests/test_sm90_fp8_fp4_g128.py +++ b/tests/test_sm90_fp8_fp4_g128.py @@ -13,7 +13,7 @@ - 同时覆盖 contiguous / masked 两条常用调用路径 用法： -python /sgl-workspace/sglang/test_sm90_fp8_fp4_g128.py --repo-root /root/work/DeepGEMM +python /sgl-workspace/DeepGEMM/tests/test_sm90_fp8_fp4_g128.py --repo-root /sgl-workspace/DeepGEMM """ from __future__ import annotations From 31083e351257ed6648e056fb56aa2ecda9c4cee3 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Wed, 15 Jul 2026 09:53:59 +0000 Subject: [PATCH 43/51] Update layout.hpp --- csrc/apis/layout.hpp | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index a28468cf92..69ff8a0da5 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -36,13 +36,10 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, // Pre-transform checks check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups); - // (BF16, 1, 32/128) on SM90 SFB fast path: path-A (k128) 与 path-B fast-path (k32) - // 都支持 bf16 sfb，体积砍半。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,8) - // 构造，这里跳过 fp32 only 的 transpose，直接复用 align 路径。 + // SM90: BF16 SFB supports gran_k {32, 128}; expects MN-major TMA-aligned strides. if (sf.scalar_type() == torch::kBFloat16 and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 9) return get_mn_major_tma_aligned_tensor(sf); - // (UInt8/E8M0, 1, 32) on SM90 SFB fast path：path-B 专用，每元素 1B 即 fp32 的 - // 8 位指数。tensor 已由调用方按 MN-major + tma_aligned_mn=align(N,16) 构造。 + // SM90: UInt8 E8M0 SFB supports gran_k=32; expects MN-major TMA-aligned strides. if (sf.scalar_type() == torch::kUInt8 and gran_mn == 1 and gran_k == 32 and arch_major == 9) return get_mn_major_tma_aligned_tensor(sf); From 43f71205bc2b840f4a90e589463e77330c638007 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Wed, 15 Jul 2026 09:59:37 +0000 Subject: [PATCH 44/51] Update layout.hpp --- csrc/utils/layout.hpp | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 17b08e4831..045e3d85f9 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include @@ -92,19 +90,15 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf, // Always do shape checks const auto sf_dtype = sf.scalar_type(); - // BF16 / E8M0 (UInt8) SFB fast path 也允许；其它情况维持 fp32/int 限制。 + // Supported SF storage formats. INT is fixed to 4-way packed UE8M0. DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt or sf_dtype == torch::kBFloat16 or sf_dtype == torch::kUInt8); DG_HOST_ASSERT(sf.dim() == static_cast(num_groups.has_value()) + 2); if (num_groups.has_value()) DG_HOST_ASSERT(sf.size(-3) == num_groups.value()); DG_HOST_ASSERT(sf.size(-2) == ceil_div(mn, gran_mn)); - const char* scale_b_lut_env = std::getenv("DG_W4_SCALE_B_LUT"); - const bool scale_b_lut = sf_dtype == torch::kInt and - scale_b_lut_env != nullptr and scale_b_lut_env[0] != '\0' and scale_b_lut_env[0] != '0'; - const int expected_sf_k = scale_b_lut ? - ceil_div(k, gran_k) * 2 : - ceil_div(k, gran_k * (sf_dtype == torch::kInt ? 4 : 1)); + const int sf_pack_factor = sf_dtype == torch::kInt ? 4 : 1; + const int expected_sf_k = ceil_div(k, gran_k * sf_pack_factor); DG_HOST_ASSERT(sf.size(-1) == expected_sf_k); // TMA stride checks: TMA aligned and MN-major From 01a855ed58d4a49c4c71d394eef838b94b615337 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Wed, 15 Jul 2026 10:00:31 +0000 Subject: [PATCH 45/51] Updata sm90_fp8_fp4_gemm_1d2d_rs.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 74 ++++--------------- 1 file changed, 15 insertions(+), 59 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index dad8c18d91..1d6d6bbbf6 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -29,50 +29,10 @@ class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime Date: Wed, 15 Jul 2026 10:01:02 +0000 Subject: [PATCH 46/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 524 ++---------------- 1 file changed, 43 insertions(+), 481 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index febe4bbb26..dff5a1064d 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -1,7 +1,5 @@ #pragma once -#include - #include #include "../../jit/compiler.hpp" @@ -635,36 +633,6 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( } } DG_HOST_ASSERT(found_layout); - const auto env_enabled = [](const char* name) { - const char* value = std::getenv(name); - return value != nullptr and value[0] != '\0' and value[0] != '0'; - }; - const auto env_disabled = [](const char* name) { - const char* value = std::getenv(name); - return value != nullptr and value[0] == '0'; - }; - const auto env_int = [](const char* name, int default_value) { - const char* value = std::getenv(name); - return (value != nullptr and value[0] != '\0') ? std::atoi(value) : default_value; - }; - // BM=64 fast-path 是否启用（函数作用域统一判据，供三处共用：layout 选择 / - // bm32_skew_layout(stages) / bm32_skew_fast_path(device 三件套总闸)）。 - // 两个触发源： - // (1) DG_W4_PATHB_BM64_FASTPATH=1：手动强制（对照实验用）。 - // (2) DG_W4_PATHB_BM64_AUTO=1（默认开）+ 高置信子集 max_hint>=128 且 - // active<=8：自动触发，抓“少数活跃组 + 大 max_m”（MTP verify spec-len - // 命中形态）。详见 layout 选择处 bm64_fastpath 的长注释。 - // (3) DG_W4_PATHB_BM64_FORCE_ALL=1：无条件全部走 BM=64 fast-path - // （用于在线业务峰值吞吐 A/B 对照，绕开 hint 判据；hint 在生产 cuda graph - // 下恒为 None 时，是验证“假设拿到 hint 全命中”的速度上界）。 - // 注意：必须同时让 stages / fast-path 总闸认账，否则 layout.block_m=64 却 - // 走通用慢路径会退化 ~2x（见 DG_W4_PATHB_BM64 非 fastpath 的坑）。 - const bool bm64_fastpath_enabled = - env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 or - env_int("DG_W4_PATHB_BM64_FASTPATH", 0) != 0 or - (env_int("DG_W4_PATHB_BM64_AUTO", 1) != 0 and - masked_m_max_hint.value_or(0) >= 128 and - active_groups_hint.value_or(static_cast(desc.num_groups)) <= 8); // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 @@ -740,267 +708,45 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( layout.block_m = best.first; layout.block_n = best.second; - // RS masked W4 empirical fallback: - // pick BM close to expected_m to avoid over-computing promotion work, - // and use BN=256 for fewer CTAs while staying within Hopper TMA limits. - // bm_select_m 用 max(expected_m, masked_m_max_hint) —— 当调用方传入 - // hot group 的真实大小时，host 据此选 BM，避免 BM 与 hot group 严重 - // 失配；fast-path gating（k32 quad_reduce / direct_load 等）下面仍按 - // expected_m 判断，保留 small-m 路径的优化。 + // RS masked W4 fallback: + // pick BM close to max(expected_m, masked_m_max_hint) to avoid + // over-computing promotion work on hot groups. Path-B per-32 scale keeps + // the proven BM=32 fast path for hot groups; small/fan-out cases use the + // BM ladder and simple scheduler where applicable. const int bm_select_m = std::max( expected_m, masked_m_max_hint.value_or(expected_m)); if (desc.gemm_type == GemmType::MGroupedMasked) { - // 公共形状判定：DSV4 MTP/speculative verify 形状下 - // host scheduler 才会做 hard-code 干预，避开通用 shape。 - // A1+A2: 放开到 N∈{4096,7168}, K∈{2048,3072,4096,7168} 以覆盖 wide_n / wide_k 系列。 - // K=3072 引入动机：DSV4 gateup/down_proj 还有一组 n=7168+k=3072 的真实 - // shape，K=3072 满足 BLOCK_K=128 整除 (3072/128=24) + gran_k_b=32 整除 - // (3072/32=96) ⇒ device fast-path (quad_reduce/direct_load/compact_sched) - // 物理可命中。原守护把 k=3072 拒在外面，强制走 path-B 通用 cache_sfb_k32 - // 路径，W4 GB/s 只有 ~600（vs fast-path 命中时 941+ GB/s）。 - // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： - // 把守护放宽到 num_groups>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}。 - // 动机：DSV4 down_proj 真实形状是 g24, n=6144, k=7168，原守护把它 - // 挡在 fast-path 外，强制走 path-B 通用 cache_sfb_k32 路径，W4 GB/s - // 只有 627（vs fast-path 命中时 1500-3000+ GB/s），speedup 0.43-0.56x。 - // device fast-path 三件套（quad_reduce / direct_load / compact_sched） - // 实际 BN-agnostic，仅硬约束 BM==32 + gran_k_b==32，冗余专家场景 - // kNumGroups<=36 也可按同一路径处理。 - // 实测：g24+n=6144+k=7168 整张表 0.43-0.56x → 0.71-0.81x； - // g24+n=7168+k=3072 整张表 0.30-0.54x → 0.45-0.79x； - // 默认 dsv4 形状 (g32+n∈{4096,7168}+k∈{2048,4096,7168}) 不变。 - // 开启后所有依赖 dsv4_shape 的路径（BN256 门槛 / small_hot - // / large_bm / cluster_n / bm64_hint）都同步放开，已验证无退化。 - // 设 DG_W4_PATHB_FAST_PATH_RELAX=0 可显式回退到原守护做对照。 - const bool fast_path_relax = - env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0; + // DSV4 / speculative-verify shape set. This is the former relaxed + // shape set, now fixed as the only specialized masked path. const bool dsv4_shape = - (static_cast(desc.num_groups) >= 8 and - (static_cast(desc.n) == 4096 or - static_cast(desc.n) == 7168) and - (static_cast(desc.k) == 2048 or - static_cast(desc.k) == 3072 or - static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168)) - or - (fast_path_relax and - static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) >= 8 and (static_cast(desc.n) == 4096 or static_cast(desc.n) == 6144 or static_cast(desc.n) == 7168) and (static_cast(desc.k) == 2048 or static_cast(desc.k) == 3072 or static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168)); + static_cast(desc.k) == 7168); if (gran_k_b == 32) { - // path-B：device fast-path（BM=32 + quad_reduce + direct_load） - // 是 gran_k_b=32 下的唯一高效形状。BN 选取根据 hint 大小： - // * hint <= 32（小 hot）：BM=32 BN=128（短 wgmma 流水线 + 高 stage 数） - // * hint > 32（大 hot）：BM=32 BN=256（grid 减半，scheduler 调度密度提升） - // - // 历史教训（skewed masked benchmark 实测）： - // * 试图基于 (active, max_m) 把"path-B 输给 FP8"的 case 切到 - // BM 阶梯 + BN=256 fan-out → device kernel 落入 path-B 通用 - // cache_sfb_k32 路径，比 fast-path 慢 1.5–3x（如 eight_hot_64 - // 190→298us、dense_tail 368→1011us）。 - // * 结论：path-B 必须保持 BM=32 + fast-path 守护命中，否则崩盘。 - // BN=128↔256 两挡都在 device fast-path 守护内（参见 device - // kernel kK32QuadReduce 路径，BN-agnostic：compute_n_0/1 - // 寻址含 m_offset = local_idx * WAVE_BLOCK_M， - // load_k32_quad_scale_b 按 compute_n 单独寻址）。 - // - // hint<=16（纯 fan-out）保留原 BM 阶梯走 BN=256 的 path-B 通用 - // 路径——这部分 device 守护本来就不命中。 - // real_hot_present 判定：caller 是否真的传入 max_hint > 16？ - // * has_value=true：按真实 max_hint 判断（>16 才认为 hot 存在）。 - // * has_value=false（caller 没传 hint）： - // - expected_m > 16：保留旧行为（保守认为 hot 存在）， - // 走 BM=32 fast-path（m>=1024 时命中 fast-path 守护）。 - // - expected_m <= 16：**新行为**，认为 hot 不存在，让 case - // 落到 expected_m candidate 路径（BM∈{8,16,32} 自动选 + - // small_m_simple_sched 命中）。 - // 动机：DSV4 EP 业务真实 shape 大量出现 expected_m=1/2/3 + - // max_hint=None + m=256 的情形。旧逻辑强制 BM=32，但 - // bm32_skew_fast_path 守护要 m>=1024，m=256 时 fast-path 不 - // 命中反而走通用 cache_sfb_k32 路径，padding 浪费 - // (32-expected_m)/32 ≈ 94%。落到 expected_m candidate 后 - // device 端 BLOCK_M<=16 命中 kUseSmallMSimpleSched，物理零浪费。 const bool real_hot_present = masked_m_max_hint.has_value() ? (masked_m_max_hint.value() > 16) : (expected_m > 16); - // fuse_decode 路径与 fast-path 互斥：env 开启时强制走通用 BM - // 阶梯（device 端 kFuseScaleBDecode 走 path-B 通用 cache_sfb_k32， - // 而 fast-path layout 假设 quad_reduce + direct_load）。 - const bool fuse_decode_hint = - env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; - // BM=64 实验性路径（默认关，env DG_W4_PATHB_BM64=1）： - // 想法：path-B swap_AB 下 wgmma 物理 N = BLOCK_M。BM=64 → wgmma - // 64x64x32，单次 wgmma 内 b_decoded 服务 64 cols 而非 32， - // W4 解码 (prmt+lop3) 成本天然摊薄 2x，弥补 fast-path - // (quad_reduce + direct_load) 关闭带来的 sfb load + fmul - // overhead。同时 grid 减半。 - // 代价：device 端 kK32QuadReduce 假设 BM=32 (compute_n_0/1 只覆盖 - // 32 cols)，BM=64 必须走通用 cache_sfb_k32 + fp32 sfb 路径。 - // 实测结果（DG_W4_PATHB_BM64=1 vs 默认 fast-path）： - // gateup_dense_tail_hot64 356us → 748us (2.1x 退化) - // gateup_mtp_dp4 158us → 275us (1.74x 退化) - // gateup_eight_hot_64 188us → 226us (1.20x 退化) - // gateup_one_hot_32 31us → 47us (1.52x 退化) - // 绝大多数 case 退化 1.2-2.1x，无一受益。 - // 退化根因：cache_sfb_k32 阶段 SMEM build LUT + register fmul + barrier - // 比 fast-path 直读 gmem + quad_reduce 慢得多；BM=64 grid 砍半 - // 反而让单 SM 排队工作量翻倍，stages 不够掩 K-loop（GB/s 从 - // 1818 跌到 452）。decode 复用收益被 sfb pipeline 重建吃光。 - // 保留 env 入口仅做对照实验，生产环境永不开启。 - const bool bm64_hint = - dsv4_shape and env_int("DG_W4_PATHB_BM64", 0) != 0; - // BM=64 fast-path 实验路径（默认关，env DG_W4_PATHB_BM64_FASTPATH=1）： - // 与 DG_W4_PATHB_BM64 不同——后者刻意让 device 落回通用 - // cache_sfb_k32 路径；本开关**保持** fast-path 三件套 - // (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched) 开启， - // 仅把 host BM 从 32 抬到 64，验证“weight 解码次数减半 + grid 减半” - // 能否吃掉 max_m>=64 差场景。 - // 依赖：device 端 swap_AB 下 WGMMA N=BLOCK_M=64 指令已存在 - // (FP8MMASelectorRS<64> = MMA_64x64x32_RS_TN)，scale 寻址 BM-agnostic。 - // BM=64 强制 BN=128，避免 A-smem 叠加 BN=256 把 stages 压崩。 - // - // 实测结论（DG_W4_PATHB_BM64_FASTPATH=1 vs 默认 BM=32，ncu 对照， - // g24 N=6144 K=7168 / N=7168 K=3072 skew benchmark）：**全面退化，已证伪**。 - // uniform_17 784us/0.75x → 1391us/0.43x (max_m=17) - // uniform_32 775us/0.77x → 1387us/0.42x (max_m=32) - // dense_tail_hot17 767/0.77x → 1381/0.43x - // uniform_64 1541us/0.38x → 1390us/0.43x (唯一微升) - // one_hot_512 542us/0.58x → 542us/0.58x (完全无变化) - // ncu 四项硬指标 BM32 vs BM64 **逐字节相同**： - // Registers/Thread 168=168, Dyn SMEM 149984=149984, - // Grid (78,1,1)=(78,1,1), Achieved Occupancy 14.0%≈13.9%。 - // 证伪根因（结构性，非实现 bug；正确性 OK diff=0、无 spill）： - // 1. persistent kernel grid 恒 = num_sms，BLOCK_M 不影响 grid。 - // 2. swap_AB 下 WGMMA 形状 / kNumAccum / WAVE_BLOCK_M 全由 - // BLOCK_N 决定，BLOCK_M 加倍不动寄存器、不动 smem。 - // 3. M-tile 数 = ceil(max_m/BM)。DSV4 真实 max_m 多 <=32， - // BM 32→64 时 tile 数不减，但每 tile WGMMA 覆盖 64 行、 - // ~47 行是 padding 空算 → 工作量翻倍、收益为零 → 耗时近翻倍。 - // 4. 仅 max_m=64 时 tile 2→1 抵消 padding 才微升；max_m 更大 - // (one_hot_512 active=1) 瓶颈不在 tile 数，故毫无变化。 - // 结论：当前 DSV4 负载 (max_m 集中 17-64) 下 BM=64 是负优化。 - // 保留 env 入口仅做对照实验，生产环境永不开启。 - // 真正受益方向见下方思路 2（active=1 大 max_m 时跨 m_block 复用 - // weight decode），不受 padding 问题影响。 - // BM=64 fast-path **自动触发规则**（默认开，env 可关）： - // 前述 DG_W4_PATHB_BM64_FASTPATH 手动实验证明，BM 32→64 让 - // M-tile 数 ceil(m/32)→ceil(m/64) 减半，仅当“每活跃组平均 - // token 数”够大（每组 m>=64，padding 占比可忽略）才净赚；组小则 - // tile 数不减、padding 翻倍纯亏。判别量本应是 sum_m/active，但 - // host API 只有 max_hint / active（拿不到 sum_m，且不能同步读 - // device 上的 masked_m）。 - // 折中：用高置信子集 `max_hint>=128 且 active<=8`（少数活跃组 + - // 大 max_m，正是 MTP verify 的 spec-len 命中形态）触发，刻意避开 - // max_m=64 / 多活跃组的模糊带（uniform_64 赢、dense_tail_hot64 - // 亏，二者 host 信号 (max_hint=64,active=24) 撞车无法区分）。 - // 实测（g24 N6144/7168 与 g32 N4096，BM64+pair vs 默认 quad）两族 - // shape 一致：命中子集 one_hot_128/384、mtp_384/512_hot、 - // one_hot_512、mtp_768_hot512 等普降 5~26%，被避开的 - // dense_tail/uniform/eight_hot 无一受损。 - // BM=64 下 quad 4 累加器会 spill 爆炸，故命中时**必须**配 - // pair_reduce（见下方 k32_pair_reduce 联动）。 - const bool bm64_fastpath = dsv4_shape and bm64_fastpath_enabled; - if (real_hot_present and not fuse_decode_hint and - not bm64_hint and - env_int("DG_W4_PATHB_FAST_PATH", 1) != 0) { - layout.block_m = bm64_fastpath ? 64 : 32; - // H1: 同时满足 max_m > 32 + active*max_m >= 1024 时走 BN=256， - // grid 减半。门槛公式来自 skewed masked benchmark 实测： - // * eight_hot_64 (8×64=512)、four_hot_64 (4×64=256)、 - // one_hot_384 (1×384=384) 在 BN=256 下 grid 太稀， - // 单 tile WAVE_WGMMA=2 的两次 wgmma 反而拖慢 8-9%。 - // * mtp_dp4 (8×214=1712)、dense_tail_hot64 a32 - // (32×64=2048)、mtp_512_hot384 (8×384=3072) 等 - // 总工作量充足的 case 才能从 BN=256 grid 减半中受益。 - // 需 dsv4_shape 才能命中 device fast-path 守护 - // (N=4096, K∈{4096,2048}, groups=32)，否则 BN=256 落到 - // path-B 通用 cache_sfb_k32 路径反而崩盘。 + if (real_hot_present) { + layout.block_m = 32; const int hint_m = masked_m_max_hint.value_or(0); - // 缺省 active 估计为 num_groups 的一半（保守，避免误升 BN=256） const int hint_active = active_groups_hint.value_or( static_cast(desc.num_groups) / 2); - // 总工作量估计（FLOPS 代理）：active × max_m × N × K。 - // 用 int64 防溢出，N/K 单位 elements。 - const int64_t hint_workload = - static_cast(hint_active) * hint_m * - static_cast(desc.n) * - static_cast(desc.k); - // 基础门槛：g32 + N=4096 + K=4096 实测 hint_active*hint_m >= 1024 - // 且 hint_m > 32 时 BN=256 受益。 const bool bn256_baseline = hint_m > 32 and static_cast(hint_active) * hint_m >= 1024; - // 大 N×K 形态扩展（实验性，env DG_W4_PATHB_BN256_BIG=1 默认关）： - // 动机：DSV4 down_proj real shape g24, N=6144, K=7168 下， - // uniform_17/24/32 + dense_tail_hot17 全被 hint_m > 32 - // 卡住走 BN=128 (~941 GB/s)，理论上 grid 4-9 wave 已脱离 - // "grid 太稀"危险区，单 tile work 翻倍应可吞。 - // 阈值原本设计：active>=8 + workload>=1.5e10。 - // 实测结果（DG_W4_PATHB_BN256_BIG=1 vs 默认 BN=128）： - // uniform_17 711us → 797us (-12%, 941→838 GB/s) - // uniform_24 716us → 799us (-12%) - // uniform_32 705us → 790us (-12%) - // dense_tail_hot17 701us → 786us (-12%) - // 退化根因：K=7168 + BN=256 + BM=32 + stages=8 ⇒ SMEM B-tile - // 占用翻倍 (256 cols × 32 K × 8 stages = 64KB)，超出 228KB - // 上限被 JIT 自动降到 stages=4，K-loop 掩蔽崩盘 (GB/s - // 941→842 整齐下跌 11%)。BN=256 grid 减半的收益 << stages - // 减半的带宽损失。 - // 保留 env 入口仅做对照实验，生产环境永不开启。 - constexpr int64_t kBN256BigWorkloadThreshold = 15'000'000'000LL; - const bool bn256_big_shape = - env_int("DG_W4_PATHB_BN256_BIG", 0) != 0 and - hint_workload >= kBN256BigWorkloadThreshold and - static_cast(hint_active) >= 8; const bool bn256_eligible = dsv4_shape and masked_m_max_hint.has_value() and - (bn256_baseline or bn256_big_shape) and - env_int("DG_W4_PATHB_BN256", 1) != 0; - // BM=64 fast-path 强制 BN=128：A-smem 已随 BM 翻倍，再叠加 - // BN=256 会把 B-tile + A-tile 占用一起推高，stages 被压崩。 - layout.block_n = (bm64_fastpath or not bn256_eligible) ? 128 : 256; - // H2: cluster_n=2（A 多播）实测**全军退化**，默认关闭。 - // 退化数据（DG_W4_PATHB_CLUSTER_N=1 vs 默认）： - // BN=256 路径 ~9-10% 退化（mtp_512: 201→220, dense_tail_hot64: - // 354→388, mtp_dp4: 158→172） - // BN=128 路径 ~3-12% 退化（four_hot_64: 95→106, - // longtail_128: 157→169） - // 退化根因： - // 1. compact_masked_sched 按 (group, n_block) 调度，cluster=2 - // 把相邻 n_block 绑定到同 cluster，但 masked 下相邻 - // n_block 可能属于不同 group，A tile 复用率被打断。 - // 2. cluster barrier 同步 + cluster launch stagger 对 small - // grid（28-56 个 tile）开销大于多播收益。 - // 3. W4 的 sfb 不参与 multicast，同步开销均摊到两个 CTA。 - // 保留 env 入口仅做对照实验，生产环境永不开启。 - if (dsv4_shape and - env_int("DG_W4_PATHB_CLUSTER_N", 0) != 0 and - static_cast(desc.n) % - (static_cast(layout.block_n) * 2) == 0 and - desc.num_sms % 2 == 0) { - layout.cluster_n = 2; - } - } else if (bm64_hint and real_hot_present) { - // BM=64 实验路径：仅 dsv4_shape + DG_W4_PATHB_BM64=1 时进入。 - // 走 path-B 通用 cache_sfb_k32 路径，依赖下面 fast-path 守护 - // 自动关闭 (kK32QuadReduce / kScaleBDirectLoad / kCompactMaskedSched - // 通过 bm32_skew_fast_path 的 BM==32 检查失败而被关闭)。 - layout.block_m = 64; - // BN=128 默认：BM=64 + BN=128 → wgmma 64x64x32 × 2 wgmma per - // K_block，b_decoded 在每次 wgmma 内部服务 64 cols。 - // env DG_W4_PATHB_BM64_BN=256 时尝试更大 grid 减半。 - const int bm64_bn_override = env_int("DG_W4_PATHB_BM64_BN", 0); - layout.block_n = (bm64_bn_override == 256) ? 256 : 128; + bn256_baseline; + layout.block_n = bn256_eligible ? 256 : 128; } else { - // 纯 fan-out（hint<=16）：base 阶梯，BN=256 if (bm_select_m <= 8) layout.block_m = 8; else if (bm_select_m <= 16) layout.block_m = 16; else if (bm_select_m <= 32) layout.block_m = 32; @@ -1015,39 +761,17 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( else if (bm_select_m <= 64) layout.block_m = 64; layout.block_n = 256; - // DSV4 + 大 BM：BN=256→128 减小单 tile promote/store 链长， - // 让 grid 上 n-tile 翻倍提升 SM 间负载均衡。 - if (layout.block_m >= 64 and dsv4_shape and - env_int("DG_W4_LARGE_BM_BN128", 1) != 0) { + if (layout.block_m >= 64 and dsv4_shape) { layout.block_n = 128; } - // DSV4 + small hot (bm_select∈(32,64])：BM=64 padding 浪费太大， - // 改 BM=32 BN=128 兼顾 hot/small group。 - if (bm_select_m > 32 and bm_select_m <= 64 and dsv4_shape and - env_int("DG_W4_SMALL_HOT_BM32", 1) != 0) { + if (bm_select_m > 32 and bm_select_m <= 64 and dsv4_shape) { layout.block_m = 32; layout.block_n = 128; } } - // 历史经验注记： - // * 曾尝试 path-B BM=32→64 升级（hot hint），device kernel 的 - // kK32QuadReduce 在 BM>=64 + swap_AB 下 sfb 寻址 / 4-wgmma fence - // 跟 BM=32 紧耦合，dp4 退化到 0.37x。 - // * 已知天花板（DSV4 skew, groups=32, N=4096）：进一步提升须从 - // device 侧扩展 BM=64 quad_reduce，工作量大且收益不确定。 } } if (not block_m_override and not block_n_override) { - // DG_W4_PATHB_BM64_FORCE_ALL=1 时，把 BLOCK_N 强制锁到 128（与 BM=64 - // fast-path 联动）。动机：BM=64 fast-path 下 final_accum 寄存器量 - // ∝ BLOCK_N/64，BN=256 → 4 份 tile 累加器逼出 168reg/thread、occupancy - // 14%；BN=128 砍半 final_accum，腾出寄存器预算给更多 warp，对大 batch / - // 高 active 场景 occupancy 上限更高。属于在线吞吐对照实验的联动开关， - // 不影响 hint 判据自身路径。 - if (env_int("DG_W4_PATHB_BM64_FORCE_ALL", 0) != 0 and - bm64_fastpath_enabled and layout.block_m == 64) { - layout.block_n = 128; - } DG_HOST_ASSERT(layout.block_m == 8 or layout.block_m == 16 or layout.block_m == 32 or layout.block_m == 64 or layout.block_m == 128); DG_HOST_ASSERT(layout.block_n == 64 or layout.block_n == 128 or layout.block_n == 256); @@ -1124,36 +848,19 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) chosen_stages = std::min(chosen_stages, 6); const bool bm32_skew_layout = gran_k_b == 32 and - (config.layout.block_m == 32 or - (config.layout.block_m == 64 and bm64_fastpath_enabled)) and + config.layout.block_m == 32 and (config.layout.block_n == 128 or config.layout.block_n == 256) and static_cast(desc.m) >= 1024 and - ((static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 36 and - (static_cast(desc.n) == 4096 or - static_cast(desc.n) == 7168) and - (static_cast(desc.k) == 2048 or - static_cast(desc.k) == 3072 or - static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168)) or - // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）： - // 放开到 g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168} - (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and - static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 36 and - (static_cast(desc.n) == 4096 or - static_cast(desc.n) == 6144 or - static_cast(desc.n) == 7168) and - (static_cast(desc.k) == 2048 or - static_cast(desc.k) == 3072 or - static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168))); - const int bm32_skew_stage_override = env_int("DG_W4_BM32_SKEW_STAGES", 0); - if (bm32_skew_layout and bm32_skew_stage_override > 0) { - DG_HOST_ASSERT(bm32_skew_stage_override >= 3 and bm32_skew_stage_override <= kW4DefaultMaxStages); - DG_HOST_ASSERT(fits(bm32_skew_stage_override)); - chosen_stages = bm32_skew_stage_override; - } else if (bm32_skew_layout and fits(kW4DefaultMaxStages)) { + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168); + if (bm32_skew_layout and fits(kW4DefaultMaxStages)) { chosen_stages = kW4DefaultMaxStages; } DG_HOST_ASSERT(chosen_stages >= 3); @@ -1164,12 +871,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( // R2b-A swap_ab maps original N onto WGMMA M. Use enough math warpgroups // to cover the 64-row WGMMA-M strips selected by BLOCK_N. DG_HOST_ASSERT(config.layout.block_n == 64 or config.layout.block_n == 128 or config.layout.block_n == 256); - int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; - const int rs_num_math_threads_override = env_int("DG_W4_RS_MATH_THREADS", 0); - if (rs_num_math_threads_override > 0) { - DG_HOST_ASSERT(rs_num_math_threads_override == 128 or rs_num_math_threads_override == 256); - rs_num_math_threads = rs_num_math_threads_override; - } + const int rs_num_math_threads = config.layout.block_n <= 64 ? 128 : 256; config.launch_config.num_math_threads = rs_num_math_threads; config.launch_config.num_threads = config.launch_config.num_tma_threads + rs_num_math_threads; @@ -1191,108 +893,23 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( config.storage_config.swizzle_cd_mode); const bool bm32_skew_fast_path = gran_k_b == 32 and - (config.layout.block_m == 32 or - (config.layout.block_m == 64 and bm64_fastpath_enabled)) and + config.layout.block_m == 32 and (config.layout.block_n == 128 or config.layout.block_n == 256) and static_cast(desc.m) >= 1024 and - ((static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 36 and - (static_cast(desc.n) == 4096 or - static_cast(desc.n) == 7168) and - (static_cast(desc.k) == 2048 or - static_cast(desc.k) == 3072 or - static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168)) or - // FAST_PATH_RELAX 扩展（DG_W4_PATHB_FAST_PATH_RELAX=1 **默认开启**）， - // 与 bm32_skew_layout / dsv4_shape 同步生效。device fast-path 三件套 - // (quad_reduce / direct_load / compact_sched) 仅硬约束 BM==32，g<=36 - // 冗余专家场景按同一路径处理， - // 与 N/K 数值无关，物理可命中。 - (env_int("DG_W4_PATHB_FAST_PATH_RELAX", 1) != 0 and - static_cast(desc.num_groups) >= 8 and - static_cast(desc.num_groups) <= 36 and - (static_cast(desc.n) == 4096 or - static_cast(desc.n) == 6144 or - static_cast(desc.n) == 7168) and - (static_cast(desc.k) == 2048 or - static_cast(desc.k) == 3072 or - static_cast(desc.k) == 4096 or - static_cast(desc.k) == 7168))); - // Fused-decode 路径（path-B 通用 cache_sfb_k32 + LUT decode）： - // * scale 在 cache 阶段编进 LUT，wgmma 后省一次 fmul（按 quad 算 = 4 次） - // * sfb 4x packed UE8M0：每 4 个 e8m0 打成 1 个 int32，体积 = fp32 的 1/4 - // 与 fast-path（kK32QuadReduce + kScaleBDirectLoad）**互斥**（cuh 强制 - // not kScaleBDirectLoad），开启后强制走通用路径。需要 sfb 是 packed int 布局。 - // 默认关，由 DG_W4_PATHB_FUSE_DECODE=1 显式开启。 - const bool fuse_scale_b_decode = - gran_k_b == 32 and - sfb.scalar_type() == torch::kInt and - get_major_type_ab(sfb) == cute::UMMA::Major::MN and - env_int("DG_W4_PATHB_FUSE_DECODE", 0) != 0; - const bool scale_b_packed_ue8m0 = fuse_scale_b_decode; - // 思路2 性能探针：把 scale_a 从 promote 内移除（device kScaleAStub 置 scale_a=1 - // 并跳过 SFA 的 TMA load），用于测 fuse_scale_b_decode 单累加器 + scale_a 后移到 - // silu_and_mul 的速度上界。开启后结果数值错误，仅做性能对照。默认关。 - const bool scale_a_stub = env_int("DG_W4_SCALE_A_STUB", 0) != 0; - // 思路2 串行依赖实验：双缓冲预解码（staged_pair_a_regs，issue k 时预取 k+2）， - // 让 decode 与 WGMMA 多重叠一级，代价 +8 reg/线程。仅在 fuse_scale_b_decode 开启 - // 时有意义，用于测"解 issue 串行依赖"能否补回 fuse 路径在大 M 的退化。默认关。 - const bool fuse_predecode_pair = - fuse_scale_b_decode and env_int("DG_W4_FUSE_PREDECODE_PAIR", 0) != 0; - // fuse_scale_b_decode 一旦开启，所有 fast-path / direct-load / e8m0 / bf16 - // 互斥关闭——device 那边 static_assert 会拒绝同时开启。 + static_cast(desc.num_groups) >= 8 and + static_cast(desc.num_groups) <= 36 and + (static_cast(desc.n) == 4096 or + static_cast(desc.n) == 6144 or + static_cast(desc.n) == 7168) and + (static_cast(desc.k) == 2048 or + static_cast(desc.k) == 3072 or + static_cast(desc.k) == 4096 or + static_cast(desc.k) == 7168); const bool scale_b_direct_load = - gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and - not fuse_scale_b_decode; + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path); const bool k32_quad_reduce = - gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path) and - not fuse_scale_b_decode; - // 杠杆3 / BM=64 联动：把 4 累加器 quad-reduce（峰值 64 float，在 - // __launch_bounds__(384,1) 的 168reg 硬上限下触发 local spill）退化为已存在的 - // 2 累加器串行 pair-reduce（峰值 32 float）降低寄存器峰值活跃量。仅改单线程内 - // 累加器生命周期、不动线程模型/barrier，数值 bit 级等价。 - // 两种启用： - // (A) BM=32 下**单独**开 pair（DG_W4_K32_PAIR_REDUCE=1）：已证伪——ncu spill - // 223680→134016(-40%) 但 wall-clock 反而 +10~25%（串行累加拉长依赖链， - // spill 不在 BM=32 关键路径上）。默认关。 - // (B) BM=64 fast-path 命中时**必须**配 pair（自动，下面 block_m==64 分支）： - // BM=64 让 kNumAccum 32→quad 4 份 = 128 fp32，spill 爆炸退化 ~2x；pair - // 2 份 = 64 fp32 受控。BM=64 靠 M-tile 数减半换吞吐，pair 是其必要前提。 - // 实测高置信子集（max_hint>=128 且 active<=8）净降 5~26%。 - const bool k32_pair_reduce = - k32_quad_reduce and - ((config.layout.block_m == 64 and bm64_fastpath_enabled) or - env_enabled("DG_W4_K32_PAIR_REDUCE")); - // 杠杆3 续（已证伪，默认关）：pair-reduce 之上把长驻 final_accum 从 fp32 改 bf16 - // 存储（nv_bfloat162）再减半长驻寄存器、继续压 spill（ncu 134016→6480，-97% vs - // quad）。精度无损（W4 diff 0.0001/0.0000）。但 wall-clock 仍慢于 quad baseline - // （随 pair-reduce 一起退化，叠加 bf16 promote 反而更长依赖链）。默认关，需 - // DG_W4_K32_BF16_FINAL_ACCUM=1 显式开启复现。 - const bool k32_bf16_final_accum = - k32_pair_reduce and env_enabled("DG_W4_K32_BF16_FINAL_ACCUM"); - // 杠杆3 续2（已证伪，默认关）：把长驻 final_accum 整体搬进 smem_d（kDirectStore - // + kFinalAccumScratch），final_accum_regs 缩成 1 元素。实测 spill 虽归零，但 - // (1) 寄存器仍顶 168（launch_bounds 硬顶，省下的 reg 没还给 occupancy，Block - // Limit Registers 仍为 1）；(2) DirectStore 逐元素非合并写 gmem 严重拖累大 - // max_m：one_hot_256 0.77x→0.50x、one_hot_512 0.49x→0.29x。净负收益，仅保留 - // 开关供复现，默认关。需 DG_W4_K32_FINAL_ACCUM_SMEM=1 显式开启。 - const bool k32_final_accum_smem = - k32_pair_reduce and env_enabled("DG_W4_K32_FINAL_ACCUM_SMEM"); - // RS baseline for direct E8M0 B scales: keep scale products in split form by default. - // The exponent-adjust path is experimental and can regress some small-M shapes. - const bool scale_b_pow2_promote = - k32_quad_reduce and env_enabled("DG_W4_SCALE_B_POW2_PROMOTE"); - const bool k32_quad_split_promote = - k32_quad_reduce and not env_disabled("DG_W4_K32_QUAD_SPLIT_PROMOTE"); - const bool k32_quad_scale_b_inline = - k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_INLINE"); - const bool k32_quad_scale_b_prefetch = - k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_PREFETCH"); - const bool k32_quad_scale_b_vec4 = - k32_quad_reduce and env_enabled("DG_W4_K32_QUAD_SCALE_B_VEC4"); - const bool k32_quad_pair4x2_promote = - k32_quad_reduce and not k32_quad_split_promote and - env_enabled("DG_W4_K32_QUAD_PAIR4X2_PROMOTE"); + gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path); + const bool k32_quad_split_promote = k32_quad_reduce; // small_m_simple_sched：device 端 kUseSmallMSimpleSched 仅编译期检查 // BLOCK_M<=16 + GroupedMasked + multicast=1，与 N/K 数值无关。 // 默认守护 (k>=4096 + n<=4096) 来自历史 g32 + n=4096 dsv4 形状的保守覆盖。 @@ -1311,59 +928,17 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( static_cast(desc.k) == 3072 or static_cast(desc.k) == 4096 or static_cast(desc.k) == 7168))); - const bool compact_masked_sched = - bm32_skew_fast_path and not env_disabled("DG_W4_COMPACT_MASKED_SCHED"); - // compact_masked_sched 按 m_max 降序遍历 active group（实验性扩展，默认关）： - // 想法：compact_masked_sched 默认按 group_idx 升序遍历 active group。 - // 当各 group masked_m 不均时，long-tail group 决定 last-wave 收敛节奏。 - // 开启后 inner loop 每步在剩余 active mask 内 O(active) 扫一次找 group_m - // 最大者，让 wave 0 优先吃重 group，wave 末尾留小 group 收尾。 - // 实测结果（DG_W4_PATHB_REORDER_BY_MM=1 vs 默认 fast-path）： - // gateup_dense_tail_hot64 356us → 356us (0%) - // gateup_dense_tail_hot128 399us → 396us (+0.8%) - // gateup_dense_tail_hot214 401us → 400us (-0.2%) - // gateup_mtp_dp4 157us → 157us (0%) - // down_dense_tail_hot64 186us → 185us (+0.5%) - // down_mtp_512_hot384_a8 98us → 100us (-2%) - // down_dense_tail_hot214 206us → 206us (0%) - // gateup_one_hot_32 31us → 31us (+1%, 边缘 case) - // 全军 ±2% 噪声振荡，无任何 case 出现 5-10% 实质收益。 - // 失败根因：compact_masked_sched 已经做了 (group, m_block, n_block) 三维紧凑 - // 展开，dense_tail_hot64 总 tile = 32 group × 2 m_blocks × 32 n_blocks = - // 2048 ≈ 15.5 wave × 132 SM。当 total_tiles >> num_sms 时，wave 边界由 - // ceil(total/132) 决定，与 group 顺序无关；重排只改变"哪些 SM 在 wave 0 - // 拿到 hot group"，last-wave 余数 = (total mod 132) 不变。 - // 边缘 case (one_hot_32 total=32 tile = 1 wave × 32 SM) 的 +1-2% 收益 - // 是 SM 提前进入 idle 的小幅噪声，不构成可推广收益。 - // 保留 env 入口仅做对照实验，生产环境永不开启。 - // 仅在 compact_masked_sched 已开启时生效（即 bm32_skew_fast_path）， - // 否则即便 env 开启也不会传到 device。 - const bool reorder_masked_by_max_m = - compact_masked_sched and env_int("DG_W4_PATHB_REORDER_BY_MM", 0) != 0; - // bf16 SFB：体积砍半，**默认开启**。两条 fast-path： - // - path-B (gran_k_b==32 + direct-load)：load_sfb / load_k32_quad_scale_b 直读 gmem。 - // - path-A (gran_k_b==128)：cooperative prefetch 经 smem，smem 也存 bf16。 - // 仍要避开 packed-UE8M0 / fused-decode 等假设 fp32 位级布局的分支。 - // 显式 `DG_W4_SCALE_B_BF16=0` 时回退 fp32；用户传的 sfb 实际是 fp32 也自然回退。 + const bool compact_masked_sched = bm32_skew_fast_path; const bool scale_b_bf16 = - ((scale_b_direct_load and gran_k_b == 32 and not k32_quad_scale_b_prefetch) or + ((scale_b_direct_load and gran_k_b == 32) or gran_k_b == 128) and - not env_disabled("DG_W4_SCALE_B_BF16") and sfb.scalar_type() == torch::kBFloat16; - // E8M0 SFB（仅 path-B fast-path）：每元素 1B = fp32 的 8 位指数，体积再砍 2x。 - // 解码 `__uint_as_float(uint32(e) << 23)` 零误差。**默认开启**：当用户传入 - // uint8 sfb 时自动启用；显式 `DG_W4_SCALE_B_E8M0=0` 时回退。 - // 与 bf16 互斥（因 sfb 实际 dtype 已不同），不会同时命中。 const bool scale_b_e8m0 = scale_b_direct_load and gran_k_b == 32 and - not k32_quad_scale_b_prefetch and - (not env_disabled("DG_W4_SCALE_B_E8M0") or - env_enabled("DG_W4_SCALE_B_E8M0_ONLY")) and sfb.scalar_type() == torch::kUInt8; DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat or (scale_b_bf16 and sfb.scalar_type() == torch::kBFloat16) or - (scale_b_e8m0 and sfb.scalar_type() == torch::kUInt8) or - (fuse_scale_b_decode and sfb.scalar_type() == torch::kInt)); + (scale_b_e8m0 and sfb.scalar_type() == torch::kUInt8)); const SM90FP8FP4Gemm1D2DRSRuntime::Args& rs_args = { .gemm_desc = desc, .gemm_config = config, @@ -1372,27 +947,14 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( config.layout.get_cluster_size()), .major_sfb = get_major_type_ab(sfb), .scale_b_direct_load = scale_b_direct_load, - .scale_b_pow2_promote = scale_b_pow2_promote, .k32_quad_reduce = k32_quad_reduce, - .k32_pair_reduce = k32_pair_reduce, - .k32_bf16_final_accum = k32_bf16_final_accum, - .k32_final_accum_smem = k32_final_accum_smem, .k32_quad_split_promote = k32_quad_split_promote, - .k32_quad_scale_b_inline = k32_quad_scale_b_inline, - .k32_quad_scale_b_prefetch = k32_quad_scale_b_prefetch, - .k32_quad_scale_b_vec4 = k32_quad_scale_b_vec4, - .k32_quad_pair4x2_promote = k32_quad_pair4x2_promote, .small_m_simple_sched = small_m_simple_sched, .compact_masked_sched = compact_masked_sched, - .fuse_scale_b_decode = fuse_scale_b_decode, - .scale_a_stub = scale_a_stub, - .fuse_predecode_pair = fuse_predecode_pair, - .scale_b_packed_ue8m0 = scale_b_packed_ue8m0, .scale_b_gran_k = static_cast(gran_k_b), .b_is_int4_sym = b_is_int4_sym, .scale_b_bf16 = scale_b_bf16, .scale_b_e8m0 = scale_b_e8m0, - .reorder_masked_by_max_m = reorder_masked_by_max_m, .gmem_b_ptr = b.first.data_ptr(), .gmem_d_ptr = d.data_ptr(), .sfb = sfb.data_ptr(), From 9e94a795b5905c99b2e3b445af306ffcf18b007c Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Wed, 15 Jul 2026 10:01:55 +0000 Subject: [PATCH 47/51] Update test_sm90_fp8_fp4 and int4_a8 .py --- tests/test_sm90_fp8_fp4.py | 79 +++++++++++++++----------------------- tests/test_sm90_int4_a8.py | 44 +++++++++++---------- 2 files changed, 54 insertions(+), 69 deletions(-) diff --git a/tests/test_sm90_fp8_fp4.py b/tests/test_sm90_fp8_fp4.py index 9f2f687d89..6b79ecacdb 100644 --- a/tests/test_sm90_fp8_fp4.py +++ b/tests/test_sm90_fp8_fp4.py @@ -1,6 +1,6 @@ +import argparse import sys import time -import os from pathlib import Path import torch @@ -215,22 +215,21 @@ def _masked_benchmark_case( a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) a_sf = torch.empty((groups, max_m, k // a_gran_k), device="cuda", dtype=torch.float) b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) - use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) - b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) - b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + b_sf_k = k // b_gran_k + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) for group_id in range(groups): a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k ) b_fp4[group_id], b_sf[group_id] = per_token_cast_to_fp4( - b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k ) a = (a_data, a_sf) b_w4 = (b_fp4, b_sf) assert a[1].shape == (groups, max_m, k // a_gran_k) assert b_w4[1].shape == (groups, n, b_sf_k) - if b_gran_k == 128 and not use_packed_b_sf: + if b_gran_k == 128: assert b_w4[1].dtype == torch.float assert b_w4[1].shape == (groups, n, k // 128) @@ -242,7 +241,7 @@ def _masked_benchmark_case( for group_id in range(groups): valid_m = int(masked_m[group_id].item()) b_dequant = cast_back_from_fp4( - b_w4[0][group_id], b_w4[1][group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + b_w4[0][group_id], b_w4[1][group_id], gran_k=b_gran_k ) b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( b_dequant, use_ue8m0=False, gran_k=a_gran_k @@ -342,51 +341,21 @@ def _masked_skew_benchmark_case( a_data = torch.empty((groups, max_m, k), device="cuda", dtype=torch.float8_e4m3fn) a_sf = torch.empty((groups, max_m, k // a_gran_k), device="cuda", dtype=torch.float) b_fp4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) - use_packed_b_sf = bool(int(os.getenv("DG_W4_FUSE_SCALE_B_DECODE", "0"))) - b_sf_k = k // (b_gran_k * (4 if use_packed_b_sf else 1)) - block_m_override = int(os.getenv("DG_W4_BLOCK_M_OVERRIDE", "0")) or None - block_n_override = int(os.getenv("DG_W4_BLOCK_N_OVERRIDE", "0")) or None + b_sf_k = k // b_gran_k bm32_skew_fast_path = ( b_gran_k == 32 and masked_m_max_hint is not None and masked_m_max_hint > 16 - and os.getenv("DG_W4_PATHB_FUSE_DECODE", "0") == "0" - and os.getenv("DG_W4_PATHB_FAST_PATH", "1") != "0" - and os.getenv("DG_W4_PATHB_BM64", "0") == "0" - and (block_m_override is None or block_m_override == 32) - and (block_n_override is None or block_n_override in (128, 256)) and max_m >= 1024 - and groups == 32 - and n in (4096, 7168) + and 8 <= groups <= 36 + and n in (4096, 6144, 7168) and k in (2048, 3072, 4096, 7168) ) scale_b_direct_load = b_gran_k == 32 and ( expected_m <= 16 or bm32_skew_fast_path ) - scale_b_dtype_fast_path = ( - scale_b_direct_load - and os.getenv("DG_W4_K32_QUAD_SCALE_B_PREFETCH", "0") == "0" - and os.getenv("DG_W4_SCALE_B_POW2_PROMOTE", "0") == "0" - ) - # bf16 sfb：path-A (gran_k_b=128) 与 path-B fast-path (gran_k_b=32) 都支持， - # 体积砍半。按 MN-major + tma_aligned_mn=align(N, 8) 直接构造，避开 host 端 - # fp32-only 的 transpose 路径。**默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 - use_bf16_b_sf = ( - ((b_gran_k == 32 and scale_b_dtype_fast_path) or b_gran_k == 128) - and not use_packed_b_sf - and bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) - ) - # E8M0 sfb（uint8）：仅 path-B (gran_k_b=32)。每元素 = fp32 pow2 scale 的指数位。 - # per_token_cast_to_fp4(..., use_ue8m0=True) 已保证 sf 是严格 pow2，因此抽指数无损。 - # **默认开启**，DG_W4_SCALE_B_E8M0=0 时回退。优先级高于 bf16（互斥）。 - use_e8m0_b_sf = ( - b_gran_k == 32 - and scale_b_dtype_fast_path - and not use_packed_b_sf - and bool(int(os.getenv("DG_W4_SCALE_B_E8M0", "1"))) - ) - if use_e8m0_b_sf: - use_bf16_b_sf = False # e8m0 优先 + use_e8m0_b_sf = b_gran_k == 32 and scale_b_direct_load + use_bf16_b_sf = b_gran_k == 128 if use_e8m0_b_sf: # uint8: tma_aligned_mn = ceil(N, 16)，要求 N % 16 == 0。 assert n % 16 == 0 @@ -410,14 +379,14 @@ def _masked_skew_benchmark_case( ) b_sf_fp32 = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) else: - b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.int if use_packed_b_sf else torch.float) + b_sf = torch.empty((groups, n, b_sf_k), device="cuda", dtype=torch.float) b_sf_fp32 = b_sf for group_id in range(groups): a_data[group_id], a_sf[group_id] = per_token_cast_to_fp8( a_ref_src[group_id], use_ue8m0=False, gran_k=a_gran_k ) b_fp4[group_id], b_sf_fp32[group_id] = per_token_cast_to_fp4( - b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + b_ref_src[group_id], use_ue8m0=True, gran_k=b_gran_k ) if use_e8m0_b_sf: # b_sf_fp32 已是严格 pow2（per_token_cast_to_fp4 use_ue8m0=True）。 @@ -453,7 +422,7 @@ def _masked_skew_benchmark_case( b_sf_for_ref = b_sf_fp32 if (use_e8m0_b_sf or use_bf16_b_sf) else b_w4[1] for group_id, valid_m in enumerate(masked_m_values): b_dequant = cast_back_from_fp4( - b_w4[0][group_id], b_sf_for_ref[group_id], gran_k=b_gran_k, use_packed_ue8m0=use_packed_b_sf + b_w4[0][group_id], b_sf_for_ref[group_id], gran_k=b_gran_k ) b_fp8_data[group_id], b_fp8_sf[group_id] = per_token_cast_to_fp8( b_dequant, use_ue8m0=False, gran_k=a_gran_k @@ -908,13 +877,25 @@ def long_tail_values( # _print_skew_table(sorted(rows, key=lambda row: float(row["speedup"]))[:12]) +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="SM90 FP8xFP4 accuracy and benchmark cases") + parser.add_argument( + "--case", + choices=("masked", "masked-skew", "masked-direct-fp32-scale", "contiguous"), + default="masked", + help="Benchmark case to run", + ) + return parser.parse_args() + + if __name__ == "__main__": + args = _parse_args() start_time = time.time() - # if os.getenv("DG_W4_CONTIGUOUS_DIRECT_FP32_SCALE", "0") not in ("", "0"): - # test_sm90_fp8_fp4_contiguous() - if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + if args.case == "contiguous": + test_sm90_fp8_fp4_contiguous() + elif args.case == "masked-skew": test_sm90_fp8_fp4_masked_skew_cases() - elif os.getenv("DG_W4_MASKED_DIRECT_FP32_SCALE", "0") not in ("", "0"): + elif args.case == "masked-direct-fp32-scale": test_sm90_fp8_fp4_masked_direct_fp32_scale() else: test_sm90_fp8_fp4_masked() diff --git a/tests/test_sm90_int4_a8.py b/tests/test_sm90_int4_a8.py index a9d654ed4f..52412b9ae8 100644 --- a/tests/test_sm90_int4_a8.py +++ b/tests/test_sm90_int4_a8.py @@ -31,7 +31,7 @@ from __future__ import annotations -import os +import argparse import sys import time from pathlib import Path @@ -303,26 +303,18 @@ def _build_int4_a8_skew_inputs(masked_m_values, max_m, n, k, *, gran_k=128): b_ref = torch.randn((groups, n, k), device="cuda", dtype=torch.bfloat16) b_int4 = torch.empty((groups, n, k // 2), device="cuda", dtype=torch.int8) - # path-A bf16 sfb：与 fp8_fp4 测试同款，按 MN-major + tma_aligned_mn=align(N,8) 直构造。 - # **默认开启**，DG_W4_SCALE_B_BF16=0 时回退 fp32。 - use_bf16_b_sf = bool(int(os.getenv("DG_W4_SCALE_B_BF16", "1"))) - if use_bf16_b_sf: - assert n % 8 == 0 - tma_aligned_n = (n + 7) // 8 * 8 - b_int4_sf = torch.empty_strided( - (groups, n, k // gran_k), - (tma_aligned_n * (k // gran_k), 1, tma_aligned_n), - device="cuda", - dtype=torch.bfloat16, - ) - b_int4_sf_fp32 = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) - else: - b_int4_sf = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) - b_int4_sf_fp32 = b_int4_sf + assert n % 8 == 0 + tma_aligned_n = (n + 7) // 8 * 8 + b_int4_sf = torch.empty_strided( + (groups, n, k // gran_k), + (tma_aligned_n * (k // gran_k), 1, tma_aligned_n), + device="cuda", + dtype=torch.bfloat16, + ) + b_int4_sf_fp32 = torch.empty((groups, n, k // gran_k), device="cuda", dtype=torch.float) for g in range(groups): b_int4[g], b_int4_sf_fp32[g] = per_token_cast_to_int4_sym(b_ref[g], gran_k=gran_k) - if use_bf16_b_sf: - b_int4_sf.copy_(b_int4_sf_fp32.to(torch.bfloat16)) + b_int4_sf.copy_(b_int4_sf_fp32.to(torch.bfloat16)) group_idx_a = torch.arange(k, device="cuda") // gran_k a_dequant = a_fp8.float() * a_sf[..., group_idx_a] @@ -514,10 +506,22 @@ def skew_values(total: int, hot: int, active: int = 8) -> list[int]: raise AssertionError(f"{fail_count} skew cases failed cos_abs > {COS_ABS_THRESHOLD}") +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="SM90 INT4-A8 masked accuracy and benchmark cases") + parser.add_argument( + "--case", + choices=("accuracy", "skew"), + default="accuracy", + help="Benchmark case to run", + ) + return parser.parse_args() + + if __name__ == "__main__": + args = _parse_args() start_time = time.time() _require_sm90() - if os.getenv("DG_W4_MASKED_SKEW_CASES", "0") not in ("", "0"): + if args.case == "skew": test_int4_a8_masked_skew_cases() else: test_int4_a8_masked_accuracy() From 7f0be399bbb311dcad3224bd5c3fc2534b2c907c Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Thu, 16 Jul 2026 06:56:19 +0000 Subject: [PATCH 48/51] Update sm90_fp8_fp4_gemm_1d2d_rs.hpp --- csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp index 1d6d6bbbf6..9da6bdc2fe 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d_rs.hpp @@ -37,10 +37,9 @@ class SM90FP8FP4Gemm1D2DRSRuntime final: public LaunchRuntime Date: Thu, 16 Jul 2026 06:57:00 +0000 Subject: [PATCH 49/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d.hpp | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp index dff5a1064d..0c211740ae 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_fp4_gemm_1d2d.hpp @@ -541,10 +541,10 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( const std::optional& masked_m_max_hint = std::nullopt, // active_groups_hint: caller passes count of groups with masked_m > 0 // (i.e. (masked_m != 0).sum()). Combined with masked_m_max_hint this - // is enough to estimate "工作量分布" and decide fast-path 是否合适： - // * 单热点 / 极少活跃 group → fast-path (BM=32 BN=128) 大胜 - // * 大量活跃 group + 高 max_m → fan-out (BM 阶梯, BN=256) 取胜 - // 不传时退化为旧行为（仅看 max_hint）。 + // is enough to estimate workload distribution and decide fast-path suitability: + // * Single-hot / few active groups → fast-path (BM=32 BN=128) dominates + // * Many active groups + high max_m → fan-out (BM ladder, BN=256) wins + // Falls back to old behavior (max_hint only) when not provided. const std::optional& active_groups_hint = std::nullopt) { DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); const int gran_k_a_requested = gran_k_a_override.value_or(gran_k); @@ -633,24 +633,24 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( } } DG_HOST_ASSERT(found_layout); - // W4 masked 启发式：以 weight HBM 带宽为主要瓶颈，需要足够的 pipeline stages - // 来隐藏 TMA B 的延迟。在 expected_m 较小时（典型 MoE 场景），优先选择能让 - // stages 数最深的 (BM, BN) 组合，并兼顾 wave 利用率。 + // W4 masked heuristic: weight HBM bandwidth is the main bottleneck; sufficient pipeline + // stages are needed to hide TMA B latency. For small expected_m (typical MoE), prioritize + // (BM, BN) combinations that maximize stages, while also considering wave utilization. // - // 参数空间：BM ∈ {8, 16, 32, 64, 128}（BM<64 用于 masked small-M）， - // BN ∈ {64, 128, 256}。 + // Parameter space: BM ∈ {8, 16, 32, 64, 128} (BM<64 for masked small-M), + // BN ∈ {64, 128, 256}. // - // 选择目标：在 (waves <= ceil_div(total_tiles, num_sms)) 的前提下最大化 stages， - // 其次最大化 last_wave 利用率，最后倾向更小的 per-stage（即更小 BN）。 + // Selection criteria: maximize stages under (waves <= ceil_div(total_tiles, num_sms)), + // then maximize last_wave utilization, then prefer smaller per-stage (smaller BN). if (not block_m_override and not block_n_override) { const int num_sms = desc.num_sms; const int block_k = layout.block_k; const int shape_k_scales_b = ceil_div(static_cast(k), gran_k_b); auto eval_layout = [&](int bm, int bn) -> std::tuple { - // 返回 (sat_stages, -waves, last_wave_util, -per_stage)，越大越好。 - // stages 在 ~6 之上对 TMA 隐藏几乎饱和，因此用饱和 stage 数比较，避免 - // 小 BN 因 stages=8 击败 wave 利用率更高的候选。 + // Returns (sat_stages, -waves, last_wave_util, -per_stage); higher is better. + // Stages saturate TMA hiding around ~6, so use saturated stage count to avoid + // small BN with stages=8 defeating candidates with better wave utilization. const int tiles = ceil_div(expected_m, bm) * ceil_div(static_cast(n), bn) * num_groups; const int waves = ceil_div(tiles, num_sms); const int last = tiles - (waves - 1) * num_sms; @@ -692,10 +692,10 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( for (const auto& cand : w4_candidates) { const int bm = cand.first; const int bn = cand.second; - // 1D2D 内核 unroll 要求 + // 1D2D kernel unroll requirements if (bn > block_k and (bn % (bn - block_k) != 0 and block_k % (bn - block_k) != 0)) continue; - // masked 路径 multicast 合法性：当前固定 cluster_m=1, cluster_n=1，恒满足 + // Masked multicast validity: currently fixed cluster_m=1, cluster_n=1, always satisfied const auto score = eval_layout(bm, bn); if (std::get<0>(score) < 3) continue; @@ -754,7 +754,7 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( layout.block_n = 256; } } else { - // path-A：cooperative prefetch + sfb→smem，BM 阶梯有效。 + // Path-A: cooperative prefetch + sfb→smem, BM ladder effective. if (bm_select_m <= 8) layout.block_m = 8; else if (bm_select_m <= 16) layout.block_m = 16; else if (bm_select_m <= 32) layout.block_m = 32; @@ -910,12 +910,12 @@ static void sm90_m_grouped_fp8_fp4_gemm_masked_1d1d_fused( const bool k32_quad_reduce = gran_k_b == 32 and (expected_m <= 16 or bm32_skew_fast_path); const bool k32_quad_split_promote = k32_quad_reduce; - // small_m_simple_sched：device 端 kUseSmallMSimpleSched 仅编译期检查 - // BLOCK_M<=16 + GroupedMasked + multicast=1，与 N/K 数值无关。 - // 默认守护 (k>=4096 + n<=4096) 来自历史 g32 + n=4096 dsv4 形状的保守覆盖。 - // 放开到 RELAX 形状集 (g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}) - // 让 DSV4 EP 业务真实 shape (g24 + n∈{6144,7168} + k∈{3072,7168} + expected_m=1/2/3) - // 也能命中 simple_sched，避开通用 masked scheduler 的额外开销。 + // small_m_simple_sched: device-side kUseSmallMSimpleSched is compile-time only. + // BLOCK_M<=16 + GroupedMasked + multicast=1, independent of N/K values. + // Default guard (k>=4096 + n<=4096) from historical g32 + n=4096 DSV4 shapes. + // Extended to RELAX shape set (g>=8 + n∈{4096,6144,7168} + k∈{2048,3072,4096,7168}) + // to allow DSV4 EP real shapes (g24 + n∈{6144,7168} + k∈{3072,7168} + expected_m=1/2/3) + // to hit simple_sched and avoid generic masked scheduler overhead. const bool small_m_simple_sched = gran_k_b == 32 and expected_m <= 16 and ((static_cast(desc.k) >= 4096 and static_cast(desc.n) <= 4096) or From 655d3cd1a696ccdfe45becd99acede0180b29dcc Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Thu, 16 Jul 2026 06:57:20 +0000 Subject: [PATCH 50/51] Update smxx_layout.hpp --- csrc/jit_kernels/impls/smxx_layout.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index ef022e6c07..c4b96e555d 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -107,7 +107,7 @@ static std::tuple preprocess_sf(const to if (require_float) { DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); } else { - // 放过 bf16 / E8M0 (UInt8) sfb（path-B fast-path，体积砍半 / 砍 4 倍）。 + // Allow BF16/UInt8 E8M0 SFB for compressed scale paths. DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat or sf.scalar_type() == torch::kInt or sf.scalar_type() == torch::kBFloat16 or @@ -130,8 +130,8 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn) return (dim == 2) ? batched_sf.squeeze(0) : batched_sf; - // bf16 / E8M0 (UInt8) sfb 必须由调用方按 MN-major + tma_aligned_mn 直接构造， - // 不能落入下面的 fp32-only TransposeFP32Runtime。 + // BF16/UInt8 E8M0 SFB must be pre-constructed with MN-major + TMA-aligned strides. + // They bypass the fp32-only transpose path above. DG_HOST_ASSERT(sf.scalar_type() != torch::kBFloat16 and sf.scalar_type() != torch::kUInt8); const auto out = torch::empty_strided({num_groups, mn, sf_k}, From 5e895256f53dfb219fc653ca1e1b1e6c817ea358 Mon Sep 17 00:00:00 2001 From: "shiyang.814" Date: Thu, 16 Jul 2026 06:57:40 +0000 Subject: [PATCH 51/51] Update sm90_fp8_fp4_gemm_1d2d.hpp --- .../impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh | 207 +++--------------- 1 file changed, 25 insertions(+), 182 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh index ab68ef4013..06d5914bf2 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_fp4_gemm_1d2d_rs.cuh @@ -627,71 +627,16 @@ template -CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, kLaunchBoundsMinBlocks) void + bool kScaleBE8M0 = false> +CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layout, nv_bfloat16* gmem_d_ptr, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, @@ -704,21 +649,10 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling"); DG_STATIC_ASSERT(kScaleBGranK == 32 or kScaleBGranK == 128, "Only support per-32 or per-128-channel FP4 scaling"); - DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or (kFuseScaleBDecode and kScaleBGranK == 32), - "Packed UE8M0 SFB is only supported by fused per-32 scale_b decode"); - DG_STATIC_ASSERT(kScaleBGranK == 128 or not kOverlapPromote, - "DG_W4_OVERLAP_PROMOTE does not support per-32 FP4 scaling"); - DG_STATIC_ASSERT(kScaleBGranK == 128 or kScaleKGroup == 1, - "DG_W4_SCALE_K_GROUP does not support per-32 FP4 scaling"); - DG_STATIC_ASSERT(kScaleKGroup == 1 or kScaleKGroup == 2 or kScaleKGroup == 4, - "DG_W4_SCALE_K_GROUP only supports 1/2/4"); DG_STATIC_ASSERT(not kBIsInt4Sym, "RS-mode FP8xFP4 kernel does not support INT4-sym B"); DG_STATIC_ASSERT(not (kScaleBBF16 and kScaleBE8M0), "Scale-B cannot be both BF16 and E8M0"); DG_STATIC_ASSERT((not kScaleBBF16 and not kScaleBE8M0) or (kScaleBDirectLoad and kScaleBGranK == 32), "Compressed Scale-B dtypes are only supported by direct-load per-32 path"); - DG_STATIC_ASSERT(kFuseScaleBDecodeAssumeExp == 0 or kFuseScaleBDecodeAssumeExp == 5 or - kFuseScaleBDecodeAssumeExp == 6, - "DG_W4_FUSE_SCALE_B_DECODE_ASSUME_EXP only supports 0/5/6"); DG_STATIC_ASSERT( math::constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (math::constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block"); @@ -759,9 +693,7 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo // SFB cache aliases smem_d when it fits. Small-M tiles may not have enough // smem_d capacity, so they fall back to a separate SFB region. const uint32_t smem_sfb_bytes = kScaleBGranK == 32 ? - math::align((BLOCK_K / kScaleBGranK) * - (kFuseScaleBDecode ? math::ceil_div(BLOCK_N, 4u) : BLOCK_N) * - sizeof(float), 16u) : + math::align((BLOCK_K / kScaleBGranK) * BLOCK_N * sizeof(float), 16u) : math::align(shape_k_scales_b * BLOCK_N * sizeof(float), 16u); // NOTES: Make sure we have enough shared memory for WGMMA padding static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); @@ -774,14 +706,8 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; - constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; - constexpr bool kParallelNWavesEnabled = - kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; - DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, - "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); - constexpr uint32_t kEmptyBarrierMathWarps = - kParallelNWavesEnabled ? kBaseWGMMAStoreThreads / 32 : kNumMathThreads / 32; + constexpr uint32_t kEmptyBarrierMathWarps = kNumMathThreads / 32; // Prefetch TMA descriptors at the very beginning if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { @@ -926,25 +852,21 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo tma::copy(&tensor_map_a, &full_barrier, smem_a[stage_idx], k_idx, scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx), num_tma_multicast_a, batch_idx); - if constexpr (not kScaleAStub) { - tma::copy(&tensor_map_sfa, &full_barrier, - smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales_a, 1, k_block_idx), - num_tma_multicast_a); - } - - if constexpr (not kWeightStub) { - // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier. - const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; - tma::copy(&tensor_map_b, &full_barrier, - reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), - k_idx_packed, - scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), - num_tma_multicast_b, batch_idx); - } + tma::copy(&tensor_map_sfa, &full_barrier, + smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.template get_global_idx(shape_k_scales_a, 1, k_block_idx), + num_tma_multicast_a); + + // Issue TMA B (packed FP4 bytes loaded as raw uint8 via FP8 alias) on the same barrier. + const uint32_t k_idx_packed = k_block_idx * BLOCK_K_PACKED; + tma::copy(&tensor_map_b, &full_barrier, + reinterpret_cast<__nv_fp8_e4m3*>(smem_b_packed[stage_idx]), + k_idx_packed, + scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx), + num_tma_multicast_b, batch_idx); constexpr uint32_t kExpectedTxBytes = SMEM_A_TMA_SIZE_PER_STAGE + - (kWeightStub ? 0 : SMEM_B_PACKED_SIZE_PER_STAGE) + - (kScaleAStub ? 0 : SMEM_SFA_TMA_SIZE_PER_STAGE); + SMEM_B_PACKED_SIZE_PER_STAGE + + SMEM_SFA_TMA_SIZE_PER_STAGE; full_barrier.arrive_and_expect_tx(kExpectedTxBytes); } } @@ -1045,62 +967,11 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo } auto cache_sfb_k32 = [&](uint32_t k_block_idx) { - if constexpr (kScaleBGranK == 32 and not kScaleBStub and not kFuseScaleBDecodeStub and not kScaleBDirectLoad) { + if constexpr (kScaleBGranK == 32 and not kScaleBDirectLoad) { const uint32_t n_block_base = n_block_idx * BLOCK_N; const uint32_t scale_k_base = k_block_idx * (BLOCK_K / kScaleBGranK); constexpr uint32_t kScaleRows = BLOCK_K / kScaleBGranK; - if constexpr (kFuseScaleBDecode) { - DG_STATIC_ASSERT(not kScaleBPackedUE8M0 or kMajorSFB == cute::UMMA::Major::MN, - "Packed UE8M0 SFB path expects MN-major transformed scale layout"); - constexpr uint32_t kVec = 4; - DG_STATIC_ASSERT(BLOCK_N % kVec == 0, - "BLOCK_N must be a multiple of 4 for packed K/32 SFB exp cache"); - constexpr uint32_t kVecsPerRow = BLOCK_N / kVec; - const uint32_t total_vecs = kScaleRows * kVecsPerRow; - for (uint32_t i = threadIdx.x; i < total_vecs; i += kNumMathThreads) { - const uint32_t k_off = i / kVecsPerRow; - const uint32_t n_off = (i % kVecsPerRow) * kVec; - const uint32_t n_idx = n_block_base + n_off; - uint32_t packed_offsets = 0; - #pragma unroll - for (uint32_t j = 0; j < kVec; ++ j) { - uint32_t exp_offset = 6; - if constexpr (kScaleBPackedUE8M0) { - if (n_idx + j < shape_n) { - const uint32_t packed_shape_k_scales_b = math::ceil_div(shape_k_scales_b, 4u); - const uint32_t packed_k_idx = (scale_k_base + k_off) / 4u; - const uint32_t byte_idx = (scale_k_base + k_off) % 4u; - const uint32_t* sfb_packed = reinterpret_cast(sfb); - const uint32_t packed_scale = sfb_packed[ - current_group_idx * aligned_shape_n_sfb * packed_shape_k_scales_b + - packed_k_idx * aligned_shape_n_sfb + n_idx + j]; - const uint32_t e8m0_exp = (packed_scale >> (byte_idx * 8)) & 0xffu; - exp_offset = (e8m0_exp > 121u) ? (e8m0_exp - 121u) : 0u; - } - } else { - float val = 1.0f; - if (n_idx + j < shape_n) { - if constexpr (kMajorSFB == cute::UMMA::Major::MN) { - const float* sfb_base = sfb + - current_group_idx * aligned_shape_n_sfb * shape_k_scales_b; - const float* ptr = sfb_base + - (scale_k_base + k_off) * aligned_shape_n_sfb + n_idx + j; - val = *ptr; - } else { - const float* ptr = sfb + - current_group_idx * shape_n * shape_k_scales_b + - (n_idx + j) * shape_k_scales_b + scale_k_base + k_off; - val = *ptr; - } - } - const int32_t exp_shift = fp4_rs_detail::pow2_scale_to_exp_shift(val); - exp_offset = static_cast(exp_shift + 6) & 0xffu; - } - packed_offsets |= exp_offset << (j * 8); - } - ptx::st_shared(smem_sfb_exp + k_off * kVecsPerRow + n_off / kVec, packed_offsets); - } - } else if constexpr (kMajorSFB == cute::UMMA::Major::MN) { + if constexpr (kMajorSFB == cute::UMMA::Major::MN) { constexpr uint32_t kVec = 4; DG_STATIC_ASSERT(BLOCK_N % kVec == 0, "BLOCK_N must be a multiple of 4 for vectorized K/32 SFB load"); @@ -1216,61 +1087,33 @@ sm90_fp8_fp4_gemm_1d2d_rs_impl(int8_t* gmem_b_ptr, float* sfb, int* grouped_layo constexpr uint32_t WAVE_BLOCK_M = BLOCK_N <= WGMMA::M ? BLOCK_N : WGMMA::M * 2; DG_STATIC_ASSERT(BLOCK_N % WAVE_BLOCK_M == 0, "Invalid block sizes"); constexpr uint32_t WAVE_WGMMA = BLOCK_N / WAVE_BLOCK_M; - constexpr uint32_t kWGsPerNWave = WAVE_BLOCK_M / WGMMA::M; - constexpr bool kParallelNWavesEnabled = - kParallelNWaves and WAVE_WGMMA == 2 and kWGsPerNWave == 2 and kNumMathThreads == 512; - constexpr bool kUseScaleKGroup = (kScaleKGroup > 1 and (WAVE_WGMMA == 1 or WAVE_WGMMA == 2)); - constexpr bool kUseScaleKGroupExact = - (kScaleKGroupExact and kUseScaleKGroup and kScaleKGroup == 2 and not kFusedPromote); DG_STATIC_ASSERT(kNumMathThreads % 128 == 0, "RS-mode math threads must be whole warpgroups"); - DG_STATIC_ASSERT(not kParallelNWaves or kParallelNWavesEnabled, - "DG_W4_PARALLEL_N_WAVES only supports BN256 with 512 math threads"); - const uint32_t wave_group_idx = kParallelNWavesEnabled ? math_wg_idx / kWGsPerNWave : 0; - const uint32_t wave_mwg_idx = kParallelNWavesEnabled ? math_wg_idx % kWGsPerNWave : math_wg_idx; - const uint32_t wave_warp_idx = wave_mwg_idx * 4 + warp_in_wg; + const uint32_t wave_warp_idx = math_wg_idx * 4 + warp_in_wg; const uint32_t r_0 = wave_warp_idx * 16 + row_idx; const uint32_t r_1 = r_0 + 8; using final_accum_t = typename fp4_rs_detail::FinalAccumStorage::type; - constexpr uint32_t kNumAccumSets = kUseScaleKGroup ? WAVE_WGMMA : 1; constexpr uint32_t kFinalAccumStride = kBF16FinalAccum ? WGMMA::kNumAccum / 2 : WGMMA::kNumAccum; constexpr uint32_t kNumFinalAccumRegs = kFinalAccumStride * WAVE_WGMMA; - constexpr bool kUseFinalAccumScratch = - kFinalAccumScratch and kDirectStore and kGemmType == GemmType::MGroupedMasked and - BLOCK_M < WGMMA::M and not kUseScaleKGroup; constexpr uint32_t kBaseWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); - constexpr uint32_t kNumWGMMAStoreThreads = - kParallelNWavesEnabled ? kBaseWGMMAStoreThreads * WAVE_WGMMA : kBaseWGMMAStoreThreads; - constexpr uint32_t kFinalAccumScratchBytes = - kNumMathThreads * kNumFinalAccumRegs * sizeof(final_accum_t); - float accum_storage[WGMMA::kNumAccum * kNumAccumSets]; - final_accum_t final_accum_regs[kUseFinalAccumScratch ? 1 : kNumFinalAccumRegs]; + float accum_storage[WGMMA::kNumAccum]; + final_accum_t final_accum_regs[kNumFinalAccumRegs]; final_accum_t* final_accum = final_accum_regs; - if constexpr (kUseFinalAccumScratch) { - DG_STATIC_ASSERT(kFinalAccumScratchBytes <= SMEM_D_SIZE, - "DG_W4_FINAL_ACCUM_SCRATCH needs more smem_d scratch space"); - final_accum = reinterpret_cast(smem_d) + - (warp_idx * 32 + lane_idx) * kNumFinalAccumRegs; - } #pragma unroll for (uint32_t i = 0; i < kNumFinalAccumRegs; ++ i) fp4_rs_detail::final_accum_init(final_accum, i); // Pick threads whose WGMMA results are to be stored in shared memory DG_STATIC_ASSERT(BLOCK_N >= 64, "RS-mode swap_ab requires at least one 64-row compute tile"); - const bool do_wgmma_store = warp_idx < kNumWGMMAStoreThreads / 32; + const bool do_wgmma_store = warp_idx < kBaseWGMMAStoreThreads / 32; // Empty barrier arrival auto empty_barrier_arrive_stage = [&](uint32_t target_stage) { - if constexpr (kParallelNWavesEnabled) - cutlass::arch::NamedBarrier::sync(kNumMathThreads, 2); if constexpr (kNumTMAMulticast == 1) { - (lane_idx == 0 and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? - empty_barriers[target_stage]->arrive() : void(); + (lane_idx == 0) ? empty_barriers[target_stage]->arrive() : void(); } else { auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); - (lane_idx < kNumTMAMulticast and (not kParallelNWavesEnabled or wave_group_idx == 0)) ? - empty_barriers[target_stage]->arrive(target_cta) : void(); + (lane_idx < kNumTMAMulticast) ? empty_barriers[target_stage]->arrive(target_cta) : void(); } }; auto empty_barrier_arrive = [&]() {