submission 780604

  #!POPCORN leaderboard matmul_v2
  #!POPCORN gpu B200
   
- # FP8 matmul via torch._scaled_mm with per-row / per-col scaling.
- # Caches quantized inputs + scales by pointer; bot benchmark reuses A/B across
- # iterations so first call pays quantization cost (~300us), subsequent are
- # ~60us (vs cuBLAS 115us).
- # Output always bf16 internally (fp16 output has a bug in scaled_mm with
- # rowwise scaling), then cast to requested dtype.
+ # matmul_v8_clc: v8_layoutA + Cluster Launch Control dynamic tile assignment.
+ #
+ # Goal: eliminate the 0.65-wave tail effect (35% of SMs idle on the last partial
+ # wave) by using Blackwell's CLC hardware to dynamically reassign canceled
+ # clusters' work to clusters that finish early.
+ #
+ # Architecture:
+ #   - Launch grid = (cluster_n * CTA_GROUP, cluster_m, 1) so each CTA's
+ #     (blockIdx.x, blockIdx.y) maps to (n-tile, m-pair) directly.
+ #   - Add a 5th warp role: the "scheduler" warp (warp 0).
+ #       * NUM_CLC_STAGES = 2 SMEM response slots + corresponding mbarriers.
+ #       * At kernel entry: prime both stages with try_cancel.
+ #       * Loop: wait clc_empty[s], refill stage s with try_cancel.
+ #   - Producer/consumer/epilogue warps: while(true) { do work for current tile;
+ #     wait clc_full[s]; read (m_pair, n) from response; arrive clc_empty[s];
+ #     advance s; if not valid break; }
+ #   - The first "tile" each cluster handles is its own (blockIdx.x/CTA_GROUP,
+ #     blockIdx.y) — no try_cancel on that one. Subsequent tiles come from CLC.
+  
+ import os
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":32768:8"
+  
  from task import input_t, output_t
  import torch
   
  ⋯ 6 unchanged lines
   
   
  _CUDA_SRC = r"""
+ #include <cstdint>
+ #include <cstring>
+ #include <cudaTypedefs.h>
+ #include <torch/library.h>
  #include <ATen/ATen.h>
- #include <torch/torch.h>
- void matmul_cublas(const torch::Tensor& A, const torch::Tensor& B, torch::Tensor& out) {
-     at::matmul_out(out, A, B);
+ #include <cuda_fp16.h>
+  
+ constexpr int WARP_SIZE = 32;
+  
+ template <typename T>
+ __device__ inline
+ T warp_uniform(T x) { return __shfl_sync(0xFFFF'FFFF, x, 0); }
+  
+ __device__ inline
+ uint32_t elect_sync() {
+   uint32_t pred = 0;
+   asm volatile(
+     "{\n\t"
+     ".reg .pred %%px;\n\t"
+     "elect.sync _|%%px, %1;\n\t"
+     "@%%px mov.s32 %0, 1;\n\t"
+     "}"
+     : "+r"(pred)
+     : "r"(0xFFFFFFFF)
+   );
+   return pred;
  }
- """
- _CPP_SRC = "void matmul_cublas(const torch::Tensor&, const torch::Tensor&, torch::Tensor&);"
- _mod = load_inline(
-     name="matmul_fp8_cublas_fallback",
-     cpp_sources=_CPP_SRC, cuda_sources=_CUDA_SRC,
-     functions=["matmul_cublas"],
-     extra_cuda_cflags=["-O3", "-arch=sm_100"], extra_cflags=["-O3"], verbose=False)
   
+ __device__ inline
+ void mbarrier_init(int mbar_addr, int count) {
+   asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
+ }
   
- _FP8_MAX = 448.0
- _CACHE: dict = {}
+ __device__ inline
+ void mbarrier_wait(int mbar_addr, int phase) {
+   uint32_t ticks = 0x989680;
+   asm volatile(
+     "{\n\t"
+     ".reg .pred P1;\n\t"
+     "LAB_WAIT:\n\t"
+     "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
+     "@P1 bra.uni DONE;\n\t"
+     "bra.uni LAB_WAIT;\n\t"
+     "DONE:\n\t"
+     "}"
+     :: "r"(mbar_addr), "r"(phase), "r"(ticks)
+   );
+ }
   
