submission 212722

- import torch
- from torch.utils.cpp_extension import load_inline
- from task import input_t, output_t
- 
- 
- cpp_src = """
- #include <torch/extension.h>
- 
- torch::Tensor cuda_nvfp4_dual_gemm_cublaslt(
-     torch::Tensor A,
-     torch::Tensor B1,
-     torch::Tensor B2,
-     torch::Tensor SFA,
-     torch::Tensor SFB1,
-     torch::Tensor SFB2,
-     torch::Tensor SFA_perm,
-     torch::Tensor SFB1_perm,
-     torch::Tensor SFB2_perm,
-     torch::Tensor C);
- """
- 
- 
- cuda_src = """
- #include <torch/extension.h>
- #include <cublasLt.h>
- #include <cuda_runtime.h>
- #include <cuda_fp16.h>
- #include <cuda_fp8.h>
- #include <cuda_fp4.h>
- #include <stdexcept>
- 
- namespace {
- 
- inline cublasLtHandle_t get_handle() {
-     static cublasLtHandle_t h = [] {
-         cublasLtHandle_t t;
-         cublasLtCreate(&t);
-         return t;
-     }();
-     return h;
- }
- 
- inline void check(cublasStatus_t s, const char* m) {
-     if (s != CUBLAS_STATUS_SUCCESS) throw std::runtime_error(m);
- }
- 
- // SiLU activation kernel: out = silu(acc1) * acc2
- // where silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
- __global__ void silu_mul_kernel(
-     const half* __restrict__ acc1,
-     const half* __restrict__ acc2,
-     half* __restrict__ out,
-     int64_t size)
- {
-     int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-     if (idx < size) {
-         float x = __half2float(acc1[idx]);
-         float y = __half2float(acc2[idx]);
-         // silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
-         float silu_x = x / (1.0f + expf(-x));
-         out[idx] = __float2half(silu_x * y);
-     }
- }
- 
- void run_gemm(
-     cublasLtHandle_t handle,
-     torch::Tensor A,
-     torch::Tensor B,
-     torch::Tensor SFA_perm,
-     torch::Tensor SFB_perm,
-     torch::Tensor C,
-     int64_t M, int64_t K, int64_t N)
- {
-     const int64_t lda = K;
-     const int64_t ldb = K;
-     const int64_t ldc = N;
- 
-     cublasLtMatmulDesc_t opDesc;
-     check(cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F), "desc");
- 
-     cublasOperation_t transA = CUBLAS_OP_N;
-     cublasOperation_t transB = CUBLAS_OP_T;
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA,
-                                          &transA, sizeof(transA)), "ta");
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB,
-                                          &transB, sizeof(transB)), "tb");
- 
-     cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE,
-                                          &scale_mode, sizeof(scale_mode)), "asm");
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE,
-                                          &scale_mode, sizeof(scale_mode)), "bsm");
- 
-     const void* a_scale_ptr = SFA_perm.data_ptr();
-     const void* b_scale_ptr = SFB_perm.data_ptr();
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
-                                          &a_scale_ptr, sizeof(a_scale_ptr)), "asp");
-     check(cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
-                                          &b_scale_ptr, sizeof(b_scale_ptr)), "bsp");
- 
-     cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc, Ddesc;
-     check(cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_4F_E2M1, M, K, lda), "Adesc");
-     check(cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_4F_E2M1, N, K, ldb), "Bdesc");
-     check(cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_16F, M, N, ldc), "Cdesc");
-     check(cublasLtMatrixLayoutCreate(&Ddesc, CUDA_R_16F, M, N, ldc), "Ddesc");
- 
-     cublasLtOrder_t order = CUBLASLT_ORDER_ROW;
-     check(cublasLtMatrixLayoutSetAttribute(Adesc, CUBLASLT_MATRIX_LAYOUT_ORDER,
-                                            &order, sizeof(order)), "orderA");
-     check(cublasLtMatrixLayoutSetAttribute(Bdesc, CUBLASLT_MATRIX_LAYOUT_ORDER,
-                                            &order, sizeof(order)), "orderB");
-     check(cublasLtMatrixLayoutSetAttribute(Cdesc, CUBLASLT_MATRIX_LAYOUT_ORDER,
-                                            &order, sizeof(order)), "orderC");
-     check(cublasLtMatrixLayoutSetAttribute(Ddesc, CUBLASLT_MATRIX_LAYOUT_ORDER,
-                                            &order, sizeof(order)), "orderD");
- 
-     cublasLtMatmulPreference_t pref;
-     check(cublasLtMatmulPreferenceCreate(&pref), "pref");
-     size_t ws = 0;
-     check(cublasLtMatmulPreferenceSetAttribute(
-               pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws, sizeof(ws)),
-           "pref_ws");
- 
-     cublasLtMatmulHeuristicResult_t hres{};
-     int ret = 0;
-     check(cublasLtMatmulAlgoGetHeuristic(
-               handle, opDesc, Adesc, Bdesc, Cdesc, Ddesc, pref, 1, &hres, &ret),
-           "heuristic");
-     TORCH_CHECK(ret > 0, "no algo");
- 
-     const float alpha = 1.0f;
-     const float beta = 0.0f;
- 
-     check(cublasLtMatmul(
-               handle, opDesc, &alpha,
-               A.data_ptr(), Adesc,
-               B.data_ptr(), Bdesc,
-               &beta,
-               C.data_ptr(), Cdesc,
-               C.data_ptr(), Ddesc,
-               &hres.algo,
-               nullptr, 0,
-               nullptr),
-           "matmul");
- 
-     cublasLtMatmulPreferenceDestroy(pref);
-     cublasLtMatrixLayoutDestroy(Adesc);
-     cublasLtMatrixLayoutDestroy(Bdesc);
-     cublasLtMatrixLayoutDestroy(Cdesc);
-     cublasLtMatrixLayoutDestroy(Ddesc);
-     cublasLtMatmulDescDestroy(opDesc);
- }
- 
- } // namespace
- 
- torch::Tensor cuda_nvfp4_dual_gemm_cublaslt(
-     torch::Tensor A,
-     torch::Tensor B1,
-     torch::Tensor B2,
-     torch::Tensor SFA,
-     torch::Tensor SFB1,
-     torch::Tensor SFB2,
-     torch::Tensor SFA_perm,
-     torch::Tensor SFB1_perm,
-     torch::Tensor SFB2_perm,
-     torch::Tensor C)
- {
-     const int64_t M = A.size(0);
-     const int64_t K = A.size(1) * 2;  // FP4 is packed 2 per byte
-     const int64_t N = B1.size(0);
- 
-     cublasLtHandle_t handle = get_handle();
- 
-     // Allocate temporary buffer for second GEMM result
-     auto acc2 = torch::empty_like(C);
- 
-     // Perform first GEMM: acc1 = A @ B1^T (stored in C temporarily)
-     run_gemm(handle, A, B1, SFA_perm, SFB1_perm, C, M, K, N);
- 
-     // Perform second GEMM: acc2 = A @ B2^T
-     run_gemm(handle, A, B2, SFA_perm, SFB2_perm, acc2, M, K, N);
- 
-     // Apply SiLU and multiply: C = silu(acc1) * acc2
-     int64_t size = M * N;
-     int threads = 256;
-     int blocks = (size + threads - 1) / threads;
- 
-     silu_mul_kernel<<<blocks, threads>>>(
-         reinterpret_cast<const half*>(C.data_ptr()),
-         reinterpret_cast<const half*>(acc2.data_ptr()),
-         reinterpret_cast<half*>(C.data_ptr()),
-         size
-     );
- 
-     return C;
- }
- """
- 
- 
- nvfp4_dual_gemm_module = load_inline(
-     name="nvfp4_dual_gemm_cublaslt",
-     cpp_sources=[cpp_src],
-     cuda_sources=[cuda_src],
-     functions=["cuda_nvfp4_dual_gemm_cublaslt"],
-     extra_cuda_cflags=[
-         "-std=c++17",
-         "-gencode=arch=compute_100a,code=sm_100a",
-         "--ptxas-options=--gpu-name=sm_100a",
-         "-O3",
-         "-w",
-         "-allow-unsupported-compiler",
-     ],
-     extra_ldflags=["-lcuda"],
-     verbose=False,
- )
- 
- 
- def custom_kernel(data: input_t) -> output_t:
-     a, b1, b2, sfa, sfb1, sfb2, sfa_perm, sfb1_perm, sfb2_perm, c = data
-     return nvfp4_dual_gemm_module.cuda_nvfp4_dual_gemm_cublaslt(
-         a, b1, b2, sfa, sfb1, sfb2, sfa_perm, sfb1_perm, sfb2_perm, c
-     )
+ #!POPCORN leaderboard nvfp4_dual_gemm
+ #!POPCORN gpu NVIDIA
+  
+ import torch
+ from task import input_t, output_t
+ from torch.utils.cpp_extension import load_inline
+  
+ # Common CUDA source shared by both kernels
+ CUDA_SRC_COMMON = """
+ #include <cudaTypedefs.h>
+ #include <cuda_fp16.h>
+ #include <cuda_fp8.h>
+  
+ #include <torch/library.h>
+ #include <ATen/core/Tensor.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 mbarrier_arrive(int mbar_addr) {
+     asm volatile("mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" :: "r"(mbar_addr) : "memory");
+ }
+  
+ __device__ inline
+ void mbarrier_expect_tx(int mbar_addr, int size) {
+     asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
+                         :: "r"(mbar_addr), "r"(size) : "memory");
+ }
+  
+ __device__ inline
+ void tma_copy(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_load_3d(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");
+ }
+  
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
+ __device__ inline
+ void issue_tma_interleaved(
+     int smem, int stage_id, int iter_k,
+     const CUtensorMap *A_tmap, const CUtensorMap *B1_tmap, const CUtensorMap *B2_tmap,
+     const char *SFA_ptr, const char *SFB1_ptr, const char *SFB2_ptr,
+     int off_m, int off_n, int K,
+     int mbar_addr, uint64_t cache_A, uint64_t cache_B
+ ) {
+     constexpr int A_size = BLOCK_M * BLOCK_K / 2;
+     constexpr int B_size = BLOCK_N * BLOCK_K / 2;
+     constexpr int SF_size = 128 * BLOCK_K / 16;
+     constexpr int STAGE_SIZE = A_size + B_size * 2 + SF_size * 3;
+  
+     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 + SF_size;
+     const int SFB2_smem = SFB1_smem + SF_size;
+  
+     const int off_k = iter_k * BLOCK_K;
+  
+     tma_load_3d(A_smem, A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
+     tma_load_3d(B1_smem, B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
+     tma_load_3d(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_copy(SFA_smem, SFA_src, SF_size, mbar_addr, cache_A);
+     tma_copy(SFB1_smem, SFB1_src, SF_size, mbar_addr, cache_B);
+     tma_copy(SFB2_smem, SFB2_src, SF_size, mbar_addr, cache_B);
+  
+     mbarrier_expect_tx(mbar_addr, STAGE_SIZE);
+ }
+  
+ __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(uint64_t a_desc, uint64_t b_desc, uint32_t i_desc,
+     int scale_A_tmem, int scale_B_tmem, int enable_input_d, int d_tmem) {
+     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)
+     );
+ }
+  
+ __device__ inline
+ void tcgen05_commit(int mbar_addr) {
+     asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
+                         :: "r"(mbar_addr) : "memory");
+ }
+  
+ 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";
+     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));
+ }
+  
+ __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_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); }
+  
+ __device__ inline
+ float silu(float x) {
+     return x / (1.0f + expf(-x));
+ }
+  
+ 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);
+ }
+  
+ // Interleaved dual GEMM kernel
+ template <int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
+ __global__
+ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
+ void dual_gemm_silu_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, int K
+ ) {
+     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;
+     const int num_iters = K / BLOCK_K;
+  
+     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 SF_size = 128 * BLOCK_K / 16;
+     constexpr int STAGE_SIZE = A_size + B_size * 2 + SF_size * 3;
+  
+     constexpr int SILU_OFFSET = NUM_STAGES * STAGE_SIZE;
+     half *silu_smem = reinterpret_cast<half *>(smem_ptr + SILU_OFFSET);
+  
+     #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 done_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
+  
+     constexpr int GEMM1_D_TMEM = 0;
+     constexpr int GEMM2_D_TMEM = BLOCK_N;
+     constexpr int SFA_tmem = BLOCK_N * 2;
+     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 * 4));
+     }
+     __syncthreads();
+  
+     uint64_t cache_A = (M > N) ? EVICT_FIRST : EVICT_LAST;
+     uint64_t cache_B = (M > N) ? EVICT_LAST : EVICT_FIRST;
+  
+     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 uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)BLOCK_N >> 3U << 17U) | ((uint32_t)128 >> 7U << 27U);
+  
+     if (warp_id == NUM_WARPS - 2 && elect_sync()) {
+         // TMA warp
+         for (int iter_k = 0; iter_k < NUM_STAGES && iter_k < num_iters; iter_k++) {
+             issue_tma_interleaved<BLOCK_M, BLOCK_N, BLOCK_K>(
+                 smem, iter_k, iter_k,
+                 &A_tmap, &B1_tmap, &B2_tmap,
+                 SFA_ptr, SFB1_ptr, SFB2_ptr,
+                 off_m, off_n, K, tma_mbar_addr + iter_k * 8, cache_A, cache_B);
+         }
+  
+         for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
+             const int stage_id = iter_k % NUM_STAGES;
+             mbarrier_wait(mma_mbar_addr + stage_id * 8, (iter_k / NUM_STAGES - 1) % 2);
+             issue_tma_interleaved<BLOCK_M, BLOCK_N, BLOCK_K>(
+                 smem, stage_id, iter_k,
+                 &A_tmap, &B1_tmap, &B2_tmap,
+                 SFA_ptr, SFB1_ptr, SFB2_ptr,
+                 off_m, off_n, K, tma_mbar_addr + stage_id * 8, cache_A, cache_B);
+         }
+     }
+     else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
+         // MMA warp
+         for (int iter_k = 0; iter_k < num_iters; iter_k++) {
+             const int stage_id = iter_k % NUM_STAGES;
+             mbarrier_wait(tma_mbar_addr + stage_id * 8, (iter_k / NUM_STAGES) % 2);
+  
+             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 + SF_size;
+             const int SFB2_smem = SFB1_smem + SF_size;
+  
+             constexpr uint64_t SF_desc_base = make_desc_SF(0);
+             const uint64_t SFA_desc = SF_desc_base + ((uint64_t)SFA_smem >> 4ULL);
+             const uint64_t SFB1_desc = SF_desc_base + ((uint64_t)SFB1_smem >> 4ULL);
+             const uint64_t SFB2_desc = SF_desc_base + ((uint64_t)SFB2_smem >> 4ULL);
+  
+             for (int k = 0; k < BLOCK_K / MMA_K; k++) {
+                 tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc + (uint64_t)k * (512ULL >> 4ULL));
+                 tcgen05_cp_nvfp4(SFB1_tmem + k * 4, SFB1_desc + (uint64_t)k * (512ULL >> 4ULL));
+                 tcgen05_cp_nvfp4(SFB2_tmem + k * 4, SFB2_desc + (uint64_t)k * (512ULL >> 4ULL));
+             }
+  
+             // GEMM1 MMA
+             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);
+                     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 enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
+                     tcgen05_mma_nvfp4(a_desc, b1_desc, i_desc, scale_A_tmem, scale_B1_tmem, enable_input_d, GEMM1_D_TMEM);
+                 }
+  
+             // GEMM2 MMA
+             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 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_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;
+                     tcgen05_mma_nvfp4(a_desc, b2_desc, i_desc, scale_A_tmem, scale_B2_tmem, enable_input_d, GEMM2_D_TMEM);
+                 }
+  
+             tcgen05_commit(mma_mbar_addr + stage_id * 8);
+         }
+         tcgen05_commit(done_mbar_addr);
+     }
+     else if (tid < BLOCK_M) {
+         // Epilogue warps
+         constexpr int WIDTH = 64;
+  
+         mbarrier_wait(done_mbar_addr, 0);
+         asm volatile("tcgen05.fence::after_thread_sync;");
+  
+         // Read GEMM1 result and compute SiLU
+         for (int n = 0; n < BLOCK_N / WIDTH; n++) {
+             float tmp[WIDTH];
+             tcgen05_ld_32x32bx64(tmp, warp_id * 32, GEMM1_D_TMEM + n * WIDTH);
+             asm volatile("tcgen05.wait::ld.sync.aligned;");
+  
+             for (int i = 0; i < WIDTH; i++) {
+                 silu_smem[tid * BLOCK_N + n * WIDTH + i] = __float2half(silu(tmp[i]));
+             }
+         }
+  
+         asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
+  
+         // Read GEMM2 result and compute final output with coalesced stores
+         for (int m = 0; m < 32 / 16; m++) {
+             float tmp[BLOCK_N / 2];
+             if constexpr (BLOCK_N == 128)
+                 tcgen05_ld_16x256bx16(tmp, warp_id * 32 + m * 16, GEMM2_D_TMEM);
+             else
+                 tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, GEMM2_D_TMEM);
+             asm volatile("tcgen05.wait::ld.sync.aligned;");
+  
+             for (int i = 0; i < BLOCK_N / 8; i++) {
+                 const int row0 = warp_id * 32 + m * 16 + lane_id / 4;
+                 const int row1 = row0 + 8;
+                 const int col0 = i * 8 + (lane_id % 4) * 2;
+                 const int col1 = col0 + 1;
+  
+                 const float s00 = __half2float(silu_smem[row0 * BLOCK_N + col0]);
+                 const float s01 = __half2float(silu_smem[row0 * BLOCK_N + col1]);
+                 const float s10 = __half2float(silu_smem[row1 * BLOCK_N + col0]);
+                 const float s11 = __half2float(silu_smem[row1 * BLOCK_N + col1]);
+  
+                 const float v00 = tmp[i * 4 + 0] * s00;
+                 const float v01 = tmp[i * 4 + 1] * s01;
+                 const float v10 = tmp[i * 4 + 2] * s10;
+                 const float v11 = tmp[i * 4 + 3] * s11;
+  
+                 const int out_row0 = off_m + row0;
+                 const int out_row1 = off_m + row1;
+                 const int out_col = off_n + col0;
+  
+                 reinterpret_cast<half2 *>(C_ptr + out_row0 * N + out_col)[0] = __float22half2_rn({v00, v01});
+                 reinterpret_cast<half2 *>(C_ptr + out_row1 * N + out_col)[0] = __float22half2_rn({v10, v11});
+             }
+         }
+  
+         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));
+     }
+ }
+ """
+  
+ # BLOCK_N=64 kernel for m=256 cases (5 stages)
+ CUDA_SRC_N64 = """
+ at::Tensor dual_gemm_silu_n64(
+     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);
+     const int K = A.size(1) * 2;
+  
+     constexpr int BLOCK_M = 128;
+     constexpr int BLOCK_N = 64;
+     constexpr int BLOCK_K = 256;
+     constexpr int NUM_STAGES = 5;  // 5 stages × 38KB = 190KB + 16KB silu = 206KB
+  
+     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_sz = BLOCK_M * BLOCK_K / 2;
+     int B_sz = BLOCK_N * BLOCK_K / 2;
+     int SF_sz = 128 * BLOCK_K / 16;
+     int stage_size = A_sz + B_sz * 2 + SF_sz * 3;
+     int pipeline_size = stage_size * NUM_STAGES;
+     int silu_size = BLOCK_M * BLOCK_N * sizeof(half);
+     int smem_size = pipeline_size + silu_size;
+  
+     auto kernel = dual_gemm_silu_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+     if (smem_size > 48'000)
+         cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  
+     kernel<<<grid, tb_size, smem_size>>>(
+         A_tmap, B1_tmap, B2_tmap,
+         SFA_ptr, SFB1_ptr, SFB2_ptr,
+         C_ptr, M, N, K
+     );
+  
+     return C;
+ }
+  
+ TORCH_LIBRARY(dual_gemm_n64, m) {
+     m.def("dual_gemm_silu(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
+     m.impl("dual_gemm_silu", &dual_gemm_silu_n64);
+ }
+ """
+  
+ # BLOCK_N=128 kernel for m=512 cases (3 stages)
+ CUDA_SRC_N128 = """
+ at::Tensor dual_gemm_silu_n128(
+     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);
+     const int K = A.size(1) * 2;
+  
+     constexpr int BLOCK_M = 128;
+     constexpr int BLOCK_N = 128;
+     constexpr int BLOCK_K = 256;
+     constexpr int NUM_STAGES = 3;  // 3 stages × 54KB = 162KB + 32KB silu = 194KB
+  
+     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_sz = BLOCK_M * BLOCK_K / 2;
+     int B_sz = BLOCK_N * BLOCK_K / 2;
+     int SF_sz = 128 * BLOCK_K / 16;
+     int stage_size = A_sz + B_sz * 2 + SF_sz * 3;
+     int pipeline_size = stage_size * NUM_STAGES;
+     int silu_size = BLOCK_M * BLOCK_N * sizeof(half);
+     int smem_size = pipeline_size + silu_size;
+  
+     auto kernel = dual_gemm_silu_kernel<BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+     if (smem_size > 48'000)
+         cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  
+     kernel<<<grid, tb_size, smem_size>>>(
+         A_tmap, B1_tmap, B2_tmap,
+         SFA_ptr, SFB1_ptr, SFB2_ptr,
+         C_ptr, M, N, K
+     );
+  
+     return C;
+ }
+  
+ TORCH_LIBRARY(dual_gemm_n128, m) {
+     m.def("dual_gemm_silu(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
+     m.impl("dual_gemm_silu", &dual_gemm_silu_n128);
+ }
+ """
+  
+ # Compile both kernels
+ for name, src in [("dual_gemm_n64", CUDA_SRC_N64), ("dual_gemm_n128", CUDA_SRC_N128)]:
+     load_inline(
+         name,
+         cpp_sources="",
+         cuda_sources=CUDA_SRC_COMMON + 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_silu_n64 = torch.ops.dual_gemm_n64.dual_gemm_silu
+ dual_gemm_silu_n128 = torch.ops.dual_gemm_n128.dual_gemm_silu
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     a, b1, b2 = data[0], data[1], data[2]
+     sfa_perm, sfb1_perm, sfb2_perm = data[6], data[7], data[8]
+     c = data[9]
+  
+     M = a.shape[0]
+     if M == 256:
+         # Use BLOCK_N=64 kernel (5 stages)
+         return dual_gemm_silu_n64(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)
+     else:
+         # Use BLOCK_N=128 kernel (3 stages) for m=512 cases
+         return dual_gemm_silu_n128(a, b1, b2, sfa_perm, sfb1_perm, sfb2_perm, c)