Skip to content

feat: support mxfp4 × mxfp4 gemm on sm100 #348

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
z52527 wants to merge 12 commits into
base: nv_dev
Choose a base branch
from
+1,386 −9
Open

feat: support mxfp4 × mxfp4 gemm on sm100 #348

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
7d27b8a
FP4: port MXF4 kernel + wrappers onto main's refactored APIs
z52527 May 29, 2026
7c2d685
FP4: progress port towards runtime — IMA fixed, smem cap, multicast c…
z52527 May 29, 2026
6539975
FP4: layout override to fea-fp4 tested config + test int8 packing
z52527 May 29, 2026
24eb9c2
FP4: tests/test_fp4.py — use FP32 output (kernel's only working epilo…
z52527 Jun 1, 2026
ae98142
FP4: tests/test_fp4.py — full suite passes on synced branch (10/10)
z52527 Jun 1, 2026
dfdc5ae
FP4: port fea-fp4's full get_best_fp4_config heuristic — perf parity
z52527 Jun 1, 2026
31aef98
FP4: rewrite tests/test_fp4.py to mirror test_fp8_fp4.py style
z52527 Jun 1, 2026
8d841f6
FP4: align .cuh comments with sm100_fp8_fp4_gemm_1d1d.cuh style
z52527 Jun 1, 2026
ae4e77d
FP4: address codex review (PDL sync, strict FP32 D, prune unused fields)
z52527 Jun 1, 2026
959e308
FP4: remove .item() host sync in m-grouped contiguous swap_ab gating
z52527 Jun 1, 2026
e315c8b
FP4: address codex re-audit (K-major guard, k%8 assert, dedup C copy,…
z52527 Jun 1, 2026
fc974c3
FP4: clarify swap_ab gating comment in contiguous wrapper
z52527 Jun 1, 2026
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
64 changes: 55 additions & 9 deletions csrc/apis/gemm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -73,7 +74,12 @@ static void fp8_fp4_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& 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))
Expand All @@ -93,8 +99,17 @@ static void fp8_fp4_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& 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) {
// 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 {
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");
}
Expand Down Expand Up @@ -166,7 +181,12 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
const auto [m_, n_] = get_shape<2>(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 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
Expand Down Expand Up @@ -197,9 +217,22 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
num_groups, m, n, k, major_a, major_b, major_sfb,
compiled_dims, use_psum_layout, expected_m_for_psum_layout);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
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,
compiled_dims, use_psum_layout, expected_m_for_psum_layout);
if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) {
DG_HOST_ASSERT(not use_psum_layout);
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);
} 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,
compiled_dims, use_psum_layout, expected_m_for_psum_layout);
}
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
Expand Down Expand Up @@ -244,7 +277,12 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair<torch::Tensor, torc
DG_HOST_ASSERT(num_groups == num_groups_ and num_groups == num_groups__ and num_groups == num_groups___);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
DG_HOST_ASSERT(expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
// 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
Expand All @@ -260,9 +298,17 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair<torch::Tensor, torc
sm90_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, major_sfb, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, gran_k_a, gran_k_b,
if (a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 8 == 0);
sm100_m_grouped_fp4_gemm_masked_1d1d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m,
major_a, major_b, compiled_dims);
} else {
sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, gran_k_a, gran_k_b,
major_a, major_b, compiled_dims);
}
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
Expand Down
Loading