submission 367633

#!POPCORN leaderboard nvfp4_dual_gemm
#!POPCORN gpu NVIDIA

import torch
from torch.utils.cpp_extension import load_inline

input_t = tuple
output_t = torch.Tensor

# Optimized fused dual GEMM kernel v5
# Key optimizations:
# 1. Larger BLOCK_N=128 with 4-stage pipeline for better compute intensity
# 2. Vectorized N-major epilogue using half2 writes for better memory bandwidth
# 3. Interleaved MMA operations for dual GEMM
# 4. Optimized scale factor loading
CUDA_SRC = r"""
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <torch/library.h>
#include <ATen/core/Tensor.h>
#include <ATen/ATen.h>

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

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

__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__
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 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));
}

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

__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_mma_nvfp4(
  int d_tmem, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
  int scale_A_tmem, int scale_B_tmem, int enable_input_d
) {
  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 _32x32b[] = ".32x32b";
  static constexpr char _16x128b[] = ".16x128b";
  static constexpr char _16x256b[] = ".16x256b";
};

struct NUM {
  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";
};

template <const char *SHAPE, const char *NUM>
__device__ inline 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>
__device__ inline 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>
__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_32x32bx32(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_32x32b, NUM::x32>(tmp, row, col);
}
__device__ inline void tcgen05_ld_32x32bx64(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_32x32b, NUM::x64>(tmp, row, col);
}

// N-major load functions for vectorized epilogue
__device__ inline void tcgen05_ld_16x256bx4(float *tmp, int row, int col) {
  tcgen05_ld_16regs<SHAPE::_16x256b, NUM::x4>(tmp, row, col);
}
__device__ inline void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}
__device__ inline void tcgen05_ld_16x256bx16(float *tmp, int row, int col) {
  tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(tmp, row, col);
}

void check_cu(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(
  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(err);
}

// Optimized fused dual GEMM kernel v5: C = silu(A @ B1) * (A @ B2)
// Key optimizations:
// 1. Larger BLOCK_N for better compute intensity
// 2. Vectorized N-major epilogue with half2 writes
// 3. Interleaved MMA operations for dual GEMM
// 4. Optimized pipeline stages based on problem size
template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void dual_gemm_kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B1_tmap,
  const __grid_constant__ CUtensorMap B2_tmap,
  const char *SFA_ptr, const char *SFB1_ptr, const char *SFB2_ptr,
  half *C_ptr, int M, int N
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;
  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  const int grid_n = N / BLOCK_N;
  const int bid_m = bid / grid_n;
  const int bid_n = bid % grid_n;
  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;

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

  // Shared memory layout per stage: A, B1, B2, SFA, SFB1, SFB2
  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 + 2 * B_size + SFA_size + 2 * SFB_size;

  #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;

  // TMEM layout: [0, BLOCK_N) for GEMM1, [BLOCK_N, 2*BLOCK_N) for GEMM2
  constexpr int TMEM_GEMM1 = 0;
  constexpr int TMEM_GEMM2 = BLOCK_N;
  constexpr int SFA_tmem = 2 * BLOCK_N;
  constexpr int SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
  constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);

  if (warp_id == 0 && elect_sync()) {
    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) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
                :: "r"(smem), "r"(BLOCK_N * 2 * 2));
  }
  __syncthreads();

  const int num_iters = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    // TMA warp - dynamic cache policies based on M vs N
    uint64_t cache_A = (M > N) ? EVICT_FIRST : EVICT_LAST;
    uint64_t cache_B = (M > N) ? EVICT_LAST : EVICT_FIRST;

    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B1_smem = A_smem + A_size;
      const int B2_smem = B1_smem + B_size;
      const int SFA_smem = B2_smem + B_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_size;

      const int off_k = iter_k * BLOCK_K;

      // Load A (shared between both GEMMs)
      tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
      // Load B1 and B2
      tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
      tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      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 *SFB1_src = SFB1_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;

      tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
      tma_gmem2smem(SFB1_smem, SFB1_src, SFB_size, mbar_addr, cache_B);
      tma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_B);

      constexpr int expected_bytes = A_size + 2 * B_size + SFA_size + 2 * SFB_size;
      asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                  :: "r"(mbar_addr), "r"(expected_bytes) : "memory");
    };

    for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++)
      issue_tma(iter_k, iter_k);

    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;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
  }
  else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    // MMA warp - interleaved GEMM1 and GEMM2 for better tensor core utilization
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U)
                              | ((uint32_t)BLOCK_N >> 3U << 17U)
                              | ((uint32_t)128 >> 7U << 27U);

    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;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B1_smem = A_smem + A_size;
      const int B2_smem = B1_smem + B_size;
      const int SFA_smem = B2_smem + B_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_size;

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

      constexpr uint64_t SF_desc = make_desc_SF(0);
      const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB1_desc = SF_desc + ((uint64_t)SFB1_smem >> 4ULL);
      const uint64_t SFB2_desc = SF_desc + ((uint64_t)SFB2_smem >> 4ULL);

      // Copy scale factors to TMEM
      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb1_desc = SFB1_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb2_desc = SFB2_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
        tcgen05_cp_nvfp4(SFB1_tmem + k * 4, sfb1_desc);
        tcgen05_cp_nvfp4(SFB2_tmem + k * 4, sfb2_desc);
      }

      // Interleaved MMA: alternate between GEMM1 and GEMM2 for better utilization
      for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b1_desc = make_desc_AB(B1_smem + k1 * BLOCK_N * 128 + k2 * 32);
          uint64_t b2_desc = make_desc_AB(B2_smem + k1 * BLOCK_N * 128 + k2 * 32);

          int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
          const int scale_B1_tmem = SFB1_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
          const int scale_B2_tmem = SFB2_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);

          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;

          // Issue both MMAs back-to-back for better pipelining
          tcgen05_mma_nvfp4(TMEM_GEMM1, a_desc, b1_desc, i_desc, scale_A_tmem, scale_B1_tmem, enable_input_d);
          tcgen05_mma_nvfp4(TMEM_GEMM2, a_desc, b2_desc, i_desc, scale_A_tmem, scale_B2_tmem, enable_input_d);
        }
      }

      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                  :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
    }

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                :: "r"(mainloop_mbar_addr) : "memory");
  }
  else if (tid < BLOCK_M) {
    // Epilogue - fused silu_mul with vectorized N-major output: C = silu(GEMM1) * GEMM2
    mbarrier_wait(mainloop_mbar_addr, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");

    // Use N-major epilogue with half2 writes for better memory bandwidth
    // Each iteration processes 16 rows x BLOCK_N columns
    for (int m = 0; m < 32 / 16; m++) {
      float tmp1[BLOCK_N / 2], tmp2[BLOCK_N / 2];

      // Load GEMM1 result using N-major layout
      if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp1, warp_id * 32 + m * 16, TMEM_GEMM1);
      else if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp1, warp_id * 32 + m * 16, TMEM_GEMM1);
      else if constexpr (BLOCK_N == 32) tcgen05_ld_16x256bx4(tmp1, warp_id * 32 + m * 16, TMEM_GEMM1);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      // Load GEMM2 result
      if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp2, warp_id * 32 + m * 16, TMEM_GEMM2);
      else if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp2, warp_id * 32 + m * 16, TMEM_GEMM2);
      else if constexpr (BLOCK_N == 32) tcgen05_ld_16x256bx4(tmp2, warp_id * 32 + m * 16, TMEM_GEMM2);
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      // Fused silu_mul with vectorized half2 writes
      #pragma unroll
      for (int i = 0; i < BLOCK_N / 8; i++) {
        const int row = off_m + warp_id * 32 + m * 16 + lane_id / 4;
        const int col = off_n + i * 8 + (lane_id % 4) * 2;

        // Apply silu to GEMM1 results and multiply with GEMM2 results
        float x0 = tmp1[i * 4 + 0];
        float x1 = tmp1[i * 4 + 1];
        float x2 = tmp1[i * 4 + 2];
        float x3 = tmp1[i * 4 + 3];

        float silu0 = x0 / (1.0f + expf(-x0));
        float silu1 = x1 / (1.0f + expf(-x1));
        float silu2 = x2 / (1.0f + expf(-x2));
        float silu3 = x3 / (1.0f + expf(-x3));

        float out0 = silu0 * tmp2[i * 4 + 0];
        float out1 = silu1 * tmp2[i * 4 + 1];
        float out2 = silu2 * tmp2[i * 4 + 2];
        float out3 = silu3 * tmp2[i * 4 + 3];

        // Vectorized half2 writes for better memory bandwidth
        reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] = __float22half2_rn({out0, out1});
        reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] = __float22half2_rn({out2, out3});
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
    if (warp_id == 0)
      asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2 * 2));
  }
}

template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
void dual_gemm_launch(
  const at::Tensor& A, const at::Tensor& B1, const at::Tensor& B2,
  const at::Tensor& SFA, const at::Tensor& SFB1, const at::Tensor& SFB2,
  at::Tensor& C
) {
  const int M = A.size(0);
  const int N = B1.size(0);

  auto A_ptr = reinterpret_cast<const char *>(A.data_ptr());
  auto B1_ptr = reinterpret_cast<const char *>(B1.data_ptr());
  auto B2_ptr = reinterpret_cast<const char *>(B2.data_ptr());
  auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
  auto SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());
  auto SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());
  auto C_ptr = reinterpret_cast<half *>(C.data_ptr());

  CUtensorMap A_tmap, B1_tmap, B2_tmap;
  init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
  init_AB_tmap(&B1_tmap, B1_ptr, N, K, BLOCK_N, BLOCK_K);
  init_AB_tmap(&B2_tmap, B2_ptr, N, K, BLOCK_N, BLOCK_K);

  int grid = (M / BLOCK_M) * (N / BLOCK_N);
  int tb_size = BLOCK_M + 2 * WARP_SIZE;

  int A_size = BLOCK_M * BLOCK_K / 2;
  int B_size = BLOCK_N * BLOCK_K / 2;
  int SFA_size = 128 * BLOCK_K / 16;
  int SFB_size = 128 * BLOCK_K / 16;
  int smem_size = (A_size + 2 * B_size + SFA_size + 2 * SFB_size) * NUM_STAGES;

  auto this_kernel = dual_gemm_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
  if (smem_size > 48000)
    cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N);
}

void dual_gemm(
  const at::Tensor& A, const at::Tensor& B1, const at::Tensor& B2,
  const at::Tensor& SFA, const at::Tensor& SFB1, const at::Tensor& SFB2,
  at::Tensor& C
) {
  const int K = A.size(1) * 2;
  const int M = A.size(0);
  const int N = B1.size(0);

  // Optimized configurations for v5:
  // - Use BLOCK_N=128 with 4 stages for larger M and N (more compute per block)
  // - Use BLOCK_N=64 with 5 stages for smaller M (better parallelism)
  // The key insight is that BLOCK_N=128 only helps when M >= 512 because:
  //   - Larger tiles reduce total threadblocks, which hurts SM utilization for small M
  //   - M=256 with BLOCK_N=128 gives only (M/128)*(N/128) = 2*32 = 64 blocks
  //   - M=256 with BLOCK_N=64 gives (M/128)*(N/64) = 2*64 = 128 blocks (better utilization)
  if (K == 7168) {
    if (M >= 512 && N >= 4096)
      dual_gemm_launch<7168, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, C);
    else
      dual_gemm_launch<7168, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
  }
  else if (K == 4096) {
    if (M >= 512 && N >= 4096)
      dual_gemm_launch<4096, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, C);
    else
      dual_gemm_launch<4096, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
  }
  else if (K == 2048) dual_gemm_launch<2048, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
  else if (K == 2304) dual_gemm_launch<2304, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
  else if (K == 1536) dual_gemm_launch<1536, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
  else if (K == 512)  dual_gemm_launch<512, 128, 64, 256, 2>(A, B1, B2, SFA, SFB1, SFB2, C);
  else if (K == 256)  dual_gemm_launch<256, 128, 64, 256, 1>(A, B1, B2, SFA, SFB1, SFB2, C);
  else                dual_gemm_launch<7168, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, C);
}

TORCH_LIBRARY(fused_dual_gemm, m) {
  m.def("dual_gemm(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> ()");
  m.impl("dual_gemm", &dual_gemm);
}
"""

