Skip to content

Performance Parity for H100_mma_ABt and H100_mma Kernels #98

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

Merged
DanFu09 merged 7 commits into from
Jun 16, 2025
Merged

Performance Parity for H100_mma_ABt and H100_mma Kernels #98

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5d8dddf
Performance on MMA_ABt +60 TFLOPS by using larger wgmma instructions
Mar 5, 2025
ace5ba2
MMA_ABt works on non-uniform sizes
Mar 5, 2025
d8afe81
Merge branch 'HazyResearch:main' into mma-abt-performance
SohamGovande Jun 15, 2025
f2b081a
.
SohamGovande Jun 15, 2025
3fc1a35
.
SohamGovande Jun 15, 2025
d8c6b60
.
SohamGovande Jun 15, 2025
138f687
.
SohamGovande Jun 15, 2025
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
22 changes: 11 additions & 11 deletions kernels/matmul/H100_mma_ABt/matmul.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,14 @@ struct matmul_template {
zero(args.state.accum[n]);
}
__device__ static void compute(consumer_compute_args<layout> args) {
for(int n = 0; n < N_BLOCK; n++) {
warpgroup::mma_ABt(
args.state.accum[n],
args.input.a[warpgroup::groupid()],
args.input.b[n]
);
}
using wide_rt = rt_fl<16, 64*N_BLOCK>;
using tall_st = st_bf<64*N_BLOCK, 64>;
// dispatch the largest possible tensor core instruction to maximize TFLOPS (64x16x256 on M_BLOCK=2, N_BLOCK=4)
warpgroup::mma_ABt(
reinterpret_cast<wide_rt&>(args.state.accum),
args.input.a[warpgroup::groupid()],
reinterpret_cast<tall_st&>(args.input.b)
);
warpgroup::mma_async_wait();
if(laneid() == 0) arrive(args.inputs_finished);
}
Expand Down Expand Up @@ -128,7 +129,7 @@ void inner_run(bf16 *d_A, bf16 *d_B, bf16 *d_C, size_t M, size_t N, size_t K, di
using globals = typename mmt::layout::globals;
// printf("M: %d, N: %d, K: %d\n", M, N, K);
global_layout Ag{d_A, nullptr, nullptr, M, K};
global_layout Bg{d_B, nullptr, nullptr, K, N};
global_layout Bg{d_B, nullptr, nullptr, N, K};
global_layout Cg{d_C, nullptr, nullptr, M, N};
globals G{Ag, Bg, Cg};
prototype::lcf::kernel<mmt><<<grid, block, MAX_SHARED_MEMORY-1024>>>(G);
Expand Down Expand Up @@ -280,8 +281,7 @@ int run_benchmark(size_t M, size_t N, size_t K) {
}

int main() {
int N;
N = 4096;
run_benchmark<matmul_template<2,4,8>>(N, N, N);
int M = 2048, N = 4096, K = 8192;
run_benchmark<matmul_template<2,4,8>>(M, N, K);
return 0;
}