submission 487322

#!POPCORN leaderboard nvfp4_group_gemm
#!POPCORN gpu NVIDIA

from __future__ import annotations

from functools import lru_cache
from typing import cast

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t

BENCH_NOTES = """\
g: 8; k: [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168]; m: [80, 176, 128, 72, 64, 248, 96, 160]; n: [4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096]; seed: 1111
 60.2 ± 0.05 us

g: 8; k: [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048]; m: [40, 76, 168, 72, 164, 148, 196, 160]; n: [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168]; seed: 1111
 52.4 ± 0.05 us

g: 2; k: [4096, 4096]; m: [192, 320]; n: [3072, 3072]; seed: 1111
 20.6 ± 0.02 us

g: 2; k: [1536, 1536]; m: [128, 384]; n: [4096, 4096]; seed: 1111
 16.0 ± 0.02 us
"""

CUDA_SRC = """
#include <vector>
#include <unordered_map>
#include <cstdint>
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>

static inline int ceil_div(int a, int b) { return (a + b - 1) / b; }

#define CUDA_CHECK(expr)                                                       \\
  do {                                                                         \\
    cudaError_t _err = (expr);                                                 \\
    TORCH_CHECK(_err == cudaSuccess, "CUDA error: ", cudaGetErrorString(_err)); \\
  } while (0)

static inline uint64_t hash_combine_u64(uint64_t h, uint64_t x) {
  h ^= x;
  h *= 1099511628211ULL;
  return h;
}

constexpr int WARP_SIZE = 32;
constexpr int MMA_K = 64;

constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
constexpr uint64_t EVICT_LAST = 0x14F0000000000000;

struct WorkItem {
  int problem_idx;
  int tile_m;
  int tile_n;
};

struct __align__(128) ProblemInfo {
  CUtensorMap A_tmap;
  CUtensorMap B_tmap;
  CUtensorMap B_tmap_256;
  const char* SFA_ptr;
  const char* SFB_ptr;
  half* C_ptr;
  int M, N, K;
  int64_t Cs0, Cs1;
};

__device__ inline
constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }

__device__
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;
}

__device__ inline void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

__device__ inline void mbarrier_arrive_expect_tx(int mbar_addr, int size) {
  asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" 
              :: "r"(mbar_addr), "r"(size) : "memory");
}

__device__ 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 LAB_WAIT;\\n\\t"
    "}"
    :: "r"(mbar_addr), "r"(phase), "r"(ticks)
  );
}

template <int CTA_GROUP = 1>
__device__ inline void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z,
                                         int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::%7.L2::cache_hint "
    "[%0], [%1, {%2, %3, %4}], [%5], %6;"
    :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy), "n"(CTA_GROUP)
    : "memory"
  );
}

__device__ inline void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
    "[%0], [%1], %2, [%3], %4;"
    :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy)
    : "memory"
  );
}

__device__ inline void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

__device__ inline void tcgen05_commit(int mbar_addr) {
  asm volatile(
    "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\\n"
    :: "r"(mbar_addr) : "memory"
  );
}

__device__ inline void tcgen05_mma_nvfp4(
  uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
  int scale_A_tmem, int scale_B_tmem, int enable_input_d
) {
  const int d_tmem = 0;
  asm volatile(
    "{\\n\\t"
    ".reg .pred p;\\n\\t"
    "setp.ne.b32 p, %6, 0;\\n\\t"
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\\n\\t"
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

struct SHAPE {
  static constexpr char _16x256b[] = ".16x256b";
};
struct NUM {
  static constexpr char x16[] = ".x16";
};

template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_64regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%65%66.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
              "  %8,  %9, %10, %11, %12, %13, %14, %15, "
              " %16, %17, %18, %19, %20, %21, %22, %23, "
              " %24, %25, %26, %27, %28, %29, %30, %31, "
              " %32, %33, %34, %35, %36, %37, %38, %39, "
              " %40, %41, %42, %43, %44, %45, %46, %47, "
              " %48, %49, %50, %51, %52, %53, %54, %55, "
              " %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
              : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
                "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
                "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
                "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31]),
                "=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
                "=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
                "=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
                "=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

__device__ inline void tcgen05_ld_16x256b_x16(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(tmp, row, col);
}

__device__ __forceinline__ void tcgen05_dealloc_cols_cta1(uint32_t tmem, int count) {
  asm volatile(
    "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\\n"
    :: "r"(tmem), "r"(count)
    : "memory"
  );
}

constexpr int TMA_BLOCK_N = 128;
constexpr int TMA_BLOCK_K = 256;
constexpr int TMA_NUM_WARPS = 8;
constexpr int MMA_M = 128;

constexpr int LOW_M_THRESHOLD = 96;

constexpr int TMA_WARP = 4;
constexpr int MMA_WARP = 5;

template <int BLOCK_M, int BLOCK_N, bool LOW_M_EPILOGUE>
__device__ __forceinline__ void epilogue_store(
    const ProblemInfo& prob,
    int m_offset,
    int n_offset,
    int tid,
    int warp_id,
    int lane_id
) {
  if (tid >= BLOCK_M) return;

  const int M = prob.M;
  const int N = prob.N;
  half* C_ptr = prob.C_ptr;
  const int64_t Cs0 = prob.Cs0;
  const int64_t Cs1 = prob.Cs1;

  const bool full_n = (n_offset + BLOCK_N <= N);
  const bool full_m = (m_offset + BLOCK_M <= M);
  const bool full_tile = full_n && full_m;
  const bool contiguous = (Cs1 == 1);

  const int warp_row_base = m_offset + warp_id * 32;
  if (LOW_M_EPILOGUE && warp_row_base >= M) return;

  int m_iters = 2;
  if (LOW_M_EPILOGUE) {
    const int remaining = M - warp_row_base;
    m_iters = (remaining <= 16) ? 1 : 2;
  }

  const int lane_row = lane_id >> 2;
  const int lane_col = (lane_id & 3) * 2;
  constexpr int HALF_N = 128;
  const int halves = BLOCK_N / HALF_N;

  for (int m = 0; m < m_iters; ++m) {
    for (int half_idx = 0; half_idx < halves; ++half_idx) {
      float tmp[HALF_N / 2];
      const int col_base = half_idx * HALF_N;
      tcgen05_ld_16x256b_x16(tmp, warp_id * 32 + m * 16, col_base);
      asm volatile("tcgen05.wait::ld.sync.aligned;\\n");

    const int row0 = warp_row_base + m * 16 + lane_row;
    const int row1 = row0 + 8;

    if (contiguous) {
      if (full_tile) {
        half2* row0_ptr = reinterpret_cast<half2*>(C_ptr + row0 * Cs0 + n_offset + col_base);
        half2* row1_ptr = reinterpret_cast<half2*>(C_ptr + row1 * Cs0 + n_offset + col_base);
        #pragma unroll
        for (int i = 0; i < HALF_N / 8; i++) {
          const int idx = i * 4;
          const int col = i * 8 + lane_col;
          const int h2_idx = col >> 1;
          row0_ptr[h2_idx] = __halves2half2(__float2half_rn(tmp[idx + 0]), __float2half_rn(tmp[idx + 1]));
          row1_ptr[h2_idx] = __halves2half2(__float2half_rn(tmp[idx + 2]), __float2half_rn(tmp[idx + 3]));
        }
        continue;
      }

      const bool row0_in = row0 < M;
      const bool row1_in = row1 < M;
      if (full_n) {
        half2* row0_ptr = row0_in ? reinterpret_cast<half2*>(C_ptr + row0 * Cs0 + n_offset + col_base) : nullptr;
        half2* row1_ptr = row1_in ? reinterpret_cast<half2*>(C_ptr + row1 * Cs0 + n_offset + col_base) : nullptr;
        #pragma unroll
        for (int i = 0; i < HALF_N / 8; i++) {
          const int idx = i * 4;
          const int col = i * 8 + lane_col;
          const int h2_idx = col >> 1;
          if (row0_in) {
            row0_ptr[h2_idx] = __halves2half2(__float2half_rn(tmp[idx + 0]), __float2half_rn(tmp[idx + 1]));
          }
          if (row1_in) {
            row1_ptr[h2_idx] = __halves2half2(__float2half_rn(tmp[idx + 2]), __float2half_rn(tmp[idx + 3]));
          }
        }
      } else {
        #pragma unroll
        for (int i = 0; i < HALF_N / 8; i++) {
          const int idx = i * 4;
          const int col = n_offset + col_base + i * 8 + lane_col;
          if (col < N) {
            const half h00 = __float2half_rn(tmp[idx + 0]);
            const half h01 = __float2half_rn(tmp[idx + 1]);
            const half h10 = __float2half_rn(tmp[idx + 2]);
            const half h11 = __float2half_rn(tmp[idx + 3]);
            if (row0_in) {
              if (col + 1 < N) {
                reinterpret_cast<half2*>(C_ptr + row0 * Cs0 + col)[0] = __halves2half2(h00, h01);
              } else {
                C_ptr[row0 * Cs0 + col] = h00;
              }
            }
            if (row1_in) {
              if (col + 1 < N) {
                reinterpret_cast<half2*>(C_ptr + row1 * Cs0 + col)[0] = __halves2half2(h10, h11);
              } else {
                C_ptr[row1 * Cs0 + col] = h10;
              }
            }
          }
        }
      }
    } else {
      const bool row0_in = row0 < M;
      const bool row1_in = row1 < M;
      #pragma unroll
      for (int i = 0; i < HALF_N / 8; i++) {
        const int idx = i * 4;
        const int col = n_offset + col_base + i * 8 + lane_col;
        if (col < N) {
          const half h00 = __float2half_rn(tmp[idx + 0]);
          const half h01 = __float2half_rn(tmp[idx + 1]);
          const half h10 = __float2half_rn(tmp[idx + 2]);
          const half h11 = __float2half_rn(tmp[idx + 3]);
          if (row0_in) {
            C_ptr[row0 * Cs0 + col * Cs1] = h00;
            if (col + 1 < N) C_ptr[row0 * Cs0 + (col + 1) * Cs1] = h01;
          }
          if (row1_in) {
            C_ptr[row1 * Cs0 + col * Cs1] = h10;
            if (col + 1 < N) C_ptr[row1 * Cs0 + (col + 1) * Cs1] = h11;
          }
        }
      }
    }
    }
  }
}

template <bool PERSISTENT, int BLOCK_M, int BLOCK_N, int NS>
__global__ __launch_bounds__(TMA_NUM_WARPS * WARP_SIZE)
void grouped_gemm_kernel_v4(
  const ProblemInfo* __restrict__ global_probs,
  const WorkItem* __restrict__ work_items,
  int num_items,
  int* __restrict__ work_counter
) {
  constexpr int TMA_A_SMEM_BYTES = BLOCK_M * (TMA_BLOCK_K / 2);
  constexpr int TMA_B_SMEM_BYTES = BLOCK_N * (TMA_BLOCK_K / 2);
  constexpr int TMA_SFA_SMEM_BYTES = MMA_M * (TMA_BLOCK_K / 16);
  constexpr int TMA_SFB_SMEM_BYTES = BLOCK_N * (TMA_BLOCK_K / 16);
  constexpr int STAGE_SIZE = TMA_A_SMEM_BYTES + TMA_B_SMEM_BYTES + TMA_SFA_SMEM_BYTES + TMA_SFB_SMEM_BYTES;
  const int tid = threadIdx.x;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem_base = static_cast<int>(__cvta_generic_to_shared(smem_ptr));

  constexpr int B_off = TMA_A_SMEM_BYTES;
  constexpr int SFA_off = B_off + TMA_B_SMEM_BYTES;
  constexpr int SFB_off = SFA_off + TMA_SFA_SMEM_BYTES;

  const int mbar_base = smem_base + STAGE_SIZE * NS;
  
  constexpr int TMEM_COLS = BLOCK_N * 2;
  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (TMA_BLOCK_K / MMA_K);
  
  constexpr uint32_t idesc = (1U << 7U) | (1U << 10U)
                           | ((uint32_t)BLOCK_N >> 3U << 17U)
                           | ((uint32_t)MMA_M >> 7U << 27U);

  if (warp_id == 0) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem_base), "r"(TMEM_COLS));
  }
  else if (warp_id == 1 && elect_sync()) {
    for (int i = 0; i < num_items && i < 8; ++i) {
      const ProblemInfo* prob = &global_probs[work_items[i].problem_idx];
      asm volatile("prefetch.tensormap [%0];" :: "l"(&prob->A_tmap) : "memory");
      if constexpr (BLOCK_N == 256) {
        asm volatile("prefetch.tensormap [%0];" :: "l"(&prob->B_tmap_256) : "memory");
      } else {
        asm volatile("prefetch.tensormap [%0];" :: "l"(&prob->B_tmap) : "memory");
      }
    }
  }
  __syncthreads();

  __shared__ int shared_work_idx;

  int work_idx = blockIdx.x;

  if (work_idx < num_items) {
    if (tid == 0) {
      for (int i = 0; i < NS; ++i) {
        mbarrier_init(mbar_base + i * 8, 1);
        mbarrier_init(mbar_base + (NS + i) * 8, 1);
      }
      asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
    }
    __syncthreads();
  }

  while (work_idx < num_items) {
    const WorkItem& work = work_items[work_idx];
    const ProblemInfo& prob = global_probs[work.problem_idx];

    const int m_offset = work.tile_m * BLOCK_M;
    const int n_offset = work.tile_n * BLOCK_N;
    const int K = prob.K;
    const int num_k_iters = K / TMA_BLOCK_K;

    if (warp_id == TMA_WARP && elect_sync()) {
      constexpr uint64_t cache_A = EVICT_LAST;
      constexpr uint64_t cache_B = EVICT_FIRST;

      auto issue_tma = [&](int k_iter, int stage) {
        const int mbar_addr = mbar_base + stage * 8;
        const int stage_base = smem_base + stage * STAGE_SIZE;
        const int off_k = k_iter * TMA_BLOCK_K;

        tma_3d_gmem2smem<1>(stage_base, &prob.A_tmap, 0, m_offset, off_k / 256, mbar_addr, cache_A);
        if constexpr (BLOCK_N == 256) {
          tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap_256, 0, n_offset, off_k / 256, mbar_addr, cache_B);
        } else {
          tma_3d_gmem2smem<1>(stage_base + B_off, &prob.B_tmap, 0, n_offset, off_k / 256, mbar_addr, cache_B);
        }

        const int rest_k = K / 16 / 4;
        const int k_blk = off_k / (16 * 4);
        const char* SFA_src = prob.SFA_ptr + ((m_offset / 128) * rest_k + k_blk) * 512;
        tma_gmem2smem(stage_base + SFA_off, SFA_src, TMA_SFA_SMEM_BYTES, mbar_addr, cache_A);
        if constexpr (BLOCK_N == 256) {
          constexpr int SFB_HALF_BYTES = 128 * (TMA_BLOCK_K / 16);
          const char* SFB_src0 = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
          const char* SFB_src1 = prob.SFB_ptr + (((n_offset / 128) + 1) * rest_k + k_blk) * 512;
          tma_gmem2smem(stage_base + SFB_off, SFB_src0, SFB_HALF_BYTES, mbar_addr, cache_B);
          tma_gmem2smem(stage_base + SFB_off + SFB_HALF_BYTES, SFB_src1, SFB_HALF_BYTES, mbar_addr, cache_B);
        } else {
          const char* SFB_src = prob.SFB_ptr + ((n_offset / 128) * rest_k + k_blk) * 512;
          tma_gmem2smem(stage_base + SFB_off, SFB_src, TMA_SFB_SMEM_BYTES, mbar_addr, cache_B);
        }

        mbarrier_arrive_expect_tx(mbar_addr, STAGE_SIZE);
      };

      for (int k_iter = 0; k_iter < NS && k_iter < num_k_iters; k_iter++) {
        issue_tma(k_iter, k_iter);
      }

      for (int k_iter = NS; k_iter < num_k_iters; k_iter++) {
        const int stage = k_iter % NS;
        const int mma_phase = (k_iter / NS - 1) % 2;
        mbarrier_wait(mbar_base + (NS + stage) * 8, mma_phase);
        issue_tma(k_iter, stage);
      }
    }

    else if (warp_id == MMA_WARP && elect_sync()) {
      auto make_desc_AB = [](int addr) -> uint64_t {
        const int SBO = 8 * 128;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
      };
      auto make_desc_SF = [](int addr) -> uint64_t {
        const int SBO = 8 * 16;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
      };

      for (int k_iter = 0; k_iter < num_k_iters; k_iter++) {
        const int stage = k_iter % NS;
        const int tma_phase = (k_iter / NS) % 2;
        mbarrier_wait(mbar_base + stage * 8, tma_phase);

        const int stage_base = smem_base + stage * STAGE_SIZE;

        const uint64_t SF_desc = make_desc_SF(0);
        const uint64_t SFA_desc = SF_desc + ((uint64_t)(stage_base + SFA_off) >> 4ULL);
        const uint64_t SFB_desc = SF_desc + ((uint64_t)(stage_base + SFB_off) >> 4ULL);

        for (int k = 0; k < TMA_BLOCK_K / MMA_K; k++) {
          tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc + (uint64_t)k * (512ULL >> 4ULL));
          if constexpr (BLOCK_N == 256) {
            constexpr uint64_t SFB_HALF_DESC = (uint64_t)(128 * (TMA_BLOCK_K / 16)) >> 4ULL;
            tcgen05_cp_nvfp4(SFB_tmem + k * 8, SFB_desc + (uint64_t)k * (512ULL >> 4ULL));
            tcgen05_cp_nvfp4(SFB_tmem + k * 8 + 4, SFB_desc + SFB_HALF_DESC + (uint64_t)k * (512ULL >> 4ULL));
          } else {
            tcgen05_cp_nvfp4(SFB_tmem + k * 4, SFB_desc + (uint64_t)k * (512ULL >> 4ULL));
          }
        }

        for (int k1 = 0; k1 < TMA_BLOCK_K / 256; k1++) {
          for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
            uint64_t a_desc = make_desc_AB(stage_base + k1 * BLOCK_M * 128 + k2 * 32);
            uint64_t b_desc = make_desc_AB(stage_base + B_off + k1 * BLOCK_N * 128 + k2 * 32);

            int k_sf = k1 * 4 + k2;
            const int scale_A_tmem = SFA_tmem + k_sf * 4 + (work.tile_m % (MMA_M / BLOCK_M)) * (BLOCK_M / 32);
            int scale_B_tmem;
            if constexpr (BLOCK_N == 256) {
              scale_B_tmem = SFB_tmem + k_sf * 8;
            } else {
              scale_B_tmem = SFB_tmem + k_sf * 4;
            }

            const int enable_input_d = (k_iter == 0 && k1 == 0 && k2 == 0) ? 0 : 1;
            tcgen05_mma_nvfp4(a_desc, b_desc, idesc, scale_A_tmem, scale_B_tmem, enable_input_d);
          }
        }

        tcgen05_commit(mbar_base + (NS + stage) * 8);
      }

      const int last_stage = (num_k_iters - 1) % NS;
      const int last_phase = ((num_k_iters - 1) / NS) % 2;
      mbarrier_wait(mbar_base + (NS + last_stage) * 8, last_phase);
    }

    __syncthreads();
    asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");

    const bool low_m = (prob.M <= LOW_M_THRESHOLD);
    if (low_m) {
      epilogue_store<BLOCK_M, BLOCK_N, true>(prob, m_offset, n_offset, tid, warp_id, lane_id);
    } else {
      epilogue_store<BLOCK_M, BLOCK_N, false>(prob, m_offset, n_offset, tid, warp_id, lane_id);
    }

    if constexpr (PERSISTENT) {
      if (warp_id == TMA_WARP && elect_sync()) {
        shared_work_idx = atomicAdd(work_counter, 1);
        if (shared_work_idx < num_items) {
          for (int i = 0; i < NS; ++i) {
            mbarrier_init(mbar_base + i * 8, 1);
            mbarrier_init(mbar_base + (NS + i) * 8, 1);
          }
          asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
        }
      }
    }

    __syncthreads();

    if constexpr (PERSISTENT) {
      work_idx = shared_work_idx;
    } else {
      break;
    }
  }

  if (warp_id == 0) {
    tcgen05_dealloc_cols_cta1(0, TMEM_COLS);
  }
}

void init_AB_tmap_u4(
  CUtensorMap *tmap,
  const void *ptr,
  uint64_t global_height, uint64_t global_width,
  uint32_t shared_height, uint32_t shared_width
) {
  TORCH_CHECK(ptr != nullptr, "ptr is null");
  TORCH_CHECK(((uintptr_t)ptr % 16) == 0, "ptr must be 16-byte aligned");
  TORCH_CHECK(global_width >= 256 && (global_width % 256) == 0, "K must be multiple of 256");

  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
  uint64_t globalStrides[rank-1] = {global_width / 2, 128};
  uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
  uint32_t elementStrides[rank]  = {1, 1, 1};

  // Cache cuTensorMap templates by shape.
  struct ShapeKey { uint64_t gh, gw; uint32_t sh, sw; };
  struct ShapeHash {
    size_t operator()(const ShapeKey& k) const noexcept {
      uint64_t h = k.gh;
      h ^= (k.gw + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2));
      h ^= ((uint64_t)k.sh << 32) ^ (uint64_t)k.sw;
      return (size_t)h;
    }
  };
  struct ShapeEq {
    bool operator()(const ShapeKey& a, const ShapeKey& b) const noexcept {
      return a.gh == b.gh && a.gw == b.gw && a.sh == b.sh && a.sw == b.sw;
    }
  };
  struct PtrKey { uint64_t gh, gw; uint32_t sh, sw; const void* ptr; };
  struct PtrHash {
    size_t operator()(const PtrKey& k) const noexcept {
      uint64_t h = k.gh;
      h ^= (k.gw + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2));
      h ^= ((uint64_t)k.sh << 32) ^ (uint64_t)k.sw;
      h ^= ((uint64_t)k.ptr >> 4);
      return (size_t)h;
    }
  };
  struct PtrEq {
    bool operator()(const PtrKey& a, const PtrKey& b) const noexcept {
      return a.gh == b.gh && a.gw == b.gw && a.sh == b.sh && a.sw == b.sw && a.ptr == b.ptr;
    }
  };

  static thread_local std::unordered_map<ShapeKey, CUtensorMap, ShapeHash, ShapeEq> tmpl_cache;
  static thread_local std::unordered_map<PtrKey, CUtensorMap, PtrHash, PtrEq> ptr_cache;

  PtrKey pkey{global_height, global_width, shared_height, shared_width, ptr};
  auto pit = ptr_cache.find(pkey);
  if (pit != ptr_cache.end()) { *tmap = pit->second; return; }

  ShapeKey skey{global_height, global_width, shared_height, shared_width};
  auto sit = tmpl_cache.find(skey);
  if (sit == tmpl_cache.end()) {
    CUtensorMap tmp;
    auto err = cuTensorMapEncodeTiled(
      &tmp, CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
      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
    );
    TORCH_CHECK(err == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed");
    sit = tmpl_cache.emplace(skey, tmp).first;
  }

  CUtensorMap tmp = sit->second;
  auto err = cuTensorMapReplaceAddress(&tmp, (void*)ptr);
  TORCH_CHECK(err == CUDA_SUCCESS, "cuTensorMapReplaceAddress failed");
  ptr_cache.emplace(pkey, tmp);
  *tmap = tmp;
}

std::vector<at::Tensor> group_gemm(
    std::vector<at::Tensor> A_list,
    std::vector<at::Tensor> B_list,
    std::vector<at::Tensor> C_list,
    std::vector<at::Tensor> sfa_list,
    std::vector<at::Tensor> sfb_list,
    at::Tensor sizes_cpu
) {
    int64_t G = A_list.size();
    auto dev = A_list[0].device();
    c10::cuda::CUDAGuard device_guard(dev);
    auto sizes_accessor = sizes_cpu.accessor<int64_t, 2>();

    static bool attrs_set = false;
    if (!attrs_set) {
      constexpr int NS_DEEP = 6;
      constexpr int NS_WIDE = 4;
      constexpr int STAGE_SIZE_128 = 128 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
      constexpr int MBAR_BYTES_DEEP = ((2 * NS_DEEP * 8 + 63) & ~63);
      constexpr int SMEM_SIZE_128 = STAGE_SIZE_128 * NS_DEEP + MBAR_BYTES_DEEP;
      constexpr int STAGE_SIZE_256 = 128 * 128 + 256 * 128 + 128 * 16 + 256 * 16;
      constexpr int MBAR_BYTES_WIDE = ((2 * NS_WIDE * 8 + 63) & ~63);
      constexpr int SMEM_SIZE_256 = STAGE_SIZE_256 * NS_WIDE + MBAR_BYTES_WIDE;
      CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_128));
      CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_128));
      CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<true, 128, 256, NS_WIDE>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_256));
      CUDA_CHECK(cudaFuncSetAttribute((grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE>), cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_SIZE_256));
      attrs_set = true;
    }

    std::vector<ProblemInfo> problem_infos(G);
    static thread_local std::vector<WorkItem> cached_work_items_128;
    static thread_local std::vector<WorkItem> cached_work_items_256;
    static thread_local uint64_t cached_work_hash = 0;
    static thread_local bool cached_work_valid = false;

    std::vector<uint8_t> active(G, 0);
    std::vector<uint8_t> use_256(G, 0);
    std::vector<int> num_tiles_m(G, 0);
    std::vector<int> num_tiles_n(G, 0);
    std::vector<int64_t> Ms(G, 0), Ns(G, 0), Ks(G, 0);

    uint64_t work_hash = 1469598103934665603ULL;
    for (int64_t i = 0; i < G; i++) {
        const int64_t M = sizes_accessor[i][0], N = sizes_accessor[i][1], K = sizes_accessor[i][2];
        Ms[(size_t)i] = M; Ns[(size_t)i] = N; Ks[(size_t)i] = K;
        if (A_list[i].stride(1) != 1 || B_list[i].stride(1) != 1) {
          work_hash = hash_combine_u64(work_hash, 0);
          continue;
        }
        active[(size_t)i] = 1;
        bool is_256 = (N >= 4096) && (K >= 2048) && ((N & 255) == 0);
        use_256[(size_t)i] = is_256;
        int block_m = 128;
        int block_n = is_256 ? 256 : 128;
        num_tiles_m[(size_t)i] = ceil_div((int)M, block_m);
        num_tiles_n[(size_t)i] = ceil_div((int)N, block_n);
        
        work_hash = hash_combine_u64(work_hash, (uint64_t)M);
        work_hash = hash_combine_u64(work_hash, (uint64_t)N);
        work_hash = hash_combine_u64(work_hash, (uint64_t)is_256);
    }

    if (!cached_work_valid || cached_work_hash != work_hash) {
      cached_work_items_128.clear();
      cached_work_items_256.clear();
      cached_work_items_128.reserve(G * 32);
      cached_work_items_256.reserve(G * 32);
      for (int64_t i = 0; i < G; i++) {
        if (!active[(size_t)i]) continue;
        for (int tm = 0; tm < num_tiles_m[(size_t)i]; tm++) {
          for (int tn = 0; tn < num_tiles_n[(size_t)i]; tn++) {
            if (use_256[(size_t)i]) {
              cached_work_items_256.push_back({(int)i, tm, tn});
            } else {
              cached_work_items_128.push_back({(int)i, tm, tn});
            }
          }
        }
      }
      cached_work_hash = work_hash;
      cached_work_valid = true;
    }

    uint64_t probs_hash = 1469598103934665603ULL;
    for (int64_t i = 0; i < G; i++) {
        if (!active[(size_t)i]) continue;
        const int64_t M = Ms[(size_t)i], N = Ns[(size_t)i], K = Ks[(size_t)i];

        ProblemInfo& p = problem_infos[i];
        p.M = M; p.N = N; p.K = K;
        p.Cs0 = C_list[i].stride(0); p.Cs1 = C_list[i].stride(1);
        p.C_ptr = (half*)C_list[i].data_ptr();
        p.SFA_ptr = (const char*)sfa_list[i].data_ptr();
        p.SFB_ptr = (const char*)sfb_list[i].data_ptr();

        init_AB_tmap_u4(&p.A_tmap, A_list[i].data_ptr(), A_list[i].size(0), K, 128, TMA_BLOCK_K);
        init_AB_tmap_u4(&p.B_tmap, B_list[i].data_ptr(), B_list[i].size(0), K, 128, TMA_BLOCK_K);
        if (use_256[(size_t)i]) {
          init_AB_tmap_u4(&p.B_tmap_256, B_list[i].data_ptr(), B_list[i].size(0), K, 256, TMA_BLOCK_K);
        } else {
          p.B_tmap_256 = p.B_tmap;
        }

        probs_hash = hash_combine_u64(probs_hash, (uint64_t)i);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)M);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)N);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)K);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)A_list[i].data_ptr());
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)B_list[i].data_ptr());
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.C_ptr);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.SFA_ptr);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)(uintptr_t)p.SFB_ptr);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)p.Cs0);
        probs_hash = hash_combine_u64(probs_hash, (uint64_t)p.Cs1);
    }

    if (cached_work_items_128.empty() && cached_work_items_256.empty()) return C_list;

    auto options = at::TensorOptions().dtype(at::kByte).device(dev);
    static thread_local at::Tensor d_probs_cache;
    static thread_local at::Tensor d_work_cache_128;
    static thread_local at::Tensor d_work_cache_256;
    static thread_local at::Tensor d_counter_cache;
    static thread_local uint64_t last_probs_hash = 0;
    static thread_local uint64_t last_work_hash = 0;
    static thread_local bool last_hash_valid = false;

    const int64_t probs_bytes = (int64_t)(G * sizeof(ProblemInfo));
    const int64_t work_bytes_128 = (int64_t)(cached_work_items_128.size() * sizeof(WorkItem));
    const int64_t work_bytes_256 = (int64_t)(cached_work_items_256.size() * sizeof(WorkItem));

    bool probs_realloc = false;
    if (!d_probs_cache.defined() || d_probs_cache.device() != dev || d_probs_cache.scalar_type() != at::kByte || d_probs_cache.numel() < probs_bytes) {
      d_probs_cache = at::empty({probs_bytes}, options);
      probs_realloc = true;
    }
    if (work_bytes_128 > 0 && (!d_work_cache_128.defined() || d_work_cache_128.device() != dev || d_work_cache_128.scalar_type() != at::kByte || d_work_cache_128.numel() < work_bytes_128)) {
      d_work_cache_128 = at::empty({work_bytes_128}, options);
    }
    if (work_bytes_256 > 0 && (!d_work_cache_256.defined() || d_work_cache_256.device() != dev || d_work_cache_256.scalar_type() != at::kByte || d_work_cache_256.numel() < work_bytes_256)) {
      d_work_cache_256 = at::empty({work_bytes_256}, options);
    }

    if (probs_realloc || !last_hash_valid || last_probs_hash != probs_hash) {
      CUDA_CHECK(cudaMemcpyAsync(d_probs_cache.data_ptr(), problem_infos.data(), G * sizeof(ProblemInfo), cudaMemcpyHostToDevice));
      last_probs_hash = probs_hash;
    }
    if (!last_hash_valid || last_work_hash != work_hash) {
      if (work_bytes_128 > 0) {
        CUDA_CHECK(cudaMemcpyAsync(d_work_cache_128.data_ptr(), cached_work_items_128.data(), work_bytes_128, cudaMemcpyHostToDevice));
      }
      if (work_bytes_256 > 0) {
        CUDA_CHECK(cudaMemcpyAsync(d_work_cache_256.data_ptr(), cached_work_items_256.data(), work_bytes_256, cudaMemcpyHostToDevice));
      }
      last_work_hash = work_hash;
    }
    last_hash_valid = true;

    constexpr int MAX_CTAS = 264;
    if (!d_counter_cache.defined() || d_counter_cache.device() != dev || d_counter_cache.scalar_type() != at::kInt || d_counter_cache.numel() != 1) {
      d_counter_cache = at::empty({1}, options.dtype(at::kInt));
    }

    constexpr int NS_DEEP = 6;
    constexpr int NS_WIDE = 4;

    if (!cached_work_items_128.empty()) {
      int num_items_128 = (int)cached_work_items_128.size();
      constexpr int STAGE_SIZE_128 = 128 * 128 + 128 * 128 + 128 * 16 + 128 * 16;
      constexpr int MBAR_128 = ((2 * NS_DEEP * 8 + 63) & ~63);
      constexpr int SMEM_SIZE_128 = STAGE_SIZE_128 * NS_DEEP + MBAR_128;
      if (num_items_128 > MAX_CTAS) {
        { static const int mc = MAX_CTAS; CUDA_CHECK(cudaMemcpyAsync(d_counter_cache.data_ptr(), &mc, sizeof(int), cudaMemcpyHostToDevice)); }
        grouped_gemm_kernel_v4<true, 128, 128, NS_DEEP><<<MAX_CTAS, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
            (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128, (int*)d_counter_cache.data_ptr());
      } else {
        grouped_gemm_kernel_v4<false, 128, 128, NS_DEEP><<<num_items_128, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_128>>>(
            (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_128.data_ptr(), num_items_128, nullptr);
      }
    }

    if (!cached_work_items_256.empty()) {
      int num_items_256 = (int)cached_work_items_256.size();
      constexpr int STAGE_SIZE_256 = 128 * 128 + 256 * 128 + 128 * 16 + 256 * 16;
      constexpr int MBAR_256 = ((2 * NS_WIDE * 8 + 63) & ~63);
      constexpr int SMEM_SIZE_256 = STAGE_SIZE_256 * NS_WIDE + MBAR_256;
      if (num_items_256 > MAX_CTAS) {
        { static const int mc = MAX_CTAS; CUDA_CHECK(cudaMemcpyAsync(d_counter_cache.data_ptr(), &mc, sizeof(int), cudaMemcpyHostToDevice)); }
        grouped_gemm_kernel_v4<true, 128, 256, NS_WIDE><<<MAX_CTAS, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_256>>>(
            (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_256.data_ptr(), num_items_256, (int*)d_counter_cache.data_ptr());
      } else {
        grouped_gemm_kernel_v4<false, 128, 256, NS_WIDE><<<num_items_256, TMA_NUM_WARPS * WARP_SIZE, SMEM_SIZE_256>>>(
            (ProblemInfo*)d_probs_cache.data_ptr(), (WorkItem*)d_work_cache_256.data_ptr(), num_items_256, nullptr);
      }
    }
    CUDA_CHECK(cudaGetLastError());
    return C_list;
}

TORCH_LIBRARY(my_module, m) {
    m.def("group_gemm(Tensor[] a, Tensor[] b, Tensor[] c, Tensor[] sfa, Tensor[] sfb, Tensor sizes) -> Tensor[]");
    m.impl("group_gemm", &group_gemm);
}
"""

