From 7d27b8a071b5ffd4ba6c3763a27b9ce25ac9d0cf Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Fri, 29 May 2026 02:46:48 -0700 Subject: [PATCH 01/12] FP4: port MXF4 kernel + wrappers onto main's refactored APIs Migrate fea-fp4's standalone MXF4 (SM100_MMA_MXF4_SS) GEMM on top of main's post-PR-304 refactor. Keeps the FP4-specialized hardware path distinct from main's MXF8F6F4 unified kernel. - csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp: rebuilt on GemmDesc + nested GemmConfig {layout, storage_config, pipeline_config, launch_config}. Forces a_dtype=b_dtype=kPackedFP4. Three entry points: dense, m-grouped contiguous, m-grouped masked. BLOCK_K and SHAPE_K convert main's byte/FP4-count to my kernel's int32-count units. - deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh: scheduler -> sched::Scheduler with shape_k ctor arg; KGroupedIndexType -> sched::IndexType; helpers via math:: / ptx:: / mma::sm100:: namespaces; tma_copy signature (multicast moved to runtime); make_runtime_instr_desc_with_sf_id takes sfb_id. swap_ab epilogue / MXF4 MMA path preserved. - csrc/apis/gemm.hpp: dispatch FP4xFP4 -> sm100_fp4_gemm_1d1d (MXF4 path); FP8xFP4 and other mixed cases unchanged -> main's sm100_fp8_fp4_gemm_1d1d. Status: host build clean, JIT NVCC clean (was 19 errors before). Runtime IMA on first launch -- next: localize via compute-sanitizer (TMA descriptor K-unit or SF tensor layout suspected). tests/test_fp4.py copied as-is from fea-fp4; needs recipe=(1,1,32) update for FP4 SF granularity. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/apis/gemm.hpp | 30 +- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 299 +++++++ .../deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh | 814 ++++++++++++++++++ tests/test_fp4.py | 749 ++++++++++++++++ 4 files changed, 1886 insertions(+), 6 deletions(-) create mode 100644 csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh create mode 100644 tests/test_fp4.py diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42622df7d8..f3476a31cf 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -7,6 +7,7 @@ #include "../jit_kernels/impls/sm90_fp8_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_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" #endif @@ -93,8 +94,13 @@ static void fp8_fp4_gemm_nt(const std::pair& a, sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims); } } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) { - sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, gran_k_a, gran_k_b, + if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) { + sm100_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims); + } else { + sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, gran_k_a, gran_k_b, + major_a, major_b, compiled_dims); + } } else { DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types"); } @@ -197,9 +203,15 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "../heuristics/common.hpp" +#include "../heuristics/sm100.hpp" + +#include "runtime_utils.hpp" + +namespace deep_gemm { + +// FP4xFP4 (MXFP4) GEMM via SM100_MMA_MXF4_SS instruction. +// Distinct from main's sm100_fp8_fp4_gemm_1d1d (MXF8F6F4 path) — this is the +// FP4-specialized hardware path, used when both A and B are kPackedFP4. +class SM100FP4Gemm1D1DRuntime final: public LaunchRuntime { +public: + struct Args { + GemmDesc gemm_desc; + GemmConfig gemm_config; + LaunchArgs launch_args; + int num_last_stages; + + void* grouped_layout; + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_sfa; + CUtensorMap tensor_map_sfb; + CUtensorMap tensor_map_c; + CUtensorMap tensor_map_d; + }; + + 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(&sm100_fp4_gemm_1d1d_impl< + {}, {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, + {}, {}, {} + >); +}}; +)", + to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b), + 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), + // FP4: my .cuh expects SHAPE_K in int32 count (= FP4_count / 8). + // main's desc.k is FP4 logical count. + get_compiled_dim(args.gemm_desc.k / 8, 'k', args.gemm_desc.compiled_dims), + args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, + // FP4: BLOCK_K in int32 count (main's heuristic gives bytes for int8 pack). + args.gemm_config.layout.block_k / 4, + args.gemm_desc.num_groups, + 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_last_stages, + args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads, + args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, + args.gemm_config.launch_config.num_sms, + args.gemm_config.layout.swap_ab, + to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, + to_string(args.gemm_desc.cd_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + // FP4: my .cuh expects shape_k in int32 count (= FP4_count / 8) at runtime as well. + // Note: must be non-const (launch_kernel takes &args, can't bind const* to void*). + int shape_k_int32 = args.gemm_desc.k / 8; + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, shape_k_int32, + args.tensor_map_a, args.tensor_map_b, + args.tensor_map_sfa, args.tensor_map_sfb, + args.tensor_map_c, args.tensor_map_d)); + } +}; + +// Helper: compute num_last_stages from k / block_k / num_stages. +// desc.k is FP4 logical count; main's heuristic block_k is in bytes (int8 elem of packed FP4 tensor). +// 1 byte = 2 FP4, so K_per_block_fp4 = block_k_bytes * 2. +static int compute_num_last_stages_fp4(int k, int block_k_bytes, int num_stages) { + const int num_k_blocks = ceil_div(k, block_k_bytes * 2); + const int rem = num_k_blocks % num_stages; + return rem == 0 ? num_stages : rem; +} + +// Build a GemmDesc forcing a_dtype=b_dtype=kPackedFP4 (since this is the FP4xFP4 wrapper). +static GemmDesc make_fp4_desc(GemmType gemm_type, int m, int n, int k, int num_groups, + const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, + const at::ScalarType& cd_dtype, bool with_accumulation, + const std::string& compiled_dims, + int expected_m = 0, int expected_num_groups = 1) { + return GemmDesc { + .gemm_type = gemm_type, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = kPackedFP4, .b_dtype = kPackedFP4, + .cd_dtype = cd_dtype, + .major_a = major_a, .major_b = major_b, + .with_accumulation = with_accumulation, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = expected_m > 0 ? expected_m : m, + .expected_n = n, .expected_k = k, + .expected_num_groups = expected_num_groups + }; +} + +static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const std::optional& c, + const torch::Tensor& d, + const int& m, const int& n, const int& k, + const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, + const std::string& compiled_dims) { + constexpr int gran_k = 32; + const auto desc = make_fp4_desc(GemmType::Normal, m, n, k, 1, + major_a, major_b, d.scalar_type(), + c.has_value(), compiled_dims); + const auto config = get_best_config(desc); + + const auto cd = c.value_or(d); + const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(get_non_contiguous_dim(major_b))), 1, + config.storage_config.swizzle_b_mode); + 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 auto tensor_map_c = make_tma_cd_desc(cd, m, n, + config.storage_config.store_block_m, + config.storage_config.store_block_n, + static_cast(cd.stride(-2)), 1, + config.storage_config.swizzle_cd_mode); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, gran_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, + config.layout.block_n, gran_k, 1, 0); + + if (c.has_value()) { + if (c->data_ptr() == d.data_ptr()) { + DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides()); + } else { + d.copy_(c.value()); + } + } + + const SM100FP4Gemm1D1DRuntime::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()), + .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), + .grouped_layout = nullptr, + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_c = tensor_map_c, + .tensor_map_d = tensor_map_d + }; + const auto code = SM100FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm100_fp4_gemm_1d1d", code); + SM100FP4Gemm1D1DRuntime::launch(runtime, args); +} + +static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, + const torch::Tensor& grouped_layout, + const int& num_groups, const int& m, const int& n, const int& k, + const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, + const std::string& compiled_dims) { + constexpr int gran_k = 32; + const auto desc = make_fp4_desc(GemmType::MGroupedContiguous, m, n, k, num_groups, + major_a, major_b, d.scalar_type(), + false, compiled_dims); + const auto config = get_best_config(desc); + + const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, + config.storage_config.swizzle_b_mode); + 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 auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, gran_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, + config.layout.block_n, gran_k, num_groups, 0); + + const SM100FP4Gemm1D1DRuntime::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()), + .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), + .grouped_layout = grouped_layout.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_c = tensor_map_d, + .tensor_map_d = tensor_map_d + }; + const auto code = SM100FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm100_m_grouped_fp4_gemm_contiguous_1d1d", code); + SM100FP4Gemm1D1DRuntime::launch(runtime, args); +} + +static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, + const torch::Tensor& masked_m, + const int& num_groups, const int& m, const int& n, const int& k, + const int& expected_m, + const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, + const std::string& compiled_dims) { + constexpr int gran_k = 32; + const auto desc = make_fp4_desc(GemmType::MGroupedMasked, m, n, k, num_groups, + major_a, major_b, d.scalar_type(), + false, compiled_dims, + /*expected_m=*/expected_m, /*expected_num_groups=*/num_groups); + const auto config = get_best_config(desc); + + const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(get_non_contiguous_dim(major_a))), num_groups, + config.storage_config.swizzle_a_mode); + const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, + config.storage_config.swizzle_b_mode); + 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 auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, + config.layout.block_m, gran_k, num_groups, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, + config.layout.block_n, gran_k, num_groups, 0); + + const SM100FP4Gemm1D1DRuntime::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()), + .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), + .grouped_layout = masked_m.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_c = tensor_map_d, + .tensor_map_d = tensor_map_d + }; + const auto code = SM100FP4Gemm1D1DRuntime::generate(args); + const auto runtime = compiler->build("sm100_m_grouped_fp4_gemm_masked_1d1d", code); + SM100FP4Gemm1D1DRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh new file mode 100644 index 0000000000..ef75ace86a --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh @@ -0,0 +1,814 @@ +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +using namespace deep_gemm::mma::sm100; +using namespace deep_gemm::math; +using namespace deep_gemm::ptx; +using namespace deep_gemm::utils; + +// E2M1 FP4 到 float 的转换函数 +// E2M1 格式: 4 bits = SEEM (S=符号 1bit, E=指数 2bits, M=尾数 1bit) +__device__ __forceinline__ float fp4_e2m1_to_float(uint32_t fp4_bits) { + constexpr float E2M1_LUT[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, // 正数 (S=0) + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f // 负数 (S=1) + }; + return E2M1_LUT[fp4_bits & 0xF]; +} + +// UE8M0 scale 1.0f packed as 4 bytes per 32-bit TMEM word. +// SM100 MXF4 block-scaled MMA UE8M0 scale factor. +// The bias is determined empirically; see DG_SF_BYTE env var testing. +__device__ __forceinline__ uint32_t pack_ue8m0_scale_factor_word(uint8_t byte_val) { + return uint32_t(byte_val) | (uint32_t(byte_val) << 8) | (uint32_t(byte_val) << 16) | (uint32_t(byte_val) << 24); +} +__device__ __forceinline__ uint32_t pack_ue8m0_2x_scale_factor_one_word() { + return pack_ue8m0_scale_factor_word(0x7Fu); // Will be tested with different values +} + +// Swizzle-aware shared memory index for reading TMA-loaded data. +// TMA stores data with bank-group XOR swizzle: physical_bank = logical_bank ^ (row % num_banks). +// swizzle_mode: kSwizzleAMode or kSwizzleBMode (bytes, e.g. 128) +// row: M or N row index, k: K column index, block_k: elements per row +template +__device__ __forceinline__ uint32_t swizzled_smem_k_major_idx(uint32_t row, uint32_t k, uint32_t block_k) { + constexpr uint32_t kElemBytes = sizeof(uint32_t); + constexpr uint32_t kBankBytes = 16; + constexpr uint32_t kElemsPerBank = kBankBytes / kElemBytes; // 4 + constexpr uint32_t kNumBanks = swizzle_mode / kBankBytes; // e.g. 8 for 128B + uint32_t bank = k / kElemsPerBank; + uint32_t in_bank = k % kElemsPerBank; + uint32_t swizzled_bank = bank ^ (row % kNumBanks); + return row * block_k + swizzled_bank * kElemsPerBank + in_bank; +} + +// SM100 FP4 GEMM 1D1D kernel实现 +// 支持 MXF4 block-scaled 矩阵乘法 +template +__global__ void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) +sm100_fp4_gemm_1d1d_impl(int* grouped_layout, + 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_sfa, + const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, + const __grid_constant__ cute::TmaDescriptor tensor_map_c, + const __grid_constant__ cute::TmaDescriptor tensor_map_d) { + + // ========== 基础配置和类型定义 ========== + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::conditional_t; + + if constexpr (kWithAccumulation) + DG_STATIC_ASSERT(cute::is_same_v, "Invalid C/D data dtype"); + + // ========== 核心配置参数 ========== + constexpr uint32_t LAYOUT_AD_M = 128; + constexpr uint32_t kNumMWaves = BLOCK_M / LAYOUT_AD_M; + constexpr uint32_t kNumTMAStoreStages = 2; + + // ========== MXF4 配置 ========== + constexpr uint32_t kNumSFAStagesPerLoad = 1; + constexpr uint32_t kNumSFBStagesPerLoad = 1; + constexpr uint32_t kNumUTCCPAlignedElems = 128; + constexpr uint32_t FP4_ELEMS_PER_INT32 = 8; + constexpr uint32_t MXF4_VS = 32; + constexpr uint32_t BLOCK_K_FP4 = BLOCK_K * FP4_ELEMS_PER_INT32; + constexpr uint32_t UMMA_K_FP4 = 64; + constexpr uint32_t SF_K_PER_STAGE = BLOCK_K_FP4 / MXF4_VS; + constexpr uint32_t SF_PACKED_K_PER_STAGE = SF_K_PER_STAGE / 4; + + DG_STATIC_ASSERT(BLOCK_M % LAYOUT_AD_M == 0 and 2 % kNumMWaves == 0, "Invalid block M"); + DG_STATIC_ASSERT(BLOCK_K == 16 or BLOCK_K == 32, "FP4 BLOCK_K must be 16 or 32 int32"); + + // ========== 动态形状处理 ========== + 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; + // shape_k debug disabled + // FP4 packed: shape_k 是 int32 个数，每个 int32 有 8 个 FP4。 + // 1 个 scale 覆盖 MXF4_VS=32 个 FP4 (=4 个 int32)， + // 每个 uint32 打包 4 个 scale → 1 个 packed group = 4*32 = 128 FP4。 + const uint32_t total_scales_k = + ceil_div(shape_k * FP4_ELEMS_PER_INT32, + MXF4_VS); // MXF4_VS 是 uint32_t，OK + + const uint32_t total_packed_k = + ceil_div(total_scales_k, + uint32_t(4)); // 把 4 也变成 uint32_t + + const uint32_t shape_sfa_k = total_packed_k; + const uint32_t shape_sfb_k = total_packed_k; + + // ========== 线程和warp信息 ========== + bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = get_lane_idx(); + + // ========== 共享内存分配 ========== + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + // ========== 块大小计算 ========== + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / (kIsMulticastOnA ? kNumMulticast: 1); + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N / (kIsMulticastOnA ? 1 : kNumMulticast); + // Swap-AB epilogue: STORE_BLOCK_M = 16 (fine-grained M slices to skip padding rows), + // STORE_BLOCK_N = BLOCK_N (write entire N at once). + // Non-swap (existing): STORE_BLOCK_M = BLOCK_M, STORE_BLOCK_N derived from swizzle. + constexpr uint32_t STORE_BLOCK_M = kSwapAB ? 16u : cute::min(BLOCK_M, LAYOUT_AD_M); + constexpr uint32_t STORE_BLOCK_N = kSwapAB ? BLOCK_N : kSwizzleCDMode / sizeof(cd_dtype_t); + + DG_STATIC_ASSERT(not kIsMulticastOnA or kNumMulticast == 1, "FP4 only supports B-multicast (2CTA along M)"); + DG_STATIC_ASSERT(LOAD_BLOCK_M == BLOCK_M and BLOCK_M % LAYOUT_AD_M == 0, "Only support tensor memory layout A/D"); + DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); + // Swap-AB requires BLOCK_N = LAYOUT_AD_M (= 128) so UMMA_M after swap stays = 128. + DG_STATIC_ASSERT(not kSwapAB or BLOCK_N == LAYOUT_AD_M, "kSwapAB requires BLOCK_N = LAYOUT_AD_M"); + // Swap-AB initial implementation: no multicast (cluster_n=1, cluster_m=1) for simplicity. + DG_STATIC_ASSERT(not kSwapAB or kNumMulticast == 1, "kSwapAB initial impl: no multicast"); + + // ========== 共享内存大小计算 ========== + // Swap-AB: per-stage SMEM = STORE_BLOCK_M (small) × STORE_BLOCK_N (= BLOCK_N) × sizeof(D) + // Non-swap: per-stage SMEM = STORE_BLOCK_M × kSwizzleCDMode (= STORE_BLOCK_N * sizeof(D)) + constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = kSwapAB + ? STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t) + : STORE_BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; + constexpr uint32_t SMEM_A_PACKED_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(uint32_t); + constexpr uint32_t SMEM_B_PACKED_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(uint32_t); + constexpr uint32_t SF_BLOCK_M = constexpr_align(BLOCK_M, kNumUTCCPAlignedElems); + constexpr uint32_t SF_BLOCK_N = constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); + + DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0, "Shared memory must be aligned to 1024 bytes"); + DG_STATIC_ASSERT(kNumTMAStoreStages >= 1, "Invalid number of TMA stages"); + + // ========== 张量内存配置 ========== + constexpr uint32_t kNumSFATmemCols = (SF_BLOCK_M / 32) * SF_PACKED_K_PER_STAGE; + constexpr uint32_t kNumSFBTmemCols = (SF_BLOCK_N / 32) * SF_PACKED_K_PER_STAGE; + constexpr uint32_t kNumEpilogueStages = (2 * kNumMWaves * BLOCK_N + kNumSFATmemCols + kNumSFBTmemCols) > 512 ? 1 : 2; + constexpr uint32_t kNumAccumTmemCols = kNumEpilogueStages * kNumMWaves * BLOCK_N; + constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + + // ========== TMA描述符预取 ========== + if (threadIdx.x == 0) { + 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_sfb); + cute::prefetch_tma_descriptor(&tensor_map_d); + if constexpr (kWithAccumulation) + cute::prefetch_tma_descriptor(&tensor_map_c); + } + + // ========== 共享内存指针设置 ========== + cd_dtype_t* smem_cd[kNumTMAStoreStages]; + uint32_t* smem_sfa[kNumStages]; + uint32_t* smem_sfb[kNumStages]; + uint32_t* smem_a_packed[kNumStages]; + uint32_t* smem_b_packed[kNumStages]; + + #pragma unroll + for (uint32_t i = 0; i < kNumTMAStoreStages; ++ i) + smem_cd[i] = reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); + + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + smem_a_packed[i] = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_PACKED_SIZE_PER_STAGE); + smem_b_packed[i] = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_PACKED_SIZE_PER_STAGE + i * SMEM_B_PACKED_SIZE_PER_STAGE); + } + + auto sf_start_ptr = smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_PACKED_SIZE_PER_STAGE + SMEM_B_PACKED_SIZE_PER_STAGE); + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + smem_sfa[i] = reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); + smem_sfb[i] = reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); + } + + // ========== 屏障初始化 ========== + auto barrier_start_ptr = reinterpret_cast(smem_buffer + + SMEM_CD_SIZE + + kNumStages * (SMEM_A_PACKED_SIZE_PER_STAGE + SMEM_B_PACKED_SIZE_PER_STAGE) + + kNumStages * (SMEM_SFA_SIZE_PER_STAGE + SMEM_SFB_SIZE_PER_STAGE)); + auto full_barriers = PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); + auto empty_barriers = PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); + auto with_sf_full_barriers = PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 2 + i); }); + auto tmem_full_barriers = PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + i); }); + auto tmem_empty_barriers = PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + kNumEpilogueStages + i); }); + + auto tmem_ptr_in_smem = reinterpret_cast(barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2); + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + + if (threadIdx.x == 0) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(1); + with_sf_full_barriers[i]->init(kNumMulticast * 32); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { + tmem_full_barriers[i]->init(1); + tmem_empty_barriers[i]->init(kNumMulticast * kNumEpilogueThreads); + } + cutlass::arch::fence_view_async_shared(); + cutlass::arch::fence_barrier_init(); + } + else if (threadIdx.x >= 32 and threadIdx.x < 64) { + Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + kNumMulticast > 1 ? cute::cluster_sync() : __syncthreads(); + + // ========== 块调度器初始化 ========== + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); + + // ========== K维度迭代控制 ========== + struct DivisibleK {}; + struct NotDivisibleK {}; + uint32_t phase = 0; + + auto launch_k_iterations = [&](const auto& func) { + const uint32_t current_shape_k = (kGemmType == GemmType::KGroupedContiguous ? scheduler.current_shape_k : shape_k); + const uint32_t num_iterations = ceil_div(current_shape_k, kNumStages * BLOCK_K); + const uint32_t num_last_stages = ceil_div(current_shape_k, BLOCK_K) % kNumStages; + + if (num_last_stages == 0) { + for (uint32_t k_iter = 0; k_iter < num_iterations; ++ k_iter, phase ^= 1) + func(k_iter, DivisibleK{}, k_iter == num_iterations - 1, num_last_stages); + } else { + for (uint32_t k_iter = 0; k_iter < num_iterations - 1; ++ k_iter, phase ^= 1) + func(k_iter, DivisibleK{}, false, num_last_stages); + func(num_iterations - 1, NotDivisibleK{}, true, num_last_stages), phase ^= 1; + } + }; + + auto dispatch_accum_stage_idx = [&](uint32_t accum_stage_idx, const auto& func) { + DG_STATIC_ASSERT(1 <= kNumEpilogueStages and kNumEpilogueStages <= 2, "Too many epilogue stages"); + accum_stage_idx == 0 ? func(0) : func(1); + }; + + // ========== Warp dispatch (FP8-style: independent loops per warp) ========== + // Warp 0: TMA load producer + // Warp 1: MMA consumer (+ UTCCP SF copy to TMEM) + // Warp 2: SF transpose (SMEM warp transpose for UTCCP) + // Warp 3+: Epilogue + + if (warp_idx == 0) { + // ========== Warp 0: TMA load ========== + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + launch_k_iterations([&](uint32_t k_iter, auto type, bool is_last_iter, uint32_t num_last_stages) { + constexpr bool kHasDivisibleStages = cute::is_same_v; + const uint32_t kNumInnerStages = kHasDivisibleStages ? kNumStages : num_last_stages; + + #pragma unroll + for (uint32_t s = 0; s < kNumInnerStages; ++ s) { + empty_barriers[s]->wait(phase ^ 1); + + uint32_t m_idx = scheduler.template get_global_idx<(kGemmType == GemmType::MGroupedMasked), sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx); + uint32_t n_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::K), sched::IndexType::MN>(shape_n, BLOCK_N, n_block_idx, m_block_idx); + + DG_STATIC_ASSERT(kGemmType == GemmType::Normal or kGemmType == GemmType::KGroupedContiguous or kMajorA == cute::UMMA::Major::K, "Invalid major"); + uint32_t k_block_idx = k_iter * kNumStages + s; + uint32_t k_a_idx = scheduler.template get_global_idx<(kMajorA == cute::UMMA::Major::MN), sched::IndexType::K>(shape_k, BLOCK_K, k_block_idx, m_block_idx); + uint32_t k_b_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::MN), sched::IndexType::K>(shape_k, BLOCK_K, k_block_idx, m_block_idx); + + if constexpr (kNumMulticast > 1) { + m_idx += kIsMulticastOnA ? (cute::block_rank_in_cluster() * LOAD_BLOCK_M) : 0; + n_idx += kIsMulticastOnA ? 0 : (cute::block_rank_in_cluster() * LOAD_BLOCK_N); + } + + if (cute::elect_one_sync()) { + if constexpr (kMajorA == cute::UMMA::Major::K) + tma_copy(&tensor_map_a, full_barriers[s], smem_a_packed[s], k_a_idx, m_idx); + if constexpr (kMajorA == cute::UMMA::Major::MN) + tma_copy(&tensor_map_a, full_barriers[s], smem_a_packed[s], m_idx, k_a_idx); + if constexpr (kMajorB == cute::UMMA::Major::K) + tma_copy(&tensor_map_b, full_barriers[s], smem_b_packed[s], k_b_idx, n_idx); + if constexpr (kMajorB == cute::UMMA::Major::MN) + tma_copy(&tensor_map_b, full_barriers[s], smem_b_packed[s], n_idx, k_b_idx); + } + auto num_arrival_bytes = SMEM_A_PACKED_SIZE_PER_STAGE + SMEM_B_PACKED_SIZE_PER_STAGE; + + const uint32_t sfa_tma_stage = (k_iter * kNumStages + s) % kNumSFAStagesPerLoad; + if (sfa_tma_stage == 0 and cute::elect_one_sync()) { + uint32_t sf_k_base = k_block_idx / kNumSFAStagesPerLoad * SF_PACKED_K_PER_STAGE; + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + tma_copy(&tensor_map_sfa, full_barriers[s], smem_sfa[s] + pk * SF_BLOCK_M, m_block_idx * BLOCK_M, + scheduler.template get_global_idx<(kGemmType != GemmType::MGroupedContiguous), sched::IndexType::SF_K>(shape_sfa_k, 1, sf_k_base + pk)); + } + num_arrival_bytes += BLOCK_M * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); + } + const uint32_t sfb_tma_stage = (k_iter * kNumStages + s) % kNumSFBStagesPerLoad; + if (sfb_tma_stage == 0 and cute::elect_one_sync()) { + uint32_t sf_k_base = k_block_idx / kNumSFBStagesPerLoad * SF_PACKED_K_PER_STAGE; + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + tma_copy(&tensor_map_sfb, full_barriers[s], smem_sfb[s] + pk * SF_BLOCK_N, n_block_idx * BLOCK_N, + scheduler.template get_global_idx(shape_sfb_k, 1, sf_k_base + pk, m_block_idx)); + } + num_arrival_bytes += BLOCK_N * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); + } + + if (cute::elect_one_sync()) + full_barriers[s]->arrive_and_expect_tx(num_arrival_bytes); + } + + #pragma unroll + for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) { + empty_barriers[s]->wait(phase ^ 1); + if (cute::elect_one_sync()) + full_barriers[s]->arrive(); + } + }); + } + } else if (warp_idx == 1 and is_leader_cta) { + // ========== Warp 1: UTCCP SF copy + MMA ========== + constexpr uint32_t UMMA_M = LAYOUT_AD_M * (kIsMulticastOnA ? 1 : kNumMulticast); + // Swap-AB: UMMA_N becomes BLOCK_M (the original M dim now plays N role in MMA). + constexpr uint32_t UMMA_N = kSwapAB ? BLOCK_M : BLOCK_N * (kIsMulticastOnA ? kNumMulticast : 1); + constexpr uint32_t UMMA_K_INT32 = UMMA_K_FP4 / FP4_ELEMS_PER_INT32; + constexpr uint32_t NUM_K_ITERS_PER_STAGE = BLOCK_K / UMMA_K_INT32; + // After swap, the "N iters" walk over BLOCK_M (now the MMA-N axis); without swap, over BLOCK_N. + constexpr uint32_t NUM_N_ITERS = (kSwapAB ? BLOCK_M : BLOCK_N) / UMMA_N; + + // Swap-AB: pass (b_dtype, a_dtype, ..., kMajorB, kMajorA) to MMA so B occupies MMA-A slot + // and A occupies MMA-B slot. Output in TMEM is D^T (N rows, M cols). + auto instr_desc_mxf4 = kSwapAB + ? cute::UMMA::make_instr_desc_block_scaled< + cutlass::float_e2m1_t, cutlass::float_e2m1_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, kMajorB, kMajorA>() + : cute::UMMA::make_instr_desc_block_scaled< + cutlass::float_e2m1_t, cutlass::float_e2m1_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, kMajorA, kMajorB>(); + + using cute_mma_mxf4_noswap_t = cute::conditional_t, + cute::SM100_MMA_MXF4_2x1SM_SS>; + using cute_mma_mxf4_swap_t = cute::conditional_t, + cute::SM100_MMA_MXF4_2x1SM_SS>; + using cute_mma_mxf4_t = cute::conditional_t; + + DG_STATIC_ASSERT(UMMA_M == 128 or UMMA_M == 256, "MXF4 supports M=128 (1CTA) or M=256 (2CTA)"); + DG_STATIC_ASSERT((UMMA_N % 8 == 0) and (8 <= UMMA_N) and (UMMA_N <= 256), "Invalid MXF4 N-mode size"); + + using cute_utccp_t = cute::conditional_t; + auto sf_desc = make_sf_desc(nullptr); + + constexpr uint32_t SMEM_A_SIZE_PER_STAGE_PACKED = LOAD_BLOCK_M * BLOCK_K * sizeof(uint32_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE_PACKED = LOAD_BLOCK_N * BLOCK_K * sizeof(uint32_t); + auto a_desc_base = make_umma_desc(smem_a_packed[0], 0, 0); + auto b_desc_base = make_umma_desc(smem_b_packed[0], 0, 0); + + // MXF4 SF addressing constants + constexpr uint32_t kGroupsPerUmmaStep = UMMA_K_FP4 / MXF4_VS; // 2 + constexpr uint32_t kGroupsPerPacked = 4; + constexpr uint32_t kUmmaStepsPerPacked = kGroupsPerPacked / kGroupsPerUmmaStep; // 2 + constexpr uint32_t kSfaColsPerPackedGroup = SF_BLOCK_M / 32; + constexpr uint32_t kSfbColsPerPackedGroup = SF_BLOCK_N / 32; + + // Pre-compute runtime descriptors (sf_id alternates 0, 2, 0, 2) + const auto runtime_desc_sf0 = make_runtime_instr_desc_with_sf_id(instr_desc_mxf4, 0, 0); + const auto runtime_desc_sf2 = make_runtime_instr_desc_with_sf_id(instr_desc_mxf4, kGroupsPerUmmaStep, kGroupsPerUmmaStep); + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + dispatch_accum_stage_idx(scheduler.current_iter % kNumEpilogueStages, [&](uint32_t accum_stage_idx) { + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); + tcgen05_after_thread_sync(); + + auto empty_barrier_arrive = [&](uint32_t s, bool do_tmem_full_arrive) { + auto umma_arrive = [](const uint64_t* barrier) { + if constexpr (kNumMulticast == 1) { + cutlass::arch::umma_arrive(barrier); + } else { + constexpr uint16_t kCTAMask = (1 << kNumMulticast) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + } + }; + umma_arrive(reinterpret_cast(empty_barriers[s])); + if (do_tmem_full_arrive) + umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + }; + + launch_k_iterations([&](uint32_t k_iter, auto type, bool is_last_iter, uint32_t num_last_stages) { + constexpr bool kHasDivisibleStages = cute::is_same_v; + const uint32_t kNumInnerStages = kHasDivisibleStages ? kNumStages : num_last_stages; + + #pragma unroll + for (uint32_t s = 0; s < kNumInnerStages; ++ s) { + with_sf_full_barriers[s]->wait(phase); + tcgen05_after_thread_sync(); + + // UTCCP: copy SF from SMEM → TMEM + // Must stay on warp 1: SF TMEM cols are reused across stages, + // so UTCCP must be serialized with MMA on the same warp. + const uint32_t sfa_copy_stage = (k_iter * kNumStages + s) % kNumSFAStagesPerLoad; + if (sfa_copy_stage == 0 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + replace_smem_desc_addr(sf_desc, smem_sfa[s] + pk * SF_BLOCK_M + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + pk * (SF_BLOCK_M / 32) + i * 4); + } + } + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + replace_smem_desc_addr(sf_desc, smem_sfb[s] + pk * SF_BLOCK_N + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + pk * (SF_BLOCK_N / 32) + i * 4); + } + } + } + __syncwarp(); + + // MMA loop + uint32_t a_desc_stage_lo = a_desc_base.lo + s * (SMEM_A_SIZE_PER_STAGE_PACKED / 16); + uint32_t b_desc_stage_lo = b_desc_base.lo + s * (SMEM_B_SIZE_PER_STAGE_PACKED / 16); + + #pragma unroll + for (uint32_t k = 0; k < NUM_K_ITERS_PER_STAGE; ++k) { + uint32_t packed_group = k / kUmmaStepsPerPacked; + uint32_t tmem_sfa_base = kTmemStartColOfSFA + packed_group * kSfaColsPerPackedGroup; + uint32_t tmem_sfb_k = kTmemStartColOfSFB + packed_group * kSfbColsPerPackedGroup; + const auto& runtime_desc_k = (k % kUmmaStepsPerPacked == 0) ? runtime_desc_sf0 : runtime_desc_sf2; + + #pragma unroll + for (uint32_t n = 0; n < NUM_N_ITERS; ++n) { + auto b_desc = b_desc_base; + b_desc.lo = advance_umma_desc_lo( + b_desc_stage_lo, n * UMMA_N * BLOCK_K, k * UMMA_K_INT32); + + #pragma unroll + for (uint32_t w = 0; w < kNumMWaves; ++w) { + auto a_desc = a_desc_base; + a_desc.lo = advance_umma_desc_lo( + a_desc_stage_lo, w * LAYOUT_AD_M * BLOCK_K, k * UMMA_K_INT32); + + uint32_t tmem_col = accum_stage_idx * kNumMWaves * BLOCK_N + w * BLOCK_N + n * UMMA_N; + if constexpr (kSwapAB) { + // Swap-AB: B goes into MMA-A slot, A into MMA-B slot. + // SF column args also swap (SFB first, SFA second). + cute_mma_mxf4_t::fma(b_desc, a_desc, tmem_col, + k_iter > 0 or s > 0 or k > 0, + runtime_desc_k, + tmem_sfb_k, + tmem_sfa_base + w * (kNumUTCCPAlignedElems / 32)); + } else { + cute_mma_mxf4_t::fma(a_desc, b_desc, tmem_col, + k_iter > 0 or s > 0 or k > 0, + runtime_desc_k, + tmem_sfa_base + w * (kNumUTCCPAlignedElems / 32), + tmem_sfb_k); + } + } + } + } + + empty_barrier_arrive(s, is_last_iter and s == kNumInnerStages - 1); + } + + #pragma unroll + for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) { + with_sf_full_barriers[s]->wait(phase); + empty_barrier_arrive(s, false); + } + }); + }); + } + } else if (warp_idx == 2) { + // ========== Warp 2: SF transpose ========== + auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, "Invalid aligned elements"); + uint32_t values[4]; + #pragma unroll + for (uint32_t i = 0; i < 4; ++ i) + values[i] = ld_shared(smem_ptr + (i ^ (lane_idx >> 3)) * 32 + lane_idx); + __syncwarp(); + #pragma unroll + for (uint32_t i = 0; i < 4; ++ i) + st_shared(smem_ptr + lane_idx * 4 + (i ^ (lane_idx >> 3)), values[i]); + }; + auto fill_sfb_missing_k_groups = [&](uint32_t* smem_ptr) { + if constexpr (BLOCK_N < kNumUTCCPAlignedElems) { + // Zero-fill [BLOCK_N, SF_BLOCK_N) before warp-transpose: XOR pattern + // reads all 128 elements, so uninitialized positions corrupt valid data. + #pragma unroll + for (uint32_t pos = lane_idx; pos < kNumUTCCPAlignedElems; pos += 32) { + if (pos >= BLOCK_N) + st_shared(smem_ptr + pos, 0u); + } + __syncwarp(); + } + }; + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + launch_k_iterations([&](uint32_t k_iter, auto type, bool is_last_iter, uint32_t num_last_stages) { + constexpr bool kHasDivisibleStages = cute::is_same_v; + const uint32_t kNumInnerStages = kHasDivisibleStages ? kNumStages : num_last_stages; + + #pragma unroll + for (uint32_t s = 0; s < kNumInnerStages; ++ s) { + full_barriers[s]->wait(phase); + + const uint32_t sfa_ut_stage = (k_iter * kNumStages + s) % kNumSFAStagesPerLoad; + if (sfa_ut_stage == 0) { + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) + utccp_required_smem_warp_transpose(smem_sfa[s] + pk * SF_BLOCK_M + i * kNumUTCCPAlignedElems); + } + cutlass::arch::fence_view_async_shared(); + } + + const uint32_t sfb_ut_stage = (k_iter * kNumStages + s) % kNumSFBStagesPerLoad; + if (sfb_ut_stage == 0) { + #pragma unroll + for (uint32_t pk = 0; pk < SF_PACKED_K_PER_STAGE; ++ pk) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + fill_sfb_missing_k_groups(smem_sfb[s] + pk * SF_BLOCK_N + i * kNumUTCCPAlignedElems); + utccp_required_smem_warp_transpose(smem_sfb[s] + pk * SF_BLOCK_N + i * kNumUTCCPAlignedElems); + } + } + cutlass::arch::fence_view_async_shared(); + } + + with_sf_full_barriers[s]->arrive(0u); + } + + #pragma unroll + for (uint32_t s = kNumInnerStages; s < kNumStages; ++ s) { + full_barriers[s]->wait(phase); + with_sf_full_barriers[s]->arrive(0u); + } + }); + } + } else if (warp_idx >= kNumNonEpilogueThreads / 32) { + // ========== Epilogue warp组 ========== + const auto epilogue_thread_idx = threadIdx.x - kNumNonEpilogueThreads; + const auto epilogue_warp_idx = warp_idx - (kNumNonEpilogueThreads / 32); + + DG_TRAP_ONLY_DEVICE_ASSERT(ld_shared(tmem_ptr_in_smem) == 0); + + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(cd_dtype_t); + + if constexpr (kSwapAB) { + // ================================================================= + // Swap-AB epilogue: + // TMEM holds D^T (BLOCK_N rows × BLOCK_M cols). + // Per accum stage covers cols [accum_stage_idx*BLOCK_M, +BLOCK_M) of TMEM. + // STORE_BLOCK_M=16: each `s` iter covers 16 cols of TMEM (= 16 of D's M). + // num_stores = effective_m / 16 → padding cols are skipped entirely. + // STORE_BLOCK_N = BLOCK_N: each TMA store covers entire BLOCK_N at once. + // ================================================================= + constexpr uint32_t kNumSwizzleAtomRows = 8; + constexpr uint32_t STORE_BLOCK_N_ATOM = kSwizzleCDMode / sizeof(cd_dtype_t); + constexpr uint32_t kNumWarpsPerAtom = STORE_BLOCK_N_ATOM / 32; + DG_STATIC_ASSERT(STORE_BLOCK_M % kNumSwizzleAtomRows == 0, "Invalid swap-AB store_block_m"); + DG_STATIC_ASSERT(STORE_BLOCK_N % STORE_BLOCK_N_ATOM == 0, "Invalid swap-AB store_block_n"); + DG_STATIC_ASSERT(kNumEpilogueThreads == 128, "Swap-AB requires full warpgroup"); + + uint32_t tma_stage_idx = 0; + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + dispatch_accum_stage_idx(scheduler.current_iter % kNumEpilogueStages, [&](uint32_t accum_stage_idx) { + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + + tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + tcgen05_after_thread_sync(); + + // Effective M (aligned up to STORE_BLOCK_M=16): how many M-cols of D^T are valid. + const auto effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); + const uint32_t num_stores = effective_m / STORE_BLOCK_M; + + // TMEM col where this accum stage's tile starts. + const auto tmem_base_addr = accum_stage_idx * BLOCK_M; + const auto base_m_idx = scheduler.template get_global_idx<(kGemmType != GemmType::MGroupedContiguous), sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx); + const auto base_n_idx = n_block_idx * BLOCK_N; + + #pragma unroll 1 + for (uint32_t s = 0; s < num_stores; ++ s) { + // Wait if TMA store pipeline full + if (s >= kNumTMAStoreStages) { + if (epilogue_thread_idx == 0) + cute::tma_store_wait(); + cutlass::arch::NamedBarrier(kNumEpilogueThreads).sync(); + } + + // SMEM store: 4 warps cooperatively write STORE_BLOCK_M × STORE_BLOCK_N tile. + // Each warp owns STORE_BLOCK_M rows × STORE_BLOCK_N_ATOM cols. + // Within each warp, loop covers STORE_BLOCK_M / kNumSwizzleAtomRows = 2 sub-blocks of 8 rows. + #pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_M / kNumSwizzleAtomRows; ++ i) { + uint32_t tmem_addr = tmem_base_addr + + s * STORE_BLOCK_M + // M-slice (cols of TMEM) + i * kNumSwizzleAtomRows; // Sub-block within slice + + uint32_t outer_atom_offset = (epilogue_warp_idx / kNumWarpsPerAtom) * STORE_BLOCK_M * kSwizzleCDMode; + uint32_t inner_atom_offset = i * kNumSwizzleAtomRows * kSwizzleCDMode; + auto smem_base_ptr = reinterpret_cast(smem_cd[tma_stage_idx]) + + outer_atom_offset + inner_atom_offset; + + if constexpr (cute::is_same_v) { + uint32_t values[kNumSwizzleAtomRows]; + // Load 32dp × 8 cols of TMEM per warp (= 8 M-cols × 32 N-rows of D). + cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, + values[0], values[1], values[2], values[3], + values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + uint32_t col = lane_idx / 4; + #pragma unroll + for (uint32_t row = 0; row < kNumSwizzleAtomRows; ++ row) { + auto smem_ptr = smem_base_ptr + row * (kNumBankGroupBytes * 8) + + (col ^ row) * kNumBankGroupBytes + + (lane_idx % 4) * sizeof(float); + st_shared(reinterpret_cast(smem_ptr), values[row]); + } + } + } + + // Notify TMEM empty on last store + if (s == num_stores - 1) { + tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage_idx]->arrive(0u); + } + __syncwarp(); + + cute::tma_store_fence(); + cutlass::arch::NamedBarrier(kNumEpilogueThreads).sync(); + + if (epilogue_thread_idx == 0) { + uint32_t m_idx = base_m_idx + s * STORE_BLOCK_M; + #pragma unroll + for (uint32_t ai = 0; ai < STORE_BLOCK_N / STORE_BLOCK_N_ATOM; ++ ai) { + auto smem_ptr = smem_cd[tma_stage_idx] + ai * STORE_BLOCK_M * STORE_BLOCK_N_ATOM; + uint32_t n_idx = base_n_idx + ai * STORE_BLOCK_N_ATOM; + using cute_tma_t = cute::conditional_t; + cute_tma_t::copy(&tensor_map_d, smem_ptr, n_idx, m_idx); + } + cute::tma_store_arrive(); + } + + tma_stage_idx = (tma_stage_idx + 1) % kNumTMAStoreStages; + } + + // If entire tile is padding (effective_m=0, hence num_stores=0): still arrive at empty barrier + // so MMA pipeline can advance. + if (num_stores == 0) { + tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage_idx]->arrive(0u); + } + }); + } + } else { + // ===== Existing non-swap epilogue (unchanged) ===== + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + dispatch_accum_stage_idx(scheduler.current_iter % kNumEpilogueStages, [&](uint32_t accum_stage_idx) { + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + + if (epilogue_thread_idx == 0) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier(kNumEpilogueThreads).sync(); + + tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + tcgen05_after_thread_sync(); + + DG_STATIC_ASSERT(kNumEpilogueThreads == 128, "Epilogue threads not enough"); + DG_STATIC_ASSERT(BLOCK_N % STORE_BLOCK_N == 0, "Invalid block sizes"); + + // Per cute::make_tmem_warp_partitioner (copy_traits_sm100.hpp): four epilogue warps read + // four 32-row M bands; stride is 32 * TMEM::DP (datapath), not +32 on column index. + constexpr uint32_t kTmemDpStride = + static_cast(cute::TMEM::DP{}); + // M-wave 间步进：128 行 × DP步进（DP 方向，不是列方向） + constexpr uint32_t kTmemMWaveStride = 128u * kTmemDpStride; + // epilogue warp 在一个 wave 内的 band 步进（DP 方向） + constexpr uint32_t kTmemWarpRowBandStride = 32u * kTmemDpStride; + + #pragma unroll + for (uint32_t w = 0; w < kNumMWaves; ++ w) { + constexpr uint32_t kNumStores = BLOCK_N / STORE_BLOCK_N; + #pragma unroll + for (uint32_t s = 0; s < kNumStores; ++ s) { + const uint32_t iter_idx = w * kNumStores + s; + if (iter_idx >= kNumTMAStoreStages) { + if (epilogue_thread_idx == 0) + cute::tma_store_wait(); + cutlass::arch::NamedBarrier(kNumEpilogueThreads).sync(); + } + + const auto tma_stage_idx = iter_idx % kNumTMAStoreStages; + const auto m_idx = scheduler.template get_global_idx<(kGemmType != GemmType::MGroupedContiguous), sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx) + w * LAYOUT_AD_M; + const auto n_idx = n_block_idx * BLOCK_N + s * STORE_BLOCK_N; + + #pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_N / kNumElemsPerBankGroup; ++ i) { + auto bank_group_index = i + lane_idx * (kSwizzleCDMode / kNumBankGroupBytes); + constexpr bool kHasShortcut = (kSwizzleCDMode / kNumBankGroupBytes) == 8; + auto row = kHasShortcut ? (i / 8 + lane_idx) : (bank_group_index / 8); + auto col = kHasShortcut ? (i) : (bank_group_index % 8); + col ^= row % (kSwizzleCDMode / 16); + + // uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N // 列：stage + // + s * STORE_BLOCK_N // 列：N tile + // + i * kNumElemsPerBankGroup // 列：bank group + // + w * kTmemMWaveStride // DP：M-wave ← 关键修改 + // + epilogue_warp_idx * kTmemWarpRowBandStride; // DP：warp band + // uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N // stage 列偏移 + // + w * BLOCK_N // M-wave 列偏移（每个 wave 占 BLOCK_N 列） + // + s * STORE_BLOCK_N // N tile 列偏移 + // + i * kNumElemsPerBankGroup; // bank group 列偏移 + uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N + w * BLOCK_N + s * STORE_BLOCK_N + + i * kNumElemsPerBankGroup; + // Match sm100_bf16_gemm.cuh / mma.cuh: each epilogue warp reads its 32-row TMEM band (DP stride). + // tmem_addr += epilogue_warp_idx * kTmemWarpRowBandStride; + auto smem_ptr = reinterpret_cast(smem_cd[tma_stage_idx]) + + epilogue_warp_idx * 32 * kSwizzleCDMode + + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; + + uint32_t values[kNumElemsPerBankGroup]; + if constexpr (cute::is_same_v) { + DG_STATIC_ASSERT(kNumElemsPerBankGroup == 4, "Invalid type"); + cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); + cutlass::arch::fence_view_async_tmem_load(); + st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + } + } + + if (w == kNumMWaves - 1 and s == BLOCK_N / STORE_BLOCK_N - 1) { + tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage_idx]->arrive(0u); + } + __syncwarp(); + + cute::tma_store_fence(); + cutlass::arch::NamedBarrier(kNumEpilogueThreads).sync(); + if (epilogue_thread_idx == 0) { + using cute_tma_t = cute::conditional_t; + cute_tma_t::copy(&tensor_map_d, smem_cd[tma_stage_idx], n_idx, m_idx); + cute::tma_store_arrive(); + } + } + } + }); + } + } // end of if constexpr (kSwapAB) else (non-swap) + + if (epilogue_thread_idx == 0) + cute::tma_store_wait<0>(); + + if (epilogue_warp_idx == 1) + Allocator().free(0, kNumTmemCols); + } + + if constexpr (kNumMulticast > 1) + cute::cluster_sync(); +} + +}; // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/tests/test_fp4.py b/tests/test_fp4.py new file mode 100644 index 0000000000..a547075fc9 --- /dev/null +++ b/tests/test_fp4.py @@ -0,0 +1,749 @@ +""" +FP4 (E2M1) GEMM correctness test for SM100 MXF4 block-scaled kernel. + +Usage: + python tests/test_fp4.py +""" + +import torch +import random +import deep_gemm +from deep_gemm.testing import bench_kineto, count_bytes +from generators import KernelType, get_ue8m0_usage + +# ============================================================ +# E2M1 FP4 查找表 +# ============================================================ +E2M1_LUT = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, # bits 0-7 (S=0) + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0 # bits 8-15 (S=1) +], dtype=torch.float32) + +# ============================================================ +# 工具函数 +# ============================================================ + +def pack_fp4_random(m: int, k_fp4: int, device='cuda'): + """生成随机 E2M1 FP4 数据并打包为 int32。每个 int32 包含 8 个 FP4 值。""" + assert k_fp4 % 8 == 0 + raw = torch.randint(0, 16, (m, k_fp4), dtype=torch.uint8, device=device) + packed = torch.zeros(m, k_fp4 // 8, dtype=torch.int32, device=device) + for i in range(8): + packed += (raw[:, i::8].to(torch.int32) << (i * 4)) + return packed + + +def pack_fp4_constant(m: int, k_fp4: int, fp4_bits: int = 0x2, device='cuda'): + """生成常量 FP4 打包数据。fp4_bits=0x2 -> E2M1 1.0""" + assert k_fp4 % 8 == 0 + word = 0 + for i in range(8): + word |= (fp4_bits & 0xF) << (i * 4) + return torch.full((m, k_fp4 // 8), word, dtype=torch.int32, device=device) + + +def generate_mxf4_scale_factors(m, n, k_fp4, device='cuda', random_sf=False): + """为 MXF4 (VS=32) 生成 scale factor (float32 格式)。 + + Host 端 C++ API (transform_sf_into_required_layout) 会在调用 kernel 之前 + 将 float32 SF 转换为 packed UE8M0 int32: + 1. 从 float32 IEEE 754 中提取 8-bit 指数字段 (bitwise_right_shift(23)) + 2. 每 4 个 UE8M0 打包为 1 个 int32 + 3. 转置为 MN-major + TMA 对齐 + 然后 kernel 通过 TMA 加载到 SMEM, 经 UTCCP 写入 TMEM 供 MMA 使用。 + + UE8M0: value = 2^(exp - 127)。每 VS=32 个 FP4 元素共享一个 SF。 + + Args: + random_sf: 如果 True, 生成随机的 2 的幂次 SF (0.25, 0.5, 1.0, 2.0, 4.0) + """ + VS = 32 + sf_k = ((k_fp4 // VS + 3) // 4) * 4 + if random_sf: + powers = torch.randint(-2, 3, (m, sf_k), device=device).float() + sf_a = torch.pow(2.0, powers) + powers = torch.randint(-2, 3, (n, sf_k), device=device).float() + sf_b = torch.pow(2.0, powers) + else: + sf_a = torch.ones((m, sf_k), dtype=torch.float32, device=device) + sf_b = torch.ones((n, sf_k), dtype=torch.float32, device=device) + return sf_a, sf_b + + +def fp4_reference(a_packed, b_packed, m, n, sf_a=None, sf_b=None): + """CPU 端 E2M1 FP4 GEMM reference: C = A @ B^T, 支持 block-scaled SF。 + + Block-scaled MXF4: C[m,n] = sum_g SF_A[m,g] * SF_B[n,g] * dot(A_g, B_g) + 其中 g 是 VS=32 元素的组。sf_a/sf_b 为 float32 (pre-transform), 每组一个值。 + """ + VS = 32 + a_cpu = a_packed.cpu().to(torch.int64) & 0xFFFFFFFF + b_cpu = b_packed.cpu().to(torch.int64) & 0xFFFFFFFF + bits_a = torch.stack([(a_cpu >> (i*4)) & 0xF for i in range(8)], dim=-1).reshape(m, -1) + bits_b = torch.stack([(b_cpu >> (i*4)) & 0xF for i in range(8)], dim=-1).reshape(n, -1) + a_float = E2M1_LUT[bits_a.long()] # [m, k_fp4] + b_float = E2M1_LUT[bits_b.long()] # [n, k_fp4] + + if sf_a is None: + return torch.matmul(a_float, b_float.T) + + k_fp4 = a_float.shape[1] + sf_a_cpu = sf_a.cpu().float() + sf_b_cpu = sf_b.cpu().float() + c = torch.zeros(m, n, dtype=torch.float32) + num_groups = k_fp4 // VS + for g in range(num_groups): + k_start, k_end = g * VS, (g + 1) * VS + a_g = a_float[:, k_start:k_end] + b_g = b_float[:, k_start:k_end] + sf_col = g + if sf_col < sf_a_cpu.shape[1]: + sfa_g = sf_a_cpu[:, sf_col].unsqueeze(1) + sfb_g = sf_b_cpu[:, sf_col].unsqueeze(1) + else: + sfa_g = torch.ones(m, 1) + sfb_g = torch.ones(n, 1) + c += (sfa_g * sfb_g.T) * torch.matmul(a_g, b_g.T) + return c + + +def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 128)): + """调用 FP4 GEMM kernel (复用 fp8_gemm_nt 入口, int32 dtype 触发 FP4 路径)""" + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + d = torch.empty((m, n), device='cuda', dtype=torch.float32) + deep_gemm.fp8_gemm_nt((a_packed, sf_a), (b_packed, sf_b), d, c=None, + recipe=recipe, disable_ue8m0_cast=duc) + torch.cuda.synchronize() + return d + + +def pack_fp4_random_3d(num_groups: int, n: int, k_fp4: int, device='cuda'): + """生成 [G, N, K_int32] FP4 packed tensor。""" + assert k_fp4 % 8 == 0 + raw = torch.randint(0, 16, (num_groups, n, k_fp4), dtype=torch.uint8, device=device) + packed = torch.zeros(num_groups, n, k_fp4 // 8, dtype=torch.int32, device=device) + for i in range(8): + packed += (raw[:, :, i::8].to(torch.int32) << (i * 4)) + return packed + + +def generate_mxf4_sf_3d(num_groups: int, n: int, k_fp4: int, device='cuda', random_sf=False): + """生成 SFB [G, N, sf_k] for grouped contiguous.""" + VS = 32 + sf_k = ((k_fp4 // VS + 3) // 4) * 4 + if random_sf: + powers = torch.randint(-2, 3, (num_groups, n, sf_k), device=device).float() + return torch.pow(2.0, powers) + return torch.ones((num_groups, n, sf_k), dtype=torch.float32, device=device) + + +def fp4_reference_grouped(a_packed, b_packed_grouped, m_indices, + n: int, num_groups: int, + sf_a=None, sf_b_grouped=None): + """Per-row grouped FP4 reference: D[i] = A[i] @ B[m_indices[i]].T (with SF scaling). + + Padding rows (m_indices == -1) get D[i] = 0. + """ + m = a_packed.shape[0] + c = torch.zeros(m, n, dtype=torch.float32) + m_idx_cpu = m_indices.cpu() + for g in range(num_groups): + rows = (m_idx_cpu == g).nonzero(as_tuple=True)[0] + if rows.numel() == 0: + continue + a_g = a_packed[rows] + sf_a_g = sf_a[rows] if sf_a is not None else None + sf_b_g = sf_b_grouped[g] if sf_b_grouped is not None else None + c_g = fp4_reference(a_g, b_packed_grouped[g], rows.numel(), n, sf_a_g, sf_b_g) + c[rows] = c_g + return c + + +def run_kernel_grouped(a_packed, b_packed, sf_a, sf_b, m_indices, m, n, recipe=(1, 1, 128)): + """调用 m_grouped_fp8_gemm_nt_contiguous 的 FP4 路径 (int32 dtype 触发)。""" + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + d = torch.empty((m, n), device='cuda', dtype=torch.float32) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + (a_packed, sf_a), (b_packed, sf_b), d, m_indices, + recipe=recipe, disable_ue8m0_cast=duc, + ) + torch.cuda.synchronize() + return d + + +def pack_fp4_random_3d_ga(num_groups: int, max_m: int, k_fp4: int, device='cuda'): + """生成 [G, max_M, K_int32] FP4 packed tensor (A side for masked variant).""" + assert k_fp4 % 8 == 0 + raw = torch.randint(0, 16, (num_groups, max_m, k_fp4), dtype=torch.uint8, device=device) + packed = torch.zeros(num_groups, max_m, k_fp4 // 8, dtype=torch.int32, device=device) + for i in range(8): + packed += (raw[:, :, i::8].to(torch.int32) << (i * 4)) + return packed + + +def generate_mxf4_sfa_3d(num_groups: int, max_m: int, k_fp4: int, device='cuda', random_sf=False): + """生成 SFA [G, max_M, sf_k] for grouped masked.""" + VS = 32 + sf_k = ((k_fp4 // VS + 3) // 4) * 4 + if random_sf: + powers = torch.randint(-2, 3, (num_groups, max_m, sf_k), device=device).float() + return torch.pow(2.0, powers) + return torch.ones((num_groups, max_m, sf_k), dtype=torch.float32, device=device) + + +def fp4_reference_masked(a_3d, b_3d, masked_m_cpu, max_m, n, num_groups, + sf_a_3d=None, sf_b_3d=None): + """Reference for masked variant: D[g, :masked_m[g], :] = A[g, :masked_m[g], :] @ B[g, :, :].T. + + Padding rows beyond masked_m[g] in D are not checked (kernel may leave any value). + Returns d_ref of shape [G, max_M, N] with valid rows filled, padding rows zeroed. + """ + d_ref = torch.zeros(num_groups, max_m, n, dtype=torch.float32) + for g in range(num_groups): + mg = int(masked_m_cpu[g].item()) + if mg == 0: + continue + a_g = a_3d[g, :mg] + sf_a_g = sf_a_3d[g, :mg] if sf_a_3d is not None else None + sf_b_g = sf_b_3d[g] if sf_b_3d is not None else None + c_g = fp4_reference(a_g, b_3d[g], mg, n, sf_a_g, sf_b_g) + d_ref[g, :mg] = c_g + return d_ref + + +def run_kernel_grouped_masked(a_3d, b_3d, sf_a_3d, sf_b_3d, masked_m, num_groups, + max_m, n, expected_m, recipe=(1, 1, 128)): + """调用 m_grouped_fp8_gemm_nt_masked FP4 路径。""" + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + d = torch.empty((num_groups, max_m, n), device='cuda', dtype=torch.float32) + deep_gemm.m_grouped_fp8_gemm_nt_masked( + (a_3d, sf_a_3d), (b_3d, sf_b_3d), d, masked_m, expected_m, + recipe=recipe, disable_ue8m0_cast=duc, + ) + torch.cuda.synchronize() + return d + + +# ============================================================ +# 测试用例 +# ============================================================ + +def test_constant(): + """全 1.0 常量测试: C[i,j] = K (因为 1.0 * 1.0 * K 个元素)""" + print('Test: constant values (all E2M1 1.0)') + configs = [ + # (M, N, K_fp4) — K 是 FP4 元素个数 + # 单 stage (K <= 256) + (32, 64, 256), + (128, 128, 128), + (128, 256, 256), + # 多 stage (K > 256) + (128, 128, 512), + (256, 256, 512), + (128, 128, 1024), + (128, 256, 1024), + # 大 M (multi-wave: BLOCK_M=128, 所以 M>128 需要多个 wave) + (256, 128, 256), + (256, 256, 1024), + ] + all_pass = True + for m, n, k in configs: + a = pack_fp4_constant(m, k, fp4_bits=0x2) + b = pack_fp4_constant(n, k, fp4_bits=0x2) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) + d = run_kernel(a, b, sf_a, sf_b, m, n) + expected = float(k) + ok = (d.cpu() == expected).all().item() + if not ok: + all_pass = False + print(f' M={m:4d} N={n:4d} K={k:4d}: expected={expected:.0f} got={d.cpu()[0,0].item():.0f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_random(): + """随机数据测试 (SF=1.0): 对比 GPU kernel 与 CPU reference""" + print('Test: random data (vs CPU reference, SF=1.0)') + configs = [ + (32, 64, 256), + (128, 128, 128), + (128, 256, 256), + # 多 stage + (128, 128, 512), + (256, 256, 512), + (128, 128, 1024), + # 大 M + (256, 128, 256), + (256, 256, 1024), + # 较大 N + (128, 512, 256), + ] + all_pass = True + for m, n, k in configs: + a = pack_fp4_random(m, k) + b = pack_fp4_random(n, k) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) + d = run_kernel(a, b, sf_a, sf_b, m, n) + ref = fp4_reference(a, b, m, n, sf_a, sf_b) + max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() + ok = max_diff < 1.0 + if not ok: + all_pass = False + print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_value_sweep(): + """不同 FP4 值测试: 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0""" + print('Test: FP4 value sweep') + m, n, k = 128, 128, 256 + all_pass = True + for bits, val in [(0x1, 0.5), (0x2, 1.0), (0x3, 1.5), (0x4, 2.0), + (0x5, 3.0), (0x6, 4.0), (0x7, 6.0)]: + a = pack_fp4_constant(m, k, fp4_bits=bits) + b = pack_fp4_constant(n, k, fp4_bits=bits) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) + d = run_kernel(a, b, sf_a, sf_b, m, n) + expected = float(k) * val * val + actual = d.cpu()[0, 0].item() + ok = abs(actual - expected) < 1.0 + if not ok: + all_pass = False + print(f' FP4={val:4.1f} (bits=0x{bits:X}): expected={expected:.1f} got={actual:.1f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_uniform_sf(): + """Uniform SF 测试: 验证 UTCCP SF->TMEM 路径在不同 SF 值下正确""" + print('Test: uniform scale factors (UTCCP path)') + configs = [ + (128, 256, 256), + (256, 128, 256), + (128, 128, 512), + ] + all_pass = True + for m, n, k in configs: + for sf_val in [0.25, 0.5, 1.0, 2.0, 4.0]: + a = pack_fp4_random(m, k) + b = pack_fp4_random(n, k) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) + sf_a.fill_(sf_val) + sf_b.fill_(sf_val) + d = run_kernel(a, b, sf_a, sf_b, m, n) + # uniform SF: kernel result = unscaled_result * sf_a * sf_b + ref = fp4_reference(a, b, m, n, sf_a, sf_b) + max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() + ok = max_diff < 1.0 + if not ok: + all_pass = False + print(f' M={m:4d} N={n:4d} K={k:4d} SF={sf_val:5.2f}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_asymmetric_values(): + """A 和 B 使用不同 FP4 值""" + print('Test: asymmetric A/B values') + m, n, k = 128, 128, 256 + all_pass = True + cases = [ + (0x2, 0x4, 1.0, 2.0), # A=1.0, B=2.0 + (0x1, 0x6, 0.5, 4.0), # A=0.5, B=4.0 + (0x4, 0x1, 2.0, 0.5), # A=2.0, B=0.5 + ] + for bits_a, bits_b, val_a, val_b in cases: + a = pack_fp4_constant(m, k, fp4_bits=bits_a) + b = pack_fp4_constant(n, k, fp4_bits=bits_b) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) + d = run_kernel(a, b, sf_a, sf_b, m, n) + expected = float(k) * val_a * val_b + actual = d.cpu()[0, 0].item() + ok = abs(actual - expected) < 1.0 + if not ok: + all_pass = False + print(f' A={val_a}, B={val_b}: expected={expected:.1f} got={actual:.1f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_random_sf(): + """随机数据 + 随机 per-group SF (powers of 2)""" + print('Test: random data + random scale factors') + configs = [ + (128, 128, 256), + (128, 256, 256), + (256, 256, 512), + (128, 128, 512), + (128, 128, 1024), + (256, 128, 256), + (256, 256, 1024), + (128, 256, 1024), + ] + all_pass = True + for m, n, k in configs: + a = pack_fp4_random(m, k) + b = pack_fp4_random(n, k) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k, random_sf=True) + d = run_kernel(a, b, sf_a, sf_b, m, n) + ref = fp4_reference(a, b, m, n, sf_a, sf_b) + max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() + ok = max_diff < 1.0 + if not ok: + all_pass = False + print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') + return all_pass + + +def test_m_grouped_contiguous(): + """M-grouped contiguous FP4 GEMM (MoE forward shape). + + A=[M_total, K] @ B=[G, N, K].T → D=[M_total, N] + m_indices[M_total] selects which group each row belongs to (-1 = padding). + """ + print('Test: m-grouped contiguous (MoE-style)') + BLOCK_M = 128 + + # Small/debug shapes — fast CPU LUT reference, exercise padding-row case. + debug_configs = [ + (2, 128, 128, 256), + (4, 128, 128, 256), + (4, 128, 256, 512), + (4, 256, 128, 1024), + (8, 128, 256, 512), + # Uneven actual M per group, padded to BLOCK_M (exercises padding rows). + (4, 90, 128, 256), + (4, 200, 256, 512), + # Extra N×K variety to exercise more BLOCK_N choices and SF transform paths. + (4, 128, 512, 256), + (4, 128, 768, 256), + ] + + # Production MoE shapes — mirror DeepGEMM official FP8 grouped (generators.py:105). + # CPU LUT reference takes a few seconds per shape at this size; OK for nightly. + prod_configs = [ + (4, 8192, 4096, 7168), # EP4, MoE up-projection + (4, 8192, 7168, 2048), # EP4, MoE down-projection + (8, 4096, 4096, 7168), # EP8, MoE up-projection + (8, 4096, 7168, 2048), # EP8, MoE down-projection + # Extra (n, k) variants from FP8 enumerate_normal for N/K coverage: + (4, 8192, 24576, 1536), + (4, 8192, 32768, 512), + ] + + # Mirror FP8 masked-test pattern: multiple random-data iterations per shape to + # catch flaky bugs that only show up with certain RNG seeds. Debug shapes get + # fewer iters (kept fast); prod shapes get 3 iters each. + NUM_ITERS = {'debug': 2, 'prod': 3} + + all_pass = True + for label, configs in [('debug', debug_configs), ('prod', prod_configs)]: + for num_groups, m_per_group, n, k in configs: + aligned_m = ((m_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + m_total = aligned_m * num_groups + + worst_diff = 0.0 + for _ in range(NUM_ITERS[label]): + # Build A, B, m_indices fresh per iteration + a = pack_fp4_random(m_total, k) + b = pack_fp4_random_3d(num_groups, n, k) + sf_a, _ = generate_mxf4_scale_factors(m_total, n, k, random_sf=True) + sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) + m_indices = torch.empty(m_total, dtype=torch.int32, device='cuda') + for g in range(num_groups): + start = g * aligned_m + actual_end = start + m_per_group + aligned_end = start + aligned_m + m_indices[start:actual_end] = g + m_indices[actual_end:aligned_end] = -1 + + d = run_kernel_grouped(a, b, sf_a, sf_b, m_indices, m_total, n) + d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) + + ref = fp4_reference_grouped(a, b, m_indices, n, num_groups, sf_a, sf_b) + d_max = torch.abs(d.cpu().float() - ref.float()).max().item() + if d_max > worst_diff: + worst_diff = d_max + + ok = worst_diff < 1.0 + if not ok: + all_pass = False + print(f' [{label}×{NUM_ITERS[label]}] G={num_groups} m_per_group={m_per_group:5d} ' + f'(aligned={aligned_m:5d}) N={n:5d} K_fp4={k:5d}: ' + f'max_diff={worst_diff:.4f} {"PASS" if ok else "FAIL"}') + + # Perf section — bench_kineto on production shapes only. + # Filter to the FP4 GEMM kernel name; works for both dense and grouped wrappers + # because the device kernel function is sm100_fp4_gemm_1d1d_impl in both cases. + print('\nPerf: m-grouped contiguous (production MoE shapes)') + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + for num_groups, m_per_group, n, k in prod_configs: + aligned_m = ((m_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + m_total = aligned_m * num_groups + + a = pack_fp4_random(m_total, k) + b = pack_fp4_random_3d(num_groups, n, k) + sf_a, _ = generate_mxf4_scale_factors(m_total, n, k, random_sf=True) + sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) + m_indices = torch.empty(m_total, dtype=torch.int32, device='cuda') + for g in range(num_groups): + start = g * aligned_m + m_indices[start:start + m_per_group] = g + m_indices[start + m_per_group:start + aligned_m] = -1 + d = torch.empty((m_total, n), device='cuda', dtype=torch.float32) + + def fn(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + (a, sf_a), (b, sf_b), d, m_indices, + recipe=(1, 1, 128), disable_ue8m0_cast=duc, + ) + + t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) + # FLOPs: 2 * M_total * N * K_fp4 (k here is K_fp4 element count). + # Bytes: A (int32, 4B/elem), B (int32, 4B/elem), D (fp32, 4B/elem), plus SF. + tflops = 2 * m_total * n * k / t / 1e12 + gbps = count_bytes(a, b, d) / 1e9 / t + print(f' G={num_groups} m_per_group={m_per_group:5d} N={n:5d} K_fp4={k:5d}: ' + f'{t * 1e6:6.1f} us | {tflops:6.0f} TFLOPS | {gbps:5.0f} GB/s') + + return all_pass + + +def test_m_grouped_trtllm_comparable(): + """trtllm-gen apples-to-apples: DeepSeek-R1 MoE setup with 256 experts × topK=8. + + Setup mirrors gitlab/trtllm BatchedGemm baseline: + num_groups = 256 (numExperts) + total_actual_M = numTokens * topK = numTokens * 8 + Uniform routing → rows_per_expert = total_actual_M / 256 + Each expert tile padded to BLOCK_M=128 + + CAVEAT: trtllm-gen's BatchedGemmFp4LowLatency uses batch=N (grouped-along-N); + DG here uses batch=M. Memory access patterns differ but useful FLOPs compare directly. + + trtllm-gen baseline on B200 (from memory/fp4_grouped_gemm_perf_work.md): + + FC2 (N=7168, K=2048): + tokens=32 → 0.37 ms | 20.4 TFLOPS | 5.74 TB/s + tokens=64 → 0.37 ms | 40.7 TFLOPS | 5.75 TB/s + tokens=128 → 0.37 ms | 80.8 TFLOPS | 5.72 TB/s + tokens=256 → 0.37 ms | 161.9 TFLOPS | 5.78 TB/s + tokens=512 → 0.55 ms | 219.6 TFLOPS | 3.98 TB/s + tokens=1024 → 0.51 ms | 472.3 TFLOPS | 4.40 TB/s + tokens=2048 → 0.51 ms | 941.2 TFLOPS | 4.63 TB/s + + FC1 (N=4096, K=7168, fusedAct=swiglu, routeAct=tma) — note batch=N caveat: + tokens=32 → 0.76 ms | 19.9 TFLOPS | 5.59 TB/s + tokens=64 → 0.75 ms | 39.8 TFLOPS | 5.61 TB/s + tokens=128 → 0.76 ms | 78.7 TFLOPS | 5.54 TB/s + tokens=256 → 0.76 ms | 157.5 TFLOPS | 5.55 TB/s + tokens=512 → 0.78 ms | 310.1 TFLOPS | 5.48 TB/s + tokens=1024 → 0.82 ms | 590.0 TFLOPS | 5.25 TB/s + tokens=2048 → 1.19 ms | 810.9 TFLOPS | 3.65 TB/s + """ + print('Test: m-grouped contiguous (trtllm-gen comparable, 256 experts)') + BLOCK_M = 128 + NUM_GROUPS = 256 + TOP_K = 8 + + # Latency regime (trtllm-gen batch=N baseline): tokens 32..2048, per_group_M < BLOCK_M + # Throughput regime (trtllm-gen batch=M target): tokens 4096..8192, per_group_M >= BLOCK_M + # 4096: rows/grp=128 (= BLOCK_M, zero padding) — first apples-to-apples vs trtllm batch=M + # 8192: rows/grp=256 (= 2×BLOCK_M, 2 tiles per group) + token_sweep = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + shape_configs = [ + ('FC2', 7168, 2048), + ('FC1', 4096, 7168), + ] + + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + all_pass = True + + for fc_name, n, k in shape_configs: + print(f'\n--- {fc_name}: N={n}, K_fp4={k}, num_groups={NUM_GROUPS} ---') + # Allocate B/SFB once per shape (large, reuse across token sweep — independent of tokens) + b = pack_fp4_random_3d(NUM_GROUPS, n, k) + sf_b = generate_mxf4_sf_3d(NUM_GROUPS, n, k, random_sf=True) + + for num_tokens in token_sweep: + total_actual_m = num_tokens * TOP_K + rows_per_group = max(1, total_actual_m // NUM_GROUPS) + # Per-group rows aligned UP to BLOCK_M (may exceed BLOCK_M for tokens >= 4096) + aligned_per_group = ((rows_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + total_padded_m = NUM_GROUPS * aligned_per_group + + # A/SF_A/D scale with total_padded_m, so allocate per iteration. + a = pack_fp4_random(total_padded_m, k) + sf_a, _ = generate_mxf4_scale_factors(total_padded_m, n, k, random_sf=True) + d = torch.empty((total_padded_m, n), device='cuda', dtype=torch.float32) + + # m_indices: each group gets aligned_per_group rows, first rows_per_group are valid (g), + # remainder padding (-1). + m_indices = torch.full((total_padded_m,), -1, dtype=torch.int32, device='cuda') + for g in range(NUM_GROUPS): + start = g * aligned_per_group + m_indices[start:start + rows_per_group] = g + # rows [start + rows_per_group : start + aligned_per_group) stay -1 + + # Run kernel + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + (a, sf_a), (b, sf_b), d, m_indices, + recipe=(1, 1, 128), disable_ue8m0_cast=duc, + ) + torch.cuda.synchronize() + d_clean = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) + + # Correctness: only for smaller tokens (CPU LUT reference for large M is slow) + if num_tokens <= 256: + ref = fp4_reference_grouped(a, b, m_indices, n, NUM_GROUPS, sf_a, sf_b) + max_diff = torch.abs(d_clean.cpu().float() - ref.float()).max().item() + ok = max_diff < 1.0 + if not ok: + all_pass = False + correctness = f'diff={max_diff:.4f} {"✓" if ok else "✗"}' + else: + correctness = '(skip)' + + # Perf + def fn(): + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + (a, sf_a), (b, sf_b), d, m_indices, + recipe=(1, 1, 128), disable_ue8m0_cast=duc, + ) + t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) + + # Useful FLOPS = 2 * useful_M * N * K (matches trtllm-gen convention) + useful_m = NUM_GROUPS * rows_per_group + useful_tflops = 2 * useful_m * n * k / t / 1e12 + # Bytes: A loaded once (total_padded_m), B loaded once, D written once + gbps = count_bytes(a, b, d) / 1e9 / t + print(f' tokens={num_tokens:5d} useful_M={useful_m:6d} ' + f'rows/grp={rows_per_group:4d}/{aligned_per_group:4d}: ' + f'{t * 1e3:6.3f} ms | {useful_tflops:7.1f} TFLOPS | ' + f'{gbps/1000:5.2f} TB/s | {correctness}') + + return all_pass + + +def test_m_grouped_masked(): + """M-grouped masked FP4 GEMM (MoE decode shape). + + A=[G, max_M, K] @ B=[G, N, K].T → D=[G, max_M, N] + masked_m[G] indicates the number of valid rows per group (rest is padding). + """ + print('Test: m-grouped masked (MoE decode)') + BLOCK_M = 128 + + # (num_groups, max_m, expected_m, n, k_fp4) + # Small/debug shapes (CPU reference fast) + production shapes mirroring FP8 enumerate_m_grouped_masked + debug_configs = [ + (2, 128, 64, 128, 256), # 50% util + (4, 128, 32, 128, 256), # heavy padding + (4, 256, 200, 128, 256), # multi-tile per group + (4, 128, 128, 256, 512), # full + (8, 128, 64, 256, 512), + ] + prod_configs = [ + # max_m=4096 (matching FP8 enumerate_m_grouped_masked max_m), varying num_groups & m + (1, 4096, 1024, 4096, 7168), # FP8 (1, 1024) + (2, 4096, 512, 4096, 7168), # FP8 (2, 512) + (4, 4096, 256, 4096, 7168), # FP8 (4, 256) + (1, 4096, 1024, 7168, 2048), + (2, 4096, 512, 7168, 2048), + (4, 4096, 256, 7168, 2048), + ] + + # Mirror FP8 test_m_grouped_gemm_masked: 10 random-data iterations per shape on + # production shapes to catch flaky bugs (different masked_m distribution each time). + NUM_ITERS = {'debug': 3, 'prod': 10} + + all_pass = True + duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + for label, configs in [('debug', debug_configs), ('prod', prod_configs)]: + for num_groups, max_m, expected_m, n, k in configs: + worst_diff = 0.0 + for _ in range(NUM_ITERS[label]): + # Fresh masked_m + tensors per iteration (matches FP8 pattern) + masked_m_cpu = torch.tensor([ + max(1, min(max_m, int(expected_m * random.uniform(0.7, 1.3)))) + for _ in range(num_groups) + ], dtype=torch.int32) + masked_m = masked_m_cpu.cuda() + + a = pack_fp4_random_3d_ga(num_groups, max_m, k) + b = pack_fp4_random_3d(num_groups, n, k) + sf_a = generate_mxf4_sfa_3d(num_groups, max_m, k, random_sf=True) + sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) + + d = run_kernel_grouped_masked(a, b, sf_a, sf_b, masked_m, num_groups, + max_m, n, expected_m) + ref = fp4_reference_masked(a, b, masked_m_cpu, max_m, n, num_groups, sf_a, sf_b) + + # Only compare valid rows per group + for g in range(num_groups): + mg = int(masked_m_cpu[g].item()) + if mg == 0: + continue + diff = torch.abs(d[g, :mg].cpu().float() - ref[g, :mg].float()).max().item() + if diff > worst_diff: + worst_diff = diff + + ok = worst_diff < 1.0 + if not ok: + all_pass = False + print(f' [{label}×{NUM_ITERS[label]}] G={num_groups} max_m={max_m:5d} ' + f'expected_m={expected_m:5d} N={n:5d} K_fp4={k:5d}: ' + f'max_diff={worst_diff:.4f} {"PASS" if ok else "FAIL"}') + + return all_pass + + +def test_multicast(): + """大 M 测试：触发 B-multicast (M>=512, 2CTA along M, UMMA_M=256)""" + print('Test: B-multicast (M>=512, 2CTA)') + configs = [ + (512, 128, 256), + (512, 128, 512), + (512, 128, 1024), + (1024, 128, 256), + (1024, 128, 512), + ] + all_pass = True + for m, n, k in configs: + a = pack_fp4_random(m, k) + b = pack_fp4_random(n, k) + sf_a, sf_b = generate_mxf4_scale_factors(m, n, k, random_sf=True) + d = run_kernel(a, b, sf_a, sf_b, m, n) + ref = fp4_reference(a, b, m, n, sf_a, sf_b) + max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() + ok = max_diff < 1.0 + if not ok: + all_pass = False + print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') + return all_pass + + +if __name__ == '__main__': + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.manual_seed(0) + random.seed(0) + + print(f'Library: {deep_gemm.__path__}\n') + + results = [ + ('constant', test_constant()), + ('random', test_random()), + ('sweep', test_value_sweep()), + ('asymmetric', test_asymmetric_values()), + ('uniform_sf', test_uniform_sf()), + ('random_sf', test_random_sf()), + ('multicast', test_multicast()), + ('m_grouped', test_m_grouped_contiguous()), + ('m_grouped_masked', test_m_grouped_masked()), + ('trtllm_cmp', test_m_grouped_trtllm_comparable()), + ] + + print() + passed = all(r for _, r in results) + for name, ok in results: + print(f' {name}: {"PASS" if ok else "FAIL"}') + print(f'\n{"ALL FP4 TESTS PASSED" if passed else "SOME TESTS FAILED"}') + if not passed: + exit(1) From 7c2d685dbd91eb80992c02456dc5520b6120cba4 Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Fri, 29 May 2026 03:20:54 -0700 Subject: [PATCH 02/12] =?UTF-8?q?FP4:=20progress=20port=20towards=20runtim?= =?UTF-8?q?e=20=E2=80=94=20IMA=20fixed,=20smem=20cap,=20multicast=20clamp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stack of incremental fixes on top of previous port commit (7d27b8a). Build clean, JIT clean, kernel launches without IMA. NaN output remains (deeper SF/data-layout debug needed). Wrapper changes (sm100_fp4_gemm_1d1d.hpp): - view int8 (kPackedFP4) tensors as int32 before TMA descriptor creation, so TMA uses CU_TENSOR_MAP_DATA_TYPE_INT32 (matches my kernel's int32-packed smem expectation, sidesteps main's 16U4_ALIGN16B unpacked-smem path). - Convert k and block_k to int32 units (k/8, block_k/4) for TMA + kernel template instantiation. - num_stages clamp: my kernel allocates 2x SF smem/stage vs main's heuristic estimate; cap stages so total smem fits 232448-byte capacity and num_stages <= num_k_blocks (fea-fp4 invariant). - Single-CTA force (v0): override cluster_m=cluster_n=1; recompute storage_config / pipeline_config so they stay self-consistent with the new layout. main's heuristic picks cluster=2 too aggressively for shapes fea-fp4 stayed single-CTA on. Status: kernel runs, produces wrong output (suspected SF byte ordering or TMEM column placement mismatch between main's UE8M0 packed transform and my kernel's read expectation). fea-fp4 tests historically used sf=1.0 which masks SF reading bugs (1.0 multiplier regardless of byte order); my smoke test with main's generators uses varying SF, exposing the issue. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 129 +++++++++++++++--- 1 file changed, 108 insertions(+), 21 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index 308a1a0c57..2a6599e954 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -98,6 +98,45 @@ static int compute_num_last_stages_fp4(int k, int block_k_bytes, int num_stages) return rem == 0 ? num_stages : rem; } +// My MXF4 kernel allocates 2x SF smem per stage (2 packed int32 per row to cover the +// full BLOCK_K_FP4=256 of 8 SFs) than main's heuristic assumes (1 packed int32 per +// row, covering only 128 FP4). Cap num_stages so total smem fits SM100's 232448-byte +// capacity. Returns {new_num_stages, new_smem_size}. +// Note: this is a v0 workaround. Phase 2: align main's heuristic SF computation with +// my kernel's actual usage so num_stages selection is correct upstream. +static std::pair recompute_stages_for_fp4(const GemmConfig& config, int block_m, int block_n, int k_fp4) { + constexpr int smem_capacity = 232448; + constexpr int sf_block_align = 128; // kNumUTCCPAlignedElems + const int sf_block_m = (block_m + sf_block_align - 1) / sf_block_align * sf_block_align; + const int sf_block_n = (block_n + sf_block_align - 1) / sf_block_align * sf_block_align; + + // My kernel SF per stage = SF_BLOCK_MN * SF_PACKED_K_PER_STAGE * 4 bytes; + // SF_PACKED_K_PER_STAGE = (BLOCK_K * 8 / 32) / 4 = BLOCK_K / 16 (in int32 units). + // For main's block_k bytes (= int32 units * 4), SF_PACKED_K_PER_STAGE = block_k_bytes / 64. + const int sf_packed_k_per_stage = config.layout.block_k / 64; + const int my_smem_sfa_per_stage = sf_block_m * sf_packed_k_per_stage * 4; + const int my_smem_sfb_per_stage = sf_block_n * sf_packed_k_per_stage * 4; + + // Reconstruct main's smem_per_stage and extra parts. + const int main_smem_sfa_per_stage = sf_block_m * 4; + const int main_smem_sfb_per_stage = sf_block_n * 4; + const int main_smem_per_stage = config.storage_config.load_block_m * config.layout.block_k + + config.storage_config.load_block_n * config.layout.block_k + + main_smem_sfa_per_stage + main_smem_sfb_per_stage; + const int main_smem_extra = config.pipeline_config.smem_size - + config.pipeline_config.num_stages * main_smem_per_stage; + + const int my_smem_per_stage = main_smem_per_stage - main_smem_sfa_per_stage - main_smem_sfb_per_stage + + my_smem_sfa_per_stage + my_smem_sfb_per_stage; + const int max_stages = (smem_capacity - main_smem_extra) / my_smem_per_stage; + // Also clamp to num_k_blocks: my kernel assumes pipeline depth <= K iterations + // (fea-fp4 heuristic invariant). main's heuristic ignores this. + const int num_k_blocks = ceil_div(k_fp4, config.layout.block_k * 2); + const int new_num_stages = std::min({config.pipeline_config.num_stages, max_stages, num_k_blocks}); + const int new_smem_size = main_smem_extra + new_num_stages * my_smem_per_stage; + return {new_num_stages, new_smem_size}; +} + // Build a GemmDesc forcing a_dtype=b_dtype=kPackedFP4 (since this is the FP4xFP4 wrapper). static GemmDesc make_fp4_desc(GemmType gemm_type, int m, int n, int k, int num_groups, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, @@ -132,18 +171,36 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa const auto desc = make_fp4_desc(GemmType::Normal, m, n, k, 1, major_a, major_b, d.scalar_type(), c.has_value(), compiled_dims); - const auto config = get_best_config(desc); + auto config = get_best_config(desc); + // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively + // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + config.layout.cluster_m = 1; + config.layout.cluster_n = 1; + config.launch_config.num_sms_per_cluster = 1; + config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); + const auto [new_stages, new_smem] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + config.pipeline_config.num_stages = new_stages; + config.pipeline_config.smem_size = new_smem; + + // View FP4-packed-int8 tensors as int32 (8 FP4 per int32). Same memory layout, + // different dtype tag -- makes the TMA descriptor use INT32 (not 16U4_ALIGN16B + // unpacked-smem), which matches my kernel's int32-packed smem expectation. + const auto a_int32 = a.view(torch::kInt); + const auto b_int32 = b.view(torch::kInt); + const int k_int32 = k / 8; + const int block_k_int32 = config.layout.block_k / 4; const auto cd = c.value_or(d); - const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + const auto tensor_map_a = make_tma_a_desc(major_a, a_int32, m, k_int32, config.storage_config.load_block_m, - config.layout.block_k, - static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, + block_k_int32, + static_cast(a_int32.stride(get_non_contiguous_dim(major_a))), 1, config.storage_config.swizzle_a_mode); - const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + const auto tensor_map_b = make_tma_b_desc(major_b, b_int32, n, k_int32, config.storage_config.load_block_n, - config.layout.block_k, - static_cast(b.stride(get_non_contiguous_dim(major_b))), 1, + block_k_int32, + static_cast(b_int32.stride(get_non_contiguous_dim(major_b))), 1, config.storage_config.swizzle_b_mode); const auto tensor_map_d = make_tma_cd_desc(d, m, n, config.storage_config.store_block_m, @@ -199,17 +256,32 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con const auto desc = make_fp4_desc(GemmType::MGroupedContiguous, m, n, k, num_groups, major_a, major_b, d.scalar_type(), false, compiled_dims); - const auto config = get_best_config(desc); + auto config = get_best_config(desc); + // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively + // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + config.layout.cluster_m = 1; + config.layout.cluster_n = 1; + config.launch_config.num_sms_per_cluster = 1; + config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); + const auto [new_stages_g, new_smem_g] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + config.pipeline_config.num_stages = new_stages_g; + config.pipeline_config.smem_size = new_smem_g; - const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + const auto a_int32 = a.view(torch::kInt); + const auto b_int32 = b.view(torch::kInt); + const int k_int32 = k / 8; + const int block_k_int32 = config.layout.block_k / 4; + + const auto tensor_map_a = make_tma_a_desc(major_a, a_int32, m, k_int32, config.storage_config.load_block_m, - config.layout.block_k, - static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, + block_k_int32, + static_cast(a_int32.stride(get_non_contiguous_dim(major_a))), 1, config.storage_config.swizzle_a_mode); - const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + const auto tensor_map_b = make_tma_b_desc(major_b, b_int32, n, k_int32, config.storage_config.load_block_n, - config.layout.block_k, - static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, + block_k_int32, + static_cast(b_int32.stride(get_non_contiguous_dim(major_b))), num_groups, config.storage_config.swizzle_b_mode); const auto tensor_map_d = make_tma_cd_desc(d, m, n, config.storage_config.store_block_m, @@ -254,17 +326,32 @@ static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const t major_a, major_b, d.scalar_type(), false, compiled_dims, /*expected_m=*/expected_m, /*expected_num_groups=*/num_groups); - const auto config = get_best_config(desc); + auto config = get_best_config(desc); + // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively + // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + config.layout.cluster_m = 1; + config.layout.cluster_n = 1; + config.launch_config.num_sms_per_cluster = 1; + config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); + const auto [new_stages_mk, new_smem_mk] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + config.pipeline_config.num_stages = new_stages_mk; + config.pipeline_config.smem_size = new_smem_mk; + + const auto a_int32 = a.view(torch::kInt); + const auto b_int32 = b.view(torch::kInt); + const int k_int32 = k / 8; + const int block_k_int32 = config.layout.block_k / 4; - const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, + const auto tensor_map_a = make_tma_a_desc(major_a, a_int32, m, k_int32, config.storage_config.load_block_m, - config.layout.block_k, - static_cast(a.stride(get_non_contiguous_dim(major_a))), num_groups, + block_k_int32, + static_cast(a_int32.stride(get_non_contiguous_dim(major_a))), num_groups, config.storage_config.swizzle_a_mode); - const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, + const auto tensor_map_b = make_tma_b_desc(major_b, b_int32, n, k_int32, config.storage_config.load_block_n, - config.layout.block_k, - static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, + block_k_int32, + static_cast(b_int32.stride(get_non_contiguous_dim(major_b))), num_groups, config.storage_config.swizzle_b_mode); const auto tensor_map_d = make_tma_cd_desc(d, m, n, config.storage_config.store_block_m, From 65399759a54837a2db3112f986c2fbbdb707af74 Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Fri, 29 May 2026 03:48:53 -0700 Subject: [PATCH 03/12] FP4: layout override to fea-fp4 tested config + test int8 packing Continuing port towards numerical correctness. Current state: kernel builds clean, JIT clean, runs without IMA, but writes zeros / wrong values. Wrapper changes (sm100_fp4_gemm_1d1d.hpp): - Override main's layout choice to fea-fp4's tested config: block_m=128, block_n=112 (kSwizzleCDMode=32, matches fea-fp4 test coverage), single-CTA cluster. main's heuristic picks block_n=16 / cluster=2 which fea-fp4 kernel was never validated against. - Fix sf_packed_k_per_stage formula in recompute_stages_for_fp4 (was off by 4x with block_k_int32/4; correct: block_k_bytes/64). Test adaptation (tests/test_fp4.py): - pack_fp4_random / pack_fp4_constant now produce int8 (kPackedFP4) instead of int32, matching main's API. Same byte layout when viewed as int32. - fp4_reference handles int8 packing (2 FP4 per byte: low nibble + high nibble), works for both int8 and int32 packed inputs. - run_kernel uses recipe_a=(1,32), recipe_b=(1,32) for FP4 SF granularity (was (1,1,128) which fea-fp4 commented out shape check to allow). Open issue: kernel writes 0 / wrong values for both test_constant (sf=1.0) and test_random. Needs device-side printf to localize: A/B smem content vs expected layout, MMA output before vs after epilogue, SF byte ordering in TMEM after UTCCP. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 70 ++++++++++++------- tests/test_fp4.py | 56 ++++++++------- 2 files changed, 72 insertions(+), 54 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index 2a6599e954..d70ed59f28 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -106,34 +106,35 @@ static int compute_num_last_stages_fp4(int k, int block_k_bytes, int num_stages) // my kernel's actual usage so num_stages selection is correct upstream. static std::pair recompute_stages_for_fp4(const GemmConfig& config, int block_m, int block_n, int k_fp4) { constexpr int smem_capacity = 232448; - constexpr int sf_block_align = 128; // kNumUTCCPAlignedElems + constexpr int sf_block_align = 128; const int sf_block_m = (block_m + sf_block_align - 1) / sf_block_align * sf_block_align; const int sf_block_n = (block_n + sf_block_align - 1) / sf_block_align * sf_block_align; - // My kernel SF per stage = SF_BLOCK_MN * SF_PACKED_K_PER_STAGE * 4 bytes; - // SF_PACKED_K_PER_STAGE = (BLOCK_K * 8 / 32) / 4 = BLOCK_K / 16 (in int32 units). - // For main's block_k bytes (= int32 units * 4), SF_PACKED_K_PER_STAGE = block_k_bytes / 64. + // Match fea-fp4 kernel's actual smem usage: + // A per stage: load_block_m * block_k_bytes + // B per stage: load_block_n * block_k_bytes + // SFA/B per stage: sf_block_mn * sf_packed_k_per_stage * 4 + // where sf_packed_k_per_stage = block_k_bytes / 64 (since BLOCK_K_FP4 / VS / 4 = block_k_int32 / 4 = block_k_bytes / 16 / 4 / 4 ... wait simpler: block_k_int32/4) + // Kernel: BLOCK_K_FP4 = block_k_int32 * 8; SF_K_PER_STAGE = BLOCK_K_FP4 / 32; + // SF_PACKED_K_PER_STAGE = SF_K_PER_STAGE / 4 = block_k_int32 / 16. + // In bytes: block_k_bytes / 64. const int sf_packed_k_per_stage = config.layout.block_k / 64; - const int my_smem_sfa_per_stage = sf_block_m * sf_packed_k_per_stage * 4; - const int my_smem_sfb_per_stage = sf_block_n * sf_packed_k_per_stage * 4; + const int per_stage = config.storage_config.load_block_m * config.layout.block_k + + config.storage_config.load_block_n * config.layout.block_k + + sf_block_m * sf_packed_k_per_stage * 4 + + sf_block_n * sf_packed_k_per_stage * 4; - // Reconstruct main's smem_per_stage and extra parts. - const int main_smem_sfa_per_stage = sf_block_m * 4; - const int main_smem_sfb_per_stage = sf_block_n * 4; - const int main_smem_per_stage = config.storage_config.load_block_m * config.layout.block_k + - config.storage_config.load_block_n * config.layout.block_k + - main_smem_sfa_per_stage + main_smem_sfb_per_stage; - const int main_smem_extra = config.pipeline_config.smem_size - - config.pipeline_config.num_stages * main_smem_per_stage; + // Fixed extras (CD smem + barriers + tmem_ptr, conservative estimate) + // CD (non-swap): store_block_m * swizzle_cd_mode * 2 stages + const int cd_size = config.storage_config.store_block_m * config.storage_config.swizzle_cd_mode * 2; + const int barriers = 12 * 8 * 4 + 4 * 8 * 2 + 8; // ~ 416 bytes max + const int tmem_ptr = 4; + const int fixed_extras = cd_size + barriers + tmem_ptr; - const int my_smem_per_stage = main_smem_per_stage - main_smem_sfa_per_stage - main_smem_sfb_per_stage - + my_smem_sfa_per_stage + my_smem_sfb_per_stage; - const int max_stages = (smem_capacity - main_smem_extra) / my_smem_per_stage; - // Also clamp to num_k_blocks: my kernel assumes pipeline depth <= K iterations - // (fea-fp4 heuristic invariant). main's heuristic ignores this. + const int max_stages_smem = (smem_capacity - fixed_extras) / per_stage; const int num_k_blocks = ceil_div(k_fp4, config.layout.block_k * 2); - const int new_num_stages = std::min({config.pipeline_config.num_stages, max_stages, num_k_blocks}); - const int new_smem_size = main_smem_extra + new_num_stages * my_smem_per_stage; + const int new_num_stages = std::min({12, max_stages_smem, num_k_blocks}); + const int new_smem_size = fixed_extras + new_num_stages * per_stage; return {new_num_stages, new_smem_size}; } @@ -172,8 +173,13 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa major_a, major_b, d.scalar_type(), c.has_value(), compiled_dims); auto config = get_best_config(desc); - // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively - // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. + // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). + // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) + // path appears untested upstream. + config.layout.block_m = 128; + config.layout.block_n = 112; + config.layout.swap_ab = false; config.layout.cluster_m = 1; config.layout.cluster_n = 1; config.launch_config.num_sms_per_cluster = 1; @@ -257,8 +263,13 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con major_a, major_b, d.scalar_type(), false, compiled_dims); auto config = get_best_config(desc); - // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively - // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. + // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). + // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) + // path appears untested upstream. + config.layout.block_m = 128; + config.layout.block_n = 112; + config.layout.swap_ab = false; config.layout.cluster_m = 1; config.layout.cluster_n = 1; config.launch_config.num_sms_per_cluster = 1; @@ -327,8 +338,13 @@ static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const t false, compiled_dims, /*expected_m=*/expected_m, /*expected_num_groups=*/num_groups); auto config = get_best_config(desc); - // v0: force single-CTA path (main's heuristic prefers cluster=2 too aggressively - // for shapes where fea-fp4 stayed at 1; multicast path needs more validation). + // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. + // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). + // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) + // path appears untested upstream. + config.layout.block_m = 128; + config.layout.block_n = 112; + config.layout.swap_ab = false; config.layout.cluster_m = 1; config.layout.cluster_n = 1; config.launch_config.num_sms_per_cluster = 1; diff --git a/tests/test_fp4.py b/tests/test_fp4.py index a547075fc9..17b62e74ef 100644 --- a/tests/test_fp4.py +++ b/tests/test_fp4.py @@ -24,22 +24,23 @@ # ============================================================ def pack_fp4_random(m: int, k_fp4: int, device='cuda'): - """生成随机 E2M1 FP4 数据并打包为 int32。每个 int32 包含 8 个 FP4 值。""" - assert k_fp4 % 8 == 0 + """Generate random E2M1 FP4 values, pack as int8 (kPackedFP4, 2 FP4 per byte). + Memory layout identical to int32 packing (8 FP4 per int32) when viewed as int32. + """ + assert k_fp4 % 2 == 0 raw = torch.randint(0, 16, (m, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(m, k_fp4 // 8, dtype=torch.int32, device=device) - for i in range(8): - packed += (raw[:, i::8].to(torch.int32) << (i * 4)) + packed = torch.zeros(m, k_fp4 // 2, dtype=torch.int8, device=device) + # Byte layout: low nibble = even FP4, high nibble = odd FP4 (matches per_token_cast_to_fp4) + packed |= (raw[:, 0::2].to(torch.int8) & 0x0F) + packed |= (raw[:, 1::2].to(torch.int8) & 0x0F) << 4 return packed def pack_fp4_constant(m: int, k_fp4: int, fp4_bits: int = 0x2, device='cuda'): - """生成常量 FP4 打包数据。fp4_bits=0x2 -> E2M1 1.0""" - assert k_fp4 % 8 == 0 - word = 0 - for i in range(8): - word |= (fp4_bits & 0xF) << (i * 4) - return torch.full((m, k_fp4 // 8), word, dtype=torch.int32, device=device) + """Constant FP4 packed as int8. fp4_bits=0x2 -> E2M1 1.0""" + assert k_fp4 % 2 == 0 + byte = (fp4_bits & 0x0F) | ((fp4_bits & 0x0F) << 4) + return torch.full((m, k_fp4 // 2), byte, dtype=torch.int8, device=device) def generate_mxf4_scale_factors(m, n, k_fp4, device='cuda', random_sf=False): @@ -71,18 +72,17 @@ def generate_mxf4_scale_factors(m, n, k_fp4, device='cuda', random_sf=False): def fp4_reference(a_packed, b_packed, m, n, sf_a=None, sf_b=None): - """CPU 端 E2M1 FP4 GEMM reference: C = A @ B^T, 支持 block-scaled SF。 - - Block-scaled MXF4: C[m,n] = sum_g SF_A[m,g] * SF_B[n,g] * dot(A_g, B_g) - 其中 g 是 VS=32 元素的组。sf_a/sf_b 为 float32 (pre-transform), 每组一个值。 + """CPU FP4 GEMM reference: C = A @ B^T with block-scaled SF. + Works for both int8 packing (2 FP4/byte, kPackedFP4) and int32 packing (8 FP4/int32). """ VS = 32 - a_cpu = a_packed.cpu().to(torch.int64) & 0xFFFFFFFF - b_cpu = b_packed.cpu().to(torch.int64) & 0xFFFFFFFF - bits_a = torch.stack([(a_cpu >> (i*4)) & 0xF for i in range(8)], dim=-1).reshape(m, -1) - bits_b = torch.stack([(b_cpu >> (i*4)) & 0xF for i in range(8)], dim=-1).reshape(n, -1) - a_float = E2M1_LUT[bits_a.long()] # [m, k_fp4] - b_float = E2M1_LUT[bits_b.long()] # [n, k_fp4] + a_cpu = a_packed.cpu().contiguous().view(torch.uint8).reshape(m, -1) + b_cpu = b_packed.cpu().contiguous().view(torch.uint8).reshape(n, -1) + # Each byte holds 2 FP4: low nibble = even idx, high nibble = odd idx + bits_a = torch.stack([(a_cpu & 0xF).long(), ((a_cpu >> 4) & 0xF).long()], dim=-1).reshape(m, -1) + bits_b = torch.stack([(b_cpu & 0xF).long(), ((b_cpu >> 4) & 0xF).long()], dim=-1).reshape(n, -1) + a_float = E2M1_LUT[bits_a] # [m, k_fp4] + b_float = E2M1_LUT[bits_b] # [n, k_fp4] if sf_a is None: return torch.matmul(a_float, b_float.T) @@ -107,12 +107,14 @@ def fp4_reference(a_packed, b_packed, m, n, sf_a=None, sf_b=None): return c -def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 128)): - """调用 FP4 GEMM kernel (复用 fp8_gemm_nt 入口, int32 dtype 触发 FP4 路径)""" - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) - d = torch.empty((m, n), device='cuda', dtype=torch.float32) - deep_gemm.fp8_gemm_nt((a_packed, sf_a), (b_packed, sf_b), d, c=None, - recipe=recipe, disable_ue8m0_cast=duc) +def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 32)): + """Call FP4 GEMM kernel via main's fp8_fp4_gemm_nt (kPackedFP4=int8 triggers my MXF4 path). + recipe gran_k=32 matches FP4 VS=32 SF granularity. + """ + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + deep_gemm.fp8_fp4_gemm_nt((a_packed, sf_a), (b_packed, sf_b), d, c=None, + recipe_a=(1, 32), recipe_b=(1, 32), + disable_ue8m0_cast=False) torch.cuda.synchronize() return d From 24eb9c280abf4897b20dc2bd1946d84c0b4d788b Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Sun, 31 May 2026 19:20:58 -0700 Subject: [PATCH 04/12] =?UTF-8?q?FP4:=20tests/test=5Ffp4.py=20=E2=80=94=20?= =?UTF-8?q?use=20FP32=20output=20(kernel's=20only=20working=20epilogue=20p?= =?UTF-8?q?ath)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of port "writes zeros" bug found: the kernel's epilogue store loop has only an `if constexpr (cd_dtype_t == float)` branch — no `else` for bf16 — so when d.dtype is bf16, the store body is a no-op and d stays at initial value (zero). This matches the existing memory `fp4_output_dtype.md` note that bf16 epilogue branch is a missing TODO. fea-fp4's run_kernel uses `torch.float32` for d; my earlier port test was generating with `torch.bfloat16` (main's typical FP8 output), which silently hit the no-op branch. Switching back to fp32 makes everything work. Verified PASS on fea-fp4-synced after this fix: test_constant (9/9 PASS, exact integer match for sf=1.0 case) test_random (9/9 PASS, max_diff=0.0000 vs CPU reference) test_random_sf (8/8 PASS, max_diff=0.0000 with varying SF) This closes the dense-FP4 port. m-grouped (contiguous + masked) paths still need a similar walkthrough and test_m_grouped_* validation. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- tests/test_fp4.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_fp4.py b/tests/test_fp4.py index 17b62e74ef..b13798bf70 100644 --- a/tests/test_fp4.py +++ b/tests/test_fp4.py @@ -110,8 +110,9 @@ def fp4_reference(a_packed, b_packed, m, n, sf_a=None, sf_b=None): def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 32)): """Call FP4 GEMM kernel via main's fp8_fp4_gemm_nt (kPackedFP4=int8 triggers my MXF4 path). recipe gran_k=32 matches FP4 VS=32 SF granularity. + NOTE: kernel hardcodes FP32 output — bf16 epilogue branch is a missing TODO. """ - d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + d = torch.empty((m, n), device='cuda', dtype=torch.float32) deep_gemm.fp8_fp4_gemm_nt((a_packed, sf_a), (b_packed, sf_b), d, c=None, recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False) From ae98142d4d82608343768e8cc9350be1f79b6bcb Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Sun, 31 May 2026 19:54:17 -0700 Subject: [PATCH 05/12] =?UTF-8?q?FP4:=20tests/test=5Ffp4.py=20=E2=80=94=20?= =?UTF-8?q?full=20suite=20passes=20on=20synced=20branch=20(10/10)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Complete the test adaptation to main's API surface: - Switch m_grouped helpers to int8 (kPackedFP4) packing for both 2D and 3D pack_fp4_random variants. - Rename remaining deep_gemm.m_grouped_fp8_gemm_* -> m_grouped_fp8_fp4_gemm_*. - Use recipe_a=(1,32) / recipe_b=(1,32) for FP4 SF granularity throughout. csrc/apis/gemm.hpp: relax `d.scalar_type() == kBFloat16` assertion in the two m-grouped dispatch paths to also allow `kFloat` when both A/B are kPackedFP4 (my MXF4 kernel hardcodes fp32 output; bf16 epilogue branch remains TODO). Final result: `python tests/test_fp4.py` -> ALL FP4 TESTS PASSED constant, random, sweep, asymmetric, uniform_sf, random_sf, multicast, m_grouped (contiguous), m_grouped_masked, trtllm_cmp. This completes the FP4 main-sync port from fea-fp4 -> fea-fp4-synced. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/apis/gemm.hpp | 8 +++++-- tests/test_fp4.py | 52 +++++++++++++++++++++++----------------------- 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index f3476a31cf..497056eee8 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -172,7 +172,9 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair(d); DG_HOST_ASSERT(m == m_ and n == n_ and k == k_); DG_HOST_ASSERT(n > 0 and k > 0 and num_groups > 0); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + // FP4xFP4 my kernel hardcodes fp32 output (bf16 epilogue is TODO). + const bool is_fp4_fp4 = (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or (is_fp4_fp4 and d.scalar_type() == torch::kFloat)); DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); // Layout checks @@ -256,7 +258,9 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair 0 and m > 0 and n > 0 and k > 0 and num_groups > 0); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + // FP4xFP4 my kernel hardcodes fp32 output (bf16 epilogue is TODO). + const bool is_fp4_fp4_masked = (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or (is_fp4_fp4_masked and d.scalar_type() == torch::kFloat)); DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); // D must be N-major diff --git a/tests/test_fp4.py b/tests/test_fp4.py index b13798bf70..4933587fff 100644 --- a/tests/test_fp4.py +++ b/tests/test_fp4.py @@ -121,12 +121,12 @@ def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 32)): def pack_fp4_random_3d(num_groups: int, n: int, k_fp4: int, device='cuda'): - """生成 [G, N, K_int32] FP4 packed tensor。""" - assert k_fp4 % 8 == 0 + """[G, N, K_byte] FP4 packed as int8 (kPackedFP4): 2 FP4 per byte.""" + assert k_fp4 % 2 == 0 raw = torch.randint(0, 16, (num_groups, n, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(num_groups, n, k_fp4 // 8, dtype=torch.int32, device=device) - for i in range(8): - packed += (raw[:, :, i::8].to(torch.int32) << (i * 4)) + packed = torch.zeros(num_groups, n, k_fp4 // 2, dtype=torch.int8, device=device) + packed |= (raw[:, :, 0::2].to(torch.int8) & 0x0F) + packed |= (raw[:, :, 1::2].to(torch.int8) & 0x0F) << 4 return packed @@ -162,25 +162,25 @@ def fp4_reference_grouped(a_packed, b_packed_grouped, m_indices, return c -def run_kernel_grouped(a_packed, b_packed, sf_a, sf_b, m_indices, m, n, recipe=(1, 1, 128)): - """调用 m_grouped_fp8_gemm_nt_contiguous 的 FP4 路径 (int32 dtype 触发)。""" - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) +def run_kernel_grouped(a_packed, b_packed, sf_a, sf_b, m_indices, m, n, recipe=None): + """Call m_grouped_fp8_fp4_gemm_nt_contiguous with FP4 path (kPackedFP4 int8 triggers).""" d = torch.empty((m, n), device='cuda', dtype=torch.float32) - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( (a_packed, sf_a), (b_packed, sf_b), d, m_indices, - recipe=recipe, disable_ue8m0_cast=duc, + recipe_a=(1, 32), recipe_b=(1, 32), + disable_ue8m0_cast=False, ) torch.cuda.synchronize() return d def pack_fp4_random_3d_ga(num_groups: int, max_m: int, k_fp4: int, device='cuda'): - """生成 [G, max_M, K_int32] FP4 packed tensor (A side for masked variant).""" - assert k_fp4 % 8 == 0 + """[G, max_M, K_byte] FP4 packed as int8 (kPackedFP4) for masked variant.""" + assert k_fp4 % 2 == 0 raw = torch.randint(0, 16, (num_groups, max_m, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(num_groups, max_m, k_fp4 // 8, dtype=torch.int32, device=device) - for i in range(8): - packed += (raw[:, :, i::8].to(torch.int32) << (i * 4)) + packed = torch.zeros(num_groups, max_m, k_fp4 // 2, dtype=torch.int8, device=device) + packed |= (raw[:, :, 0::2].to(torch.int8) & 0x0F) + packed |= (raw[:, :, 1::2].to(torch.int8) & 0x0F) << 4 return packed @@ -215,13 +215,13 @@ def fp4_reference_masked(a_3d, b_3d, masked_m_cpu, max_m, n, num_groups, def run_kernel_grouped_masked(a_3d, b_3d, sf_a_3d, sf_b_3d, masked_m, num_groups, - max_m, n, expected_m, recipe=(1, 1, 128)): - """调用 m_grouped_fp8_gemm_nt_masked FP4 路径。""" - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) + max_m, n, expected_m, recipe=None): + """Call m_grouped_fp8_fp4_gemm_nt_masked with FP4 path (kPackedFP4 int8).""" d = torch.empty((num_groups, max_m, n), device='cuda', dtype=torch.float32) - deep_gemm.m_grouped_fp8_gemm_nt_masked( + deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( (a_3d, sf_a_3d), (b_3d, sf_b_3d), d, masked_m, expected_m, - recipe=recipe, disable_ue8m0_cast=duc, + recipe_a=(1, 32), recipe_b=(1, 32), + disable_ue8m0_cast=False, ) torch.cuda.synchronize() return d @@ -492,9 +492,9 @@ def test_m_grouped_contiguous(): d = torch.empty((m_total, n), device='cuda', dtype=torch.float32) def fn(): - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( (a, sf_a), (b, sf_b), d, m_indices, - recipe=(1, 1, 128), disable_ue8m0_cast=duc, + recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, ) t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) @@ -585,9 +585,9 @@ def test_m_grouped_trtllm_comparable(): # rows [start + rows_per_group : start + aligned_per_group) stay -1 # Run kernel - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( (a, sf_a), (b, sf_b), d, m_indices, - recipe=(1, 1, 128), disable_ue8m0_cast=duc, + recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, ) torch.cuda.synchronize() d_clean = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) @@ -605,9 +605,9 @@ def test_m_grouped_trtllm_comparable(): # Perf def fn(): - deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( (a, sf_a), (b, sf_b), d, m_indices, - recipe=(1, 1, 128), disable_ue8m0_cast=duc, + recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, ) t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) From dfdc5aeb3b3f3f2071690410332828512e95c888 Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Sun, 31 May 2026 22:47:46 -0700 Subject: [PATCH 06/12] =?UTF-8?q?FP4:=20port=20fea-fp4's=20full=20get=5Fbe?= =?UTF-8?q?st=5Ffp4=5Fconfig=20heuristic=20=E2=80=94=20perf=20parity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the v0 hardcoded layout (block_n=112, cluster=1, swap_ab=false) with a port of fea-fp4's wave-aware FP4 heuristic. Now returns main's nested Layout struct instead of fea-fp4's flat GemmConfig. pick_fp4_layout(gemm_type, m, n, k, num_groups, num_sms, expected_m_per_group): - block_m = 128 (fixed, MXF4 UMMA_M) - block_k = 128 bytes (= 32 int32 = 256 FP4 per K block) - block_n picked by wave count + composite score = est_stages^2 * bn; 2-epi-stage tiebreak when waves >= 2 (TMEM double-buffer benefit needs >= 2 tiles per SM). - cluster_m=2, cluster_n=1 (B-multicast) when m >= 512, divisible, and gemm type is Normal / KGroupedContiguous. m-grouped stays cluster=1. - swap_ab=true for MGroupedContiguous when useful_m_per_group < BLOCK_M (sparse MoE benefit; forces block_n=128 + cluster=1 per kernel asserts). Wrapper changes: - 3 entry points now call pick_fp4_layout instead of get_best_config + override. - m-grouped contiguous derives useful_per_group from (m_indices >= 0).sum() / G to drive swap_ab gating accurately (matches fea-fp4 wrapper behavior). - m-grouped masked passes expected_m as per-group hint (swap_ab itself disabled for masked per fea-fp4 v0). - recompute_stages_for_fp4: fix CD smem formula to account for swap_ab path (STORE_BLOCK_M=16 * STORE_BLOCK_N=block_n * sizeof(cd_dtype=fp32)) — was using non-swap formula and underestimating CD smem, causing IMA on the swap_ab path when num_stages was over-budgeted. Perf vs fea-fp4 (m-grouped contiguous prod shapes): G=4 m=8192 N=4096 K=7168: 4290 -> 4255 TFLOPS (-0.8%) G=4 m=8192 N=7168 K=2048: 2879 -> 2874 TFLOPS (-0.2%) G=8 m=4096 N=4096 K=7168: 4270 -> 4249 TFLOPS (-0.5%) G=8 m=4096 N=7168 K=2048: 2877 -> 2873 TFLOPS (-0.1%) Parity within noise. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 192 ++++++++++++++---- 1 file changed, 148 insertions(+), 44 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index d70ed59f28..1d3eaea0c1 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -125,8 +125,17 @@ static std::pair recompute_stages_for_fp4(const GemmConfig& config, in + sf_block_n * sf_packed_k_per_stage * 4; // Fixed extras (CD smem + barriers + tmem_ptr, conservative estimate) - // CD (non-swap): store_block_m * swizzle_cd_mode * 2 stages - const int cd_size = config.storage_config.store_block_m * config.storage_config.swizzle_cd_mode * 2; + // CD smem (must match kernel's SMEM_CD_SIZE_PER_STAGE × kNumTMAStoreStages=2): + // swap_ab: STORE_BLOCK_M(=16) * STORE_BLOCK_N(=block_n) * sizeof(cd_dtype) + // non-swap: STORE_BLOCK_M(=block_m) * kSwizzleCDMode + // cd_dtype size: assume fp32 (4 bytes); bf16 epilogue is TODO so we can't reach it here. + int cd_size; + if (config.layout.swap_ab) { + constexpr int store_block_m_swap = 16; + cd_size = store_block_m_swap * block_n * 4 * 2; // fp32, 2 stages + } else { + cd_size = config.storage_config.store_block_m * config.storage_config.swizzle_cd_mode * 2; + } const int barriers = 12 * 8 * 4 + 4 * 8 * 2 + 8; // ~ 416 bytes max const int tmem_ptr = 4; const int fixed_extras = cd_size + barriers + tmem_ptr; @@ -138,6 +147,106 @@ static std::pair recompute_stages_for_fp4(const GemmConfig& config, in return {new_num_stages, new_smem_size}; } +// Pick the FP4-optimal Layout (block_m=128 fixed, wave-aware block_n, B-multicast for +// large M, swap_ab for sparse m-grouped). Mirrors fea-fp4's `get_best_fp4_config` but +// returns main's nested Layout struct. +// +// Constants (FP4 MXF4 path): +// block_m = 128 (UMMA_M for MXF4) +// block_k = 128 bytes (= 32 int32 = 256 FP4 per K block) +// sf_pk = 2 (SF_PACKED_K_PER_STAGE; from block_k_int32/16) +// +// block_n is chosen by wave count + composite score = est_stages^2 * bn, +// with 2-epi-stage tiebreak for multi-wave cases. swap_ab is enabled for +// MGroupedContiguous when expected_m_per_group < BLOCK_M (sparse MoE). +// +// expected_m_per_group: useful per-group row count (m_indices >= 0 sum/G). +// Pass INT_MAX to disable swap_ab gating (default for Normal / non-grouped). +static Layout pick_fp4_layout(const GemmType& gemm_type, + const int& m, const int& n, const int& k, + const int& num_groups, const int& num_sms, + const int& expected_m_per_group = INT_MAX) { + constexpr int block_m = 128; + constexpr int block_k_bytes = 128; // 32 int32; matches main's block_k convention + constexpr int sf_pk = 2; // block_k_int32 / 16 = 32/16 + constexpr int sf_block_m_cols = (128 / 32) * sf_pk; // 8 + constexpr int smem_capacity = 232448; + + // BLOCK_N legality (mirrors fea-fp4's is_fp4_block_n_legal) + auto is_legal = [&](int bn) { + if (bn % 16 != 0 || bn < 16 || bn > 256) return false; + const int sf_block_n = (bn + 127) / 128 * 128; + const int sf_block_n_cols = (sf_block_n / 32) * sf_pk; + // TMEM minimum: 1 epi stage must fit + if ((1 * bn + sf_block_m_cols + sf_block_n_cols) > 512) return false; + return true; + }; + auto fits_2_epi = [&](int bn) { + const int sf_block_n = (bn + 127) / 128 * 128; + const int sf_block_n_cols = (sf_block_n / 32) * sf_pk; + return (2 * bn + sf_block_m_cols + sf_block_n_cols) <= 512; + }; + + auto num_blocks = [&](int bn) { return ceil_div(m, block_m) * ceil_div(n, bn) * num_groups; }; + auto num_waves = [&](int bn) { return ceil_div(num_blocks(bn), num_sms); }; + + int best_bn = 0, best_waves = 0, best_score = 0; + bool best_fits_2epi = false; + for (int bn = 16; bn <= 256; bn += 16) { + if (!is_legal(bn)) continue; + const int waves = num_waves(bn); + + // est_stages (conservative, no multicast): per_stage in bytes + const int sf_block_n = (bn + 127) / 128 * 128; + const int per_stage = block_m * block_k_bytes + bn * block_k_bytes + + 128 * sf_pk * 4 + sf_block_n * sf_pk * 4; + const int avail = smem_capacity - 32768 - 200; + const int est_stages = std::min(12, std::max(1, avail / per_stage)); + const int score = est_stages * est_stages * bn; + const bool cur_fits_2epi = fits_2_epi(bn); + const bool consider_epi = waves >= 2; // 2-epi benefit only when SM processes >= 2 tiles + + bool success = false; + if (best_bn == 0 || waves < best_waves) { + success = true; + } else if (waves == best_waves && bn <= n) { + if (consider_epi && cur_fits_2epi && !best_fits_2epi) success = true; + else if ((!consider_epi || cur_fits_2epi == best_fits_2epi) && score > best_score) success = true; + } + if (success) { + best_bn = bn; best_waves = waves; best_score = score; best_fits_2epi = cur_fits_2epi; + } + } + DG_HOST_ASSERT(best_bn > 0); + + // Env override (for benchmarking) + if (const auto env_bn = get_env("DG_FP4_BLOCK_N"); env_bn > 0) { + DG_HOST_ASSERT(env_bn % 16 == 0 && env_bn <= 256); + best_bn = env_bn; + } + + // Multicast: B-multicast (cluster_m=2, cluster_n=1) when M >= 512 for Normal / KGrouped. + // For m-grouped, fea-fp4 keeps cluster=1 (m_indices iteration breaks multi-CTA M-distribution). + int cluster_m = 1, cluster_n = 1; + const bool can_multicast = (m >= 512) + && (ceil_div(m, block_m) % 2 == 0) + && (num_sms % 2 == 0) + && (gemm_type == GemmType::Normal || gemm_type == GemmType::KGroupedContiguous); + if (can_multicast) { + cluster_m = 2; + } + + // swap_ab for sparse m-grouped contiguous (forces block_n=128 + cluster=1) + bool swap_ab = false; + if (gemm_type == GemmType::MGroupedContiguous && expected_m_per_group < block_m) { + swap_ab = true; + best_bn = 128; + cluster_m = cluster_n = 1; + } + + return Layout{swap_ab, block_m, best_bn, block_k_bytes, cluster_m, cluster_n}; +} + // Build a GemmDesc forcing a_dtype=b_dtype=kPackedFP4 (since this is the FP4xFP4 wrapper). static GemmDesc make_fp4_desc(GemmType gemm_type, int m, int n, int k, int num_groups, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, @@ -172,20 +281,16 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa const auto desc = make_fp4_desc(GemmType::Normal, m, n, k, 1, major_a, major_b, d.scalar_type(), c.has_value(), compiled_dims); - auto config = get_best_config(desc); - // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. - // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). - // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) - // path appears untested upstream. - config.layout.block_m = 128; - config.layout.block_n = 112; - config.layout.swap_ab = false; - config.layout.cluster_m = 1; - config.layout.cluster_n = 1; - config.launch_config.num_sms_per_cluster = 1; - config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); - config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); - const auto [new_stages, new_smem] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + const auto layout = pick_fp4_layout(GemmType::Normal, m, n, k, 1, + device_runtime->get_num_sms()); + auto config = GemmConfig{ + .layout = layout, + .storage_config = SM100ArchSpec::get_storage_config(desc, layout), + .pipeline_config = {}, + .launch_config = SM100ArchSpec::get_launch_config(desc, layout), + }; + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, layout, config.storage_config); + const auto [new_stages, new_smem] = recompute_stages_for_fp4(config, layout.block_m, layout.block_n, k); config.pipeline_config.num_stages = new_stages; config.pipeline_config.smem_size = new_smem; @@ -262,20 +367,20 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con const auto desc = make_fp4_desc(GemmType::MGroupedContiguous, m, n, k, num_groups, major_a, major_b, d.scalar_type(), false, compiled_dims); - auto config = get_best_config(desc); - // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. - // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). - // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) - // path appears untested upstream. - config.layout.block_m = 128; - config.layout.block_n = 112; - config.layout.swap_ab = false; - config.layout.cluster_m = 1; - config.layout.cluster_n = 1; - config.launch_config.num_sms_per_cluster = 1; - config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); - config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); - const auto [new_stages_g, new_smem_g] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + // For m-grouped contiguous, derive useful-per-group row count from m_indices + // (count of valid rows / num_groups). Drives swap_ab gating for sparse MoE. + const auto useful_m = (grouped_layout >= 0).sum().item(); + const int useful_per_group = num_groups > 0 ? static_cast(useful_m) / num_groups : 0; + const auto layout = pick_fp4_layout(GemmType::MGroupedContiguous, m, n, k, num_groups, + device_runtime->get_num_sms(), useful_per_group); + auto config = GemmConfig{ + .layout = layout, + .storage_config = SM100ArchSpec::get_storage_config(desc, layout), + .pipeline_config = {}, + .launch_config = SM100ArchSpec::get_launch_config(desc, layout), + }; + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, layout, config.storage_config); + const auto [new_stages_g, new_smem_g] = recompute_stages_for_fp4(config, layout.block_m, layout.block_n, k); config.pipeline_config.num_stages = new_stages_g; config.pipeline_config.smem_size = new_smem_g; @@ -337,20 +442,19 @@ static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const t major_a, major_b, d.scalar_type(), false, compiled_dims, /*expected_m=*/expected_m, /*expected_num_groups=*/num_groups); - auto config = get_best_config(desc); - // v0: match fea-fp4's tested config block_m=128, block_n=112, single-CTA. - // block_n=112 makes kSwizzleCDMode=32 (since 112*2=224, gcd(224, 128/64)=32). - // This is what fea-fp4's tests actually exercised; block_n=128 (kSwizzleCDMode=128) - // path appears untested upstream. - config.layout.block_m = 128; - config.layout.block_n = 112; - config.layout.swap_ab = false; - config.layout.cluster_m = 1; - config.layout.cluster_n = 1; - config.launch_config.num_sms_per_cluster = 1; - config.storage_config = SM100ArchSpec::get_storage_config(desc, config.layout); - config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, config.layout, config.storage_config); - const auto [new_stages_mk, new_smem_mk] = recompute_stages_for_fp4(config, config.layout.block_m, config.layout.block_n, k); + // m-grouped masked: pass expected_m as per-group hint (fea-fp4 invariant); + // swap_ab heuristic itself only activates for MGroupedContiguous per fea-fp4 v0 + // (masked + swap_ab had NaN issues unresolved upstream). + const auto layout = pick_fp4_layout(GemmType::MGroupedMasked, m, n, k, num_groups, + device_runtime->get_num_sms(), expected_m); + auto config = GemmConfig{ + .layout = layout, + .storage_config = SM100ArchSpec::get_storage_config(desc, layout), + .pipeline_config = {}, + .launch_config = SM100ArchSpec::get_launch_config(desc, layout), + }; + config.pipeline_config = SM100ArchSpec::get_pipeline_config(desc, layout, config.storage_config); + const auto [new_stages_mk, new_smem_mk] = recompute_stages_for_fp4(config, layout.block_m, layout.block_n, k); config.pipeline_config.num_stages = new_stages_mk; config.pipeline_config.smem_size = new_smem_mk; From 31aef98fe7d293455ea807888f08e3a2e37f19cd Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Sun, 31 May 2026 23:52:43 -0700 Subject: [PATCH 07/12] FP4: rewrite tests/test_fp4.py to mirror test_fp8_fp4.py style The previous test_fp4.py (752 lines) used custom pack_fp4_random + CPU LUT reference + bespoke PASS/FAIL printing, predating the QuantConfig / generators infrastructure. This rewrites it to match the official test_fp8_fp4.py: - Reuse generate_normal / generate_m_grouped_contiguous / generate_m_grouped_masked from tests/generators.py instead of pack_fp4_random + per-test fp4_reference. - QuantConfig((32, 32, True, True)) drives both quantization and reference. - assert + max_diff() threshold instead of custom PASS/FAIL prints. - '> Perf (m=..., n=..., k=..., 1D1D, layout=NT, FP32): X us | Y TFLOPS | Z GB/s' output line matches main's official format exactly. - Three test functions named to align with test_fp8_fp4.py: test_gemm, test_m_grouped_gemm_contiguous, test_m_grouped_gemm_masked. - Entry block: torch.manual_seed(0); random.seed(0); print('Library path:'); test_*(). Shape coverage: Dense: m in [128, 4096] x 7 (n, k) production combos from main's nk_list. Contig: (num_groups, m_per_group) in [(4, 8192), (8, 4096)] x 4 (n, k). Masked: (num_groups, expected_m) in [(1, 1024), (2, 512), (4, 256)] x 2 (n, k). All shapes PASS with diff < QuantConfig.max_diff() = 0.02. Perf parity check (same kernel, same heuristic, only test wrapper changed): 4, 8192, 4096, 7168: 4255 -> 4306 TFLOPS (+1.2%, noise) 4, 8192, 7168, 2048: 2874 -> 2891 TFLOPS (+0.6%, noise) 8, 4096, 4096, 7168: 4249 -> 4305 TFLOPS (+1.3%, noise) 8, 4096, 7168, 2048: 2873 -> 2849 TFLOPS (-0.8%, noise) Note: kernel hardcodes fp32 output (bf16 epilogue is a remaining TODO); the test sets out_dtype=torch.float and casts ref_d from generators' bf16 default for the diff check. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- tests/test_fp4.py | 868 +++++++--------------------------------------- 1 file changed, 126 insertions(+), 742 deletions(-) diff --git a/tests/test_fp4.py b/tests/test_fp4.py index 4933587fff..45b49536c1 100644 --- a/tests/test_fp4.py +++ b/tests/test_fp4.py @@ -1,752 +1,136 @@ -""" -FP4 (E2M1) GEMM correctness test for SM100 MXF4 block-scaled kernel. - -Usage: - python tests/test_fp4.py -""" - -import torch import random -import deep_gemm -from deep_gemm.testing import bench_kineto, count_bytes -from generators import KernelType, get_ue8m0_usage - -# ============================================================ -# E2M1 FP4 查找表 -# ============================================================ -E2M1_LUT = torch.tensor([ - 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, # bits 0-7 (S=0) - -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0 # bits 8-15 (S=1) -], dtype=torch.float32) - -# ============================================================ -# 工具函数 -# ============================================================ - -def pack_fp4_random(m: int, k_fp4: int, device='cuda'): - """Generate random E2M1 FP4 values, pack as int8 (kPackedFP4, 2 FP4 per byte). - Memory layout identical to int32 packing (8 FP4 per int32) when viewed as int32. - """ - assert k_fp4 % 2 == 0 - raw = torch.randint(0, 16, (m, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(m, k_fp4 // 2, dtype=torch.int8, device=device) - # Byte layout: low nibble = even FP4, high nibble = odd FP4 (matches per_token_cast_to_fp4) - packed |= (raw[:, 0::2].to(torch.int8) & 0x0F) - packed |= (raw[:, 1::2].to(torch.int8) & 0x0F) << 4 - return packed - - -def pack_fp4_constant(m: int, k_fp4: int, fp4_bits: int = 0x2, device='cuda'): - """Constant FP4 packed as int8. fp4_bits=0x2 -> E2M1 1.0""" - assert k_fp4 % 2 == 0 - byte = (fp4_bits & 0x0F) | ((fp4_bits & 0x0F) << 4) - return torch.full((m, k_fp4 // 2), byte, dtype=torch.int8, device=device) - - -def generate_mxf4_scale_factors(m, n, k_fp4, device='cuda', random_sf=False): - """为 MXF4 (VS=32) 生成 scale factor (float32 格式)。 - - Host 端 C++ API (transform_sf_into_required_layout) 会在调用 kernel 之前 - 将 float32 SF 转换为 packed UE8M0 int32: - 1. 从 float32 IEEE 754 中提取 8-bit 指数字段 (bitwise_right_shift(23)) - 2. 每 4 个 UE8M0 打包为 1 个 int32 - 3. 转置为 MN-major + TMA 对齐 - 然后 kernel 通过 TMA 加载到 SMEM, 经 UTCCP 写入 TMEM 供 MMA 使用。 - - UE8M0: value = 2^(exp - 127)。每 VS=32 个 FP4 元素共享一个 SF。 - - Args: - random_sf: 如果 True, 生成随机的 2 的幂次 SF (0.25, 0.5, 1.0, 2.0, 4.0) - """ - VS = 32 - sf_k = ((k_fp4 // VS + 3) // 4) * 4 - if random_sf: - powers = torch.randint(-2, 3, (m, sf_k), device=device).float() - sf_a = torch.pow(2.0, powers) - powers = torch.randint(-2, 3, (n, sf_k), device=device).float() - sf_b = torch.pow(2.0, powers) - else: - sf_a = torch.ones((m, sf_k), dtype=torch.float32, device=device) - sf_b = torch.ones((n, sf_k), dtype=torch.float32, device=device) - return sf_a, sf_b - - -def fp4_reference(a_packed, b_packed, m, n, sf_a=None, sf_b=None): - """CPU FP4 GEMM reference: C = A @ B^T with block-scaled SF. - Works for both int8 packing (2 FP4/byte, kPackedFP4) and int32 packing (8 FP4/int32). - """ - VS = 32 - a_cpu = a_packed.cpu().contiguous().view(torch.uint8).reshape(m, -1) - b_cpu = b_packed.cpu().contiguous().view(torch.uint8).reshape(n, -1) - # Each byte holds 2 FP4: low nibble = even idx, high nibble = odd idx - bits_a = torch.stack([(a_cpu & 0xF).long(), ((a_cpu >> 4) & 0xF).long()], dim=-1).reshape(m, -1) - bits_b = torch.stack([(b_cpu & 0xF).long(), ((b_cpu >> 4) & 0xF).long()], dim=-1).reshape(n, -1) - a_float = E2M1_LUT[bits_a] # [m, k_fp4] - b_float = E2M1_LUT[bits_b] # [n, k_fp4] - - if sf_a is None: - return torch.matmul(a_float, b_float.T) - - k_fp4 = a_float.shape[1] - sf_a_cpu = sf_a.cpu().float() - sf_b_cpu = sf_b.cpu().float() - c = torch.zeros(m, n, dtype=torch.float32) - num_groups = k_fp4 // VS - for g in range(num_groups): - k_start, k_end = g * VS, (g + 1) * VS - a_g = a_float[:, k_start:k_end] - b_g = b_float[:, k_start:k_end] - sf_col = g - if sf_col < sf_a_cpu.shape[1]: - sfa_g = sf_a_cpu[:, sf_col].unsqueeze(1) - sfb_g = sf_b_cpu[:, sf_col].unsqueeze(1) - else: - sfa_g = torch.ones(m, 1) - sfb_g = torch.ones(n, 1) - c += (sfa_g * sfb_g.T) * torch.matmul(a_g, b_g.T) - return c - - -def run_kernel(a_packed, b_packed, sf_a, sf_b, m, n, recipe=(1, 1, 32)): - """Call FP4 GEMM kernel via main's fp8_fp4_gemm_nt (kPackedFP4=int8 triggers my MXF4 path). - recipe gran_k=32 matches FP4 VS=32 SF granularity. - NOTE: kernel hardcodes FP32 output — bf16 epilogue branch is a missing TODO. - """ - d = torch.empty((m, n), device='cuda', dtype=torch.float32) - deep_gemm.fp8_fp4_gemm_nt((a_packed, sf_a), (b_packed, sf_b), d, c=None, - recipe_a=(1, 32), recipe_b=(1, 32), - disable_ue8m0_cast=False) - torch.cuda.synchronize() - return d - - -def pack_fp4_random_3d(num_groups: int, n: int, k_fp4: int, device='cuda'): - """[G, N, K_byte] FP4 packed as int8 (kPackedFP4): 2 FP4 per byte.""" - assert k_fp4 % 2 == 0 - raw = torch.randint(0, 16, (num_groups, n, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(num_groups, n, k_fp4 // 2, dtype=torch.int8, device=device) - packed |= (raw[:, :, 0::2].to(torch.int8) & 0x0F) - packed |= (raw[:, :, 1::2].to(torch.int8) & 0x0F) << 4 - return packed - - -def generate_mxf4_sf_3d(num_groups: int, n: int, k_fp4: int, device='cuda', random_sf=False): - """生成 SFB [G, N, sf_k] for grouped contiguous.""" - VS = 32 - sf_k = ((k_fp4 // VS + 3) // 4) * 4 - if random_sf: - powers = torch.randint(-2, 3, (num_groups, n, sf_k), device=device).float() - return torch.pow(2.0, powers) - return torch.ones((num_groups, n, sf_k), dtype=torch.float32, device=device) - - -def fp4_reference_grouped(a_packed, b_packed_grouped, m_indices, - n: int, num_groups: int, - sf_a=None, sf_b_grouped=None): - """Per-row grouped FP4 reference: D[i] = A[i] @ B[m_indices[i]].T (with SF scaling). - - Padding rows (m_indices == -1) get D[i] = 0. - """ - m = a_packed.shape[0] - c = torch.zeros(m, n, dtype=torch.float32) - m_idx_cpu = m_indices.cpu() - for g in range(num_groups): - rows = (m_idx_cpu == g).nonzero(as_tuple=True)[0] - if rows.numel() == 0: - continue - a_g = a_packed[rows] - sf_a_g = sf_a[rows] if sf_a is not None else None - sf_b_g = sf_b_grouped[g] if sf_b_grouped is not None else None - c_g = fp4_reference(a_g, b_packed_grouped[g], rows.numel(), n, sf_a_g, sf_b_g) - c[rows] = c_g - return c - - -def run_kernel_grouped(a_packed, b_packed, sf_a, sf_b, m_indices, m, n, recipe=None): - """Call m_grouped_fp8_fp4_gemm_nt_contiguous with FP4 path (kPackedFP4 int8 triggers).""" - d = torch.empty((m, n), device='cuda', dtype=torch.float32) - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - (a_packed, sf_a), (b_packed, sf_b), d, m_indices, - recipe_a=(1, 32), recipe_b=(1, 32), - disable_ue8m0_cast=False, - ) - torch.cuda.synchronize() - return d - - -def pack_fp4_random_3d_ga(num_groups: int, max_m: int, k_fp4: int, device='cuda'): - """[G, max_M, K_byte] FP4 packed as int8 (kPackedFP4) for masked variant.""" - assert k_fp4 % 2 == 0 - raw = torch.randint(0, 16, (num_groups, max_m, k_fp4), dtype=torch.uint8, device=device) - packed = torch.zeros(num_groups, max_m, k_fp4 // 2, dtype=torch.int8, device=device) - packed |= (raw[:, :, 0::2].to(torch.int8) & 0x0F) - packed |= (raw[:, :, 1::2].to(torch.int8) & 0x0F) << 4 - return packed - - -def generate_mxf4_sfa_3d(num_groups: int, max_m: int, k_fp4: int, device='cuda', random_sf=False): - """生成 SFA [G, max_M, sf_k] for grouped masked.""" - VS = 32 - sf_k = ((k_fp4 // VS + 3) // 4) * 4 - if random_sf: - powers = torch.randint(-2, 3, (num_groups, max_m, sf_k), device=device).float() - return torch.pow(2.0, powers) - return torch.ones((num_groups, max_m, sf_k), dtype=torch.float32, device=device) - - -def fp4_reference_masked(a_3d, b_3d, masked_m_cpu, max_m, n, num_groups, - sf_a_3d=None, sf_b_3d=None): - """Reference for masked variant: D[g, :masked_m[g], :] = A[g, :masked_m[g], :] @ B[g, :, :].T. - - Padding rows beyond masked_m[g] in D are not checked (kernel may leave any value). - Returns d_ref of shape [G, max_M, N] with valid rows filled, padding rows zeroed. - """ - d_ref = torch.zeros(num_groups, max_m, n, dtype=torch.float32) - for g in range(num_groups): - mg = int(masked_m_cpu[g].item()) - if mg == 0: - continue - a_g = a_3d[g, :mg] - sf_a_g = sf_a_3d[g, :mg] if sf_a_3d is not None else None - sf_b_g = sf_b_3d[g] if sf_b_3d is not None else None - c_g = fp4_reference(a_g, b_3d[g], mg, n, sf_a_g, sf_b_g) - d_ref[g, :mg] = c_g - return d_ref - - -def run_kernel_grouped_masked(a_3d, b_3d, sf_a_3d, sf_b_3d, masked_m, num_groups, - max_m, n, expected_m, recipe=None): - """Call m_grouped_fp8_fp4_gemm_nt_masked with FP4 path (kPackedFP4 int8).""" - d = torch.empty((num_groups, max_m, n), device='cuda', dtype=torch.float32) - deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( - (a_3d, sf_a_3d), (b_3d, sf_b_3d), d, masked_m, expected_m, - recipe_a=(1, 32), recipe_b=(1, 32), - disable_ue8m0_cast=False, - ) - torch.cuda.synchronize() - return d - - -# ============================================================ -# 测试用例 -# ============================================================ - -def test_constant(): - """全 1.0 常量测试: C[i,j] = K (因为 1.0 * 1.0 * K 个元素)""" - print('Test: constant values (all E2M1 1.0)') - configs = [ - # (M, N, K_fp4) — K 是 FP4 元素个数 - # 单 stage (K <= 256) - (32, 64, 256), - (128, 128, 128), - (128, 256, 256), - # 多 stage (K > 256) - (128, 128, 512), - (256, 256, 512), - (128, 128, 1024), - (128, 256, 1024), - # 大 M (multi-wave: BLOCK_M=128, 所以 M>128 需要多个 wave) - (256, 128, 256), - (256, 256, 1024), - ] - all_pass = True - for m, n, k in configs: - a = pack_fp4_constant(m, k, fp4_bits=0x2) - b = pack_fp4_constant(n, k, fp4_bits=0x2) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) - d = run_kernel(a, b, sf_a, sf_b, m, n) - expected = float(k) - ok = (d.cpu() == expected).all().item() - if not ok: - all_pass = False - print(f' M={m:4d} N={n:4d} K={k:4d}: expected={expected:.0f} got={d.cpu()[0,0].item():.0f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_random(): - """随机数据测试 (SF=1.0): 对比 GPU kernel 与 CPU reference""" - print('Test: random data (vs CPU reference, SF=1.0)') - configs = [ - (32, 64, 256), - (128, 128, 128), - (128, 256, 256), - # 多 stage - (128, 128, 512), - (256, 256, 512), - (128, 128, 1024), - # 大 M - (256, 128, 256), - (256, 256, 1024), - # 较大 N - (128, 512, 256), - ] - all_pass = True - for m, n, k in configs: - a = pack_fp4_random(m, k) - b = pack_fp4_random(n, k) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) - d = run_kernel(a, b, sf_a, sf_b, m, n) - ref = fp4_reference(a, b, m, n, sf_a, sf_b) - max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() - ok = max_diff < 1.0 - if not ok: - all_pass = False - print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_value_sweep(): - """不同 FP4 值测试: 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0""" - print('Test: FP4 value sweep') - m, n, k = 128, 128, 256 - all_pass = True - for bits, val in [(0x1, 0.5), (0x2, 1.0), (0x3, 1.5), (0x4, 2.0), - (0x5, 3.0), (0x6, 4.0), (0x7, 6.0)]: - a = pack_fp4_constant(m, k, fp4_bits=bits) - b = pack_fp4_constant(n, k, fp4_bits=bits) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) - d = run_kernel(a, b, sf_a, sf_b, m, n) - expected = float(k) * val * val - actual = d.cpu()[0, 0].item() - ok = abs(actual - expected) < 1.0 - if not ok: - all_pass = False - print(f' FP4={val:4.1f} (bits=0x{bits:X}): expected={expected:.1f} got={actual:.1f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_uniform_sf(): - """Uniform SF 测试: 验证 UTCCP SF->TMEM 路径在不同 SF 值下正确""" - print('Test: uniform scale factors (UTCCP path)') - configs = [ - (128, 256, 256), - (256, 128, 256), - (128, 128, 512), - ] - all_pass = True - for m, n, k in configs: - for sf_val in [0.25, 0.5, 1.0, 2.0, 4.0]: - a = pack_fp4_random(m, k) - b = pack_fp4_random(n, k) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) - sf_a.fill_(sf_val) - sf_b.fill_(sf_val) - d = run_kernel(a, b, sf_a, sf_b, m, n) - # uniform SF: kernel result = unscaled_result * sf_a * sf_b - ref = fp4_reference(a, b, m, n, sf_a, sf_b) - max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() - ok = max_diff < 1.0 - if not ok: - all_pass = False - print(f' M={m:4d} N={n:4d} K={k:4d} SF={sf_val:5.2f}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_asymmetric_values(): - """A 和 B 使用不同 FP4 值""" - print('Test: asymmetric A/B values') - m, n, k = 128, 128, 256 - all_pass = True - cases = [ - (0x2, 0x4, 1.0, 2.0), # A=1.0, B=2.0 - (0x1, 0x6, 0.5, 4.0), # A=0.5, B=4.0 - (0x4, 0x1, 2.0, 0.5), # A=2.0, B=0.5 - ] - for bits_a, bits_b, val_a, val_b in cases: - a = pack_fp4_constant(m, k, fp4_bits=bits_a) - b = pack_fp4_constant(n, k, fp4_bits=bits_b) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k) - d = run_kernel(a, b, sf_a, sf_b, m, n) - expected = float(k) * val_a * val_b - actual = d.cpu()[0, 0].item() - ok = abs(actual - expected) < 1.0 - if not ok: - all_pass = False - print(f' A={val_a}, B={val_b}: expected={expected:.1f} got={actual:.1f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_random_sf(): - """随机数据 + 随机 per-group SF (powers of 2)""" - print('Test: random data + random scale factors') - configs = [ - (128, 128, 256), - (128, 256, 256), - (256, 256, 512), - (128, 128, 512), - (128, 128, 1024), - (256, 128, 256), - (256, 256, 1024), - (128, 256, 1024), - ] - all_pass = True - for m, n, k in configs: - a = pack_fp4_random(m, k) - b = pack_fp4_random(n, k) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k, random_sf=True) - d = run_kernel(a, b, sf_a, sf_b, m, n) - ref = fp4_reference(a, b, m, n, sf_a, sf_b) - max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() - ok = max_diff < 1.0 - if not ok: - all_pass = False - print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') - return all_pass - - -def test_m_grouped_contiguous(): - """M-grouped contiguous FP4 GEMM (MoE forward shape). - - A=[M_total, K] @ B=[G, N, K].T → D=[M_total, N] - m_indices[M_total] selects which group each row belongs to (-1 = padding). - """ - print('Test: m-grouped contiguous (MoE-style)') - BLOCK_M = 128 - - # Small/debug shapes — fast CPU LUT reference, exercise padding-row case. - debug_configs = [ - (2, 128, 128, 256), - (4, 128, 128, 256), - (4, 128, 256, 512), - (4, 256, 128, 1024), - (8, 128, 256, 512), - # Uneven actual M per group, padded to BLOCK_M (exercises padding rows). - (4, 90, 128, 256), - (4, 200, 256, 512), - # Extra N×K variety to exercise more BLOCK_N choices and SF transform paths. - (4, 128, 512, 256), - (4, 128, 768, 256), - ] - - # Production MoE shapes — mirror DeepGEMM official FP8 grouped (generators.py:105). - # CPU LUT reference takes a few seconds per shape at this size; OK for nightly. - prod_configs = [ - (4, 8192, 4096, 7168), # EP4, MoE up-projection - (4, 8192, 7168, 2048), # EP4, MoE down-projection - (8, 4096, 4096, 7168), # EP8, MoE up-projection - (8, 4096, 7168, 2048), # EP8, MoE down-projection - # Extra (n, k) variants from FP8 enumerate_normal for N/K coverage: - (4, 8192, 24576, 1536), - (4, 8192, 32768, 512), - ] - - # Mirror FP8 masked-test pattern: multiple random-data iterations per shape to - # catch flaky bugs that only show up with certain RNG seeds. Debug shapes get - # fewer iters (kept fast); prod shapes get 3 iters each. - NUM_ITERS = {'debug': 2, 'prod': 3} - - all_pass = True - for label, configs in [('debug', debug_configs), ('prod', prod_configs)]: - for num_groups, m_per_group, n, k in configs: - aligned_m = ((m_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M - m_total = aligned_m * num_groups - - worst_diff = 0.0 - for _ in range(NUM_ITERS[label]): - # Build A, B, m_indices fresh per iteration - a = pack_fp4_random(m_total, k) - b = pack_fp4_random_3d(num_groups, n, k) - sf_a, _ = generate_mxf4_scale_factors(m_total, n, k, random_sf=True) - sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) - m_indices = torch.empty(m_total, dtype=torch.int32, device='cuda') - for g in range(num_groups): - start = g * aligned_m - actual_end = start + m_per_group - aligned_end = start + aligned_m - m_indices[start:actual_end] = g - m_indices[actual_end:aligned_end] = -1 - - d = run_kernel_grouped(a, b, sf_a, sf_b, m_indices, m_total, n) - d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) - - ref = fp4_reference_grouped(a, b, m_indices, n, num_groups, sf_a, sf_b) - d_max = torch.abs(d.cpu().float() - ref.float()).max().item() - if d_max > worst_diff: - worst_diff = d_max - - ok = worst_diff < 1.0 - if not ok: - all_pass = False - print(f' [{label}×{NUM_ITERS[label]}] G={num_groups} m_per_group={m_per_group:5d} ' - f'(aligned={aligned_m:5d}) N={n:5d} K_fp4={k:5d}: ' - f'max_diff={worst_diff:.4f} {"PASS" if ok else "FAIL"}') - - # Perf section — bench_kineto on production shapes only. - # Filter to the FP4 GEMM kernel name; works for both dense and grouped wrappers - # because the device kernel function is sm100_fp4_gemm_1d1d_impl in both cases. - print('\nPerf: m-grouped contiguous (production MoE shapes)') - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) - for num_groups, m_per_group, n, k in prod_configs: - aligned_m = ((m_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M - m_total = aligned_m * num_groups - - a = pack_fp4_random(m_total, k) - b = pack_fp4_random_3d(num_groups, n, k) - sf_a, _ = generate_mxf4_scale_factors(m_total, n, k, random_sf=True) - sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) - m_indices = torch.empty(m_total, dtype=torch.int32, device='cuda') - for g in range(num_groups): - start = g * aligned_m - m_indices[start:start + m_per_group] = g - m_indices[start + m_per_group:start + aligned_m] = -1 - d = torch.empty((m_total, n), device='cuda', dtype=torch.float32) - - def fn(): - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - (a, sf_a), (b, sf_b), d, m_indices, - recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, - ) - - t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) - # FLOPs: 2 * M_total * N * K_fp4 (k here is K_fp4 element count). - # Bytes: A (int32, 4B/elem), B (int32, 4B/elem), D (fp32, 4B/elem), plus SF. - tflops = 2 * m_total * n * k / t / 1e12 - gbps = count_bytes(a, b, d) / 1e9 / t - print(f' G={num_groups} m_per_group={m_per_group:5d} N={n:5d} K_fp4={k:5d}: ' - f'{t * 1e6:6.1f} us | {tflops:6.0f} TFLOPS | {gbps:5.0f} GB/s') - - return all_pass - - -def test_m_grouped_trtllm_comparable(): - """trtllm-gen apples-to-apples: DeepSeek-R1 MoE setup with 256 experts × topK=8. - - Setup mirrors gitlab/trtllm BatchedGemm baseline: - num_groups = 256 (numExperts) - total_actual_M = numTokens * topK = numTokens * 8 - Uniform routing → rows_per_expert = total_actual_M / 256 - Each expert tile padded to BLOCK_M=128 - - CAVEAT: trtllm-gen's BatchedGemmFp4LowLatency uses batch=N (grouped-along-N); - DG here uses batch=M. Memory access patterns differ but useful FLOPs compare directly. - - trtllm-gen baseline on B200 (from memory/fp4_grouped_gemm_perf_work.md): - - FC2 (N=7168, K=2048): - tokens=32 → 0.37 ms | 20.4 TFLOPS | 5.74 TB/s - tokens=64 → 0.37 ms | 40.7 TFLOPS | 5.75 TB/s - tokens=128 → 0.37 ms | 80.8 TFLOPS | 5.72 TB/s - tokens=256 → 0.37 ms | 161.9 TFLOPS | 5.78 TB/s - tokens=512 → 0.55 ms | 219.6 TFLOPS | 3.98 TB/s - tokens=1024 → 0.51 ms | 472.3 TFLOPS | 4.40 TB/s - tokens=2048 → 0.51 ms | 941.2 TFLOPS | 4.63 TB/s - - FC1 (N=4096, K=7168, fusedAct=swiglu, routeAct=tma) — note batch=N caveat: - tokens=32 → 0.76 ms | 19.9 TFLOPS | 5.59 TB/s - tokens=64 → 0.75 ms | 39.8 TFLOPS | 5.61 TB/s - tokens=128 → 0.76 ms | 78.7 TFLOPS | 5.54 TB/s - tokens=256 → 0.76 ms | 157.5 TFLOPS | 5.55 TB/s - tokens=512 → 0.78 ms | 310.1 TFLOPS | 5.48 TB/s - tokens=1024 → 0.82 ms | 590.0 TFLOPS | 5.25 TB/s - tokens=2048 → 1.19 ms | 810.9 TFLOPS | 3.65 TB/s - """ - print('Test: m-grouped contiguous (trtllm-gen comparable, 256 experts)') - BLOCK_M = 128 - NUM_GROUPS = 256 - TOP_K = 8 - - # Latency regime (trtllm-gen batch=N baseline): tokens 32..2048, per_group_M < BLOCK_M - # Throughput regime (trtllm-gen batch=M target): tokens 4096..8192, per_group_M >= BLOCK_M - # 4096: rows/grp=128 (= BLOCK_M, zero padding) — first apples-to-apples vs trtllm batch=M - # 8192: rows/grp=256 (= 2×BLOCK_M, 2 tiles per group) - token_sweep = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] - shape_configs = [ - ('FC2', 7168, 2048), - ('FC1', 4096, 7168), - ] - - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) - all_pass = True - - for fc_name, n, k in shape_configs: - print(f'\n--- {fc_name}: N={n}, K_fp4={k}, num_groups={NUM_GROUPS} ---') - # Allocate B/SFB once per shape (large, reuse across token sweep — independent of tokens) - b = pack_fp4_random_3d(NUM_GROUPS, n, k) - sf_b = generate_mxf4_sf_3d(NUM_GROUPS, n, k, random_sf=True) - - for num_tokens in token_sweep: - total_actual_m = num_tokens * TOP_K - rows_per_group = max(1, total_actual_m // NUM_GROUPS) - # Per-group rows aligned UP to BLOCK_M (may exceed BLOCK_M for tokens >= 4096) - aligned_per_group = ((rows_per_group + BLOCK_M - 1) // BLOCK_M) * BLOCK_M - total_padded_m = NUM_GROUPS * aligned_per_group - - # A/SF_A/D scale with total_padded_m, so allocate per iteration. - a = pack_fp4_random(total_padded_m, k) - sf_a, _ = generate_mxf4_scale_factors(total_padded_m, n, k, random_sf=True) - d = torch.empty((total_padded_m, n), device='cuda', dtype=torch.float32) - - # m_indices: each group gets aligned_per_group rows, first rows_per_group are valid (g), - # remainder padding (-1). - m_indices = torch.full((total_padded_m,), -1, dtype=torch.int32, device='cuda') - for g in range(NUM_GROUPS): - start = g * aligned_per_group - m_indices[start:start + rows_per_group] = g - # rows [start + rows_per_group : start + aligned_per_group) stay -1 +import torch - # Run kernel +import deep_gemm +from deep_gemm.testing import bench_kineto, calc_diff, count_bytes + +from generators import ( + KernelType, MajorTypeAB, QuantConfig, + generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, +) + + +# FP4xFP4 (MXFP4): both operands packed FP4 (E2M1) with VS=32 UE8M0 scales. +# Routes through deep_gemm.fp8_fp4_gemm_nt / m_grouped_fp8_fp4_gemm_* APIs and +# dispatches to the SM100_MMA_MXF4_SS path in sm100_fp4_gemm_1d1d.cuh. +# NOTE: the kernel currently only supports FP32 output (bf16 epilogue is a TODO). +FP4_FP4 = QuantConfig((32, 32, True, True)) + + +def test_gemm() -> None: + print('Testing GEMM:') + nk_list = [(2112, 7168), (576, 7168), (24576, 1536), (32768, 512), + (7168, 16384), (4096, 7168), (7168, 2048)] + m_list = [128, 4096] + kernel_type = KernelType.Kernel1D1D + out_dtype = torch.float + recipe, recipe_a, recipe_b = FP4_FP4.get_recipes() + + for m in m_list: + for n, k in nk_list: + a, b, c, d, ref_d = generate_normal( + m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, + accumulate=False, out_dtype=out_dtype, kernel_type=kernel_type, + use_ue8m0=True, quant_config=FP4_FP4) + deep_gemm.fp8_fp4_gemm_nt( + a, b, d, c=c, disable_ue8m0_cast=False, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + diff = calc_diff(d, ref_d) + assert diff < FP4_FP4.max_diff(), \ + f'{m=}, {n=}, {k=}, {diff:.5f}' + + t = bench_kineto( + lambda: deep_gemm.fp8_fp4_gemm_nt( + a, b, d, c=c, disable_ue8m0_cast=False, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b), + 'sm100_fp4_gemm', suppress_kineto_output=True) + print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, 1D1D, layout=NT, FP32): ' + f'{t * 1e6:6.1f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + + +def test_m_grouped_gemm_contiguous() -> None: + print('Testing m-grouped contiguous GEMM:') + m_group_list = [(4, 8192), (8, 4096)] + n_k_list = [(4096, 7168), (7168, 2048), (24576, 1536), (32768, 512)] + out_dtype = torch.float + recipe, recipe_a, recipe_b = FP4_FP4.get_recipes() + + for num_groups, expected_m_per_group in m_group_list: + for n, k in n_k_list: + m, a, b, m_indices, d_bf16, _ref_bf16 = generate_m_grouped_contiguous( + num_groups, expected_m_per_group, n, k, + MajorTypeAB.KMajor, MajorTypeAB.KMajor, + use_ue8m0=True, quant_config=FP4_FP4) + # Re-generate with bf16 disabled-cast path to get a proper fp32 reference. + # generate_m_grouped_contiguous hardcodes bf16 d/ref; we allocate fp32 ourselves + # and cast the ref for the diff check. + d = torch.empty_like(d_bf16, dtype=out_dtype) + ref_d = _ref_bf16.to(out_dtype) deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - (a, sf_a), (b, sf_b), d, m_indices, - recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, - ) - torch.cuda.synchronize() - d_clean = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) - - # Correctness: only for smaller tokens (CPU LUT reference for large M is slow) - if num_tokens <= 256: - ref = fp4_reference_grouped(a, b, m_indices, n, NUM_GROUPS, sf_a, sf_b) - max_diff = torch.abs(d_clean.cpu().float() - ref.float()).max().item() - ok = max_diff < 1.0 - if not ok: - all_pass = False - correctness = f'diff={max_diff:.4f} {"✓" if ok else "✗"}' - else: - correctness = '(skip)' - - # Perf - def fn(): - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - (a, sf_a), (b, sf_b), d, m_indices, - recipe_a=(1, 32), recipe_b=(1, 32), disable_ue8m0_cast=False, - ) - t = bench_kineto(fn, 'sm100_fp4_gemm', suppress_kineto_output=True) - - # Useful FLOPS = 2 * useful_M * N * K (matches trtllm-gen convention) - useful_m = NUM_GROUPS * rows_per_group - useful_tflops = 2 * useful_m * n * k / t / 1e12 - # Bytes: A loaded once (total_padded_m), B loaded once, D written once - gbps = count_bytes(a, b, d) / 1e9 / t - print(f' tokens={num_tokens:5d} useful_M={useful_m:6d} ' - f'rows/grp={rows_per_group:4d}/{aligned_per_group:4d}: ' - f'{t * 1e3:6.3f} ms | {useful_tflops:7.1f} TFLOPS | ' - f'{gbps/1000:5.2f} TB/s | {correctness}') - - return all_pass - - -def test_m_grouped_masked(): - """M-grouped masked FP4 GEMM (MoE decode shape). - - A=[G, max_M, K] @ B=[G, N, K].T → D=[G, max_M, N] - masked_m[G] indicates the number of valid rows per group (rest is padding). - """ - print('Test: m-grouped masked (MoE decode)') - BLOCK_M = 128 - - # (num_groups, max_m, expected_m, n, k_fp4) - # Small/debug shapes (CPU reference fast) + production shapes mirroring FP8 enumerate_m_grouped_masked - debug_configs = [ - (2, 128, 64, 128, 256), # 50% util - (4, 128, 32, 128, 256), # heavy padding - (4, 256, 200, 128, 256), # multi-tile per group - (4, 128, 128, 256, 512), # full - (8, 128, 64, 256, 512), - ] - prod_configs = [ - # max_m=4096 (matching FP8 enumerate_m_grouped_masked max_m), varying num_groups & m - (1, 4096, 1024, 4096, 7168), # FP8 (1, 1024) - (2, 4096, 512, 4096, 7168), # FP8 (2, 512) - (4, 4096, 256, 4096, 7168), # FP8 (4, 256) - (1, 4096, 1024, 7168, 2048), - (2, 4096, 512, 7168, 2048), - (4, 4096, 256, 7168, 2048), - ] - - # Mirror FP8 test_m_grouped_gemm_masked: 10 random-data iterations per shape on - # production shapes to catch flaky bugs (different masked_m distribution each time). - NUM_ITERS = {'debug': 3, 'prod': 10} - - all_pass = True - duc = not get_ue8m0_usage(KernelType.Kernel1D1D) - for label, configs in [('debug', debug_configs), ('prod', prod_configs)]: - for num_groups, max_m, expected_m, n, k in configs: - worst_diff = 0.0 - for _ in range(NUM_ITERS[label]): - # Fresh masked_m + tensors per iteration (matches FP8 pattern) - masked_m_cpu = torch.tensor([ - max(1, min(max_m, int(expected_m * random.uniform(0.7, 1.3)))) - for _ in range(num_groups) - ], dtype=torch.int32) - masked_m = masked_m_cpu.cuda() - - a = pack_fp4_random_3d_ga(num_groups, max_m, k) - b = pack_fp4_random_3d(num_groups, n, k) - sf_a = generate_mxf4_sfa_3d(num_groups, max_m, k, random_sf=True) - sf_b = generate_mxf4_sf_3d(num_groups, n, k, random_sf=True) - - d = run_kernel_grouped_masked(a, b, sf_a, sf_b, masked_m, num_groups, - max_m, n, expected_m) - ref = fp4_reference_masked(a, b, masked_m_cpu, max_m, n, num_groups, sf_a, sf_b) - - # Only compare valid rows per group - for g in range(num_groups): - mg = int(masked_m_cpu[g].item()) - if mg == 0: - continue - diff = torch.abs(d[g, :mg].cpu().float() - ref[g, :mg].float()).max().item() - if diff > worst_diff: - worst_diff = diff - - ok = worst_diff < 1.0 - if not ok: - all_pass = False - print(f' [{label}×{NUM_ITERS[label]}] G={num_groups} max_m={max_m:5d} ' - f'expected_m={expected_m:5d} N={n:5d} K_fp4={k:5d}: ' - f'max_diff={worst_diff:.4f} {"PASS" if ok else "FAIL"}') - - return all_pass - - -def test_multicast(): - """大 M 测试：触发 B-multicast (M>=512, 2CTA along M, UMMA_M=256)""" - print('Test: B-multicast (M>=512, 2CTA)') - configs = [ - (512, 128, 256), - (512, 128, 512), - (512, 128, 1024), - (1024, 128, 256), - (1024, 128, 512), - ] - all_pass = True - for m, n, k in configs: - a = pack_fp4_random(m, k) - b = pack_fp4_random(n, k) - sf_a, sf_b = generate_mxf4_scale_factors(m, n, k, random_sf=True) - d = run_kernel(a, b, sf_a, sf_b, m, n) - ref = fp4_reference(a, b, m, n, sf_a, sf_b) - max_diff = torch.abs(d.cpu().float() - ref.float()).max().item() - ok = max_diff < 1.0 - if not ok: - all_pass = False - print(f' M={m:4d} N={n:4d} K={k:4d}: max_diff={max_diff:.4f} {"PASS" if ok else "FAIL"}') - return all_pass + a, b, d, m_indices, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + disable_ue8m0_cast=False, use_psum_layout=False) + diff = calc_diff(d, ref_d) + assert diff < FP4_FP4.max_diff(), \ + f'{num_groups=}, {m=}, {n=}, {k=}, {diff:.5f}' + + t = bench_kineto( + lambda: deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + a, b, d, m_indices, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + disable_ue8m0_cast=False, use_psum_layout=False), + 'sm100_fp4_gemm', suppress_kineto_output=True) + print(f' > Perf (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:5}, ' + f'n={n:5}, k={k:5}, 1D1D, FP32): ' + f'{t * 1e6:6.1f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + + +def test_m_grouped_gemm_masked() -> None: + print('Testing m-grouped masked GEMM:') + max_m = 4096 + m_group_list = [(1, 1024), (2, 512), (4, 256)] + n_k_list = [(4096, 7168), (7168, 2048)] + out_dtype = torch.float + recipe, recipe_a, recipe_b = FP4_FP4.get_recipes() + + for num_groups, expected_m_per_group in m_group_list: + for n, k in n_k_list: + a, b, masked_m, _psum_m, d_bf16, _ref_bf16 = generate_m_grouped_masked( + num_groups, max_m, expected_m_per_group, n, k, + use_ue8m0=True, quant_config=FP4_FP4) + d = torch.empty_like(d_bf16, dtype=out_dtype) + ref_d = _ref_bf16.to(out_dtype) + deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( + a, b, d, masked_m, expected_m_per_group, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + disable_ue8m0_cast=False) + # Diff over the valid (non-padding) rows per group only. + for g in range(num_groups): + mm = int(masked_m[g].item()) + diff = calc_diff(d[g, :mm], ref_d[g, :mm]) + assert diff < FP4_FP4.max_diff(), \ + f'{num_groups=}, group={g}, masked_m={mm}, {n=}, {k=}, {diff:.5f}' + + t = bench_kineto( + lambda: deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( + a, b, d, masked_m, expected_m_per_group, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + disable_ue8m0_cast=False), + 'sm100_fp4_gemm', suppress_kineto_output=True) + print(f' > Perf (num_groups={num_groups}, max_m={max_m}, expected_m_per_group={expected_m_per_group:5}, ' + f'n={n:5}, k={k:5}, 1D1D, FP32): ' + f'{t * 1e6:6.1f} us | {2 * num_groups * expected_m_per_group * n * k / t / 1e12:4.0f} TFLOPS') if __name__ == '__main__': - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True torch.manual_seed(0) random.seed(0) - print(f'Library: {deep_gemm.__path__}\n') - - results = [ - ('constant', test_constant()), - ('random', test_random()), - ('sweep', test_value_sweep()), - ('asymmetric', test_asymmetric_values()), - ('uniform_sf', test_uniform_sf()), - ('random_sf', test_random_sf()), - ('multicast', test_multicast()), - ('m_grouped', test_m_grouped_contiguous()), - ('m_grouped_masked', test_m_grouped_masked()), - ('trtllm_cmp', test_m_grouped_trtllm_comparable()), - ] + print('Library path:') + print(f' > {deep_gemm.__path__}\n') - print() - passed = all(r for _, r in results) - for name, ok in results: - print(f' {name}: {"PASS" if ok else "FAIL"}') - print(f'\n{"ALL FP4 TESTS PASSED" if passed else "SOME TESTS FAILED"}') - if not passed: - exit(1) + test_gemm() + test_m_grouped_gemm_contiguous() + test_m_grouped_gemm_masked() From 8d841f664fa83a6802136653adb90afb13caea18 Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Mon, 1 Jun 2026 00:01:26 -0700 Subject: [PATCH 08/12] FP4: align .cuh comments with sm100_fp8_fp4_gemm_1d1d.cuh style MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Comment hygiene pass against the official FP8 kernel's style: - Removed 4 unused helper functions (fp4_e2m1_to_float, pack_ue8m0_*, swizzled_smem_k_major_idx) — leftover debug code, never called. - Replaced 18 '// ========== Chinese label ==========' section banners with short English labels matching the FP8 kernel ('// MMA configs', '// SF configs', '// Shared memory sizes', '// Block scheduler', etc.). - Translated the M-wave / warp-row band stride explanation: kTmemDpStride / kTmemMWaveStride / kTmemWarpRowBandStride were only referenced from commented-out alternative tmem_addr formulas, so all three constants and the associated dead-code blocks are removed. - Removed verbose SF-packing math docstrings (3 lines) and inline type-cast notes — the code is self-evident. Net: -73 lines (-100 deletions / +27 insertions). 0 Chinese characters left in any branch file. test_fp4.py still PASS (exit 0). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- .../deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh | 127 ++++-------------- 1 file changed, 27 insertions(+), 100 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh index ef75ace86a..0d59d9bf52 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh @@ -21,44 +21,6 @@ using namespace deep_gemm::math; using namespace deep_gemm::ptx; using namespace deep_gemm::utils; -// E2M1 FP4 到 float 的转换函数 -// E2M1 格式: 4 bits = SEEM (S=符号 1bit, E=指数 2bits, M=尾数 1bit) -__device__ __forceinline__ float fp4_e2m1_to_float(uint32_t fp4_bits) { - constexpr float E2M1_LUT[16] = { - 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, // 正数 (S=0) - -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f // 负数 (S=1) - }; - return E2M1_LUT[fp4_bits & 0xF]; -} - -// UE8M0 scale 1.0f packed as 4 bytes per 32-bit TMEM word. -// SM100 MXF4 block-scaled MMA UE8M0 scale factor. -// The bias is determined empirically; see DG_SF_BYTE env var testing. -__device__ __forceinline__ uint32_t pack_ue8m0_scale_factor_word(uint8_t byte_val) { - return uint32_t(byte_val) | (uint32_t(byte_val) << 8) | (uint32_t(byte_val) << 16) | (uint32_t(byte_val) << 24); -} -__device__ __forceinline__ uint32_t pack_ue8m0_2x_scale_factor_one_word() { - return pack_ue8m0_scale_factor_word(0x7Fu); // Will be tested with different values -} - -// Swizzle-aware shared memory index for reading TMA-loaded data. -// TMA stores data with bank-group XOR swizzle: physical_bank = logical_bank ^ (row % num_banks). -// swizzle_mode: kSwizzleAMode or kSwizzleBMode (bytes, e.g. 128) -// row: M or N row index, k: K column index, block_k: elements per row -template -__device__ __forceinline__ uint32_t swizzled_smem_k_major_idx(uint32_t row, uint32_t k, uint32_t block_k) { - constexpr uint32_t kElemBytes = sizeof(uint32_t); - constexpr uint32_t kBankBytes = 16; - constexpr uint32_t kElemsPerBank = kBankBytes / kElemBytes; // 4 - constexpr uint32_t kNumBanks = swizzle_mode / kBankBytes; // e.g. 8 for 128B - uint32_t bank = k / kElemsPerBank; - uint32_t in_bank = k % kElemsPerBank; - uint32_t swizzled_bank = bank ^ (row % kNumBanks); - return row * block_k + swizzled_bank * kElemsPerBank + in_bank; -} - -// SM100 FP4 GEMM 1D1D kernel实现 -// 支持 MXF4 block-scaled 矩阵乘法 template ; if constexpr (kWithAccumulation) DG_STATIC_ASSERT(cute::is_same_v, "Invalid C/D data dtype"); - // ========== 核心配置参数 ========== + // MMA configs constexpr uint32_t LAYOUT_AD_M = 128; constexpr uint32_t kNumMWaves = BLOCK_M / LAYOUT_AD_M; constexpr uint32_t kNumTMAStoreStages = 2; - - // ========== MXF4 配置 ========== + + // MXFP4 / SF configs constexpr uint32_t kNumSFAStagesPerLoad = 1; constexpr uint32_t kNumSFBStagesPerLoad = 1; constexpr uint32_t kNumUTCCPAlignedElems = 128; @@ -102,43 +62,32 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, constexpr uint32_t UMMA_K_FP4 = 64; constexpr uint32_t SF_K_PER_STAGE = BLOCK_K_FP4 / MXF4_VS; constexpr uint32_t SF_PACKED_K_PER_STAGE = SF_K_PER_STAGE / 4; - + DG_STATIC_ASSERT(BLOCK_M % LAYOUT_AD_M == 0 and 2 % kNumMWaves == 0, "Invalid block M"); DG_STATIC_ASSERT(BLOCK_K == 16 or BLOCK_K == 32, "FP4 BLOCK_K must be 16 or 32 int32"); - // ========== 动态形状处理 ========== + // 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; - // shape_k debug disabled - // FP4 packed: shape_k 是 int32 个数，每个 int32 有 8 个 FP4。 - // 1 个 scale 覆盖 MXF4_VS=32 个 FP4 (=4 个 int32)， - // 每个 uint32 打包 4 个 scale → 1 个 packed group = 4*32 = 128 FP4。 - const uint32_t total_scales_k = - ceil_div(shape_k * FP4_ELEMS_PER_INT32, - MXF4_VS); // MXF4_VS 是 uint32_t，OK - - const uint32_t total_packed_k = - ceil_div(total_scales_k, - uint32_t(4)); // 把 4 也变成 uint32_t - + const uint32_t total_scales_k = ceil_div(shape_k * FP4_ELEMS_PER_INT32, MXF4_VS); + const uint32_t total_packed_k = ceil_div(total_scales_k, uint32_t(4)); const uint32_t shape_sfa_k = total_packed_k; const uint32_t shape_sfb_k = total_packed_k; - // ========== 线程和warp信息 ========== + // Utils bool is_leader_cta = cute::block_rank_in_cluster() == 0; const auto warp_idx = cutlass::canonical_warp_idx_sync(); const auto lane_idx = get_lane_idx(); - // ========== 共享内存分配 ========== + // Align to 1024 bytes for swizzle-128B extern __shared__ __align__(1024) uint8_t smem_buffer[]; - // ========== 块大小计算 ========== + // Load/store block sizes constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / (kIsMulticastOnA ? kNumMulticast: 1); constexpr uint32_t LOAD_BLOCK_N = BLOCK_N / (kIsMulticastOnA ? 1 : kNumMulticast); - // Swap-AB epilogue: STORE_BLOCK_M = 16 (fine-grained M slices to skip padding rows), - // STORE_BLOCK_N = BLOCK_N (write entire N at once). - // Non-swap (existing): STORE_BLOCK_M = BLOCK_M, STORE_BLOCK_N derived from swizzle. + // Swap-AB: STORE_BLOCK_M=16 (fine-grained M slices to skip padding rows), STORE_BLOCK_N=BLOCK_N. + // Non-swap: STORE_BLOCK_M=BLOCK_M, STORE_BLOCK_N derived from swizzle. constexpr uint32_t STORE_BLOCK_M = kSwapAB ? 16u : cute::min(BLOCK_M, LAYOUT_AD_M); constexpr uint32_t STORE_BLOCK_N = kSwapAB ? BLOCK_N : kSwizzleCDMode / sizeof(cd_dtype_t); @@ -147,12 +96,9 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); // Swap-AB requires BLOCK_N = LAYOUT_AD_M (= 128) so UMMA_M after swap stays = 128. DG_STATIC_ASSERT(not kSwapAB or BLOCK_N == LAYOUT_AD_M, "kSwapAB requires BLOCK_N = LAYOUT_AD_M"); - // Swap-AB initial implementation: no multicast (cluster_n=1, cluster_m=1) for simplicity. DG_STATIC_ASSERT(not kSwapAB or kNumMulticast == 1, "kSwapAB initial impl: no multicast"); - // ========== 共享内存大小计算 ========== - // Swap-AB: per-stage SMEM = STORE_BLOCK_M (small) × STORE_BLOCK_N (= BLOCK_N) × sizeof(D) - // Non-swap: per-stage SMEM = STORE_BLOCK_M × kSwizzleCDMode (= STORE_BLOCK_N * sizeof(D)) + // Shared memory sizes constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = kSwapAB ? STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t) : STORE_BLOCK_M * kSwizzleCDMode; @@ -163,11 +109,11 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, constexpr uint32_t SF_BLOCK_N = constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * SF_PACKED_K_PER_STAGE * sizeof(uint32_t); - + DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0, "Shared memory must be aligned to 1024 bytes"); DG_STATIC_ASSERT(kNumTMAStoreStages >= 1, "Invalid number of TMA stages"); - // ========== 张量内存配置 ========== + // Tensor memory size and offsets constexpr uint32_t kNumSFATmemCols = (SF_BLOCK_M / 32) * SF_PACKED_K_PER_STAGE; constexpr uint32_t kNumSFBTmemCols = (SF_BLOCK_N / 32) * SF_PACKED_K_PER_STAGE; constexpr uint32_t kNumEpilogueStages = (2 * kNumMWaves * BLOCK_N + kNumSFATmemCols + kNumSFBTmemCols) > 512 ? 1 : 2; @@ -176,7 +122,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; - // ========== TMA描述符预取 ========== + // Prefetch TMA descriptors at the very beginning if (threadIdx.x == 0) { cute::prefetch_tma_descriptor(&tensor_map_a); cute::prefetch_tma_descriptor(&tensor_map_b); @@ -187,23 +133,24 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, cute::prefetch_tma_descriptor(&tensor_map_c); } - // ========== 共享内存指针设置 ========== + // D/A/B shared memory cd_dtype_t* smem_cd[kNumTMAStoreStages]; uint32_t* smem_sfa[kNumStages]; uint32_t* smem_sfb[kNumStages]; uint32_t* smem_a_packed[kNumStages]; uint32_t* smem_b_packed[kNumStages]; - + #pragma unroll for (uint32_t i = 0; i < kNumTMAStoreStages; ++ i) smem_cd[i] = reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); - + #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { smem_a_packed[i] = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_PACKED_SIZE_PER_STAGE); smem_b_packed[i] = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_PACKED_SIZE_PER_STAGE + i * SMEM_B_PACKED_SIZE_PER_STAGE); } - + + // SFA/SFB shared memory auto sf_start_ptr = smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_PACKED_SIZE_PER_STAGE + SMEM_B_PACKED_SIZE_PER_STAGE); #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { @@ -211,7 +158,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, smem_sfb[i] = reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); } - // ========== 屏障初始化 ========== + // Barriers and tensor memory pointer auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_PACKED_SIZE_PER_STAGE + SMEM_B_PACKED_SIZE_PER_STAGE) + @@ -246,15 +193,15 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, } kNumMulticast > 1 ? cute::cluster_sync() : __syncthreads(); - // ========== 块调度器初始化 ========== + // Block scheduler uint32_t m_block_idx, n_block_idx; auto scheduler = sched::Scheduler(shape_m, shape_n, shape_k, grouped_layout); - // ========== K维度迭代控制 ========== + // K-loop driver struct DivisibleK {}; struct NotDivisibleK {}; uint32_t phase = 0; - + auto launch_k_iterations = [&](const auto& func) { const uint32_t current_shape_k = (kGemmType == GemmType::KGroupedContiguous ? scheduler.current_shape_k : shape_k); const uint32_t num_iterations = ceil_div(current_shape_k, kNumStages * BLOCK_K); @@ -586,7 +533,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, }); } } else if (warp_idx >= kNumNonEpilogueThreads / 32) { - // ========== Epilogue warp组 ========== + // Epilogue warps const auto epilogue_thread_idx = threadIdx.x - kNumNonEpilogueThreads; const auto epilogue_warp_idx = warp_idx - (kNumNonEpilogueThreads / 32); @@ -719,15 +666,6 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, DG_STATIC_ASSERT(kNumEpilogueThreads == 128, "Epilogue threads not enough"); DG_STATIC_ASSERT(BLOCK_N % STORE_BLOCK_N == 0, "Invalid block sizes"); - // Per cute::make_tmem_warp_partitioner (copy_traits_sm100.hpp): four epilogue warps read - // four 32-row M bands; stride is 32 * TMEM::DP (datapath), not +32 on column index. - constexpr uint32_t kTmemDpStride = - static_cast(cute::TMEM::DP{}); - // M-wave 间步进：128 行 × DP步进（DP 方向，不是列方向） - constexpr uint32_t kTmemMWaveStride = 128u * kTmemDpStride; - // epilogue warp 在一个 wave 内的 band 步进（DP 方向） - constexpr uint32_t kTmemWarpRowBandStride = 32u * kTmemDpStride; - #pragma unroll for (uint32_t w = 0; w < kNumMWaves; ++ w) { constexpr uint32_t kNumStores = BLOCK_N / STORE_BLOCK_N; @@ -752,19 +690,8 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, auto col = kHasShortcut ? (i) : (bank_group_index % 8); col ^= row % (kSwizzleCDMode / 16); - // uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N // 列：stage - // + s * STORE_BLOCK_N // 列：N tile - // + i * kNumElemsPerBankGroup // 列：bank group - // + w * kTmemMWaveStride // DP：M-wave ← 关键修改 - // + epilogue_warp_idx * kTmemWarpRowBandStride; // DP：warp band - // uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N // stage 列偏移 - // + w * BLOCK_N // M-wave 列偏移（每个 wave 占 BLOCK_N 列） - // + s * STORE_BLOCK_N // N tile 列偏移 - // + i * kNumElemsPerBankGroup; // bank group 列偏移 uint32_t tmem_addr = accum_stage_idx * kNumMWaves * BLOCK_N + w * BLOCK_N + s * STORE_BLOCK_N - + i * kNumElemsPerBankGroup; - // Match sm100_bf16_gemm.cuh / mma.cuh: each epilogue warp reads its 32-row TMEM band (DP stride). - // tmem_addr += epilogue_warp_idx * kTmemWarpRowBandStride; + + i * kNumElemsPerBankGroup; auto smem_ptr = reinterpret_cast(smem_cd[tma_stage_idx]) + epilogue_warp_idx * 32 * kSwizzleCDMode + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; From ae4e77d8c7316101241bb64f7b219d278e5ca22a Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Mon, 1 Jun 2026 00:37:47 -0700 Subject: [PATCH 09/12] FP4: address codex review (PDL sync, strict FP32 D, prune unused fields) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply the P0/P1/P2 items from FP4_PR_REVIEW_NOTES.md that don't change the swap_ab gating contract: P0-1. Add cudaGridDependencySynchronize() before block scheduling in sm100_fp4_gemm_1d1d.cuh. Mirrors sm100_fp8_fp4_gemm_1d1d.cuh — required because the API transforms SFA/SFB in a prior kernel; without the wait the GEMM can race the producer. P0-2. Tighten the D dtype assertions in csrc/apis/gemm.hpp at all three FP4xFP4 dispatch sites (fp8_fp4_gemm_nt, m_grouped_*_contiguous, m_grouped_*_masked): when both a and b are kPackedFP4, require d.scalar_type() == kFloat. The kernel has no bf16 epilogue store path; previously the assertion still accepted bf16 and silently produced garbage output. P1-2. Drop trailing whitespace at the `} else if` join in the kernel (`git diff --check origin/main..HEAD` is now clean). P2. Remove unused template/runtime plumbing: - kNumLastStages template parameter (kernel body recomputes the value from runtime shape_k; the compile-time arg was never referenced). - tensor_map_c kernel arg, Args field, wrapper-side make_tma_cd_desc call, and the local `cd` alias. Accumulation is folded into D pre-launch by `d.copy_(c.value())`, so the kernel never loads C. - compute_num_last_stages_fp4() helper (callers gone). P2 (style). Rewrite remaining comments to drop personal/migration wording ("my kernel", "fea-fp4", "Phase 2", "NaN issues unresolved", "wait simpler", exploratory math notes). Keep only invariants needed to read the code. Swap_ab gating in pick_fp4_layout still uses `(grouped_layout >= 0).sum().item()` for MGroupedContiguous (P0-3 deferred). That host sync only fires when swap_ab might apply; documented as a known limitation for CUDA-graph callers. Build clean. test_fp4.py PASS (exit 0). Manual bf16 input check now raises the strict-dtype assertion instead of returning zeroed D. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/apis/gemm.hpp | 25 ++-- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 118 ++++++------------ .../deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh | 11 +- 3 files changed, 63 insertions(+), 91 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 497056eee8..6b3f12e192 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -74,7 +74,12 @@ static void fp8_fp4_gemm_nt(const std::pair& a, const auto [n , k_] = check_ab_fp8_fp4(b.first, major_b, arch_major); const auto [m_, n_] = get_shape<2>(d); DG_HOST_ASSERT(m == m_ and n == n_ and k == k_); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + // FP4xFP4 path has no bf16 epilogue store: require fp32 D. + if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) { + DG_HOST_ASSERT(d.scalar_type() == torch::kFloat); + } else { + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + } // Early return for trivial cases if (early_return(m, n, k, d, c)) @@ -172,9 +177,12 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair(d); DG_HOST_ASSERT(m == m_ and n == n_ and k == k_); DG_HOST_ASSERT(n > 0 and k > 0 and num_groups > 0); - // FP4xFP4 my kernel hardcodes fp32 output (bf16 epilogue is TODO). - const bool is_fp4_fp4 = (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or (is_fp4_fp4 and d.scalar_type() == torch::kFloat)); + // FP4xFP4 path has no bf16 epilogue store: require fp32 D. + if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) { + DG_HOST_ASSERT(d.scalar_type() == torch::kFloat); + } else { + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + } DG_HOST_ASSERT(grouped_layout.scalar_type() == torch::kInt); // Layout checks @@ -258,9 +266,12 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair 0 and m > 0 and n > 0 and k > 0 and num_groups > 0); - // FP4xFP4 my kernel hardcodes fp32 output (bf16 epilogue is TODO). - const bool is_fp4_fp4_masked = (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or (is_fp4_fp4_masked and d.scalar_type() == torch::kFloat)); + // FP4xFP4 path has no bf16 epilogue store: require fp32 D. + if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) { + DG_HOST_ASSERT(d.scalar_type() == torch::kFloat); + } else { + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + } DG_HOST_ASSERT(masked_m.scalar_type() == torch::kInt); // D must be N-major diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index 1d3eaea0c1..298922a997 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -15,23 +15,23 @@ namespace deep_gemm { -// FP4xFP4 (MXFP4) GEMM via SM100_MMA_MXF4_SS instruction. -// Distinct from main's sm100_fp8_fp4_gemm_1d1d (MXF8F6F4 path) — this is the -// FP4-specialized hardware path, used when both A and B are kPackedFP4. +// FP4xFP4 (MXFP4) GEMM via SM100_MMA_MXF4_SS. Distinct from the MXF8F6F4 +// path in sm100_fp8_fp4_gemm_1d1d.hpp; used when both A and B are kPackedFP4. +// GemmDesc carries logical FP4 K; this kernel template takes K and BLOCK_K in +// packed int32 units (8 FP4 / int32, 4 bytes). The wrapper converts at JIT +// instantiation and runtime launch. class SM100FP4Gemm1D1DRuntime final: public LaunchRuntime { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; LaunchArgs launch_args; - int num_last_stages; void* grouped_layout; CUtensorMap tensor_map_a; CUtensorMap tensor_map_b; CUtensorMap tensor_map_sfa; CUtensorMap tensor_map_sfb; - CUtensorMap tensor_map_c; CUtensorMap tensor_map_d; }; @@ -48,7 +48,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, {}, {}, - {}, {}, + {}, {}, {}, {}, {}, {}, @@ -60,15 +60,14 @@ static void __instantiate_kernel() {{ to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b), 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), - // FP4: my .cuh expects SHAPE_K in int32 count (= FP4_count / 8). - // main's desc.k is FP4 logical count. + // SHAPE_K in packed int32 (= FP4_count / 8) get_compiled_dim(args.gemm_desc.k / 8, 'k', args.gemm_desc.compiled_dims), args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, - // FP4: BLOCK_K in int32 count (main's heuristic gives bytes for int8 pack). + // BLOCK_K in packed int32 (Layout.block_k is in bytes for int8-packed FP4) args.gemm_config.layout.block_k / 4, args.gemm_desc.num_groups, 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_last_stages, + args.gemm_config.pipeline_config.num_stages, args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, @@ -78,65 +77,46 @@ static void __instantiate_kernel() {{ } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - // FP4: my .cuh expects shape_k in int32 count (= FP4_count / 8) at runtime as well. - // Note: must be non-const (launch_kernel takes &args, can't bind const* to void*). + // shape_k in packed int32 (must be a non-const local — launch_kernel takes + // &args internally and a const& cannot bind to void*). int shape_k_int32 = args.gemm_desc.k / 8; DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, shape_k_int32, args.tensor_map_a, args.tensor_map_b, args.tensor_map_sfa, args.tensor_map_sfb, - args.tensor_map_c, args.tensor_map_d)); + args.tensor_map_d)); } }; -// Helper: compute num_last_stages from k / block_k / num_stages. -// desc.k is FP4 logical count; main's heuristic block_k is in bytes (int8 elem of packed FP4 tensor). -// 1 byte = 2 FP4, so K_per_block_fp4 = block_k_bytes * 2. -static int compute_num_last_stages_fp4(int k, int block_k_bytes, int num_stages) { - const int num_k_blocks = ceil_div(k, block_k_bytes * 2); - const int rem = num_k_blocks % num_stages; - return rem == 0 ? num_stages : rem; -} - -// My MXF4 kernel allocates 2x SF smem per stage (2 packed int32 per row to cover the -// full BLOCK_K_FP4=256 of 8 SFs) than main's heuristic assumes (1 packed int32 per -// row, covering only 128 FP4). Cap num_stages so total smem fits SM100's 232448-byte -// capacity. Returns {new_num_stages, new_smem_size}. -// Note: this is a v0 workaround. Phase 2: align main's heuristic SF computation with -// my kernel's actual usage so num_stages selection is correct upstream. +// Cap num_stages so SM100's 232448-byte smem capacity fits the FP4 SF smem +// footprint (sf_packed_k_per_stage = block_k_bytes / 64), which the shared +// SM100ArchSpec heuristic underestimates by 2x. static std::pair recompute_stages_for_fp4(const GemmConfig& config, int block_m, int block_n, int k_fp4) { constexpr int smem_capacity = 232448; constexpr int sf_block_align = 128; const int sf_block_m = (block_m + sf_block_align - 1) / sf_block_align * sf_block_align; const int sf_block_n = (block_n + sf_block_align - 1) / sf_block_align * sf_block_align; - // Match fea-fp4 kernel's actual smem usage: - // A per stage: load_block_m * block_k_bytes - // B per stage: load_block_n * block_k_bytes - // SFA/B per stage: sf_block_mn * sf_packed_k_per_stage * 4 - // where sf_packed_k_per_stage = block_k_bytes / 64 (since BLOCK_K_FP4 / VS / 4 = block_k_int32 / 4 = block_k_bytes / 16 / 4 / 4 ... wait simpler: block_k_int32/4) - // Kernel: BLOCK_K_FP4 = block_k_int32 * 8; SF_K_PER_STAGE = BLOCK_K_FP4 / 32; - // SF_PACKED_K_PER_STAGE = SF_K_PER_STAGE / 4 = block_k_int32 / 16. - // In bytes: block_k_bytes / 64. + // Per-stage smem footprint (matches the kernel's actual allocation): + // A: load_block_m * block_k_bytes, B: load_block_n * block_k_bytes + // SFA/B: sf_block_mn * sf_packed_k_per_stage * 4 + // sf_packed_k_per_stage = block_k_bytes / 64 (BLOCK_K_FP4 / VS / 4). const int sf_packed_k_per_stage = config.layout.block_k / 64; const int per_stage = config.storage_config.load_block_m * config.layout.block_k + config.storage_config.load_block_n * config.layout.block_k + sf_block_m * sf_packed_k_per_stage * 4 + sf_block_n * sf_packed_k_per_stage * 4; - // Fixed extras (CD smem + barriers + tmem_ptr, conservative estimate) - // CD smem (must match kernel's SMEM_CD_SIZE_PER_STAGE × kNumTMAStoreStages=2): - // swap_ab: STORE_BLOCK_M(=16) * STORE_BLOCK_N(=block_n) * sizeof(cd_dtype) - // non-swap: STORE_BLOCK_M(=block_m) * kSwizzleCDMode - // cd_dtype size: assume fp32 (4 bytes); bf16 epilogue is TODO so we can't reach it here. + // CD smem (matches kernel's SMEM_CD_SIZE_PER_STAGE × kNumTMAStoreStages=2); + // FP4xFP4 epilogue is fp32-only (4-byte cd elements). int cd_size; if (config.layout.swap_ab) { constexpr int store_block_m_swap = 16; - cd_size = store_block_m_swap * block_n * 4 * 2; // fp32, 2 stages + cd_size = store_block_m_swap * block_n * 4 * 2; } else { cd_size = config.storage_config.store_block_m * config.storage_config.swizzle_cd_mode * 2; } - const int barriers = 12 * 8 * 4 + 4 * 8 * 2 + 8; // ~ 416 bytes max + const int barriers = 12 * 8 * 4 + 4 * 8 * 2 + 8; const int tmem_ptr = 4; const int fixed_extras = cd_size + barriers + tmem_ptr; @@ -147,21 +127,13 @@ static std::pair recompute_stages_for_fp4(const GemmConfig& config, in return {new_num_stages, new_smem_size}; } -// Pick the FP4-optimal Layout (block_m=128 fixed, wave-aware block_n, B-multicast for -// large M, swap_ab for sparse m-grouped). Mirrors fea-fp4's `get_best_fp4_config` but -// returns main's nested Layout struct. -// -// Constants (FP4 MXF4 path): -// block_m = 128 (UMMA_M for MXF4) -// block_k = 128 bytes (= 32 int32 = 256 FP4 per K block) -// sf_pk = 2 (SF_PACKED_K_PER_STAGE; from block_k_int32/16) -// -// block_n is chosen by wave count + composite score = est_stages^2 * bn, -// with 2-epi-stage tiebreak for multi-wave cases. swap_ab is enabled for -// MGroupedContiguous when expected_m_per_group < BLOCK_M (sparse MoE). -// -// expected_m_per_group: useful per-group row count (m_indices >= 0 sum/G). -// Pass INT_MAX to disable swap_ab gating (default for Normal / non-grouped). +// FP4xFP4 Layout selector: +// block_m = 128 (UMMA_M for MXF4); block_k = 128 bytes (32 packed int32 = 256 FP4). +// block_n: pick by wave count + composite score = est_stages^2 * bn, with a +// 2-epi-stage tiebreak when waves >= 2 (TMEM double-buffer needs >= 2 tiles/SM). +// cluster_m=2, cluster_n=1 (B-multicast) when M >= 512 for Normal/KGroupedContiguous. +// swap_ab: enabled for MGroupedContiguous when expected_m_per_group < BLOCK_M +// (sparse MoE); forces block_n=128 + cluster=1 per kernel asserts. static Layout pick_fp4_layout(const GemmType& gemm_type, const int& m, const int& n, const int& k, const int& num_groups, const int& num_sms, @@ -172,7 +144,8 @@ static Layout pick_fp4_layout(const GemmType& gemm_type, constexpr int sf_block_m_cols = (128 / 32) * sf_pk; // 8 constexpr int smem_capacity = 232448; - // BLOCK_N legality (mirrors fea-fp4's is_fp4_block_n_legal) + // BLOCK_N legality: only TMEM column budget; B-tensor OOB is handled by TMA + // store-drop, and SFB smem padding is zero-filled before warp transpose. auto is_legal = [&](int bn) { if (bn % 16 != 0 || bn < 16 || bn > 256) return false; const int sf_block_n = (bn + 127) / 128 * 128; @@ -226,7 +199,8 @@ static Layout pick_fp4_layout(const GemmType& gemm_type, } // Multicast: B-multicast (cluster_m=2, cluster_n=1) when M >= 512 for Normal / KGrouped. - // For m-grouped, fea-fp4 keeps cluster=1 (m_indices iteration breaks multi-CTA M-distribution). + // m-grouped types use cluster=1: multi-CTA M-distribution is incompatible + // with m_indices iteration in the kernel scheduler. int cluster_m = 1, cluster_n = 1; const bool can_multicast = (m >= 512) && (ceil_div(m, block_m) % 2 == 0) @@ -294,15 +268,14 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa config.pipeline_config.num_stages = new_stages; config.pipeline_config.smem_size = new_smem; - // View FP4-packed-int8 tensors as int32 (8 FP4 per int32). Same memory layout, - // different dtype tag -- makes the TMA descriptor use INT32 (not 16U4_ALIGN16B - // unpacked-smem), which matches my kernel's int32-packed smem expectation. + // View FP4-packed-int8 tensors as int32 (8 FP4 per int32). Same memory + // layout, different dtype tag — the TMA descriptor then uses INT32 (not + // 16U4_ALIGN16B unpacked-smem), matching the kernel's int32-packed smem. const auto a_int32 = a.view(torch::kInt); const auto b_int32 = b.view(torch::kInt); const int k_int32 = k / 8; const int block_k_int32 = config.layout.block_k / 4; - const auto cd = c.value_or(d); const auto tensor_map_a = make_tma_a_desc(major_a, a_int32, m, k_int32, config.storage_config.load_block_m, block_k_int32, @@ -318,16 +291,12 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa config.storage_config.store_block_n, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - const auto tensor_map_c = make_tma_cd_desc(cd, m, n, - config.storage_config.store_block_m, - config.storage_config.store_block_n, - static_cast(cd.stride(-2)), 1, - config.storage_config.swizzle_cd_mode); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, config.layout.block_m, gran_k, 1, 0); const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, config.layout.block_n, gran_k, 1, 0); + // Pre-merge C into D for accumulation; the kernel has no separate C path. if (c.has_value()) { if (c->data_ptr() == d.data_ptr()) { DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides()); @@ -342,13 +311,11 @@ static void sm100_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, config.pipeline_config.smem_size, config.layout.get_cluster_size()), - .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, - .tensor_map_c = tensor_map_c, .tensor_map_d = tensor_map_d }; const auto code = SM100FP4Gemm1D1DRuntime::generate(args); @@ -415,13 +382,11 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, config.pipeline_config.smem_size, config.layout.get_cluster_size()), - .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), .grouped_layout = grouped_layout.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, - .tensor_map_c = tensor_map_d, .tensor_map_d = tensor_map_d }; const auto code = SM100FP4Gemm1D1DRuntime::generate(args); @@ -442,9 +407,8 @@ static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const t major_a, major_b, d.scalar_type(), false, compiled_dims, /*expected_m=*/expected_m, /*expected_num_groups=*/num_groups); - // m-grouped masked: pass expected_m as per-group hint (fea-fp4 invariant); - // swap_ab heuristic itself only activates for MGroupedContiguous per fea-fp4 v0 - // (masked + swap_ab had NaN issues unresolved upstream). + // m-grouped masked: pass expected_m as per-group hint; swap_ab gating + // inside pick_fp4_layout only activates for MGroupedContiguous. const auto layout = pick_fp4_layout(GemmType::MGroupedMasked, m, n, k, num_groups, device_runtime->get_num_sms(), expected_m); auto config = GemmConfig{ @@ -489,13 +453,11 @@ static void sm100_m_grouped_fp4_gemm_masked_1d1d(const torch::Tensor& a, const t .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, config.pipeline_config.smem_size, config.layout.get_cluster_size()), - .num_last_stages = compute_num_last_stages_fp4(k, config.layout.block_k, config.pipeline_config.num_stages), .grouped_layout = masked_m.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, - .tensor_map_c = tensor_map_d, .tensor_map_d = tensor_map_d }; const auto code = SM100FP4Gemm1D1DRuntime::generate(args); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh index 0d59d9bf52..d1d0d407e0 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh @@ -26,7 +26,7 @@ template ; @@ -129,8 +128,6 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, cute::prefetch_tma_descriptor(&tensor_map_sfa); cute::prefetch_tma_descriptor(&tensor_map_sfb); cute::prefetch_tma_descriptor(&tensor_map_d); - if constexpr (kWithAccumulation) - cute::prefetch_tma_descriptor(&tensor_map_c); } // D/A/B shared memory @@ -187,12 +184,14 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, } cutlass::arch::fence_view_async_shared(); cutlass::arch::fence_barrier_init(); - } - else if (threadIdx.x >= 32 and threadIdx.x < 64) { + } else if (threadIdx.x >= 32 and threadIdx.x < 64) { Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); } kNumMulticast > 1 ? cute::cluster_sync() : __syncthreads(); + // 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); From 959e308af8078e59d8eabccd7142915d55e6ebba Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Mon, 1 Jun 2026 00:47:31 -0700 Subject: [PATCH 10/12] FP4: remove .item() host sync in m-grouped contiguous swap_ab gating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the per-launch device→host sync `(grouped_layout >= 0).sum().item() / num_groups` with a host-side hint plumbed through the dispatch: - sm100_m_grouped_fp4_gemm_contiguous_1d1d gains an `expected_m_per_group` int parameter, consumed directly by pick_fp4_layout for swap_ab gating. - csrc/apis/gemm.hpp dispatch passes `expected_m_for_psum_layout.value_or(m / num_groups)`: - Caller-known sparse MoE workloads (DeepSeek-style routed expert shapes) pass the actual useful per-group row count via the existing `expected_m_for_psum_layout` keyword and keep the swap_ab speedup. - Dense callers get `m / num_groups`, which matches the actual per-group count when m_indices has no padding. No Python API surface change; the optional `expected_m_for_psum_layout` argument was already available. Removes the only host sync inside FP4 grouped dispatch — kernel is now safe to capture in a CUDA graph. Build clean. test_fp4.py PASS (exit 0). Production-shape perf unchanged (swap_ab didn't fire for these dense shapes either way). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/apis/gemm.hpp | 8 +++++++- csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 10 +++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 6b3f12e192..ca24a53fef 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -215,8 +215,14 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair 0 ? m / num_groups : 0); sm100_m_grouped_fp4_gemm_contiguous_1d1d(a.first, sfa, b.first, sfb, d, grouped_layout, - num_groups, m, n, k, major_a, major_b, compiled_dims); + num_groups, m, n, k, expected_m_per_group, + major_a, major_b, compiled_dims); } else { sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(a.first, sfa, b.first, sfb, d, grouped_layout, num_groups, m, n, k, gran_k_a, gran_k_b, major_a, major_b, diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index 298922a997..e6d959773b 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -328,18 +328,18 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con const torch::Tensor& d, const torch::Tensor& grouped_layout, const int& num_groups, const int& m, const int& n, const int& k, + const int& expected_m_per_group, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const std::string& compiled_dims) { constexpr int gran_k = 32; const auto desc = make_fp4_desc(GemmType::MGroupedContiguous, m, n, k, num_groups, major_a, major_b, d.scalar_type(), false, compiled_dims); - // For m-grouped contiguous, derive useful-per-group row count from m_indices - // (count of valid rows / num_groups). Drives swap_ab gating for sparse MoE. - const auto useful_m = (grouped_layout >= 0).sum().item(); - const int useful_per_group = num_groups > 0 ? static_cast(useful_m) / num_groups : 0; + // Caller-provided useful-per-group hint drives swap_ab gating. Dense layouts + // (no -1 padding) get m/num_groups; sparse MoE callers pass the actual + // per-group useful row count. No device tensor inspection on the host. const auto layout = pick_fp4_layout(GemmType::MGroupedContiguous, m, n, k, num_groups, - device_runtime->get_num_sms(), useful_per_group); + device_runtime->get_num_sms(), expected_m_per_group); auto config = GemmConfig{ .layout = layout, .storage_config = SM100ArchSpec::get_storage_config(desc, layout), From e315c8b3d73e563910af252b7a2725206632354b Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Mon, 1 Jun 2026 01:05:16 -0700 Subject: [PATCH 11/12] FP4: address codex re-audit (K-major guard, k%8 assert, dedup C copy, style) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-audit follow-ups from FP4_PR_REVIEW_NOTES.md: 1. Guard the FP4xFP4 dispatch against MN-major / transposed inputs. `a.view(torch::kInt)` requires `stride(-1) == 1`, which the API's nn/tn/tt aliases break. Add explicit major checks at all three FP4xFP4 dispatch sites in csrc/apis/gemm.hpp: DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == ...); 2. Assert `k % 8 == 0` at the same dispatch sites. The wrapper divides logical FP4 K by 8 to get packed int32 K; non-8-multiple K would be silently truncated. 3. Drop the redundant `c -> d` copy in sm100_fp4_gemm_1d1d. early_return() in csrc/apis/gemm.hpp already performs the merge before dispatch. 4. Stop overloading expected_m_for_psum_layout as a sparse-MoE hint — the shared layout-check code asserts it is unset when use_psum_layout is false, so the .value_or() never fired. The FP4xFP4 contiguous path now uses m / num_groups unconditionally and the comment reflects that. 5. Comment cleanup: drop "==========" banners and "Existing non-swap epilogue (unchanged)" wording in the .cuh, drop the "hardcodes" / "TODO" / "proper fp32 reference" notes in tests/test_fp4.py and the "matches main's block_k convention" line in the .hpp. Build clean, full test_fp4.py PASS (exit 0, 0 FAIL). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/apis/gemm.hpp | 17 +++++++--- .../jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 13 ++------ .../deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh | 31 ++++++++----------- tests/test_fp4.py | 10 +++--- 4 files changed, 32 insertions(+), 39 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index ca24a53fef..31a6f5bc99 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -100,6 +100,10 @@ static void fp8_fp4_gemm_nt(const std::pair& a, } } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) { if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) { + // FP4xFP4 path reinterprets int8 packed tensors as int32 for TMA; requires + // contiguous K-major layout and packed-int32-aligned K. + DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); + DG_HOST_ASSERT(k % 8 == 0); sm100_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims); } else { @@ -215,11 +219,12 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair 0 ? m / num_groups : 0); + DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); + DG_HOST_ASSERT(k % 8 == 0); + // swap_ab gating uses m/num_groups as the per-group estimate (accurate + // for dense layouts; sparse MoE without a hint conservatively skips + // swap_ab). + const int expected_m_per_group = num_groups > 0 ? m / num_groups : 0; sm100_m_grouped_fp4_gemm_contiguous_1d1d(a.first, sfa, b.first, sfb, d, grouped_layout, num_groups, m, n, k, expected_m_per_group, major_a, major_b, compiled_dims); @@ -294,6 +299,8 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pairdata_ptr() == d.data_ptr()) { - DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides()); - } else { - d.copy_(c.value()); - } - } - + // C is merged into D by the API's early_return() pre-launch; the kernel + // has no separate C load path. const SM100FP4Gemm1D1DRuntime::Args args = { .gemm_desc = desc, .gemm_config = config, diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh index d1d0d407e0..742ef4d40d 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_gemm_1d1d.cuh @@ -221,14 +221,13 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, accum_stage_idx == 0 ? func(0) : func(1); }; - // ========== Warp dispatch (FP8-style: independent loops per warp) ========== - // Warp 0: TMA load producer - // Warp 1: MMA consumer (+ UTCCP SF copy to TMEM) - // Warp 2: SF transpose (SMEM warp transpose for UTCCP) - // Warp 3+: Epilogue - + // Dispatch warps into different roles: + // warp 0 : TMA load producer + // warp 1 : MMA consumer + UTCCP SF copy to TMEM + // warp 2 : SF SMEM warp transpose for UTCCP + // warp 3+ : Epilogue if (warp_idx == 0) { - // ========== Warp 0: TMA load ========== + // TMA load warp while (scheduler.get_next_block(m_block_idx, n_block_idx)) { launch_k_iterations([&](uint32_t k_iter, auto type, bool is_last_iter, uint32_t num_last_stages) { constexpr bool kHasDivisibleStages = cute::is_same_v; @@ -297,7 +296,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, }); } } else if (warp_idx == 1 and is_leader_cta) { - // ========== Warp 1: UTCCP SF copy + MMA ========== + // MMA + UTCCP SF copy warp constexpr uint32_t UMMA_M = LAYOUT_AD_M * (kIsMulticastOnA ? 1 : kNumMulticast); // Swap-AB: UMMA_N becomes BLOCK_M (the original M dim now plays N role in MMA). constexpr uint32_t UMMA_N = kSwapAB ? BLOCK_M : BLOCK_N * (kIsMulticastOnA ? kNumMulticast : 1); @@ -463,7 +462,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, }); } } else if (warp_idx == 2) { - // ========== Warp 2: SF transpose ========== + // SF transpose warp auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, "Invalid aligned elements"); uint32_t values[4]; @@ -542,14 +541,10 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(cd_dtype_t); if constexpr (kSwapAB) { - // ================================================================= - // Swap-AB epilogue: - // TMEM holds D^T (BLOCK_N rows × BLOCK_M cols). - // Per accum stage covers cols [accum_stage_idx*BLOCK_M, +BLOCK_M) of TMEM. - // STORE_BLOCK_M=16: each `s` iter covers 16 cols of TMEM (= 16 of D's M). - // num_stores = effective_m / 16 → padding cols are skipped entirely. - // STORE_BLOCK_N = BLOCK_N: each TMA store covers entire BLOCK_N at once. - // ================================================================= + // Swap-AB epilogue: TMEM holds D^T (BLOCK_N rows x BLOCK_M cols); each + // accum stage covers BLOCK_M TMEM cols. STORE_BLOCK_M=16 lets the loop + // skip 16-row M slices that are entirely padding (effective_m / 16), + // while STORE_BLOCK_N=BLOCK_N stores the full N tile per iteration. constexpr uint32_t kNumSwizzleAtomRows = 8; constexpr uint32_t STORE_BLOCK_N_ATOM = kSwizzleCDMode / sizeof(cd_dtype_t); constexpr uint32_t kNumWarpsPerAtom = STORE_BLOCK_N_ATOM / 32; @@ -650,7 +645,7 @@ sm100_fp4_gemm_1d1d_impl(int* grouped_layout, }); } } else { - // ===== Existing non-swap epilogue (unchanged) ===== + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { dispatch_accum_stage_idx(scheduler.current_iter % kNumEpilogueStages, [&](uint32_t accum_stage_idx) { auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; diff --git a/tests/test_fp4.py b/tests/test_fp4.py index 45b49536c1..cb3d136a4a 100644 --- a/tests/test_fp4.py +++ b/tests/test_fp4.py @@ -11,9 +11,8 @@ # FP4xFP4 (MXFP4): both operands packed FP4 (E2M1) with VS=32 UE8M0 scales. -# Routes through deep_gemm.fp8_fp4_gemm_nt / m_grouped_fp8_fp4_gemm_* APIs and -# dispatches to the SM100_MMA_MXF4_SS path in sm100_fp4_gemm_1d1d.cuh. -# NOTE: the kernel currently only supports FP32 output (bf16 epilogue is a TODO). +# Dispatches to the SM100_MMA_MXF4_SS path; the API rejects bf16 D so we +# always allocate fp32 D below. FP4_FP4 = QuantConfig((32, 32, True, True)) @@ -62,9 +61,8 @@ def test_m_grouped_gemm_contiguous() -> None: num_groups, expected_m_per_group, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, use_ue8m0=True, quant_config=FP4_FP4) - # Re-generate with bf16 disabled-cast path to get a proper fp32 reference. - # generate_m_grouped_contiguous hardcodes bf16 d/ref; we allocate fp32 ourselves - # and cast the ref for the diff check. + # generate_m_grouped_contiguous returns bf16 d/ref; allocate fp32 d for + # the FP4xFP4 API and cast the reference for calc_diff. d = torch.empty_like(d_bf16, dtype=out_dtype) ref_d = _ref_bf16.to(out_dtype) deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( From fc974c3b91c0f2c9ad6e2d2da75865af3a6c5d94 Mon Sep 17 00:00:00 2001 From: Runchu Zhao Date: Mon, 1 Jun 2026 01:58:56 -0700 Subject: [PATCH 12/12] FP4: clarify swap_ab gating comment in contiguous wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous comment claimed sparse MoE callers pass an actual per-group row count, but no such path exists — the API passes m / num_groups unconditionally. Rewrite to describe the real behavior: dense layouts are accurate, sparse layouts with -1 padding over-estimate and conservatively disable swap_ab, which is the accepted trade-off vs a device sync to inspect m_indices. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Runchu Zhao --- csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp index 27db33d864..da4fa4fc27 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_gemm_1d1d.hpp @@ -328,9 +328,10 @@ static void sm100_m_grouped_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, con const auto desc = make_fp4_desc(GemmType::MGroupedContiguous, m, n, k, num_groups, major_a, major_b, d.scalar_type(), false, compiled_dims); - // Caller-provided useful-per-group hint drives swap_ab gating. Dense layouts - // (no -1 padding) get m/num_groups; sparse MoE callers pass the actual - // per-group useful row count. No device tensor inspection on the host. + // swap_ab is gated on expected_m_per_group. The API passes m / num_groups + // (a host-side estimate that avoids inspecting m_indices). Accurate for + // dense layouts; over-estimates for sparse MoE with -1 padding, which + // conservatively disables swap_ab. Acceptable trade-off vs a device sync. const auto layout = pick_fp4_layout(GemmType::MGroupedContiguous, m, n, k, num_groups, device_runtime->get_num_sms(), expected_m_per_group); auto config = GemmConfig{