submission 340222

# Dual GEMM: C = silu(A @ B1) * (A @ B2)


import torch
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline




CUDA_SRC_COMMON = r"""
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <torch/library.h>
#include <ATen/core/Tensor.h>


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


// https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197
constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
constexpr uint64_t EVICT_LAST = 0x14F0000000000000;

// x & 0x3'FFFFULL: keep only the lower 18 bit, 0x3'FFFF is 256KB - 1
// limit the addressable range to 256KB (max shared memory size)
// >> 4ULL: right shift 4 bits, divide by 16, because the shared memory is 16-byte aligned
// the tcgen05_mma_nvfp4 will automatically do the * 16 to get the actual shared memory byte address.
// It designed as this to save bits (save 4 bits, also force alignment)
__device__ inline
constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };


// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cute/arch/cluster_sm90.hpp#L180
__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_init_pred(int mbar_addr, int count, int pred) {
 asm volatile(
   "{\n\t"
   ".reg .pred p;\n\t"
   "setp.ne.b32 p, %2, 0;\n\t"
   "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t"
   "}"
   :: "r"(mbar_addr), "r"(count), "r"(pred)
 );
}


// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cutlass/arch/barrier.h#L408
__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) {
  // This is for SFA, SFB copy.
  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) {
  // This is for transfer data for A, B, which is represent as 3d tensor.
  // r is for 32-bit register for int values; l is for 64 bit register for pointers and uint64_t
  // memory: tells compiler this instruction touches memory.
  // cp.async.bulk: TMA copy
  // tensor.3d: operating on 3d tensor
  // shared::cta: destination is CTA-scoped shared memory; can use shared::cluster for all CTAs in the same cluster, all CTA get the same data in one load.
  // global: source is global memory
  // mbarrier::complete_tx::byte: use mbarrier to track completion, track bytes trasferred
  // cta_group::1: cta group size is 1, only this CTA track completion
  // L2::cache_hint: apply L2 cache policy hint
  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");
}


// The tcgen05.cp instruction is issued by warp 5 (the MMA warp in this code), 
// but the actual TMEM storage is managed by the TMEM controller, which organizes 
// data according to the warp structure of warps 0-3. (my understanding)
__device__ inline
void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
 // .32x128b corresponds to (32, 16) 8-bit scale -> 1 MMA for nvfp4.
 // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-cp
 // .32x128b require .warpx4, the data will be multicasted into all 4 warps.
 // row 0 - 31 replicated to row 32 - 63 and row 64 - 95 and row 96 - 127,
 // because each warp's tensor core can only access it's portion of tmem during mma.
 asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}



// setp.ne.b32 p, %6, 0;
// p = (enable_input_d != 0), p = false is overwrite mode,
// p = true is accumulate mode, when k = 0 is overwrite, when k > 0 is accumulate.
// block_scale: block wise scaling, not per element scaling
// .block16: 1 scale per 16 elements
__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)
 );
}


// see https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html
struct SHAPE {
 static constexpr char _32x32b[] = ".32x32b";  // 32x1 tile for each warp
};


struct NUM {
 static constexpr char x32[] = ".x32";
 static constexpr char x64[] = ".x64";
 static constexpr char x128[] = ".x128";
};


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


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


__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); }
__device__ inline void tcgen05_ld_32x32bx128(float *tmp, int row, int col) { tcgen05_ld_128regs<SHAPE::_32x32b, NUM::x128>(tmp, row, col); }


// Fast silu using PTX SFU with FMA for better pipelining
// silu(x) = x * sigmoid(x) = x / (1 + exp(-x)) = x / (1 + 2^(-x*log2e))
__device__ inline
float silu(float x) {
 float result;
 asm volatile(
   "{\n\t"
   ".reg .f32 neg_scaled, exp_val, denom, rcp_val;\n\t"
   "mul.f32 neg_scaled, %1, 0FBFB8AA3B;\n\t"    // -x * log2(e) in one instruction
   "ex2.approx.ftz.f32 exp_val, neg_scaled;\n\t" // exp(-x)
   "add.f32 denom, exp_val, 0F3F800000;\n\t"     // 1 + exp(-x)
   "rcp.approx.ftz.f32 rcp_val, denom;\n\t"      // 1/(1+exp(-x))
   "mul.f32 %0, %1, rcp_val;\n\t"                // x * sigmoid(x)
   "}"
   : "=f"(result)
   : "f"(x)
 );
 return result;
}


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 check_cuda(cudaError_t err) {
 if (err == cudaSuccess) return;
 TORCH_CHECK(false, cudaGetErrorString(err));
}


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};  // in bytes
 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);
}
"""