load_inline(
    "fused_dual_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"],
)

dual_gemm_cuda = torch.ops.fused_dual_gemm.dual_gemm


def custom_kernel(data: input_t) -> output_t:
    a_ref, b1_ref, b2_ref, sfa_ref, sfb1_ref, sfb2_ref, sfa_perm, sfb1_perm, sfb2_perm, c_ref = data

    m, n, l = c_ref.shape

    for l_idx in range(l):
        # Get 2D slices for this batch
        a_2d = a_ref[:, :, l_idx]
        b1_2d = b1_ref[:, :, l_idx]
        b2_2d = b2_ref[:, :, l_idx]

        # Convert permuted scale factors to kernel-expected layout
        sfa_slice = sfa_perm[:, :, :, :, :, l_idx]
        sfb1_slice = sfb1_perm[:, :, :, :, :, l_idx]
        sfb2_slice = sfb2_perm[:, :, :, :, :, l_idx]

        # Permute to [rest_m, rest_k, 32, 4, 4]
        sfa_2d = sfa_slice.permute(2, 4, 0, 1, 3).contiguous()
        sfb1_2d = sfb1_slice.permute(2, 4, 0, 1, 3).contiguous()
        sfb2_2d = sfb2_slice.permute(2, 4, 0, 1, 3).contiguous()

        c_2d = c_ref[:, :, l_idx]

        # Run fused dual GEMM
        dual_gemm_cuda(a_2d, b1_2d, b2_2d, sfa_2d, sfb1_2d, sfb2_2d, c_2d)

    return c_ref