[nihalpasham]

What This PR Does

This PR re-authors the Blackwell (sm_100a) kernels in gemm_sol_final (zero compiler changes) to beat live cuBLASLt across all three target sizes (4096³, 8192³, and 16384³) on NVIDIA B200:

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

---

1. 7-Stage Contiguous SMEM Pipeline + Register-Direct 128-Bit Epilogue (U32x4)

Previously, gemm_sol_final reserved 64 KiB of shared memory (SMEM_OUT) to stage epilogue outputs through stmatrix.m8n8.x2 + ld.shared, which capped the TMA/UMMA input ring at 4 stages (128 KiB).

BEFORE (4 SMEM stages + 64 KiB `SMEM_OUT`, 64/96 regs):
  ┌──────────────────────────┬──────────────────┐
  │ 4 × 32 KiB A/B (128 KiB) │ SMEM_OUT (64 KiB)│   tcgen05.ld ──► stmatrix ──► bar.warp.sync ──► ld.shared ──► st.global.b64
  └──────────────────────────┴──────────────────┘

AFTER (7 SMEM stages + 0 KiB `SMEM_OUT`, 60 regs):
  ┌─────────────────────────────────────────────┐
  │ 7 × 32 KiB Contiguous A/B Ring (224 KiB)    │   tcgen05.ld.16x256b.x8 ──► shfl.sync.bfly (xor 1, 2) ──► st.global.v4.b32
  └─────────────────────────────────────────────┘

• Eliminates SMEM_OUT completely (64 KiB -> 0 B) by transposing adjacent thread pairs in registers via two warp-butterfly shuffles (xor 1 and xor 2) and issuing 128-bit st.global.v4.b32 (U32x4) stores directly from registers.
• Expands the TMA $\leftrightarrow$ UMMA ring from 4 stages to 7 stages (224 KiB) while dropping register pressure from 64/96 registers down to 60 registers (0 spills).
• Signals ACCUM_EMPTY from a single elected thread per warp (lane_id == 0, arrival count 8 instead of 256) immediately after the final tcgen05.wait::ld, releasing TMEM before the final global stores complete.

---

2. Wave-Tail Tile Balancing Across 74 B200 CTA Pairs (M256×N192 & M256×N224)

On B200 (148 SMs = 74 cta_group::2 cluster pairs), uniform 256×256 tiles leave part of the GPU idle in the final wave:

Problem Size   Uniform 256×256 Waves on 74 CTA Pairs         Wave-Tail Balanced Schedule (This PR)
────────────   ───────────────────────────────────────────   ──────────────────────────────────────────────────────────
4096³          256 tiles / 74 = 3.46 waves                   Replaces last 40 `3×256` (768-col) triplets with `4×192`
               (40 of 74 CTA pairs sit idle in wave 4!)      (`M256×N192`) tiles ──► 296 = 4 × 74 work items (100% fill)
                                                             Cuts waves 3 & 4 from `1.0 + 1.0` to `0.75 + 0.75`!

8192³          1024 tiles / 74 = 13.84 waves                 Replaces last 12 `7×256` (1792-col) heptets with `8×224`
               (12 of 74 CTA pairs sit idle in wave 14)      (`M256×N224`) tiles ──► 1036 = 14 × 74 work items

16384³         4096 tiles / 74 = 55.35 waves                 Replaces last 48 `7×256` (1792-col) heptets with `8×224`
               (48 of 74 CTA pairs sit idle in wave 56!)     (`M256×N224`) tiles ──► 4144 = 56 × 74 work items

---

3. Overlapped CLC Work-Stealing & Hoisted SMEM Descriptors

• Asynchronous CLC prefetch (k_idx + 4 == k_iters): Issues clc_try_cancel_multicast 4 K-stages before the end of each tile into a 4-slot TILE_INFO / TILE_READY ring buffer, hiding the GigaThread work-stealing latency behind the last 4 UMMA stages.
• Hoisted SMEM descriptors: Precomputes a_desc0 and b_desc0 once before the tile loop and advances per stage via a_desc0 + (stage << 10) + j * 2, removing 36 64-bit shift/mask/or instructions per K-stage in MMA_WARP.
• L2 Promotion: Enables CU_TENSOR_MAP_L2_PROMOTION_L2_128B on the 2D swizzled TMA descriptors.

---

Verification

• Blackwell B200 (sm_100a) Full Correctness & Live Benchmark (cargo oxide run gemm_sol_final):
  • gemm_sol_clc_multicast_4_stage_pipeline (M256xN256) 4096³ full 16,777,216-element exact-BF16 check: PASS (0.0000 max rel err)
  • gemm_sol_clc_multicast_4_stage_pipeline_large (M512xN256) 4096³ full 16,777,216-element exact-BF16 check: PASS (0.0000 max rel err)

Problem Size	Before (main)	After (This PR)	Live cuBLASLt top-8	cuBLASLt Exhaustive*	vs. cuBLASLt top-8
4096³	1475.5 TFLOPS (93.15 us)	1857.5 TFLOPS (73.99 us)	1839.6 TFLOPS (74.81 us)	1927.0 TFLOPS	100.97% (+17.9 TFLOPS)
8192³	1772.1 TFLOPS (620.47 us)	2084.7 TFLOPS (527.42 us)	2077.2 TFLOPS (529.37 us)	2084.0 TFLOPS	100.36% (+7.5 TFLOPS)
16384³	1992.5 TFLOPS (4414.58 us)	2121.1 TFLOPS (4146.98 us)	2063.9 TFLOPS (4263.79 us)	2103.0 TFLOPS	102.77% (+57.2 TFLOPS)

Test plan

• cargo oxide run gemm_sol_final on B200 (sm_100a) (full 16,777,216-element exact-BF16 validation + 4K/8K/16K benchmark vs. live cuBLASLt)
• cargo test --manifest-path crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml
• cargo clippy --manifest-path crates/rustc-codegen-cuda/examples/gemm_sol_final/Cargo.toml --all-targets -- -D warnings
• cargo oxide fmt --check

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.