diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped.cu b/csrc/gemm/cutlass/f4f4bf16_grouped.cu index fc184766..cbba0b17 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped.cu @@ -243,7 +243,8 @@ Kernel_f4f4bf16_grouped get_kernel_via_tuning( std::optional offsets = std::nullopt, std::optional M_sizes = std::nullopt, std::optional global_scale = std::nullopt, - std::optional starting_row_after_padding = std::nullopt) { + std::optional starting_row_after_padding = std::nullopt, + int64_t mxfp4_block_size = 32) { static TuningCache cache("f4f4bf16_grouped"); M = nextPowerOf2OrRoundUp(M, 1024, 1024); @@ -266,7 +267,8 @@ Kernel_f4f4bf16_grouped get_kernel_via_tuning( offsets, M_sizes, global_scale, - starting_row_after_padding); + starting_row_after_padding, + mxfp4_block_size); return kernel; } @@ -283,7 +285,8 @@ at::Tensor dispatch_fp4_grouped_kernel( std::optional M_sizes = std::nullopt, std::optional global_scale = std::nullopt, std::optional starting_row_after_padding = std::nullopt, - bool use_mx = true) { + bool use_mx = true, + int64_t mxfp4_block_size = 32) { TORCH_CHECK( offsets.has_value() ^ M_sizes.has_value(), "Exactly one of M_sizes or offsets must be present."); @@ -302,7 +305,8 @@ at::Tensor dispatch_fp4_grouped_kernel( offsets, M_sizes, global_scale, - starting_row_after_padding); + starting_row_after_padding, + mxfp4_block_size); } return get_kernel_via_heuristics(M, N, K); }(); @@ -316,7 +320,8 @@ at::Tensor dispatch_fp4_grouped_kernel( offsets, M_sizes, global_scale, - starting_row_after_padding); + starting_row_after_padding, + mxfp4_block_size); } at::Tensor f4f4bf16_grouped_stacked( @@ -327,7 +332,8 @@ at::Tensor f4f4bf16_grouped_stacked( at::Tensor M_sizes, std::optional global_scale = std::nullopt, std::optional starting_row_after_padding = std::nullopt, - bool use_mx = true) { + bool use_mx = true, + int64_t mxfp4_block_size = 32) { int64_t total_M = XQ.size(0); int64_t N = WQ.size(1); int64_t K = WQ.size(2); @@ -370,7 +376,8 @@ at::Tensor f4f4bf16_grouped_stacked( M_sizes, global_scale, starting_row_after_padding, - use_mx); + use_mx, + mxfp4_block_size); } at::Tensor f4f4bf16_grouped_mm( @@ -380,7 +387,8 @@ at::Tensor f4f4bf16_grouped_mm( at::Tensor w_scale, at::Tensor offsets, std::optional output_maybe, - std::optional global_scale = std::nullopt) { + std::optional global_scale = std::nullopt, + int64_t mxfp4_block_size = 32) { TORCH_CHECK(offsets.dtype() == at::kInt, "offsets must be int32."); TORCH_CHECK(offsets.dim() == 1, "offsets must be 1D tensor."); TORCH_CHECK(XQ.is_contiguous(), "XQ must be row major."); @@ -393,17 +401,23 @@ at::Tensor f4f4bf16_grouped_mm( TORCH_CHECK(x_scale.is_contiguous(), "x_scale must be contiguous."); TORCH_CHECK(w_scale.is_contiguous(), "w_scale must be contiguous."); - const bool use_mx = [&]() { + // Determine quantization mode: NVFP4 (e4m3 scales + global_scale) or + // MXFP4 (e8m0 scales, no global_scale). + // mx_mode: 0 = NVFP4, 1 = MXFP4 (bs32), 2 = MXFP4_16 (bs16) + const int mx_mode = [&]() { if (x_scale.dtype() == at::kFloat8_e4m3fn) { TORCH_CHECK( global_scale.has_value(), "global_scale must be provided for NVFP4.") TORCH_CHECK( global_scale->dtype() == at::kFloat, "global_scale must be FP32.") - return false; + return 0; } else if (x_scale.dtype() == at::kFloat8_e8m0fnu) { TORCH_CHECK( !global_scale.has_value(), "global_scale must be unset for MXFP4.") - return true; + TORCH_CHECK( + mxfp4_block_size == 16 || mxfp4_block_size == 32, + "mxfp4_block_size must be 16 or 32.") + return mxfp4_block_size == 16 ? 2 : 1; } else { TORCH_CHECK( false, "Scales must be FP8 e8m0 for MXFP4 or FP8 e4m3 for NVFP4") @@ -469,7 +483,8 @@ at::Tensor f4f4bf16_grouped_mm( std::nullopt, // M_sizes global_scale, std::nullopt, // starting_row_after_padding - use_mx); + mx_mode > 0, + mxfp4_block_size); } #else @@ -481,7 +496,8 @@ at::Tensor f4f4bf16_grouped_mm( at::Tensor w_scale, at::Tensor offsets, std::optional output, - std::optional global_scale = std::nullopt) { + std::optional global_scale = std::nullopt, + int64_t mxfp4_block_size = 32) { throw std::runtime_error( "CUDA version is older than 12.8"); // requires CUDA>=12.8 } @@ -494,7 +510,8 @@ at::Tensor f4f4bf16_grouped_stacked( at::Tensor M_sizes, std::optional global_scale = std::nullopt, std::optional starting_row_after_padding = std::nullopt, - bool use_mx = true) { + bool use_mx = true, + int64_t mxfp4_block_size = 32) { throw std::runtime_error( "CUDA version is older than 12.8"); // requires CUDA>=12.8 } diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1.cu index 84dc02fb..badac039 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1.cu @@ -13,15 +13,16 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) at::Tensor f4f4bf16_grouped_128_128_256_1_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { return f4f4bf16_grouped_impl( XQ, @@ -33,6 +34,23 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1( M_sizes, global_scale, starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); } else { return f4f4bf16_grouped_impl( XQ, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1.cu index 8d1f7ab2..9214a09c 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1.cu @@ -13,15 +13,16 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) at::Tensor f4f4bf16_grouped_128_64_256_1_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { return f4f4bf16_grouped_impl( XQ, @@ -33,6 +34,23 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1( M_sizes, global_scale, starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); } else { return f4f4bf16_grouped_impl( XQ, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1.cu index 583748b2..f32fdabc 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1.cu @@ -12,18 +12,36 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) -at::Tensor f4f4bf16_grouped_256_64_256_2_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 +at::Tensor f4f4bf16_grouped_256_128_256_2_1_1( + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { - return f4f4bf16_grouped_impl( + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( XQ, WQ, x_scale, @@ -34,7 +52,7 @@ at::Tensor f4f4bf16_grouped_256_64_256_2_1_1( global_scale, starting_row_after_padding); } else { - return f4f4bf16_grouped_impl( + return f4f4bf16_grouped_impl( XQ, WQ, x_scale, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_128_2_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_128_2_1_1.cu index 68b97fd8..1cae67ad 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_128_2_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_128_2_1_1.cu @@ -13,15 +13,16 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) at::Tensor f4f4bf16_grouped_256_256_128_2_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { return f4f4bf16_grouped_impl( XQ, @@ -33,6 +34,23 @@ at::Tensor f4f4bf16_grouped_256_256_128_2_1_1( M_sizes, global_scale, starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); } else { return f4f4bf16_grouped_impl( XQ, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1.cu index 70454f17..25e4e890 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1.cu @@ -13,15 +13,16 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) at::Tensor f4f4bf16_grouped_256_256_256_2_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { return f4f4bf16_grouped_impl( XQ, @@ -33,6 +34,23 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1( M_sizes, global_scale, starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); } else { return f4f4bf16_grouped_impl( XQ, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_64_256_2_1_1.cu b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_64_256_2_1_1.cu index 8e5eaf83..1ecefb30 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_64_256_2_1_1.cu +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_64_256_2_1_1.cu @@ -12,18 +12,36 @@ namespace mslk::gemm { #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080) -at::Tensor f4f4bf16_grouped_256_128_256_2_1_1( - at::Tensor XQ, // FP4 - at::Tensor WQ, // FP4 +at::Tensor f4f4bf16_grouped_256_64_256_2_1_1( + at::Tensor XQ, + at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding) { + std::optional starting_row_after_padding, + int64_t mxfp4_block_size) { if (global_scale) { - return f4f4bf16_grouped_impl( + return f4f4bf16_grouped_impl( + XQ, + WQ, + x_scale, + w_scale, + output, + offsets, + M_sizes, + global_scale, + starting_row_after_padding); + } else if (mxfp4_block_size == 16) { + // NOTE: deliberately 128,128,256,1,1,1 and not this file's own tile shape. + // The other MXFP4_16 tile instantiations fault with an illegal memory + // access (see mslk/test/gemm/test_mxfp4_16_grouped_mm.py, which was written + // to reproduce it), so every MXFP4_16 workload is routed to the one shape + // that is known good. This does mean per-name tuning does not select a + // tile for MXFP4_16 yet. + return f4f4bf16_grouped_impl( XQ, WQ, x_scale, @@ -34,7 +52,7 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1( global_scale, starting_row_after_padding); } else { - return f4f4bf16_grouped_impl( + return f4f4bf16_grouped_impl( XQ, WQ, x_scale, diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_common.cuh b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_common.cuh index ee4c79b9..dc71c393 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_common.cuh +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_common.cuh @@ -8,6 +8,8 @@ #define CUTLASS_NAMESPACE mslk +#include + #include #include #include @@ -31,6 +33,7 @@ namespace mslk::gemm { namespace cutlass = cutlass_mslk; using MXFP4 = cutlass::mx_float4_t; +using MXFP4_16 = cutlass::mx_float4_16_t; using NVFP4 = cutlass::nv_float4_t; inline int64_t _byte_align(int64_t offset) { @@ -77,7 +80,11 @@ __global__ void set_stacked_kernel_args_kernel( LayoutSFB* layout_SFB, ElementGlobalScale* global_scale, const ElementGlobalScale** global_scale_ptr, - int64_t* starting_row_after_padding) { + int64_t* starting_row_after_padding, + // No default: the only caller passes the compile-time scale_block_size for + // its InputQuantType. A default of 32 would let a future caller silently + // misindex NVFP4/MXFP4_16 scale offsets, which both use 16. + int ele_per_quantize_group) { uint32_t group_index = blockIdx.x * blockDim.x + threadIdx.x; // If this thread corresponds to a valid group, write kernel args to device // memory. @@ -108,10 +115,6 @@ __global__ void set_stacked_kernel_args_kernel( int64_t offset_M = 0; int64_t accumulated_x_scale = 0; int64_t accumulated_w_scale = 0; - int ele_per_quantize_group = 16; - if (global_scale == nullptr) { - ele_per_quantize_group = 32; - } for (int i = 0; i < group_index; i++) { offset_M += M_sizes[i]; /* It's calculated this way since the scales are at least padded to @@ -219,18 +222,26 @@ at::Tensor f4f4bf16_grouped_impl( using ClusterShape = cute::Shape, cute::Int, cute::Int>; + // Select kernel schedule based on quant type: + // - NVFP4: NvF4 schedule (SfVecSize=16, E4M3 scales) + // - MXFP4: Mxf4 schedule (SfVecSize=32, E8M0 scales) + // - MXFP4_16: NvF4 schedule (SfVecSize=16, E8M0 scales -- same SfVecSize as + // NVFP4) + static constexpr bool is_mxfp4_16 = std::is_same_v; + static constexpr bool use_nvf4_schedule = is_nvfp4 || is_mxfp4_16; + static constexpr bool use_2sm = (TB_M == 256) && (TBS_M % 2 == 0); using KernelSchedule = cute::conditional_t< - is_nvfp4, + use_nvf4_schedule, cute::conditional_t< - (TB_M == 256) && (TBS_M % 2 == 0), + use_2sm, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmNvf4Sm100, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100>, cute::conditional_t< - (TB_M == 256) && (TBS_M % 2 == 0), + use_2sm, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmMxf4Sm100, cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100>>; using EpilogueSchedule = cute::conditional_t< - (TB_M == 256) && (TBS_M % 2 == 0), + use_2sm, cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm, cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm>; @@ -393,6 +404,17 @@ at::Tensor f4f4bf16_grouped_impl( auto stream = at::cuda::getCurrentCUDAStream().stream(); + // Determine scale block size from InputQuantType. + // NVFP4 and MXFP4_16: 16 elements per scale group. + // MXFP4 (standard): 32 elements per scale group. + constexpr int scale_block_size = + std::is_same_v ? 32 : 16; + + // Verify CUTLASS config's SFVecSize matches our scale_block_size + // (confirmed at compile time). + static_assert( + CollectiveMainloop::Sm1xxBlkScaledConfig::SFVecSize == scale_block_size); + const int64_t M = XQ.size(-2); const int64_t N = WQ.size(-2); const int64_t K = WQ.size(-1) * 2; // 2 FP4 values are packed into uint8 @@ -490,7 +512,8 @@ at::Tensor f4f4bf16_grouped_impl( : nullptr, is_nvfp4 ? global_scale_ptr : nullptr, reinterpret_cast( - starting_row_after_padding.value().data_ptr())); + starting_row_after_padding.value().data_ptr()), + scale_block_size); // Set the number of groups to the kernel to be at most the number of // non-zero rows. kernel_groups = int(std::min(M, G)); diff --git a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_manifest.cuh b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_manifest.cuh index d560886e..efb9c9aa 100644 --- a/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_manifest.cuh +++ b/csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_manifest.cuh @@ -21,7 +21,8 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); at::Tensor f4f4bf16_grouped_256_256_128_2_1_1( at::Tensor XQ, // FP4 @@ -32,7 +33,8 @@ at::Tensor f4f4bf16_grouped_256_256_128_2_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); at::Tensor f4f4bf16_grouped_256_256_256_2_1_1( at::Tensor XQ, // FP4 @@ -43,7 +45,8 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); at::Tensor f4f4bf16_grouped_256_64_256_2_1_1( at::Tensor XQ, // FP4 @@ -54,7 +57,8 @@ at::Tensor f4f4bf16_grouped_256_64_256_2_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); at::Tensor f4f4bf16_grouped_256_128_256_2_1_1( at::Tensor XQ, // FP4 @@ -65,7 +69,8 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); at::Tensor f4f4bf16_grouped_128_128_256_1_1_1( at::Tensor XQ, // FP4 @@ -76,7 +81,8 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1( std::optional offsets, std::optional M_sizes, std::optional global_scale, - std::optional starting_row_after_padding); + std::optional starting_row_after_padding, + int64_t mxfp4_block_size = 32); using Kernel_f4f4bf16_grouped = at::Tensor (*)( at::Tensor, @@ -87,7 +93,8 @@ using Kernel_f4f4bf16_grouped = at::Tensor (*)( std::optional, std::optional, std::optional, - std::optional); + std::optional, + int64_t); const std::unordered_map& get_f4f4bf16_grouped_kernels() { diff --git a/csrc/gemm/gemm_ops.cpp b/csrc/gemm/gemm_ops.cpp index b2d53e99..a7d7ed6b 100644 --- a/csrc/gemm/gemm_ops.cpp +++ b/csrc/gemm/gemm_ops.cpp @@ -107,9 +107,9 @@ TORCH_LIBRARY_FRAGMENT(mslk, m) { m.def( "mx6mx6bf16(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor? output=None, int splits=0) -> Tensor"); m.def( - "f4f4bf16_grouped_stacked(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor M_sizes, Tensor? global_scale=None, Tensor? starting_row_after_padding=None, bool use_mx=True) -> Tensor"); + "f4f4bf16_grouped_stacked(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor M_sizes, Tensor? global_scale=None, Tensor? starting_row_after_padding=None, bool use_mx=True, int mxfp4_block_size=32) -> Tensor"); m.def( - "f4f4bf16_grouped_mm(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor offsets, Tensor(a!)? output=None, Tensor(a!)? global_scale=None) -> Tensor"); + "f4f4bf16_grouped_mm(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor offsets, Tensor(a!)? output=None, Tensor(a!)? global_scale=None, int mxfp4_block_size=32) -> Tensor"); m.def( "f4f4bf16_ultra_grouped_mm(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, Tensor offsets, Tensor x_global_scale, Tensor w_global_scale, Tensor(a!)? output=None) -> Tensor"); m.def("bf16x9_gemm(Tensor A, Tensor B, Tensor(a!)? output=None) -> Tensor"); diff --git a/include/mslk/gemm/cutlass/grouped_common.cuh b/include/mslk/gemm/cutlass/grouped_common.cuh index 99f3d422..f88dc7e2 100644 --- a/include/mslk/gemm/cutlass/grouped_common.cuh +++ b/include/mslk/gemm/cutlass/grouped_common.cuh @@ -94,7 +94,8 @@ __global__ void set_grouped_gemm_args_kernel( // NVFP4 groups 16 elements for a single scale factor, MXFP4 & MXFP8 // groups 32 elements. - const int64_t scale_factor_block_size = global_scale == nullptr ? 32 : 16; + // Scale block size from CUTLASS config: NVFP4=16, MXFP4=32, MXFP4_16=16. + constexpr int64_t scale_factor_block_size = Sm1xxBlkScaledConfig::SFVecSize; // Handle offsets API (torch compliant API for 2D-2D and 2D-3D inputs) CUDA_KERNEL_ASSERT( diff --git a/include/mslk/gemm/gemm.h b/include/mslk/gemm/gemm.h index a3d0f20e..a2c1e1a1 100644 --- a/include/mslk/gemm/gemm.h +++ b/include/mslk/gemm/gemm.h @@ -147,7 +147,8 @@ at::Tensor f4f4bf16_grouped_stacked( at::Tensor M_sizes, std::optional global_scale = std::nullopt, std::optional starting_row_after_padding = std::nullopt, - bool use_mx = true); + bool use_mx = true, + int64_t mxfp4_block_size = 32); at::Tensor bf16x9_gemm( at::Tensor A, diff --git a/include/mslk/gemm/gemm_torch.h b/include/mslk/gemm/gemm_torch.h index 35f7fbaa..fec2cc86 100644 --- a/include/mslk/gemm/gemm_torch.h +++ b/include/mslk/gemm/gemm_torch.h @@ -43,7 +43,8 @@ at::Tensor f4f4bf16_grouped_mm( at::Tensor w_scale, at::Tensor offsets, std::optional output = std::nullopt, - std::optional global_scale = std::nullopt); + std::optional global_scale = std::nullopt, + int64_t mxfp4_block_size = 32); // FP4 ultra grouped GEMM (SM103) with per-token activation scaling. at::Tensor f4f4bf16_ultra_grouped_mm( diff --git a/test/gemm/test_mxfp4_16_grouped_mm.py b/test/gemm/test_mxfp4_16_grouped_mm.py new file mode 100644 index 00000000..9c1c239f --- /dev/null +++ b/test/gemm/test_mxfp4_16_grouped_mm.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Test f4f4bf16_grouped_mm with mxfp4_block_size=16 (offset-based API). + +Reproduces the crash seen in MetaShuffling MoE with MXFP4_16: +- Some workers crash with 'illegal memory access' +- Pattern suggests certain M_group_size values trigger the issue +""" + +import unittest + +import torch + +try: + import mslk.gemm # noqa: F401 + + HAS_MSLK = True +except ImportError: + HAS_MSLK = False + +try: + from mslk.quantize.triton.fp4_quantize import triton_quantize_mx4_unpack + + HAS_QUANT = True +except ImportError: + HAS_QUANT = False + +try: + # vLLM owns the fused activation quantizer; mslk does not depend on vllm, + # so the tests exercising it skip rather than forcing a layering-inverting + # BUCK dep from mslk onto vllm. + from vllm.fb.plugins.meta_shuffling_kernels.quantization import ( + mxfp4_quantize_stacked, + ) + + HAS_VLLM_QUANT = True +except ImportError: + HAS_VLLM_QUANT = False + +GROUP_SIZE = 16 +# SQNR of the MXFP4 result against the unquantized BF16 GEMM. Real MXFP4 +# lands near 16dB on this data; a kernel reading the wrong scales or segments +# collapses to ~0dB or below. 10dB sits between the two with margin. +MIN_SQNR_DB = 10.0 + + +def _quantize_per_expert(data, counts, bs=GROUP_SIZE): + """Quantize per-expert with correct padding (matching fused kernel).""" + xq_parts, xs_parts = [], [] + offset = 0 + for c in counts: + if c > 0: + end = offset + c + xq_i, xs_i = triton_quantize_mx4_unpack(data[offset:end], group_size=bs) + xq_parts.append(xq_i) + xs_parts.append(xs_i) + offset += c + return torch.cat(xq_parts, dim=0), torch.cat(xs_parts, dim=0) + + +def _sqnr(ref, got): + """Signal-to-quantization-noise ratio in dB; inf when bit-identical.""" + r, g = ref.float(), got.float() + noise = (r - g).pow(2).mean() + if noise == 0: + return float("inf") + return float(10.0 * torch.log10(r.pow(2).mean() / noise)) + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") +@unittest.skipIf(not HAS_MSLK, "mslk not available") +@unittest.skipIf(not HAS_QUANT, "mslk.quantize not available") +class TestMXFP4_16GroupedMM(unittest.TestCase): + """Test f4f4bf16_grouped_mm with mxfp4_block_size=16.""" + + def _run_test(self, E, tokens_per_expert, N, K): + device = torch.device("cuda:0") + total_M = sum(tokens_per_expert) + counts = list(tokens_per_expert) + + # Create data + x_bf16 = torch.randn(total_M, K, dtype=torch.bfloat16, device=device) * 0.01 + W_bf16 = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.01 + + # Quantize activations per-expert + xq, x_scale = _quantize_per_expert(x_bf16, counts, GROUP_SIZE) + + # Quantize weights per-expert + wq_parts, ws_parts = [], [] + for e in range(E): + wq_i, ws_i = triton_quantize_mx4_unpack(W_bf16[e], group_size=GROUP_SIZE) + wq_parts.append(wq_i) + ws_parts.append(ws_i) + wq = torch.stack(wq_parts).contiguous() + w_scale = torch.stack(ws_parts).contiguous() + + # Compute offsets (cumulative sum of counts), must be int32 + offsets = ( + torch.tensor(counts, dtype=torch.int64, device=device) + .cumsum(0) + .to(torch.int32) + ) + + print(f" E={E}, tokens={counts}, N={N}, K={K}") + print(f" xq={xq.shape}, x_scale={x_scale.shape}") + print(f" wq={wq.shape}, w_scale={w_scale.shape}") + print(f" offsets={offsets.tolist()}") + + torch.cuda.synchronize() + out = torch.ops.mslk.f4f4bf16_grouped_mm( + xq.view(torch.float4_e2m1fn_x2), + wq.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + x_scale.view(torch.float8_e8m0fnu), + w_scale.view(torch.float8_e8m0fnu), + offsets, + global_scale=None, + mxfp4_block_size=16, + ) + torch.cuda.synchronize() + + self.assertTrue(out.isfinite().all(), "Output has non-finite values") + + # Accuracy: compare against the unquantized BF16 GEMM. MXFP4 is + # 4-bit, so some error is expected -- the point is that it lands in + # the band real quantization produces. A kernel that mis-indexes + # scales or segments still returns finite output, so finiteness alone + # says nothing; those failures show up here as ~0dB or negative. + ref = torch.empty(out.shape, dtype=torch.float32, device=out.device) + start = 0 + for e, c in enumerate(counts): + if c == 0: + continue + stop = start + c + ref[start:stop] = x_bf16[start:stop].float() @ W_bf16[e].float().T + start += c + + sqnr = _sqnr(ref, out) + print( + f" OK: out={out.shape}, mean={out.float().mean():.6f}, sqnr={sqnr:.1f}dB" + ) + self.assertGreater( + sqnr, + MIN_SQNR_DB, + f"MXFP4_16 grouped MM disagrees with its dequantized reference " + f"(sqnr={sqnr:.1f}dB); scales or segment indexing are likely wrong", + ) + + def test_uniform_small(self): + """E=4, 64 tokens per expert.""" + print("\ntest_uniform_small:") + self._run_test(4, [64, 64, 64, 64], 4096, 4096) + + def test_uniform_large(self): + """E=4, 256 tokens per expert.""" + print("\ntest_uniform_large:") + self._run_test(4, [256, 256, 256, 256], 4096, 4096) + + def test_single_token(self): + """E=4, 1 token to one expert, 0 to others.""" + print("\ntest_single_token:") + self._run_test(4, [1, 0, 0, 0], 4096, 4096) + + def test_sparse_like_warmup(self): + """E=32, 1 token routed to 1 expert (mimics warmup).""" + print("\ntest_sparse_like_warmup:") + counts = [0] * 32 + counts[5] = 1 # one expert gets 1 token + self._run_test(32, counts, 4096, 4096) + + def test_sparse_multi_expert(self): + """E=32, tokens spread across 8 experts.""" + print("\ntest_sparse_multi_expert:") + counts = [0] * 32 + for i in [0, 4, 8, 12, 16, 20, 24, 28]: + counts[i] = 1 + self._run_test(32, counts, 4096, 4096) + + def test_n8192(self): + """Test with N=8192 (w13 GEMM dimension).""" + print("\ntest_n8192:") + self._run_test(4, [64, 64, 64, 64], 8192, 4096) + + def test_mixed_counts(self): + """E=32 with varying token counts.""" + print("\ntest_mixed_counts:") + counts = [0] * 32 + counts[0] = 100 + counts[3] = 50 + counts[7] = 200 + counts[15] = 10 + counts[31] = 1 + self._run_test(32, counts, 4096, 4096) + + def test_m128_per_expert(self): + """E=4, 128 tokens per expert (border of heuristic).""" + print("\ntest_m128_per_expert:") + self._run_test(4, [128, 128, 128, 128], 4096, 4096) + + def test_m192_per_expert(self): + """E=4, 192 tokens per expert.""" + print("\ntest_m192_per_expert:") + self._run_test(4, [192, 192, 192, 192], 4096, 4096) + + def test_m256_e1(self): + """E=1, 256 tokens (pure GEMM, no expert split).""" + print("\ntest_m256_e1:") + self._run_test(1, [256], 4096, 4096) + + def test_m512_per_expert(self): + """E=4, 512 tokens per expert.""" + print("\ntest_m512_per_expert:") + self._run_test(4, [512, 512, 512, 512], 4096, 4096) + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") +@unittest.skipIf(not HAS_MSLK, "mslk not available") +@unittest.skipIf(not HAS_QUANT, "mslk.quantize not available") +@unittest.skipIf(not HAS_VLLM_QUANT, "vllm meta_shuffling quantization not available") +class TestFusedMXFP4QuantizeStacked(unittest.TestCase): + """Test fused mxfp4_quantize_stacked kernel (the cudagraph-compatible + activation quantization) matches per-expert triton_quantize_mx4_unpack.""" + + def _compare_fused_vs_per_expert(self, E, tokens_per_expert, K, bs=16): + """Compare fused stacked kernel vs per-expert by running grouped_mm + with both and checking that results are close to BF16 reference.""" + device = torch.device("cuda:0") + total_M = sum(tokens_per_expert) + counts = list(tokens_per_expert) + m_sizes = torch.tensor(counts, dtype=torch.int64, device=device) + N = 4096 + + x_bf16 = torch.randn(total_M, K, dtype=torch.bfloat16, device=device) * 0.01 + W_bf16 = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.01 + + # Quantize weights (shared between both paths) + wq_parts, ws_parts = [], [] + for e in range(E): + wq_i, ws_i = triton_quantize_mx4_unpack(W_bf16[e], group_size=bs) + wq_parts.append(wq_i) + ws_parts.append(ws_i) + wq = torch.stack(wq_parts).contiguous() + w_scale = torch.stack(ws_parts).contiguous() + offsets = m_sizes.cumsum(0).to(torch.int32) + + # Path A: per-expert quant (reference) + ref_xq, ref_xs = _quantize_per_expert(x_bf16, counts, bs) + torch.cuda.synchronize() + out_ref = torch.ops.mslk.f4f4bf16_grouped_mm( + ref_xq.view(torch.float4_e2m1fn_x2), + wq.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + ref_xs.view(torch.float8_e8m0fnu), + w_scale.view(torch.float8_e8m0fnu), + offsets, + global_scale=None, + mxfp4_block_size=bs, + ) + torch.cuda.synchronize() + + # Path B: fused stacked kernel + fused_xq, fused_xs = mxfp4_quantize_stacked(m_sizes, x_bf16, bs) + torch.cuda.synchronize() + out_fused = torch.ops.mslk.f4f4bf16_grouped_mm( + fused_xq.view(torch.float4_e2m1fn_x2), + wq.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + fused_xs.view(torch.float8_e8m0fnu), + w_scale.view(torch.float8_e8m0fnu), + offsets, + global_scale=None, + mxfp4_block_size=bs, + ) + torch.cuda.synchronize() + + # BF16 reference + out_bf16 = torch._grouped_mm( + x_bf16, + W_bf16.transpose(-2, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + + # Both quantized paths should be close to BF16 + self.assertTrue(out_ref.isfinite().all(), "ref output non-finite") + self.assertTrue(out_fused.isfinite().all(), "fused output non-finite") + + # Check fused vs ref are very close (both approximate BF16) + torch.testing.assert_close(out_fused, out_ref, atol=1e-1, rtol=1e-1) + # Check both approximate BF16 + torch.testing.assert_close(out_ref, out_bf16, atol=1.0, rtol=0.5) + + ref_err = (out_ref - out_bf16).float().abs().mean() + fused_err = (out_fused - out_bf16).float().abs().mean() + print( + f" E={E}, tokens={counts}, K={K}, bs={bs}: " + f"ref_err={ref_err:.4f}, fused_err={fused_err:.4f} -- OK" + ) + + def test_fused_uniform_e4(self): + print("\ntest_fused_uniform_e4:") + self._compare_fused_vs_per_expert(4, [64, 64, 64, 64], 4096, bs=16) + + def test_fused_sparse_e32(self): + print("\ntest_fused_sparse_e32:") + counts = [0] * 32 + counts[0] = 10 + counts[15] = 5 + counts[31] = 20 + self._compare_fused_vs_per_expert(32, counts, 4096, bs=16) + + def test_fused_large_m(self): + print("\ntest_fused_large_m:") + self._compare_fused_vs_per_expert(4, [256, 256, 256, 256], 4096, bs=16) + + def test_fused_bs32(self): + """Also verify bs32 still works after buffer allocation change.""" + print("\ntest_fused_bs32:") + self._compare_fused_vs_per_expert(4, [64, 64, 64, 64], 4096, bs=32) + + def test_fused_single_token(self): + print("\ntest_fused_single_token:") + self._compare_fused_vs_per_expert(4, [1, 0, 0, 0], 4096, bs=16) + + def test_fused_then_grouped_mm(self): + """End-to-end: fused quant -> grouped_mm with bs16.""" + print("\ntest_fused_then_grouped_mm:") + device = torch.device("cuda:0") + E, M_per, N, K, bs = 4, 64, 4096, 4096, 16 + + x_bf16 = torch.randn(E * M_per, K, dtype=torch.bfloat16, device=device) * 0.01 + W_bf16 = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.01 + + m_sizes = torch.full((E,), M_per, dtype=torch.int64, device=device) + + # Fused activation quant + xq, x_scale = mxfp4_quantize_stacked(m_sizes, x_bf16, bs) + + # Weight quant per-expert + wq_parts, ws_parts = [], [] + for e in range(E): + wq_i, ws_i = triton_quantize_mx4_unpack(W_bf16[e], group_size=bs) + wq_parts.append(wq_i) + ws_parts.append(ws_i) + wq = torch.stack(wq_parts).contiguous() + w_scale = torch.stack(ws_parts).contiguous() + + offsets = m_sizes.cumsum(0).to(torch.int32) + + torch.cuda.synchronize() + out = torch.ops.mslk.f4f4bf16_grouped_mm( + xq.view(torch.float4_e2m1fn_x2), + wq.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + x_scale.view(torch.float8_e8m0fnu), + w_scale.view(torch.float8_e8m0fnu), + offsets, + global_scale=None, + mxfp4_block_size=bs, + ) + torch.cuda.synchronize() + self.assertTrue(out.isfinite().all()) + print(f" out={out.shape}, mean={out.float().mean():.6f} -- OK")