load_inline(
    "group_gemm",
    cpp_sources="",
    cuda_sources=CUDA_SRC,
    verbose=True,
    is_python_module=False,
    no_implicit_headers=True,
    extra_cuda_cflags=[
        "-O3",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--use_fast_math",
        "--expt-relaxed-constexpr",
        "--relocatable-device-code=false",
        "-lineinfo",
        "-Xptxas=-v",
    ],
    extra_ldflags=["-lcuda"],
)
group_gemm = torch.ops.my_module.group_gemm


@lru_cache(maxsize=128)
def _sizes_cpu_cached(key: tuple[tuple[int, int, int], ...]) -> torch.Tensor:
    return torch.tensor(key, dtype=torch.int64, device="cpu")


def custom_kernel(data: input_t) -> output_t:
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
    A_list = [t[0] for t in abc_tensors]
    B_list = [t[1] for t in abc_tensors]
    C_list = [t[2] for t in abc_tensors]
    sfa_list = [t[0] for t in sfasfb_reordered_tensors]
    sfb_list = [t[1] for t in sfasfb_reordered_tensors]
    key = tuple(tuple(int(v) for v in x) for x in problem_sizes)
    sizes_cpu = _sizes_cpu_cached(key)
    group_gemm(A_list, B_list, C_list, sfa_list, sfb_list, sizes_cpu)
    return cast(output_t, C_list)