submission 487950

 # AUTO-GENERATED by run.py
# Do not edit directly; edit cuda_lib/* and ptx_lib/*.cuh instead.

CUDA_SRC = r'''
// ----- ptx_common.cuh -----

#include <stdint.h>
#include <cuda_fp16.h>
#include <cudaTypedefs.h>
#include <cuda_runtime.h>
#include <ATen/ATen.h>
#include <torch/torch.h>

// Common helpers for PTX inline asm wrappers.

#if defined(__CUDA_ARCH__)
#define PTX_DEVICE __device__ inline

PTX_DEVICE uint32_t ptx_laneid() {
  uint32_t lane;
  asm volatile("mov.u32 %0, %laneid;" : "=r"(lane));
  return lane;
}

PTX_DEVICE uint32_t ptx_activemask() {
  uint32_t mask;
  asm volatile("activemask.b32 %0;" : "=r"(mask));
  return mask;
}

PTX_DEVICE bool ptx_elect_one_sync() {
  uint32_t mask = ptx_activemask();
  int leader = __ffs(mask) - 1;
  return (int)ptx_laneid() == leader;
}

PTX_DEVICE uint32_t ptx_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"
    "}\n\t"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

#else
#define PTX_DEVICE __device__ inline

PTX_DEVICE uint32_t ptx_laneid() { return 0; }
PTX_DEVICE uint32_t ptx_activemask() { return 0xFFFFFFFF; }
PTX_DEVICE bool ptx_elect_one_sync() { return true; }
PTX_DEVICE uint32_t ptx_elect_sync() { return 1; }
#endif

#ifndef PTX_NO_ELECT
#define PTX_ELECT_ONE()          \
  do {                           \
    if (!ptx_elect_one_sync()) { \
      return;                    \
    }                            \
  } while (0)
#else
#define PTX_ELECT_ONE() do { } while (0)
#endif

PTX_DEVICE void ptx_bar_sync(int bar_id, int count) {
  asm volatile("bar.sync %0, %1;" :: "r"(bar_id), "r"(count) : "memory");
}

// ----- ptx_mbarrier.cuh -----


// mbarrier helpers (CTA scope)

// NOTE: Keep gemm1 semantics (no implicit election in wrappers).
PTX_DEVICE void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

// CTA-scope arrive expect_tx (gemm1 uses CTA scope)
PTX_DEVICE void mbarrier_arrive_expect_tx_cta(int mbar_addr, int size) {
  asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
              :: "r"(mbar_addr), "r"(size) : "memory");
}

PTX_DEVICE 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");
}

PTX_DEVICE void mbarrier_wait(int mbar_addr, int phase) {
  // gemm1 uses a ticked wait loop and exits when P1 is true.
  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"
      "}\n\t"
      :: "r"(mbar_addr), "r"(phase), "r"(ticks));
}

// Explicit wait loop with ticks (as in gemm1)
PTX_DEVICE void mbarrier_wait_ticks(int mbar_addr, int phase, uint32_t ticks) {
  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"
      "}\n\t"
      :: "r"(mbar_addr), "r"(phase), "r"(ticks));
}

PTX_DEVICE void mbarrier_wait_relaxed(int mbar_addr, int phase) {
  asm volatile(
      "{\n\t"
      ".reg .pred P1;\n\t"
      "WAIT: \n\t"
      "mbarrier.try_wait.parity.relaxed.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
      "@P1 bra WAIT;\n\t"
      "}\n\t"
      :: "r"(mbar_addr), "r"(phase), "r"(0xFFFFFFFF));
}

PTX_DEVICE void mbarrier_fence_init_release() {
  asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
}

// Cluster barrier helpers (used to synchronize CTAs in a cluster after mbarrier init)
PTX_DEVICE void barrier_cluster_arrive_relaxed_aligned() {
  asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory");
}

PTX_DEVICE void barrier_cluster_wait_acquire_aligned() {
  asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
}

// ----- ptx_tcgen05_cp.cuh -----


// tcgen05.cp wrappers (CTA group 1 only for now)

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_cp_32x128b_warpx4(int taddr, uint64_t s_desc) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_cp_128x128b(int taddr, uint64_t s_desc) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("tcgen05.cp.cta_group::1.128x128b [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_cp_128x256b(int taddr, uint64_t s_desc) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("tcgen05.cp.cta_group::1.128x256b [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

// Alias used by gemm1 (NVFP4 block scaling)
template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  tcgen05_cp_32x128b_warpx4<CTA_GROUP>(taddr, s_desc);
}

// ----- ptx_tcgen05_ldst.cuh -----


// tcgen05.ld / tcgen05.st wrappers

// see https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html
struct SHAPE {
  static constexpr char _32x32b[]  = ".32x32b";
  static constexpr char _16x128b[] = ".16x128b";
  static constexpr char _16x256b[] = ".16x256b";
};

struct NUM {
  static constexpr char x1[]   = ".x1";
  static constexpr char x2[]   = ".x2";
  static constexpr char x4[]   = ".x4";
  static constexpr char x8[]   = ".x8";
  static constexpr char x16[]  = ".x16";
  static constexpr char x32[]  = ".x32";
  static constexpr char x64[]  = ".x64";
  static constexpr char x128[] = ".x128";
};

template <int NUM_REGS, const char *SHAPE, int NUM>
PTX_DEVICE void tcgen05_ld(float *tmp, int row, int col) {
  int addr = (row << 16) | col;

  if constexpr (NUM_REGS == 1)
  asm volatile("tcgen05.ld.sync.aligned%2.x%3.b32 {%0}, [%1];"
              : "=f"(tmp[0]) : "r"(addr), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 2)
  asm volatile("tcgen05.ld.sync.aligned%3.x%4.b32 {%0, %1}, [%2];"
              : "=f"(tmp[0]), "=f"(tmp[1]) : "r"(addr), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 4)
  asm volatile("tcgen05.ld.sync.aligned%5.x%6.b32 "
              "{%0, %1, %2, %3}, [%4];"
              : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3])
              : "r"(addr), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 8)
  asm volatile("tcgen05.ld.sync.aligned%9.x%10.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7}, [%8];"
              : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]), "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
              : "r"(addr), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 16)
  asm volatile("tcgen05.ld.sync.aligned%17.x%18.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
              "  %8,  %9, %10, %11, %12, %13, %14, %15}, [%16];"
              : "=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])
              : "r"(addr), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 32)
  asm volatile("tcgen05.ld.sync.aligned%33.x%34.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];"
              : "=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])
              : "r"(addr), "C"(SHAPE), "n"(NUM));
}

// Explicit tcgen05.ld variants (verbatim from gemm1)
template <const char *SHAPE, const char *NUM>
PTX_DEVICE void tcgen05_ld_16regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%17%18.b32 "
              "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
              "  %8,  %9, %10, %11, %12, %13, %14, %15}, [%16];"
              : "=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])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

template <const char *SHAPE, const char *NUM>
PTX_DEVICE void tcgen05_ld_32regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%33%34.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];"
              : "=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])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

template <const char *SHAPE, const char *NUM>
PTX_DEVICE 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));
}

template <const char *SHAPE, const char *NUM>
PTX_DEVICE void tcgen05_ld_128regs(float *tmp, int row, int col) {
  asm volatile("tcgen05.ld.sync.aligned%129%130.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, %65, %66, %67, %68, %69, %70, %71, "
              " %72, %73, %74, %75, %76, %77, %78, %79, "
              " %80, %81, %82, %83, %84, %85, %86, %87, "
              " %88, %89, %90, %91, %92, %93, %94, %95, "
              " %96, %97, %98, %99,%100,%101,%102,%103, "
              "%104,%105,%106,%107,%108,%109,%110,%111, "
              "%112,%113,%114,%115,%116,%117,%118,%119, "
              "%120,%121,%122,%123,%124,%125,%126,%127}, [%128];"
              : "=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]),
                "=f"(tmp[64]), "=f"(tmp[65]), "=f"(tmp[66]), "=f"(tmp[67]), "=f"(tmp[68]), "=f"(tmp[69]), "=f"(tmp[70]), "=f"(tmp[71]),
                "=f"(tmp[72]), "=f"(tmp[73]), "=f"(tmp[74]), "=f"(tmp[75]), "=f"(tmp[76]), "=f"(tmp[77]), "=f"(tmp[78]), "=f"(tmp[79]),
                "=f"(tmp[80]), "=f"(tmp[81]), "=f"(tmp[82]), "=f"(tmp[83]), "=f"(tmp[84]), "=f"(tmp[85]), "=f"(tmp[86]), "=f"(tmp[87]),
                "=f"(tmp[88]), "=f"(tmp[89]), "=f"(tmp[90]), "=f"(tmp[91]), "=f"(tmp[92]), "=f"(tmp[93]), "=f"(tmp[94]), "=f"(tmp[95]),
                "=f"(tmp[96]), "=f"(tmp[97]), "=f"(tmp[98]), "=f"(tmp[99]), "=f"(tmp[100]),"=f"(tmp[101]),"=f"(tmp[102]),"=f"(tmp[103]),
                "=f"(tmp[104]),"=f"(tmp[105]),"=f"(tmp[106]),"=f"(tmp[107]),"=f"(tmp[108]),"=f"(tmp[109]),"=f"(tmp[110]),"=f"(tmp[111]),
                "=f"(tmp[112]),"=f"(tmp[113]),"=f"(tmp[114]),"=f"(tmp[115]),"=f"(tmp[116]),"=f"(tmp[117]),"=f"(tmp[118]),"=f"(tmp[119]),
                "=f"(tmp[120]),"=f"(tmp[121]),"=f"(tmp[122]),"=f"(tmp[123]),"=f"(tmp[124]),"=f"(tmp[125]),"=f"(tmp[126]),"=f"(tmp[127])
              : "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}

// Limited tcgen05.st variants (b32). Use uint32_t registers.
template <int NUM_REGS, const char *SHAPE, int NUM>
PTX_DEVICE void tcgen05_st(uint32_t const* tmp, int row, int col) {
  int addr = (row << 16) | col;

  if constexpr (NUM_REGS == 1)
  asm volatile("tcgen05.st.sync.aligned%2.x%3.b32 [%0], {%1};"
              :: "r"(addr), "r"(tmp[0]), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 2)
  asm volatile("tcgen05.st.sync.aligned%3.x%4.b32 [%0], {%1, %2};"
              :: "r"(addr), "r"(tmp[0]), "r"(tmp[1]), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 4)
  asm volatile("tcgen05.st.sync.aligned%5.x%6.b32 [%0], {%1, %2, %3, %4};"
              :: "r"(addr), "r"(tmp[0]), "r"(tmp[1]), "r"(tmp[2]), "r"(tmp[3]), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 8)
  asm volatile("tcgen05.st.sync.aligned%9.x%10.b32 [%0], "
              "{%1, %2, %3, %4, %5, %6, %7, %8};"
              :: "r"(addr), "r"(tmp[0]), "r"(tmp[1]), "r"(tmp[2]), "r"(tmp[3]), "r"(tmp[4]), "r"(tmp[5]), "r"(tmp[6]), "r"(tmp[7]), "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 16)
  asm volatile("tcgen05.st.sync.aligned%17.x%18.b32 [%0], "
              "{%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};"
              :: "r"(addr),
                 "r"(tmp[ 0]), "r"(tmp[ 1]), "r"(tmp[ 2]), "r"(tmp[ 3]), "r"(tmp[ 4]), "r"(tmp[ 5]), "r"(tmp[ 6]), "r"(tmp[ 7]),
                 "r"(tmp[ 8]), "r"(tmp[ 9]), "r"(tmp[10]), "r"(tmp[11]), "r"(tmp[12]), "r"(tmp[13]), "r"(tmp[14]), "r"(tmp[15]),
                 "C"(SHAPE), "n"(NUM));
  if constexpr (NUM_REGS == 32)
  asm volatile("tcgen05.st.sync.aligned%33.x%34.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};"
              :: "r"(addr),
                 "r"(tmp[ 0]), "r"(tmp[ 1]), "r"(tmp[ 2]), "r"(tmp[ 3]), "r"(tmp[ 4]), "r"(tmp[ 5]), "r"(tmp[ 6]), "r"(tmp[ 7]),
                 "r"(tmp[ 8]), "r"(tmp[ 9]), "r"(tmp[10]), "r"(tmp[11]), "r"(tmp[12]), "r"(tmp[13]), "r"(tmp[14]), "r"(tmp[15]),
                 "r"(tmp[16]), "r"(tmp[17]), "r"(tmp[18]), "r"(tmp[19]), "r"(tmp[20]), "r"(tmp[21]), "r"(tmp[22]), "r"(tmp[23]),
                 "r"(tmp[24]), "r"(tmp[25]), "r"(tmp[26]), "r"(tmp[27]), "r"(tmp[28]), "r"(tmp[29]), "r"(tmp[30]), "r"(tmp[31]),
                 "C"(SHAPE), "n"(NUM));
}

// Convenience wrappers
PTX_DEVICE void tcgen05_ld_32x32b(float *tmp, int row, int col, int num) {
  if (num == 1)  tcgen05_ld<1,  SHAPE::_32x32b, 1>(tmp, row, col);
  if (num == 2)  tcgen05_ld<2,  SHAPE::_32x32b, 2>(tmp, row, col);
  if (num == 4)  tcgen05_ld<4,  SHAPE::_32x32b, 4>(tmp, row, col);
  if (num == 8)  tcgen05_ld<8,  SHAPE::_32x32b, 8>(tmp, row, col);
  if (num == 16) tcgen05_ld<16, SHAPE::_32x32b, 16>(tmp, row, col);
  if (num == 32) tcgen05_ld<32, SHAPE::_32x32b, 32>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x128b(float *tmp, int row, int col, int num) {
  if (num == 1)  tcgen05_ld<1,  SHAPE::_16x128b, 1>(tmp, row, col);
  if (num == 2)  tcgen05_ld<2,  SHAPE::_16x128b, 2>(tmp, row, col);
  if (num == 4)  tcgen05_ld<4,  SHAPE::_16x128b, 4>(tmp, row, col);
  if (num == 8)  tcgen05_ld<8,  SHAPE::_16x128b, 8>(tmp, row, col);
  if (num == 16) tcgen05_ld<16, SHAPE::_16x128b, 16>(tmp, row, col);
  if (num == 32) tcgen05_ld<32, SHAPE::_16x128b, 32>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x256b(float *tmp, int row, int col, int num) {
  if (num == 1)  tcgen05_ld<1,  SHAPE::_16x256b, 1>(tmp, row, col);
  if (num == 2)  tcgen05_ld<2,  SHAPE::_16x256b, 2>(tmp, row, col);
  if (num == 4)  tcgen05_ld<4,  SHAPE::_16x256b, 4>(tmp, row, col);
  if (num == 8)  tcgen05_ld<8,  SHAPE::_16x256b, 8>(tmp, row, col);
  if (num == 16) tcgen05_ld<16, SHAPE::_16x256b, 16>(tmp, row, col);
  if (num == 32) tcgen05_ld<32, SHAPE::_16x256b, 32>(tmp, row, col);
}

// Named wrappers used by gemm1
PTX_DEVICE void tcgen05_ld_32x32bx32(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_32x32b, NUM::x32>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_32x32bx64(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_32x32b, NUM::x64>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_32x32bx128(float *tmp, int row, int col) {
  tcgen05_ld_128regs<SHAPE::_32x32b, NUM::x128>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x128bx8(float *tmp, int row, int col) {
  tcgen05_ld_16regs<SHAPE::_16x128b, NUM::x8>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x128bx16(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x128b, NUM::x16>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x128bx32(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_16x128b, NUM::x32>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x256bx4(float *tmp, int row, int col) {
  tcgen05_ld_16regs<SHAPE::_16x256b, NUM::x4>(tmp, row, col);
}

PTX_DEVICE void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}

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

// ----- ptx_tcgen05_mma.cuh -----


// tcgen05.mma wrappers (CTA group 1 only for now)

struct COLLECTOR_USAGE {
  static constexpr char NONE[]      = "";
  static constexpr char A_FILL[]    = ".collector::a::fill";
  static constexpr char A_USE[]     = ".collector::a::use";
  static constexpr char A_LASTUSE[] = ".collector::a::lastuse";
  static constexpr char A_DISCARD[] = ".collector::a::discard";
};

// Block-scaled NVFP4 MMA (smem A/B descriptors, tmem scales)
template <int CTA_GROUP = 1, const char *collector_usage = COLLECTOR_USAGE::NONE>
PTX_DEVICE void tcgen05_mma_mxf4nvf4_block16(
  int d_tmem,
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t idesc,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  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%7 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(idesc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d),
       "C"(collector_usage)
  );
}

// F16/BF16 MMA, SS (A/B in smem desc), C in tmem
PTX_DEVICE void tcgen05_mma_f16_ss(
  uint32_t tmem_c,
  uint64_t desc_a,
  uint64_t desc_b,
  uint32_t idesc,
  int accumulate
) {
  PTX_ELECT_ONE();
  uint32_t mask[4] = {0, 0, 0, 0};
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %4, 0;\n\t"
    "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, {%5, %6, %7, %8}, p;\n\t"
    "}"
    :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc), "r"(accumulate),
       "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3])
  );
}

// F16/BF16 MMA, TS (A in tmem, B in smem desc), C in tmem
PTX_DEVICE void tcgen05_mma_f16_ts(
  uint32_t tmem_c,
  uint32_t tmem_a,
  uint64_t desc_b,
  uint32_t idesc,
  int accumulate
) {
  PTX_ELECT_ONE();
  uint32_t mask[4] = {0, 0, 0, 0};
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %4, 0;\n\t"
    "tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p;\n\t"
    "}"
    :: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(idesc), "r"(accumulate),
       "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3])
  );
}

// F16/BF16 MMA, WS (warp-specialized), C in tmem
PTX_DEVICE void tcgen05_mma_ws_f16_ts(
  uint32_t tmem_c,
  uint32_t tmem_a,
  uint64_t desc_b,
  uint32_t idesc,
  int accumulate
) {
  PTX_ELECT_ONE();
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %4, 0;\n\t"
    "tcgen05.mma.ws.cta_group::1.kind::f16 [%0], [%1], %2, %3, p, 0;\n\t"
    "}"
    :: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(idesc), "r"(accumulate)
  );
}

// Alias used by gemm1 (block-scaled NVFP4 MMA, d_tmem assumed 0)
PTX_DEVICE 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
) {
  tcgen05_mma_mxf4nvf4_block16<1, COLLECTOR_USAGE::NONE>(
    0, a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d
  );
}

// ----- ptx_tcgen05_sync.cuh -----


// tcgen05 commit, wait, fence wrappers (CTA group 1 only for now)

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_commit(int mbar_addr) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
              :: "r"(mbar_addr) : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_commit_mcast(int mbar_addr, uint16_t cta_mask) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
              :: "r"(mbar_addr), "h"(cta_mask) : "memory");
}

// tcgen05 tmem allocation / deallocation (CTA group 1 only for now)
template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_alloc(int smem_addr, int num_cols) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
              :: "r"(smem_addr), "r"(num_cols));
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tcgen05_dealloc(int tmem_addr, int num_cols) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
              :: "r"(tmem_addr), "r"(num_cols));
}

PTX_DEVICE void tcgen05_wait_ld() {
  asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory");
}

PTX_DEVICE void tcgen05_wait_st() {
  asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
}

PTX_DEVICE void tcgen05_fence_before_thread_sync() {
  asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory");
}

PTX_DEVICE void tcgen05_fence_after_thread_sync() {
  asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
}

// ----- ptx_tma.cuh -----


// TMA bulk tensor loads (CTA group 1 only for now)

// Bulk global->shared copy (non-tensor)
PTX_DEVICE void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
  PTX_ELECT_ONE();
  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));
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_1d_gmem2smem(int dst, const void *tmap_ptr, int x, int mbar_addr, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.1d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
              "[%0], [%1, {%2}], [%3], %4;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(mbar_addr), "l"(cache_policy)
              : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_2d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int mbar_addr, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
              "[%0], [%1, {%2, %3}], [%4], %5;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(mbar_addr), "l"(cache_policy)
              : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.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)
              : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_1d_gmem2smem_mcast(int dst, const void *tmap_ptr, int x, int mbar_addr, int16_t cta_mask, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::1.L2::cache_hint "
              "[%0], [%1, {%2}], [%3], %4, %5;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(mbar_addr), "h"(cta_mask), "l"(cache_policy)
              : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_2d_gmem2smem_mcast(int dst, const void *tmap_ptr, int x, int y, int mbar_addr, int16_t cta_mask, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::1.L2::cache_hint "
              "[%0], [%1, {%2, %3}], [%4], %5, %6;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(mbar_addr), "h"(cta_mask), "l"(cache_policy)
              : "memory");
}

template <int CTA_GROUP = 1>
PTX_DEVICE void tma_3d_gmem2smem_mcast(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, int16_t cta_mask, uint64_t cache_policy) {
  static_assert(CTA_GROUP == 1, "Only CTA_GROUP=1 supported for now");
  PTX_ELECT_ONE();
  asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.cta_group::1.L2::cache_hint "
              "[%0], [%1, {%2, %3, %4}], [%5], %6, %7;"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "h"(cta_mask), "l"(cache_policy)
              : "memory");
}

// ----- device.cuh -----
// experimental device snippet (placeholder)
__device__ __forceinline__ int exp_add(int a, int b) { return a + b; }
__device__ __forceinline__ int exp_sub(int a, int b) { return a - b; }


// Grouped GEMM kernel for NVFP4 block-scaled matrices
// Constraints: M % 128 == 0, N % 128 == 0, K % 256 == 0


constexpr int WARP_SIZE_GG = 32;
constexpr int MMA_K_GG = 64;
constexpr uint64_t EVICT_NORMAL_GG = 0x1000000000000000;

__device__ inline
constexpr uint64_t desc_encode_gg(uint64_t x) { return (x & 0x3FFFFULL) >> 4ULL; }

__device__ inline
uint32_t elect_sync_gg() {
  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;
}

// Problem info structure (passed per group)
struct GroupedGemmProblem {
  int M;
  int N;
  int K;
  int L;  // always 1
  int tiles_m;  // M / BLOCK_M
  int tiles_n;  // N / BLOCK_N
  int cumulative_tiles;  // sum of tiles for groups 0..i-1
};


template <
  int BLOCK_M,
  int BLOCK_N,
  int BLOCK_K,
  int NUM_STAGES
>
// @kernel name=kernel_grouped_gemm arch=sm_100a warp_size=WARP_SIZE_GG num_warps=NUM_WARPS cluster_ctas=1
__global__
__launch_bounds__(BLOCK_M + 2 * WARP_SIZE_GG)
void kernel_grouped_gemm(
  // Per-group tensor pointers (contiguous arrays)
  const char* const* __restrict__ A_ptrs,    // [num_groups]
  const char* const* __restrict__ B_ptrs,    // [num_groups]
  half** __restrict__ C_ptrs,                // [num_groups]
  const char* const* __restrict__ SFA_ptrs,  // [num_groups]
  const char* const* __restrict__ SFB_ptrs,  // [num_groups]
  // Per-group TMA descriptors (already initialized on host)
  const CUtensorMap* __restrict__ A_tmaps,   // [num_groups]
  const CUtensorMap* __restrict__ B_tmaps,   // [num_groups]
  // Problem info per group
  const GroupedGemmProblem* __restrict__ problems,  // [num_groups]
  int num_groups
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;  // Linear CTA index across all groups

  const int lane_id = tid % WARP_SIZE_GG;
  const int warp_id = tid / WARP_SIZE_GG;
  constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE_GG + 2;

  // Map bid to (group_idx, tile_m, tile_n) via binary search
  int group_idx = 0;
  for (int g = 0; g < num_groups; g++) {
    if (bid >= problems[g].cumulative_tiles) {
      group_idx = g;
    }
  }
  
  const GroupedGemmProblem& prob = problems[group_idx];
  const int local_bid = bid - prob.cumulative_tiles;
  const int bid_m = local_bid % prob.tiles_m;
  const int bid_n = local_bid / prob.tiles_m;

  const int M = prob.M;
  const int N = prob.N;
  const int K = prob.K;
  
  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  // Get per-group pointers
  const char* A_ptr = A_ptrs[group_idx];
  const char* B_ptr = B_ptrs[group_idx];
  half* C_ptr = C_ptrs[group_idx];
  const char* SFA_ptr = SFA_ptrs[group_idx];
  const char* SFB_ptr = SFB_ptrs[group_idx];
  const CUtensorMap* A_tmap = &A_tmaps[group_idx];
  const CUtensorMap* B_tmap = &B_tmaps[group_idx];

  // @buffer name=tmem0 space=tmem cols=BLOCK_N*2 dtype=f32
  // @buffer name=A_smem space=smem dtype=fp4 major=K swizzle=128B element_bits=4
  // @buffer name=B_smem space=smem dtype=fp4 major=K swizzle=128B element_bits=4
  // @buffer name=SFA_smem space=smem dtype=fp8 major=K swizzle=none element_bits=8
  // @buffer name=SFB_smem space=smem dtype=fp8 major=K swizzle=none element_bits=8

  // Set up smem
  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 / 2;
  constexpr int B_size = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

  // mbarriers: NUM_STAGES for TMA, NUM_STAGES for MMA, 1 for mainloop
  // @barrier name=tma_mbar scope=cta count=NUM_STAGES
  // @barrier name=mma_mbar scope=cta count=NUM_STAGES
  // @barrier name=mainloop_mbar scope=cta count=1
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K_GG);


  // Initialize mbarriers and allocate tmem
  if (warp_id == 0 && elect_sync_gg()) {
    // @op mbarrier_init bar=tma_mbar count=1 scope=cta
    // @op warp_id=0 lane_id=elect
    // @op mbarrier_init bar=mma_mbar count=1 scope=cta
    // @op warp_id=0 lane_id=elect
    // @op mbarrier_init bar=mainloop_mbar count=1 scope=cta
    // @op warp_id=0 lane_id=elect
    for (int i = 0; i < NUM_STAGES * 2 + 1; i++)
      mbarrier_init(tma_mbar_addr + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  }
  else if (warp_id == 1) {
    // @op tcgen05_alloc tmem=tmem0 cols=BLOCK_N*2 cta_group=1 scope=one_warp
    // @op warp_id=1
    tcgen05_alloc(smem, BLOCK_N * 2);
  }
  __syncthreads();

  const int num_iters = K / BLOCK_K;


  // Warp-specialization: TMA warp
  if (warp_id == NUM_WARPS - 2 && elect_sync_gg()) {
    uint64_t cache_A = EVICT_NORMAL_GG;
    uint64_t cache_B = EVICT_NORMAL_GG;

    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem_addr = smem + stage_id * STAGE_SIZE;
      const int B_smem_addr = A_smem_addr + A_size;
      const int SFA_smem_addr = B_smem_addr + B_size;
      const int SFB_smem_addr = SFA_smem_addr + SFA_size;

      const int off_k = iter_k * BLOCK_K;
      // @op tma_3d_gmem2smem bar=tma_mbar tmap=A_tmap
      // @op warp_id=NUM_WARPS-2 lane_id=elect rank=3
      tma_3d_gmem2smem(A_smem_addr, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
      // @op tma_3d_gmem2smem bar=tma_mbar tmap=B_tmap
      // @op warp_id=NUM_WARPS-2 lane_id=elect rank=3
      tma_3d_gmem2smem(B_smem_addr, B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      // Scale factors layout: [M/128, rest_k, 32, 4, 4]
      const int rest_k = K / 16 / 4;
      const char* SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char* SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      // @op tma_gmem2smem bar=tma_mbar size=SFA_size dst_align=16 src_align=16
      // @op warp_id=NUM_WARPS-2 lane_id=elect
      tma_gmem2smem(SFA_smem_addr, SFA_src, SFA_size, mbar_addr, cache_A);
      // @op tma_gmem2smem bar=tma_mbar size=SFB_size dst_align=16 src_align=16
      // @op warp_id=NUM_WARPS-2 lane_id=elect
      tma_gmem2smem(SFB_smem_addr, SFB_src, SFB_size, mbar_addr, cache_B);

      // @op mbarrier_arrive_expect_tx bar=tma_mbar size=STAGE_SIZE
      // @op warp_id=NUM_WARPS-2 lane_id=elect
      asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                  :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
    };

    // Pre-issue NUM_STAGES TMA loads
    for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++)
      issue_tma(iter_k, iter_k);

    // @loop var=iter_k iters=num_iters start=NUM_STAGES
    for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int mma_phase = (iter_k / NUM_STAGES - 1) % 2;
      // @op mbarrier_wait bar=mma_mbar phase=mma_phase
      // @op warp_id=NUM_WARPS-2 lane_id=elect
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
    // @endloop
  }


  // Warp-specialization: MMA warp
  else if (warp_id == NUM_WARPS - 1 && elect_sync_gg()) {
    constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
                              | (1U << 10U)  // btype=E2M1
                              | ((uint32_t)BLOCK_N >> 3U << 17U)
                              | ((uint32_t)128 >> 7U << 27U);

    // @loop var=iter_k iters=num_iters
    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) % 2;
      // @op mbarrier_wait bar=tma_mbar phase=tma_phase
      // @op warp_id=NUM_WARPS-1 lane_id=elect
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem_addr = smem + stage_id * STAGE_SIZE;
      const int B_smem_addr = A_smem_addr + A_size;
      const int SFA_smem_addr = B_smem_addr + B_size;
      const int SFB_smem_addr = SFA_smem_addr + SFA_size;

      // @desc name=AB_desc major=K swizzle=128B sbo=8*128 lbo=1
      auto make_desc_AB = [](int addr) -> uint64_t {
        const int SBO = 8 * 128;
        return desc_encode_gg(addr) | (desc_encode_gg(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
      };
      // @desc name=SF_desc major=K swizzle=none sbo=8*16 lbo=1
      auto make_desc_SF = [](int addr) -> uint64_t {
        const int SBO = 8 * 16;
        return desc_encode_gg(addr) | (desc_encode_gg(SBO) << 32ULL) | (1ULL << 46ULL);
      };

      constexpr uint64_t SF_desc = make_desc_SF(0);
      const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem_addr >> 4ULL);
      const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem_addr >> 4ULL);

      // @loop var=k iters=BLOCK_K/MMA_K_GG
      for (int k = 0; k < BLOCK_K / MMA_K_GG; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
        // @op tcgen05_cp tmem=tmem0 cta_group=1 issue=one_thread
        // @op shape=32x128b tile=warpx4 warp_id=NUM_WARPS-1 lane_id=elect desc=SF_desc smem_buf=SFA_smem
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
        // @op tcgen05_cp tmem=tmem0 cta_group=1 issue=one_thread
        // @op shape=32x128b tile=warpx4 warp_id=NUM_WARPS-1 lane_id=elect desc=SF_desc smem_buf=SFB_smem
        tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
      }
      // @endloop

      // @loop var=k1 iters=BLOCK_K/256
      for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
        // @loop var=k2 iters=256/MMA_K_GG
        for (int k2 = 0; k2 < 256 / MMA_K_GG; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem_addr + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b_desc = make_desc_AB(B_smem_addr + k1 * BLOCK_N * 128 + k2 * 32);

          int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4 + (off_m % 128) / 32;
          const int scale_B_tmem = SFB_tmem + k_sf * 4 + (off_n % 128) / 32;

          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          // @op tcgen05_mma tmem=tmem0 cta_group=1 issue=one_thread
          // @op shape=mxf4nvf4.block16 warp_id=NUM_WARPS-1 lane_id=elect desc_a=AB_desc desc_b=AB_desc
          tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }
        // @endloop
      }
      // @endloop

      // @op tcgen05_commit bar=mma_mbar cta_group=1
      // @op warp_id=NUM_WARPS-1 lane_id=elect
      tcgen05_commit(mma_mbar_addr + stage_id * 8);
    }
    // @endloop

    // @op tcgen05_commit bar=mainloop_mbar cta_group=1
    // @op warp_id=NUM_WARPS-1 lane_id=elect
    tcgen05_commit(mainloop_mbar_addr);
  }


  // Epilogue warps: read from tmem and write to global
  else if (tid < BLOCK_M) {
    // @op mbarrier_wait bar=mainloop_mbar phase=0
    // @op
    mbarrier_wait(mainloop_mbar_addr, 0);
    // @op tcgen05_fence_after_thread_sync
    // @op
    tcgen05_fence_after_thread_sync();

    // N-major output: C[m, n] stored as C[m * N + n]
    constexpr int WIDTH = std::min(BLOCK_N, 64);

    // @loop var=n iters=BLOCK_N/WIDTH
    for (int n = 0; n < BLOCK_N / WIDTH; n++) {
      float tmp[WIDTH];
      // @op tcgen05_ld tmem=tmem0 cta_group=1 when=WIDTH==128
      // @op shape=32x32b num=128 warp_id=warp_id lane_id=lane_id
      if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp, warp_id * 32, n * WIDTH);
      // @op tcgen05_ld tmem=tmem0 cta_group=1 when=WIDTH==64
      // @op shape=32x32b num=64 warp_id=warp_id lane_id=lane_id
      if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp, warp_id * 32, n * WIDTH);
      // @op tcgen05_ld tmem=tmem0 cta_group=1 when=WIDTH==32
      // @op shape=32x32b num=32 warp_id=warp_id lane_id=lane_id
      if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp, warp_id * 32, n * WIDTH);
      // @op tcgen05_wait_ld
      // @op
      tcgen05_wait_ld();

      // Write to N-major C: C[row, col] at C[row * N + col]
      for (int i = 0; i < WIDTH; i++) {
        const int row = off_m + tid;
        const int col = off_n + n * WIDTH + i;
        if (row < M && col < N) {
          C_ptr[row * N + col] = __float2half(tmp[i]);
        }
      }
    }
    // @endloop

    // @op ptx_bar_sync bar_id=1 count=BLOCK_M
    // @op
    asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
    if (warp_id == 0)
      // @op tcgen05_dealloc tmem=tmem0 cols=BLOCK_N*2 cta_group=1 scope=one_warp
      // @op warp_id=0
      tcgen05_dealloc(0, BLOCK_N * 2);
  }
}

// @endkernel
// ----- host.cuh -----
// experimental host snippet (placeholder)
inline int exp_host_noop(int x) { return x; }


// Host-side helpers for grouped GEMM

void check_cu_gg(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *error_msg_ptr;
  if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
    error_msg_ptr = "unable to get error string";
  TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
}

void init_AB_tmap_gg(
  CUtensorMap *tmap,
  const char *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]       = {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};

  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
    rank,
    (void *)ptr,
    globalDim,
    globalStrides,
    boxDim,
    elementStrides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu_gg(err);
}


template <
  int BLOCK_M,
  int BLOCK_N,
  int BLOCK_K,
  int NUM_STAGES
>
std::vector<at::Tensor> grouped_gemm_launch(
  const std::vector<at::Tensor>& A_tensors,
  const std::vector<at::Tensor>& B_tensors,
  const std::vector<at::Tensor>& SFA_tensors,
  const std::vector<at::Tensor>& SFB_tensors,
  std::vector<at::Tensor>& C_tensors
) {
  static_assert(BLOCK_K % 256 == 0);

  const int num_groups = A_tensors.size();
  
  // Allocate device memory for pointers and problem info
  std::vector<const char*> h_A_ptrs(num_groups);
  std::vector<const char*> h_B_ptrs(num_groups);
  std::vector<half*> h_C_ptrs(num_groups);
  std::vector<const char*> h_SFA_ptrs(num_groups);
  std::vector<const char*> h_SFB_ptrs(num_groups);
  std::vector<CUtensorMap> h_A_tmaps(num_groups);
  std::vector<CUtensorMap> h_B_tmaps(num_groups);
  std::vector<GroupedGemmProblem> h_problems(num_groups);

  int total_tiles = 0;
  for (int g = 0; g < num_groups; g++) {
    const int M = A_tensors[g].size(0);
    const int N = B_tensors[g].size(0);
    const int K = A_tensors[g].size(1) * 2;
    
    h_A_ptrs[g] = reinterpret_cast<const char*>(A_tensors[g].data_ptr());
    h_B_ptrs[g] = reinterpret_cast<const char*>(B_tensors[g].data_ptr());
    h_C_ptrs[g] = reinterpret_cast<half*>(C_tensors[g].data_ptr());
    h_SFA_ptrs[g] = reinterpret_cast<const char*>(SFA_tensors[g].data_ptr());
    h_SFB_ptrs[g] = reinterpret_cast<const char*>(SFB_tensors[g].data_ptr());
    
    // @op cute_tmap name=A_tmap rank=3
    // @op dtype=16u4_align8b interleave=none swizzle=128b l2=none oob=none
    init_AB_tmap_gg(&h_A_tmaps[g], h_A_ptrs[g], M, K, BLOCK_M, BLOCK_K);
    // @op cute_tmap name=B_tmap rank=3
    // @op dtype=16u4_align8b interleave=none swizzle=128b l2=none oob=none
    init_AB_tmap_gg(&h_B_tmaps[g], h_B_ptrs[g], N, K, BLOCK_N, BLOCK_K);
    
    h_problems[g].M = M;
    h_problems[g].N = N;
    h_problems[g].K = K;
    h_problems[g].L = 1;
    h_problems[g].tiles_m = (M + BLOCK_M - 1) / BLOCK_M;
    h_problems[g].tiles_n = (N + BLOCK_N - 1) / BLOCK_N;
    h_problems[g].cumulative_tiles = total_tiles;
    total_tiles += h_problems[g].tiles_m * h_problems[g].tiles_n;
  }
  
  // Allocate device arrays
  const char** d_A_ptrs;
  const char** d_B_ptrs;
  half** d_C_ptrs;
  const char** d_SFA_ptrs;
  const char** d_SFB_ptrs;
  CUtensorMap* d_A_tmaps;
  CUtensorMap* d_B_tmaps;
  GroupedGemmProblem* d_problems;
  
  cudaMalloc(&d_A_ptrs, num_groups * sizeof(const char*));
  cudaMalloc(&d_B_ptrs, num_groups * sizeof(const char*));
  cudaMalloc(&d_C_ptrs, num_groups * sizeof(half*));
  cudaMalloc(&d_SFA_ptrs, num_groups * sizeof(const char*));
  cudaMalloc(&d_SFB_ptrs, num_groups * sizeof(const char*));
  cudaMalloc(&d_A_tmaps, num_groups * sizeof(CUtensorMap));
  cudaMalloc(&d_B_tmaps, num_groups * sizeof(CUtensorMap));
  cudaMalloc(&d_problems, num_groups * sizeof(GroupedGemmProblem));
  
  cudaMemcpy(d_A_ptrs, h_A_ptrs.data(), num_groups * sizeof(const char*), cudaMemcpyHostToDevice);
  cudaMemcpy(d_B_ptrs, h_B_ptrs.data(), num_groups * sizeof(const char*), cudaMemcpyHostToDevice);
  cudaMemcpy(d_C_ptrs, h_C_ptrs.data(), num_groups * sizeof(half*), cudaMemcpyHostToDevice);
  cudaMemcpy(d_SFA_ptrs, h_SFA_ptrs.data(), num_groups * sizeof(const char*), cudaMemcpyHostToDevice);
  cudaMemcpy(d_SFB_ptrs, h_SFB_ptrs.data(), num_groups * sizeof(const char*), cudaMemcpyHostToDevice);
  cudaMemcpy(d_A_tmaps, h_A_tmaps.data(), num_groups * sizeof(CUtensorMap), cudaMemcpyHostToDevice);
  cudaMemcpy(d_B_tmaps, h_B_tmaps.data(), num_groups * sizeof(CUtensorMap), cudaMemcpyHostToDevice);
  cudaMemcpy(d_problems, h_problems.data(), num_groups * sizeof(GroupedGemmProblem), cudaMemcpyHostToDevice);
  
  // Launch kernel
  int tb_size = BLOCK_M + 2 * WARP_SIZE_GG;
  int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
  int SFAB_size = 128 * (BLOCK_K / 16) * 2;
  int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
  
  auto this_kernel = kernel_grouped_gemm<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
  if (smem_size > 48000)
    cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
    
  this_kernel<<<total_tiles, tb_size, smem_size>>>(
    d_A_ptrs, d_B_ptrs, d_C_ptrs, d_SFA_ptrs, d_SFB_ptrs,
    d_A_tmaps, d_B_tmaps, d_problems, num_groups
  );
  
  // Cleanup
  cudaFree(d_A_ptrs);
  cudaFree(d_B_ptrs);
  cudaFree(d_C_ptrs);
  cudaFree(d_SFA_ptrs);
  cudaFree(d_SFB_ptrs);
  cudaFree(d_A_tmaps);
  cudaFree(d_B_tmaps);
  cudaFree(d_problems);
  
  return C_tensors;
}


std::vector<at::Tensor> grouped_gemm(
  const std::vector<at::Tensor>& A_tensors,
  const std::vector<at::Tensor>& B_tensors,
  const std::vector<at::Tensor>& SFA_tensors,
  const std::vector<at::Tensor>& SFB_tensors,
  const std::vector<at::Tensor>& C_tensors
) {
  // Use fixed tile sizes: 128x128x256 with 6 stages
  std::vector<at::Tensor> C_out = C_tensors;
  return grouped_gemm_launch<128, 128, 256, 6>(
    A_tensors, B_tensors, SFA_tensors, SFB_tensors, C_out
  );
}

TORCH_LIBRARY(my_grouped_gemm, m) {
  m.def("grouped_gemm(Tensor[] A, Tensor[] B, Tensor[] SFA, Tensor[] SFB, Tensor[] C) -> Tensor[]");
  m.impl("grouped_gemm", &grouped_gemm);
}
'''

