Skip to content

Add mxfp4_block_size parameter to f4f4bf16_grouped_mm for MXFP4_16 support (#496) #496

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ghjeong12 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add mxfp4_block_size parameter to f4f4bf16_grouped_mm for MXFP4_16 support (#496) #496

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 31 additions & 14 deletions csrc/gemm/cutlass/f4f4bf16_grouped.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,8 @@ Kernel_f4f4bf16_grouped get_kernel_via_tuning(
std::optional<at::Tensor> offsets = std::nullopt,
std::optional<at::Tensor> M_sizes = std::nullopt,
std::optional<at::Tensor> global_scale = std::nullopt,
std::optional<at::Tensor> starting_row_after_padding = std::nullopt) {
std::optional<at::Tensor> starting_row_after_padding = std::nullopt,
int64_t mxfp4_block_size = 32) {
static TuningCache cache("f4f4bf16_grouped");

M = nextPowerOf2OrRoundUp(M, 1024, 1024);
Expand All @@ -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;
}

Expand All @@ -283,7 +285,8 @@ at::Tensor dispatch_fp4_grouped_kernel(
std::optional<at::Tensor> M_sizes = std::nullopt,
std::optional<at::Tensor> global_scale = std::nullopt,
std::optional<at::Tensor> 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.");
Expand All @@ -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);
}();
Expand All @@ -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(
Expand All @@ -327,7 +332,8 @@ at::Tensor f4f4bf16_grouped_stacked(
at::Tensor M_sizes,
std::optional<at::Tensor> global_scale = std::nullopt,
std::optional<at::Tensor> 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);
Expand Down Expand Up @@ -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(
Expand All @@ -380,7 +387,8 @@ at::Tensor f4f4bf16_grouped_mm(
at::Tensor w_scale,
at::Tensor offsets,
std::optional<at::Tensor> output_maybe,
std::optional<at::Tensor> global_scale = std::nullopt) {
std::optional<at::Tensor> 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.");
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -481,7 +496,8 @@ at::Tensor f4f4bf16_grouped_mm(
at::Tensor w_scale,
at::Tensor offsets,
std::optional<at::Tensor> output,
std::optional<at::Tensor> global_scale = std::nullopt) {
std::optional<at::Tensor> 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
}
Expand All @@ -494,7 +510,8 @@ at::Tensor f4f4bf16_grouped_stacked(
at::Tensor M_sizes,
std::optional<at::Tensor> global_scale = std::nullopt,
std::optional<at::Tensor> 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
}
Expand Down
24 changes: 21 additions & 3 deletions csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<at::Tensor> offsets,
std::optional<at::Tensor> M_sizes,
std::optional<at::Tensor> global_scale,
std::optional<at::Tensor> starting_row_after_padding) {
std::optional<at::Tensor> starting_row_after_padding,
int64_t mxfp4_block_size) {
if (global_scale) {
return f4f4bf16_grouped_impl<NVFP4, 128, 128, 256, 1, 1, 1>(
XQ,
Expand All @@ -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<MXFP4_16, 128, 128, 256, 1, 1, 1>(
XQ,
WQ,
x_scale,
w_scale,
output,
offsets,
M_sizes,
global_scale,
starting_row_after_padding);
} else {
return f4f4bf16_grouped_impl<MXFP4, 128, 128, 256, 1, 1, 1>(
XQ,
Expand Down
24 changes: 21 additions & 3 deletions csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<at::Tensor> offsets,
std::optional<at::Tensor> M_sizes,
std::optional<at::Tensor> global_scale,
std::optional<at::Tensor> starting_row_after_padding) {
std::optional<at::Tensor> starting_row_after_padding,
int64_t mxfp4_block_size) {
if (global_scale) {
return f4f4bf16_grouped_impl<NVFP4, 128, 64, 256, 1, 1, 1>(
XQ,
Expand All @@ -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<MXFP4_16, 128, 128, 256, 1, 1, 1>(
XQ,
WQ,
x_scale,
w_scale,
output,
offsets,
M_sizes,
global_scale,
starting_row_after_padding);
} else {
return f4f4bf16_grouped_impl<MXFP4, 128, 64, 256, 1, 1, 1>(
XQ,
Expand Down
30 changes: 24 additions & 6 deletions csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<at::Tensor> offsets,
std::optional<at::Tensor> M_sizes,
std::optional<at::Tensor> global_scale,
std::optional<at::Tensor> starting_row_after_padding) {
std::optional<at::Tensor> starting_row_after_padding,
int64_t mxfp4_block_size) {
if (global_scale) {
return f4f4bf16_grouped_impl<NVFP4, 256, 64, 256, 2, 1, 1>(
return f4f4bf16_grouped_impl<NVFP4, 256, 128, 256, 2, 1, 1>(
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<MXFP4_16, 128, 128, 256, 1, 1, 1>(
XQ,
WQ,
x_scale,
Expand All @@ -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<MXFP4, 256, 64, 256, 2, 1, 1>(
return f4f4bf16_grouped_impl<MXFP4, 256, 128, 256, 2, 1, 1>(
XQ,
WQ,
x_scale,
Expand Down
24 changes: 21 additions & 3 deletions csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_128_2_1_1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<at::Tensor> offsets,
std::optional<at::Tensor> M_sizes,
std::optional<at::Tensor> global_scale,
std::optional<at::Tensor> starting_row_after_padding) {
std::optional<at::Tensor> starting_row_after_padding,
int64_t mxfp4_block_size) {
if (global_scale) {
return f4f4bf16_grouped_impl<NVFP4, 256, 256, 128, 2, 1, 1>(
XQ,
Expand All @@ -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<MXFP4_16, 128, 128, 256, 1, 1, 1>(
XQ,
WQ,
x_scale,
w_scale,
output,
offsets,
M_sizes,
global_scale,
starting_row_after_padding);
} else {
return f4f4bf16_grouped_impl<MXFP4, 256, 256, 128, 2, 1, 1>(
XQ,
Expand Down
24 changes: 21 additions & 3 deletions csrc/gemm/cutlass/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<at::Tensor> offsets,
std::optional<at::Tensor> M_sizes,
std::optional<at::Tensor> global_scale,
std::optional<at::Tensor> starting_row_after_padding) {
std::optional<at::Tensor> starting_row_after_padding,
int64_t mxfp4_block_size) {
if (global_scale) {
return f4f4bf16_grouped_impl<NVFP4, 256, 256, 256, 2, 1, 1>(
XQ,
Expand All @@ -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<MXFP4_16, 128, 128, 256, 1, 1, 1>(
XQ,
WQ,
x_scale,
w_scale,
output,
offsets,
M_sizes,
global_scale,
starting_row_after_padding);
} else {
return f4f4bf16_grouped_impl<MXFP4, 256, 256, 256, 2, 1, 1>(
XQ,
Expand Down
Loading
Loading