+ __device__ inline
+ void mbarrier_arrive_expect_tx(int mbar_addr, int size) {
+   asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;"
+               :: "r"(mbar_addr), "r"(size) : "memory");
+ }
   
- def _prepare_fp8(A: torch.Tensor, B: torch.Tensor):
-     # Per-row scale for A (shape M), per-col scale for B (shape N).
-     max_a = A.abs().amax(dim=1).clamp(min=1e-8).float()
-     max_b = B.abs().amax(dim=0).clamp(min=1e-8).float()
-     scale_a = (max_a / _FP8_MAX).contiguous()  # (M,)
-     scale_b = (max_b / _FP8_MAX).contiguous()  # (N,)
+ // Per-CTA (cluster-aware) arrive+expect_tx: arrives on a remote cluster CTA's
+ // shared mbarrier. Used to signal CLC full barrier across cluster.
+ __device__ inline
+ void mbarrier_arrive_expect_tx_remote(int mbar_addr, int size, int dst_cta_rank) {
+   asm volatile(
+     "{\n\t"
+     ".reg .b32 remAddr32;\n\t"
+     "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
+     "mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %2;\n\t"
+     "}"
+     :: "r"(mbar_addr), "r"(dst_cta_rank), "r"(size)
+     : "memory"
+   );
+ }
   
-     A_fp8 = (A.float() / scale_a.view(-1, 1)).to(torch.float8_e4m3fn).contiguous()  # (M,K) row-major
-     # Make B into column-major fp8: compute fp8 of B^T row-major then view-transpose.
-     B_t_rm = (B.float().t().contiguous() / scale_b.view(-1, 1)).to(torch.float8_e4m3fn)  # (N,K) row-major
-     B_fp8 = B_t_rm.t()  # (K,N) column-major view
+ __device__ inline
+ void mbarrier_arrive(int mbar_addr) {
+   asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" :: "r"(mbar_addr) : "memory");
+ }
   
-     return A_fp8, B_fp8, scale_a.view(-1, 1), scale_b.view(1, -1)
+ template <int CTA_GROUP = 1>
+ __device__ inline
+ void tma_3d_g2s(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) {
+   asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.cta_group::%6 "
+               "[%0], [%1, {%2, %3, %4}], [%5];"
+               :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "n"(CTA_GROUP)
+               : "memory");
+ }
   
+ template <int CTA_GROUP = 1>
+ __device__ inline
+ void tcgen05_alloc(int smem_addr, int size) {
+   asm volatile("tcgen05.alloc.cta_group::%2.sync.aligned.shared::cta.b32 [%0], %1;"
+               :: "r"(smem_addr), "r"(size), "n"(CTA_GROUP));
+ }
   
- def custom_kernel(data: input_t) -> output_t:
-     A, B, out = data
+ template <int CTA_GROUP = 1>
+ __device__ inline
+ void tcgen05_dealloc(int taddr, int size) {
+   asm volatile("tcgen05.dealloc.cta_group::%2.sync.aligned.b32 %0, %1;"
+               :: "r"(taddr), "r"(size), "n"(CTA_GROUP));
+ }
   
-     if A.dtype == torch.bfloat16 and A.is_cuda and B.is_cuda:
-         try:
-             key = (A.data_ptr(), B.data_ptr(),
-                    tuple(A.shape), tuple(B.shape),
-                    A.dtype)
-             entry = _CACHE.get(key)
-             if entry is None:
-                 entry = _prepare_fp8(A, B)
-                 _CACHE[key] = entry
-             A_fp8, B_fp8, scale_a, scale_b = entry
+ template <int CTA_GROUP = 1>
+ __device__ inline
+ void tcgen05_mma_f16(int taddr, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc, int enable_input_d) {
+   asm volatile(
+     "{\n\t"
+     ".reg .pred p;\n\t"
+     "setp.ne.b32 p, %4, 0;\n\t"
+     "tcgen05.mma.cta_group::%5.kind::f16 [%0], %1, %2, %3, p;\n\t"
+     "}"
+     :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc), "r"(enable_input_d), "n"(CTA_GROUP)
+   );
+ }
   