# @chunk name=grouped_python_header
#!POPCORN leaderboard nvfp4_grouped_gemm
#!POPCORN gpu NVIDIA
import torch
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline
load_inline(
    'gemm_all',
    cpp_sources='',
    cuda_sources=CUDA_SRC,
    verbose=False,
    is_python_module=False,
    no_implicit_headers=False,
    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'],
)
# @chunk name=grouped_python_bindings
grouped_gemm_fn = torch.ops.my_grouped_gemm.grouped_gemm
# @chunk name=grouped_python_kernel
def custom_kernel(data: input_t) -> output_t:
    """
    Custom kernel for grouped NVFP4 GEMM.
    """
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
    
    # Prepare tensor lists
    A_list = []
    B_list = []
    SFA_list = []
    SFB_list = []
    C_list = []
    
    for i, ((a, b, c), (sfa_reordered, sfb_reordered), (m, n, k, l)) in enumerate(
        zip(abc_tensors, sfasfb_reordered_tensors, problem_sizes)
    ):
        # a shape: [m, k//2, l=1], squeeze to [m, k//2]
        # b shape: [n, k//2, l=1], squeeze to [n, k//2]
        A_list.append(a.squeeze(-1).contiguous())
        B_list.append(b.squeeze(-1).contiguous())
        C_list.append(c.squeeze(-1).contiguous())
        
        # sfa_reordered shape: [32, 4, rest_m, 4, rest_k, l]
        # Need to permute to [rest_m, rest_k, 32, 4, 4] for kernel
        # Permute: (2, 4, 0, 1, 3, 5) -> [rest_m, rest_k, 32, 4, 4, l]
        # Then squeeze last dim and flatten to contiguous
        sfa_perm = sfa_reordered.permute(2, 4, 0, 1, 3, 5).squeeze(-1).contiguous()
        sfb_perm = sfb_reordered.permute(2, 4, 0, 1, 3, 5).squeeze(-1).contiguous()
        
        # Flatten to 1D for TMA bulk copy (kernel expects [M/128, K/64, 32, 4, 4])
        SFA_list.append(sfa_perm.view(-1))
        SFB_list.append(sfb_perm.view(-1))
    
    # Call the CUDA kernel
    result = grouped_gemm_fn(A_list, B_list, SFA_list, SFB_list, C_list)
    
    # Add back the L dimension
    output = []
    for c in result:
        output.append(c.unsqueeze(-1))
    
    return output