Skip to content

perf(cuda-oxide): 7-stage SMEM pipeline and wave-tail tile balancing in gemm_sol_final #1663

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
nihalpasham merged 1 commit into from
Oct 27, 2026
+930 −989
Merged

perf(cuda-oxide): 7-stage SMEM pipeline and wave-tail tile balancing in gemm_sol_final #1663

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
29 changes: 25 additions & 4 deletions cuda-oxide/crates/rustc-codegen-cuda/examples/gemm_sol_final/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,42 @@ The goal is simple: show how close a Rust/cuda-oxide kernel can get to the
closest supported live cuBLASLt GEMM on the same GPU, without weakening the
workload or correctness check.

## Measured B200 throughput vs cuBLASLt

```text
Visual Bar Order on B200 Panel (0 – 2.2 PFLOP/s axis):

4096³ cuda-oxide ██████████████████████████████████████████ 1.857 (beats top-8!)
cuBLASLt top-8 █████████████████████████████████████████ 1.840
Exhaustive* ███████████████████████████████████████████ 1.927

8192³ cuda-oxide ███████████████████████████████████████████████ 2.085 (beats BOTH!)
cuBLASLt top-8 ███████████████████████████████████████████████ 2.080
Exhaustive* ███████████████████████████████████████████████ 2.084

16384³ cuda-oxide ████████████████████████████████████████████████ 2.121 (beats BOTH!)
cuBLASLt top-8 ██████████████████████████████████████████████ 2.065
Exhaustive* ███████████████████████████████████████████████ 2.103
```

## Final design

```text
KernelFamily: gemm_sol_final/output_tile@1
├── M256xN256 -> existing small resource envelope
└── M512xN256 -> existing large resource envelope
├── M256xN256 -> 4096 / 8192 entry
└── M512xN256 -> 16384 entry

Auto selector:
4096 / 8192 -> M256xN256
16384 -> M512xN256
```

Both entries use CLC work scheduling, a two-CTA cluster, `cta_group::2` UMMA,
four TMA/shared-memory pipeline stages, two TMEM accumulator halves,
compiler-directed K-loop unrolling, and L2-aware output-tile ordering.
a seven-stage TMA/shared-memory pipeline (`224 KiB` of contiguous A/B stages),
two TMEM accumulator halves, register-direct 128-bit epilogue stores (`U32x4`),
wave-tail tile balancing (`M256×N192` / `M256×N224`), overlapped CLC
work-stealing, compiler-directed K-loop unrolling, and L2-aware output-tile
ordering.

## What moved performance

Expand Down
Loading
Loading