-             result = torch._scaled_mm(
-                 A_fp8, B_fp8,
-                 scale_a=scale_a, scale_b=scale_b,
-                 out_dtype=torch.bfloat16,
-             )
-             if A.dtype == torch.bfloat16:
-                 out.copy_(result)
-             else:
-                 out.copy_(result.to(A.dtype))
-             return out
-         except Exception:
-             pass  # Fall through to cuBLAS on any FP8 issue
+ template <int CTA_GROUP = 1>
+ __device__ inline
+ void tcgen05_commit_mcast(int mbar_addr, int16_t cta_mask) {
+   asm volatile("tcgen05.commit.cta_group::%2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
+               :: "r"(mbar_addr), "h"(cta_mask), "n"(CTA_GROUP) : "memory");
+ }
   
-     _mod.matmul_cublas(A, B, out)
+ __device__ inline
+ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
+  
+ // CLC: issue try_cancel into the response slot, completing 16 bytes of full mbar.
+ __device__ inline
+ void clc_try_cancel(int response_addr, int mbar_addr) {
+   asm volatile(
+     "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 "
+     "[%0], [%1];"
+     :: "r"(response_addr), "r"(mbar_addr)
+     : "memory"
+   );
+ }
+  
+ // CLC: query a stored response. Returns is_canceled predicate (as int 0/1) and
+ // the canceled cluster's first CTA blockIdx (.x, .y, .z).
+ __device__ inline
+ void clc_query(int response_addr, int &valid, int &bx, int &by, int &bz) {
+   uint32_t v = 0, x = 0, y = 0, z = 0;
+   asm volatile(
+     "{\n\t"
+     ".reg .pred p1;\n\t"
+     ".reg .b128 r;\n\t"
+     "ld.shared.b128 r, [%4];\n\t"
+     "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, r;\n\t"
+     "selp.u32 %3, 1, 0, p1;\n\t"
+     "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, r;\n\t"
+     "}"
+     : "=r"(x), "=r"(y), "=r"(z), "=r"(v)
+     : "r"(response_addr)
+     : "memory"
+   );
+   valid = v; bx = x; by = y; bz = z;
+ }
+  
+ inline
+ void check_cu(CUresult err) {
+   if (err == CUDA_SUCCESS) return;
+   const char *msg;
+   if (cuGetErrorString(err, &msg) != CUDA_SUCCESS)
+     msg = "unable to get error string";
+   TORCH_CHECK(false, msg);
+ }
+  
+ inline
+ void check_cuda(cudaError_t err) {
+   if (err == cudaSuccess) return;
+   TORCH_CHECK(false, cudaGetErrorString(err));
+ }
+  
+ inline
+ void init_tmap_3d_128B_fp16(
+   CUtensorMap *tmap,
+   const __half *ptr,
+   uint64_t global_height, uint64_t global_width,
+   uint32_t shared_height, uint32_t shared_width
+ ) {
+   constexpr uint32_t rank = 3;
+   uint64_t globalDim[rank]         = {64, global_height, global_width / 64};
+   uint64_t globalStrides[rank - 1] = {global_width * sizeof(__half), 128};
+   uint32_t boxDim[rank]            = {64, shared_height, shared_width / 64};
+   uint32_t elementStrides[rank]    = {1, 1, 1};
+  
+   auto err = cuTensorMapEncodeTiled(
+     tmap, CU_TENSOR_MAP_DATA_TYPE_FLOAT16, rank, (void *)ptr,
+     globalDim, globalStrides, boxDim, elementStrides,
+     CU_TENSOR_MAP_INTERLEAVE_NONE,
+     CU_TENSOR_MAP_SWIZZLE_128B,
+     CU_TENSOR_MAP_L2_PROMOTION_NONE,
+     CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
+   );
+   check_cu(err);
+ }
+  
+ constexpr int NUM_WARPS = 7;  // 1 scheduler + 4 epilogue + 1 producer + 1 consumer
+ constexpr int TB_SIZE = NUM_WARPS * WARP_SIZE;
+  
+ constexpr int BLOCK_M = 128;
+ constexpr int BLOCK_K = 64;
+ constexpr int MMA_K = 16;
+ constexpr int NUM_CLC_STAGES = 2;
+ constexpr int CLC_RESP_BYTES = 16;
+  
+ // Warp roles (kept layout compatible with original v8: epilogue at warps 0..3)
+ constexpr int WARP_SCHED    = 4;        // CLC scheduler warp
+ constexpr int WARP_PRODUCER = 5;
+ constexpr int WARP_CONSUMER = 6;
+ // warps 0..3 are epilogue
+  
+ template <int BLOCK_N, int CTA_GROUP, int NUM_STAGES, bool EPILOGUE_CACHE_MOD>
+ __global__
+ __cluster_dims__(CTA_GROUP, 1, 1)
+ __launch_bounds__(TB_SIZE)
+ void matmul_v8_clc_kernel(
+   const __grid_constant__ CUtensorMap A_tmap,
+   const __grid_constant__ CUtensorMap B_tmap,
+   __half *C_ptr,
+   int M, int N, int K
+ ) {
+   const int tid = threadIdx.x;
+   const int warp_id = warp_uniform(tid / WARP_SIZE);
+  
+   int cta_rank;
+   asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank));
+  
+   extern __shared__ __align__(1024) char smem_ptr[];
+   const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
+   constexpr int A_size = BLOCK_M * BLOCK_K * sizeof(__half);
+   constexpr int BN_PER_CTA = BLOCK_N / CTA_GROUP;
+   constexpr int B_size = BLOCK_K * BN_PER_CTA * sizeof(__half);
+  
+   // SMEM layout (after A/B stages):
+   //   tma_mbar      [NUM_STAGES] (8B each)
+   //   mma_mbar      [NUM_STAGES] (8B each)
+   //   mainloop_mbar [2]          (8B each)
+   //   epilogue_mbar [2]          (8B each, +4B for tcgen05_alloc out-ptr after)
+   //   clc_full_mbar [NUM_CLC_STAGES] (8B each)
+   //   clc_empty_mbar[NUM_CLC_STAGES] (8B each)
+   //   clc_response  [NUM_CLC_STAGES] (16B each, 16-byte aligned)
+   const int tma_mbar_addr = smem + (A_size + B_size) * NUM_STAGES;
+   const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
+   const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
+   const int epilogue_mbar_addr = mainloop_mbar_addr + 2 * 8;
+  
+   // CLC mbarriers and response slots. tcgen05_alloc writes a 4B out-ptr at
+   // epilogue_mbar+16 — leave 8 bytes of pad before CLC region.
+   int clc_full_mbar_addr = epilogue_mbar_addr + 2 * 8 + 8;
+   int clc_empty_mbar_addr = clc_full_mbar_addr + NUM_CLC_STAGES * 8;
+   int clc_response_base = clc_empty_mbar_addr + NUM_CLC_STAGES * 8;
+   // Round up to 16-byte alignment for response slots.
+   clc_response_base = (clc_response_base + 15) & ~15;
+  
+   // Number of warps that consume each CLC response: 1 producer + 1 consumer + 4 epilogue.
+   constexpr int CLC_CONSUMERS = 6;
+  
+   if (warp_id == 0 && elect_sync()) {
+     for (int i = 0; i < NUM_STAGES; i++) {
+       mbarrier_init(tma_mbar_addr + i * 8, CTA_GROUP);
+       mbarrier_init(mma_mbar_addr + i * 8, 1);
+     }
+     for (int i = 0; i < 2; i++) {
+       mbarrier_init(mainloop_mbar_addr + i * 8, 1);
+       mbarrier_init(epilogue_mbar_addr + i * 8, 4 * CTA_GROUP * WARP_SIZE);
+     }
+     for (int i = 0; i < NUM_CLC_STAGES; i++) {
+       mbarrier_init(clc_full_mbar_addr + i * 8, 1);
+       mbarrier_init(clc_empty_mbar_addr + i * 8, CLC_CONSUMERS);
+     }
+     asm volatile("fence.mbarrier_init.release.cluster;");
+   }
+  
+   if constexpr (CTA_GROUP > 1) {
+     asm volatile("barrier.cluster.arrive.relaxed.aligned;");
+     asm volatile("barrier.cluster.wait.acquire.aligned;");
+   } else {
+     __syncthreads();
+   }
+  
+   const int num_iters = K / BLOCK_K;
+  
+   // Decode "current tile" from blockIdx (initial work) or from CLC response.
+   // Grid layout: grid_x = grid_n * CTA_GROUP, grid_y = grid_m_pair.
+   //   bid_n = blockIdx.x / CTA_GROUP, m_pair = blockIdx.y, cta_rank = blockIdx.x % CTA_GROUP.
+   // CLC response (get_first_ctaid) returns canceled CLUSTER's first CTA blockIdx,
+   // which has x = cluster_x * CTA_GROUP. So canceled bid_n = bx / CTA_GROUP.
+  
+   // ---- Scheduler warp ----
+   if (warp_id == WARP_SCHED) {
+     if (cta_rank == 0 && elect_sync()) {
+       // Initial: prime both CLC stages by issuing try_cancels.
+       for (int s = 0; s < NUM_CLC_STAGES; s++) {
+         const int mbar_s = clc_full_mbar_addr + s * 8;
+         // Arrive+expect_tx on cluster CTA 0
+         mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 0);
+         if constexpr (CTA_GROUP >= 2)
+           mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 1);
+         const int resp_s = clc_response_base + s * CLC_RESP_BYTES;
+         clc_try_cancel(resp_s, mbar_s);
+       }
+  
+       // Refill loop: wait clc_empty[s], query response, if is_canceled refill,
+       // else exit. Bound by 4096 as a safety net.
+       int sched_stage = 0;
+       int empty_phase = 0;
+       for (int i = 0; i < 4096; i++) {
+         mbarrier_wait(clc_empty_mbar_addr + sched_stage * 8, empty_phase);
+         // Query the response that consumers just drained.
+         int v, bx, by, bz;
+         clc_query(clc_response_base + sched_stage * CLC_RESP_BYTES, v, bx, by, bz);
+         if (!v) break;  // No more work to steal; consumers will all see false next.
+         const int mbar_s = clc_full_mbar_addr + sched_stage * 8;
+         mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 0);
+         if constexpr (CTA_GROUP >= 2)
+           mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 1);
+         const int resp_s = clc_response_base + sched_stage * CLC_RESP_BYTES;
+         clc_try_cancel(resp_s, mbar_s);
+         sched_stage = (sched_stage + 1) % NUM_CLC_STAGES;
+         if (sched_stage == 0) empty_phase ^= 1;
+       }
+     }
+     return;
+   }
+  
+   // ---- Producer warp ----
+   if (warp_id == WARP_PRODUCER) {
+     if (elect_sync()) {
+       int tma_stage = 0;
+       int mma_phase = 1;
+       const int tma_mbar_addr_ = tma_mbar_addr & 0xFEFFFFFF;
+  
+       int cur_bid_n = blockIdx.x / CTA_GROUP;
+       int cur_m_pair = blockIdx.y;
+       int valid = 1;
+  
+       int clc_stage = 0;
+       int clc_full_phase = 0;
+  
+       while (valid) {
+         const int bid_m = cur_m_pair * CTA_GROUP + cta_rank;
+         const int off_m = bid_m * BLOCK_M;
+         const int off_n = cur_bid_n * BLOCK_N + cta_rank * BN_PER_CTA;
+  
+         for (int iter_k = 0; iter_k < num_iters; iter_k++) {
+           const int mbar_addr = tma_mbar_addr_ + tma_stage * 8;
+           const int A_smem = smem + tma_stage * (A_size + B_size);
+           const int B_smem = A_smem + A_size;
+           const int off_k = iter_k * BLOCK_K;
+  
+           mbarrier_wait(mma_mbar_addr + tma_stage * 8, mma_phase);
+  
+           tma_3d_g2s<CTA_GROUP>(A_smem, &A_tmap, 0, off_m, off_k / 64, mbar_addr);
+           tma_3d_g2s<CTA_GROUP>(B_smem, &B_tmap, 0, off_k, off_n / 64, mbar_addr);
+           mbarrier_arrive_expect_tx(mbar_addr, A_size + B_size);
+  
+           tma_stage = (tma_stage + 1) % NUM_STAGES;
+           if (tma_stage == 0) mma_phase ^= 1;
+         }
+  
+         // Wait for next CLC response, decode it.
+         mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
+         int v, bx, by, bz;
+         clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
+         const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
+         mbarrier_arrive(empty_addr);
+         clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
+         if (clc_stage == 0) clc_full_phase ^= 1;
+  
+         valid = v;
+         cur_bid_n = bx / CTA_GROUP;
+         cur_m_pair = by;
+       }
+     }
+     return;
+   }
+  
+   // ---- Consumer (MMA) warp ----
+   if (warp_id == WARP_CONSUMER) {
+     tcgen05_alloc<CTA_GROUP>(epilogue_mbar_addr + 8 * 2, BLOCK_N * 2);
+  
+     constexpr uint32_t MMA_M = BLOCK_M * CTA_GROUP;
+     constexpr uint32_t MMA_N = BLOCK_N;
+     constexpr uint32_t i_desc = (1U << 4U)
+                               | (0U << 7U)
+                               | (0U << 10U)
+                               | (1U << 16U)
+                               | (MMA_N >> 3U << 17U)
+                               | (MMA_M >> 4U << 24U)
+                               ;
+  
+     constexpr uint64_t A_desc_base = (desc_encode(8 * 128) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
+     constexpr uint64_t B_desc_base = (desc_encode(BLOCK_K * 128) << 16ULL)
+                                    | (desc_encode(8 * 128) << 32ULL)
+                                    | (1ULL << 46ULL)
+                                    | (2ULL << 61ULL);
+  
+     if (cta_rank == 0 && elect_sync()) {
+       int tma_stage = 0;
+       int tma_phase = 0;
+       int mainloop_stage = 0;
+       int epilogue_phase = 1;
+  
+       constexpr int16_t cta_mask = (1 << CTA_GROUP) - 1;
+  
+       int valid = 1;
+       int clc_stage = 0;
+       int clc_full_phase = 0;
+  
+       while (valid) {
+         mbarrier_wait(epilogue_mbar_addr + mainloop_stage * 8, epilogue_phase);
+  
+         for (int iter_k = 0; iter_k < num_iters; iter_k++) {
+           const int A_smem = smem + tma_stage * (A_size + B_size);
+           const int B_smem = A_smem + A_size;
+           const int tmem = mainloop_stage * BLOCK_N;
+  
+           uint64_t a_desc = A_desc_base | (A_smem >> 4);
+           uint64_t b_desc = B_desc_base | (B_smem >> 4);
+  
+           mbarrier_wait(tma_mbar_addr + tma_stage * 8, tma_phase);
+           asm volatile("tcgen05.fence::after_thread_sync;");
+  
+           tcgen05_mma_f16<CTA_GROUP>(tmem, a_desc, b_desc, i_desc, iter_k);
+           for (int k = 1; k < BLOCK_K / MMA_K; k++) {
+             a_desc += (32ULL >> 4);
+             b_desc += ((MMA_K * 128ULL) >> 4);
+             tcgen05_mma_f16<CTA_GROUP>(tmem, a_desc, b_desc, i_desc, 1);
+           }
+  
+           tcgen05_commit_mcast<CTA_GROUP>(mma_mbar_addr + tma_stage * 8, cta_mask);
+  
+           tma_stage = (tma_stage + 1) % NUM_STAGES;
+           if (tma_stage == 0) tma_phase ^= 1;
+         }
+  
+         tcgen05_commit_mcast<CTA_GROUP>(mainloop_mbar_addr + mainloop_stage * 8, cta_mask);
+  
+         mainloop_stage = (mainloop_stage + 1) % 2;
+         if (mainloop_stage == 0) epilogue_phase ^= 1;
+  
+         // Read next CLC response.
+         mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
+         int v, bx, by, bz;
+         clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
+         const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
+         mbarrier_arrive(empty_addr);
+         clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
+         if (clc_stage == 0) clc_full_phase ^= 1;
+         valid = v;
+       }
+     }
+     return;
+   }
+  
+   // ---- Epilogue warps (warp_id ∈ {0, 1, 2, 3}) ----
+   {
+     int mainloop_stage = 0;
+     int mainloop_phase = 0;
+  
+     auto epilogue_sync = []() {
+       asm volatile("bar.sync %0, %1;" :: "r"(1), "r"(4 * WARP_SIZE) : "memory");
+     };
+  
+     int cur_bid_n = blockIdx.x / CTA_GROUP;
+     int cur_m_pair = blockIdx.y;
+     int valid = 1;
+  
+     int clc_stage = 0;
+     int clc_full_phase = 0;
+  
+     const int epi_warp_id = warp_id;  // 0..3 (warps 0..3 are epilogue)
+     const int lane = tid % WARP_SIZE;
+  
+     while (valid) {
+       const int bid_m = cur_m_pair * CTA_GROUP + cta_rank;
+       const int bid_n = cur_bid_n;
+  
+       if (epi_warp_id == 0)
+         mbarrier_wait(mainloop_mbar_addr + mainloop_stage * 8, mainloop_phase);
+       epilogue_sync();
+       asm volatile("tcgen05.fence::after_thread_sync;");
+  
+       constexpr int WIDTH = 16;
+       for (int n = 0; n < BLOCK_N / WIDTH; n++) {
+         const int t_row = cta_rank * 128 + epi_warp_id * 32;
+         const int t_col = mainloop_stage * BLOCK_N + n * WIDTH;
+         const int t_addr = (t_row << 16) + t_col;
+  
+         const int g_row = bid_m * BLOCK_M + epi_warp_id * 32 + lane;
+         const int g_col = bid_n * BLOCK_N + n * WIDTH;
+  
+         if constexpr (EPILOGUE_CACHE_MOD)
+           asm volatile(
+             "{\n"
+             ".reg .f32 f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15;\n"
+             ".reg .b32 b0, b1, b2, b3, b4, b5, b6, b7;\n"
+             "tcgen05.ld.sync.aligned.32x32b.x16.b32\n"
+             "  {f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15}, [%1];\n"
+             "tcgen05.wait::ld.sync.aligned;\n"
+             "cvt.rn.f16x2.f32 b0, f1, f0;\n"
+             "cvt.rn.f16x2.f32 b1, f3, f2;\n"
+             "cvt.rn.f16x2.f32 b2, f5, f4;\n"
+             "cvt.rn.f16x2.f32 b3, f7, f6;\n"
+             "cvt.rn.f16x2.f32 b4, f9, f8;\n"
+             "cvt.rn.f16x2.f32 b5, f11, f10;\n"
+             "cvt.rn.f16x2.f32 b6, f13, f12;\n"
+             "cvt.rn.f16x2.f32 b7, f15, f14;\n"
+             "st.relaxed.cta.global.L1::no_allocate.v8.b32 [%0], {b0, b1, b2, b3, b4, b5, b6, b7};\n"
+             "}"
+             :: "l"(C_ptr + g_row * N + g_col), "r"(t_addr)
+           );
+         else
+           asm volatile(
+             "{\n"
+             ".reg .f32 f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15;\n"
+             ".reg .b32 b0, b1, b2, b3, b4, b5, b6, b7;\n"
+             "tcgen05.ld.sync.aligned.32x32b.x16.b32\n"
+             "  {f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15}, [%1];\n"
+             "tcgen05.wait::ld.sync.aligned;\n"
+             "cvt.rn.f16x2.f32 b0, f1, f0;\n"
+             "cvt.rn.f16x2.f32 b1, f3, f2;\n"
+             "cvt.rn.f16x2.f32 b2, f5, f4;\n"
+             "cvt.rn.f16x2.f32 b3, f7, f6;\n"
+             "cvt.rn.f16x2.f32 b4, f9, f8;\n"
+             "cvt.rn.f16x2.f32 b5, f11, f10;\n"
+             "cvt.rn.f16x2.f32 b6, f13, f12;\n"
+             "cvt.rn.f16x2.f32 b7, f15, f14;\n"
+             "st.global.v8.b32 [%0], {b0, b1, b2, b3, b4, b5, b6, b7};\n"
+             "}"
+             :: "l"(C_ptr + g_row * N + g_col), "r"(t_addr)
+           );
+       }
+  
+       const int mbar_addr_ep = (epilogue_mbar_addr + mainloop_stage * 8) & 0xFEFFFFFF;
+       mbarrier_arrive(mbar_addr_ep);
+  
+       mainloop_stage = (mainloop_stage + 1) % 2;
+       if (mainloop_stage == 0) mainloop_phase ^= 1;
+  
+       // Wait for next CLC response — only one elect_sync lane queries; broadcast result.
+       mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
+       int v = 0, bx = 0, by = 0, bz = 0;
+       if (elect_sync()) {
+         clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
+       }
+       v = __shfl_sync(0xFFFFFFFF, v, 0);
+       bx = __shfl_sync(0xFFFFFFFF, bx, 0);
+       by = __shfl_sync(0xFFFFFFFF, by, 0);
+  
+       if (elect_sync()) {
+         const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
+         mbarrier_arrive(empty_addr);
+       }
+  
+       clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
+       if (clc_stage == 0) clc_full_phase ^= 1;
+       valid = v;
+       cur_bid_n = bx / CTA_GROUP;
+       cur_m_pair = by;
+     }
+  
+     if constexpr (CTA_GROUP > 1) {
+       asm volatile("barrier.cluster.arrive.relaxed.aligned;");
+       asm volatile("barrier.cluster.wait.acquire.aligned;");
+     } else {
+       epilogue_sync();
+     }
+  
+     if (epi_warp_id == 0)
+       tcgen05_dealloc<CTA_GROUP>(0, BLOCK_N * 2);
+   }
+ }
+  
+ template <int BLOCK_N, int CTA_GROUP, bool EPILOGUE_CACHE_MOD>
+ void matmul_v8_clc_launch(
+   const __half *A_ptr,
+   const __half *B_ptr,
+         __half *C_ptr,
+   int M, int N, int K
+ ) {
+   CUtensorMap A_tmap, B_tmap;
+   init_tmap_3d_128B_fp16(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
+   init_tmap_3d_128B_fp16(&B_tmap, B_ptr, K, N, BLOCK_K, BLOCK_N / CTA_GROUP);
+  
+   // Launch grid covers ALL cluster-tiles. With CLC, idle clusters cancel and
+   // contribute their work to busy clusters — no host-level persistent loop.
+   const int grid_n = N / BLOCK_N;
+   const int grid_m_pair = M / (BLOCK_M * CTA_GROUP);
+   const int grid_x = grid_n * CTA_GROUP;
+   const int grid_y = grid_m_pair;
+  
+   constexpr int AB_size = (BLOCK_M + BLOCK_N / CTA_GROUP) * BLOCK_K * sizeof(__half);
+   constexpr int dynamic_size = AB_size + 2 * 8;  // per-stage A/B + tma_mbar + mma_mbar
+   // Static: mainloop_mbar 2*8 + epilogue_mbar 2*8 + tcgen05_alloc out-ptr 4 + pad 4
+   //         + clc_full 2*8 + clc_empty 2*8 + clc_response 2*16 + 16-byte align pad 8
+   constexpr int static_size = 2 * 8 + 2 * 8 + 4 + 4
+                             + 8  // pad before clc
+                             + NUM_CLC_STAGES * 8 + NUM_CLC_STAGES * 8
+                             + NUM_CLC_STAGES * CLC_RESP_BYTES
+                             + 16;  // align pad
+  
+   constexpr int sm100_size = 227 * 1024;
+   constexpr int NUM_STAGES = (sm100_size - static_size) / dynamic_size;
+   constexpr int smem_size = NUM_STAGES * dynamic_size + static_size;
+  
+   auto this_kernel = matmul_v8_clc_kernel<BLOCK_N, CTA_GROUP, NUM_STAGES, EPILOGUE_CACHE_MOD>;
+   cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  
+   dim3 grid(grid_x, grid_y, 1);
+   this_kernel<<<grid, TB_SIZE, smem_size>>>(
+     A_tmap, B_tmap, C_ptr, M, N, K
+   );
+ }
+  
+ void matmul_fwd(const at::Tensor& A, const at::Tensor& B, at::Tensor& C) {
+   int M = A.size(0);
+   int K = A.size(1);
+   int N = B.size(1);
+   if (A.dtype() == at::kHalf
+       && B.dtype() == at::kHalf
+       && C.dtype() == at::kHalf
+       && A.is_contiguous() && B.is_contiguous() && C.is_contiguous()
+       && (M % 256 == 0) && (K % 64 == 0) && (N % 256 == 0)) {
+     matmul_v8_clc_launch<256, 2, true>(
+       reinterpret_cast<const __half*>(A.data_ptr()),
+       reinterpret_cast<const __half*>(B.data_ptr()),
+       reinterpret_cast<__half*>(C.data_ptr()),
+       M, N, K
+     );
+   } else {
+     at::matmul_out(C, A, B);
+   }
+ }
+ """
+  
+ _CPP_SRC = (
+     "void matmul_fwd(const at::Tensor&, const at::Tensor&, at::Tensor&);\n"
+ )
+  
+ _mod = load_inline(
+     name="matmul_v8_clc",
+     cpp_sources=_CPP_SRC,
+     cuda_sources=_CUDA_SRC,
+     functions=["matmul_fwd"],
+     extra_cuda_cflags=[
+         "-O3",
+         "-std=c++17",
+         "-gencode=arch=compute_100a,code=sm_100a",
+         "--expt-relaxed-constexpr",
+     ],
+     extra_cflags=["-O3", "-std=c++17"],
+     extra_ldflags=["-lcuda"],
+     verbose=False,
+ )
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     A, B, out = data
+     _mod.matmul_fwd(A, B, out)
      return out