CUDA_SRC_DUAL_GEMM = r"""
template <
 int K,
 int BLOCK_M,
 int BLOCK_N,
 int BLOCK_K,
 int NUM_STAGES
>
__global__
__launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
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_m = M / BLOCK_M;
 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));
 constexpr int A_size = BLOCK_M * BLOCK_K / 2;
 constexpr int B_size = BLOCK_N * BLOCK_K / 2;
 // always copy 128xBLOCK_K/16, because the tcgen05.cp that copy from smem to tmem
 // has work work with fixed (32, 16) 8 bit scale, which is 4 tmem columns, each column has 32 rows
 // and each column value has 4 8-bit scale factor as draw in https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x
 // The scale factor need to layout in that way, so our SFA and SFB both need to be 128 rows,
 // so each MMA work on (128, 4) scale factors which can be (32, 4, 4) which is (32, 16).
 // so no matter what the block_m and block_n are we need the 128 * (BLOCK_K / 16), the BLOCK_K / 16
 // includes the scale factors for all MMA in the block (e.g. BLOCK_K=256, MMA_K = 64, so it has shape
 // [128, 4 (4 MMA), 32, 4, 4])
 constexpr int SFA_size = 128 * BLOCK_K / 16;
 constexpr int SFB_size = 128 * BLOCK_K / 16;
 // For dual GEMM: A, B1, B2, SFA, SFB1, SFB2
 constexpr int STAGE_SIZE = A_size + B_size * 2 + SFA_size + SFB_size * 2;


 // set up mbarriers and tmem
 // we have NUM_STAGES mbars for TMA, NUM_STAGES mbars for MMA, 1 mbar for mainloop
 #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));
 // byte offset: each barrier takes int64_t which is 8 bytes, so *8 here.
 const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
 const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;


 // TMEM layout for dual GEMM:
 // - Acc1: columns 0 to BLOCK_N-1 (d_tmem = 0)
 // - SFA: columns BLOCK_N onwards
 // - SFB1, SFB2: after SFA
 // - Acc2: columns BLOCK_N*2 onwards (d_tmem = BLOCK_N*2)
 constexpr int SFA_tmem = BLOCK_N;
 // The columns of SFA_tmem is 4 * 4, the first 4 is 4 tmem columns
 // (128, 64 / 16) = (128, 4) per mma = (32, 4, 4) as in https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x
 // the second 4 mmas per block_k.
 constexpr int SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
 constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);


 mbarrier_init_pred(tma_mbar_addr + lane_id * 8, 1, (warp_id == 0) & (lane_id < NUM_STAGES * 2 + 1));
 // We actually not need BLOCK_N * 4 TMEM, but if we use BLOCK_N * 3 here it will error out,
 // so somehow TMEM require it to be power of 2?
 if (warp_id == 1) {
   // TMEM allocation MUST be power of 2! Tested: 384 fails, 448 fails, 512 works
   // For BLOCK_N=64: need 192 cols (Acc2 ends at 191), so 256 works
   // For BLOCK_N=128: need 384 cols (Acc2 ends at 383), so 512 needed
   asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 4));
 }
 __syncthreads();


 const int num_iters = K / BLOCK_K;


 if (warp_id == 4 && elect_sync()) {
   // TMA warp - cache eviction policy (optimized for M < N benchmark cases)
   // A tile reuse = N/BLOCK_N (high when N > M) → keep in cache longest
   // B tile reuse = M/BLOCK_M (low when M < N) → evict first
   constexpr uint64_t cache_A = EVICT_LAST;
   constexpr uint64_t cache_B = EVICT_FIRST;


   // K / 16 / 4: 16 elements share same scale factor, 4 fp8 packed to 1 value
   // it's the total number of (4 scale factors pack to one 32 bit) scale factors along K.
   constexpr int rest_k = K / 16 / 4;
   // layout of SFA IS [M/128, rest_k, 32, 4, 4]
   // so SFA_base:
   // Move ptr to the current block's position for the SFA, SFB - not including
   // the k iter offset yet, which will be added during the loop along k.
   const char *SFA_base = SFA_ptr + (off_m / 128) * rest_k * 512;
   // SFB is [N/128, rest_k, 32, 4, 4]
   // When block_n = 64, because we are always using 128 here, whcih is required by
   // tcgen05, then block 0 and block 1 both load the entire 128 n rows, and block 0
   // use the first half and block 1 use the second half.
   const char *SFB1_base = SFB1_ptr + (off_n / 128) * rest_k * 512;
   const char *SFB2_base = SFB2_ptr + (off_n / 128) * rest_k * 512;


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


     // issue TMA for A, B1, B2
     // iter_k * (BLOCK_K / 256), BLOCK_K = 256
     const int k_idx = iter_k;
     // the A_tmap has globalDim as (256, M, K/256), off_m is for row (M), k_idx is for col (K/256).
     // 0 means start from the first value in the block.
     tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, k_idx, mbar_addr, cache_A);
     tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, k_idx, mbar_addr, cache_B);
     tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, k_idx, mbar_addr, cache_B);


     // Scale factor TMA with pre-computed base
     const char *SFA_src = SFA_base + iter_k * (BLOCK_K / 64) * 512;
     const char *SFB1_src = SFB1_base + iter_k * (BLOCK_K / 64) * 512;
     const char *SFB2_src = SFB2_base + iter_k * (BLOCK_K / 64) * 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);


     // mbarrier: memory barrier
     // .arrive: signal arrival at the barrier
     // .expect_tx: expect a transaction - tells barrier how many bytes to wait for
     // .release: ensure all prior writes are visible before arriving
     // .cta: within the CTA barrier sync
     // .shared::cta: barrier is in CTA-scoped shared memory
     // .b64: barrier is 64 bit
     // _ discard the return value
     // This is a non blocking instruction that tells the barrier to expect bytes from
     // TMA transaction and send arrival signal in the following mbarrier_wait.
     asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                 :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
   };


   // prefetch: issue NUM_STAGES TMA call
   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;
     // e.g. when NUM_STAGES == 5, when iter_k = 5 - 9, wait mma_phase = 0,
     // when iter_k = 10 - 14, wait mma_phase = 1,
     // when iter_k = 15 - 19, wait mma_phase = 0, ...
     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 == 5 && elect_sync()) {
   // from https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instuction-desc-kind-mxf4-mxf4nvf4
   // see E2M1 = 1, so we do (1U << 7U) and (1U << 10U)
   // also according to the table, the block_n need >> 3U, and block_m need >>7U
   // which is assuming N must be 8X, and M must be 128X, so those 3 or 7 bits are
   // all zero, so we can save some bit for better compression.
   // MMA instruction descriptor.
   constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
                             | (1U << 10U)  // btype=E2M1
                             | ((uint32_t)BLOCK_N >> 3U << 17U)  // MMA_N
                             | ((uint32_t)128 >> 7U << 27U);     // MMA_M


   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;

     // This matches the swizzle as specified in init_AB_tmap CU_TENSOR_MAP_SWIZZLE_128B
     // which is how TMA saved data in SMEM.
     // The bit shifting follows https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-shared-memory-descriptor.
     // Need to set bits that are not 0.
     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);
     };
     // no swizzling for scale factors
     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);
     // address need to make 4 bit right shift, same as in desc_encode.
     // SFA_desc: describe the SFA in shared memory - it's address, swizzle mode, etc
     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);


     constexpr int d_tmem_acc2 = BLOCK_N * 2;
     // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x
     // 32 * 4 * 4 = 512
     constexpr uint64_t SF_stride = 512ULL >> 4ULL;


     // TMEM column addresses for scale factors.
     // Scale factors are always loaded in 128-row chunks (tcgen05.cp requirement).
     // Each (32, 4, 4) tile contains scale factors for 128 M/N rows.
     // When BLOCK_M/N < 128, one tile covers multiple MMA tiles, so we offset to
     // select the portion corresponding to this CTA's M/N range:
     //   - (bid % (128/BLOCK)): which sub-portion of the 128 rows this MMA tile uses
     //   - (BLOCK/32): TMEM columns per sub-portion (32 rows per TMEM column)
     // E.g., BLOCK_N=64: bid_n=0 uses cols 0-1 (N=0-63), bid_n=1 uses cols 2-3 (N=64-127)
     // When BLOCK_M/N=128: offset is always 0 (MMA tile uses entire scale factor tile)
     // This is the offset calculating for the row (N) direction - when BLOCK_N = 64, for bid_n = 1, 3, 5, ...
     // they will have 2 TMEM column offset in the TMEM, and the first 2 TMEM rows are wasted.
     // Loaded: Full 128 N rows of scale factors (tcgen05.cp requirement)
     // Used: Only 64 N rows (half), corresponding to this CTA's N range
     // What's loaded into TMEM (per K iteration):┌─────────────┬─────────────┐│  Cols 0-1   │  Cols 2-3   ││  N = 0-63   │  N = 64-127 ││  (64 rows)  │  (64 rows)  │└─────────────┴─────────────┘      ↑              ↑   bid_n=0        bid_n=1   uses this      uses this   (half)         (half)
     // So there's some inefficiency when BLOCK_N < 128:
     // You load 128 rows but only use BLOCK_N rows
     // The unused portion wastes TMEM space and memory bandwidth
     const int scale_A_base = SFA_tmem + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
     const int scale_B1_base = SFB1_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
     const int scale_B2_base = SFB2_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);


     const uint64_t a_desc_base = make_desc_AB(A_smem);
     const uint64_t b1_desc_base = make_desc_AB(B1_smem);
     const uint64_t b2_desc_base = make_desc_AB(B2_smem);


     // Prefetch first scale factor
     tcgen05_cp_nvfp4(SFA_tmem, SFA_desc);
     tcgen05_cp_nvfp4(SFB1_tmem, SFB1_desc);
     tcgen05_cp_nvfp4(SFB2_tmem, SFB2_desc);

     constexpr int NUM_K_ITERS = BLOCK_K / MMA_K;


     #pragma unroll
     for (int k = 0; k < NUM_K_ITERS; k++) {
       // Prefetch next scale (only if within bounds)
       if (k + 1 < NUM_K_ITERS) {
         const int next_k = k + 1;
         // Always move 4 TMEM cols to next MMA no matter what the BLOCK_N is.
         tcgen05_cp_nvfp4(SFA_tmem + (next_k << 2), SFA_desc + next_k * SF_stride);
         tcgen05_cp_nvfp4(SFB1_tmem + (next_k << 2), SFB1_desc + next_k * SF_stride);
         tcgen05_cp_nvfp4(SFB2_tmem + (next_k << 2), SFB2_desc + next_k * SF_stride);
       }


       // Offset within shared memory for this MMA_K=64 chunk
       // Each MMA_K=64 nvfp4 = 32 bytes, so offset = k * 32 = k << 5
       // we are assuming use BLOCK_K = 256 here. Otherwisae we need to do:
       // ab_off = k1 * A_k256_stride + (k2 << 5), bb_off = k1 * B_k256_stride + (k2 << 5)
       const int ab_off = k << 5;
       const int bb_off = k << 5;
       const uint64_t a_desc = a_desc_base + (ab_off >> 4);
       const uint64_t b1_desc = b1_desc_base + (bb_off >> 4);
       const uint64_t b2_desc = b2_desc_base + (bb_off >> 4);

       // jump 4 TMEM columns per MMA iteration.
       const int scale_A = scale_A_base + (k << 2);
       const int scale_B1 = scale_B1_base + (k << 2);
       const int scale_B2 = scale_B2_base + (k << 2);


       const int enable_d = (k == 0) ? iter_k : 1;
       // we don't need explicit wait on TMEM here, because the TC tracks dependencies on TMEM and stalls if data not ready.
       tcgen05_mma_nvfp4(0, a_desc, b1_desc, i_desc, scale_A, scale_B1, enable_d);
       tcgen05_mma_nvfp4(d_tmem_acc2, a_desc, b2_desc, i_desc, scale_A, scale_B2, enable_d);
     }

     // signal the mma barrier of arrival, consumer is TMA warp
     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 warps - read both accumulators and apply silu(Acc1) * Acc2
   mbarrier_wait(mainloop_mbar_addr, 0);
   // Fence before reading TMEM
   asm volatile("tcgen05.fence::after_thread_sync;");

   // The WIDTH constant controls how many TMEM columns are read per tcgen05.ld instruction, and it's capped at 64 due to register pressure constraints.
   constexpr int WIDTH = 64;


   for (int n = 0; n < BLOCK_N / WIDTH; n++) {
     float tmp1[WIDTH];
     float tmp2[WIDTH];


     // Issue BOTH TMEM reads first to overlap latency
     if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp1, warp_id * 32, n * WIDTH);
     if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);
     //synchronization barrier that stalls the thread until all previously issued tcgen05.ld instructions have finished loading data into registers.
     asm volatile("tcgen05.wait::ld.sync.aligned;");


     const int col = off_m + tid;
     const int base_row = off_n + n * WIDTH;


     #pragma unroll
     for (int i = 0; i < WIDTH; i += 2) {
       float r0 = silu(tmp1[i]) * tmp2[i];
       float r1 = silu(tmp1[i+1]) * tmp2[i+1];
       C_ptr[(base_row + i) * M + col] = __float2half(r0);
       C_ptr[(base_row + i + 1) * M + col] = __float2half(r1);
     }
   }


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


template <
 int K,
 int BLOCK_M,
 int BLOCK_N,
 int BLOCK_K,
 int NUM_STAGES
>
at::Tensor cutlass_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
) {
 static_assert(BLOCK_K % 256 == 0);


 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 AB_size = (BLOCK_M + BLOCK_N * 2) * (BLOCK_K / 2);
 int SFAB_size = 128 * (BLOCK_K / 16) * 3;
 int smem_size = (AB_size + SFAB_size) * NUM_STAGES;


 auto this_kernel = dual_gemm_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
 if (smem_size > 48'000)
   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);


 return C.view({N, M, 1}).transpose(0, 1);
}


at::Tensor 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);


#define LAUNCH_EXACT(K_, M_, N_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \
 else if (K == K_ && M == M_ && N == N_) C = cutlass_dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);


#define LAUNCH(K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \
 else if (K == K_) C = cutlass_dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);


 if (false) {}
 LAUNCH_EXACT(7168, 256, 4096, 128, 64, 256, 5)
 LAUNCH_EXACT(7168, 512, 4096, 128, 128, 256, 4)
 LAUNCH_EXACT(4096, 256, 3072, 128, 64, 256, 5)
 LAUNCH_EXACT(7168, 512, 3072, 128, 128, 256, 4)
 LAUNCH(16384, 128, 64, 256, 5)
 LAUNCH( 7168, 128, 64, 256, 5)
 LAUNCH( 4096, 128, 64, 256, 5)
 LAUNCH( 3072, 128, 64, 256, 5)
 LAUNCH( 2048, 128, 64, 256, 5)
 LAUNCH( 2304, 128, 64, 256, 5)
 LAUNCH( 1536, 128, 64, 256, 5)
 LAUNCH( 1024, 128, 64, 256, 5)
 LAUNCH(  512, 128, 64, 256, 5)
 LAUNCH(  256, 128, 64, 256, 5)


#undef LAUNCH_EXACT
#undef LAUNCH


 return C;
}


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




load_inline(
   "dual_gemm_cuda_best_v1",
   cpp_sources="",
   cuda_sources=CUDA_SRC_COMMON + CUDA_SRC_DUAL_GEMM,
   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",
       "-maxrregcount=192",
   ],
   extra_ldflags=["-lcuda"],
)


dual_gemm_cuda = torch.ops.my_module_dual_gemm.dual_gemm




def custom_kernel(data: input_t) -> output_t:
   # data contains: A, B1, B2, _, _, _, SFA, SFB1, SFB2, C
   # Dual GEMM computes: C = silu(A @ B1) * (A @ B2)
   return dual_gemm_cuda(data[0], data[1], data[2], data[6], data[7], data[8], data[9])