submission 190523

+ #!POPCORN leaderboard nvfp4_dual_gemm
+ #!POPCORN gpu NVIDIA
+  
  import torch
  from torch.utils.cpp_extension import load_inline
   
   
- cpp_src = r"""
- #include <torch/extension.h>
+ # 编译内核并缓存
+ _CUDA_SRC_COMMON = r"""
+ #include <cuda.h>
+ #include <cudaTypedefs.h>
+ #include <cuda_fp16.h>
+ #include <cuda_runtime.h>
+ #include <math.h>
+  
+ #include <torch/library.h>
+ #include <ATen/core/Tensor.h>
  #include <ATen/ATen.h>
- #include <ATen/ops/_scaled_mm.h>
   
- void silu_mul_cuda(const at::Half* g1, const at::Half* g2, at::Half* out, int64_t count, int64_t out_stride);
+ constexpr int WARP_SIZE = 32;
+ constexpr int MMA_K = 64;
   
- torch::Tensor to_blocked(torch::Tensor input) {
-     auto rows = input.size(0);
-     auto cols = input.size(1);
-     auto n_row_blocks = (rows + 127) / 128;
-     auto n_col_blocks = (cols + 3) / 4;
-     auto padded = input.contiguous();
-     auto blocks = padded.view({n_row_blocks, 128, n_col_blocks, 4}).permute({0, 2, 1, 3});
-     auto rearranged = blocks.reshape({-1, 4, 32, 4}).transpose(1, 2).reshape({-1, 32, 16});
-     return rearranged.reshape({-1}).contiguous();
+ constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
+ 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;
  }
   
- torch::Tensor fused_dual_gemm(torch::Tensor a, torch::Tensor b1, torch::Tensor b2, torch::Tensor sfa, torch::Tensor sfb1, torch::Tensor sfb2, torch::Tensor c) {
-     auto out = c.contiguous();
-     auto l = out.size(2);
-     auto out_stride = out.size(2);
-     for (int64_t l_idx = 0; l_idx < l; ++l_idx) {
-         auto a_l = a.select(2, l_idx);
-         auto b1_l = b1.select(2, l_idx);
-         auto b2_l = b2.select(2, l_idx);
-         auto scale_a = to_blocked(sfa.select(2, l_idx));
-         auto scale_b1 = to_blocked(sfb1.select(2, l_idx));
-         auto scale_b2 = to_blocked(sfb2.select(2, l_idx));
-         auto g1 = at::_scaled_mm(
-             a_l,
-             b1_l.transpose(0, 1),
-             scale_a,
-             scale_b1,
-             c10::optional<at::Tensor>(),
-             c10::optional<at::Tensor>(),
-             c10::optional<at::ScalarType>(at::kHalf),
-             false
-         );
-         auto g2 = at::_scaled_mm(
-             a_l,
-             b2_l.transpose(0, 1),
-             scale_a,
-             scale_b2,
-             c10::optional<at::Tensor>(),
-             c10::optional<at::Tensor>(),
-             c10::optional<at::ScalarType>(at::kHalf),
-             false
-         );
-         auto g1_c = g1.contiguous();
-         auto g2_c = g2.contiguous();
-         int64_t count = g1_c.numel();
-         auto out_ptr = reinterpret_cast<at::Half*>(out.data_ptr()) + l_idx;
-         silu_mul_cuda(
-             reinterpret_cast<const at::Half*>(g1_c.data_ptr()),
-             reinterpret_cast<const at::Half*>(g2_c.data_ptr()),
-             out_ptr,
-             count,
-             out_stride
-         );
-     }
-     return out;
+ __device__ inline void mbarrier_init(int mbar_addr, int count) {
+   asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
  }
   
- PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fused_dual_gemm", &fused_dual_gemm, "nvfp4 dual gemm"); }
+ __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(
+   uint64_t a_desc,
+   uint64_t b_desc,
+   uint32_t i_desc,
+   int scale_A_tmem,
+   int scale_B_tmem,
+   int enable_input_d
+ ) {
+   const int d_tmem = 0;
+   asm volatile(
+     "{\n\t"
+     ".reg .pred p;\n\t"
+     "setp.ne.b32 p, %6, 0;\n\t"
+     "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
+     "}"
+     :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
+        "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
+   );
+ }
+  
+ struct SHAPE {
+   static constexpr char _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_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));
+ }
+  
+ 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); }
+  
+ __device__ inline void tcgen05_ld_16x128bx8(float *tmp, int row, int col) { tcgen05_ld_16regs<SHAPE::_16x128b, NUM::x8>(tmp, row, col); }
+ __device__ inline void tcgen05_ld_16x128bx16(float *tmp, int row, int col) { tcgen05_ld_32regs<SHAPE::_16x128b, NUM::x16>(tmp, row, col); }
+ __device__ inline void tcgen05_ld_16x128bx32(float *tmp, int row, int col) { tcgen05_ld_64regs<SHAPE::_16x128b, NUM::x32>(tmp, row, col); }
+  
+ __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 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};
+  
+   CUresult status = 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
+   );
+   (void)status;
+ }
  """
   
+ _CUDA_SRC_V4 = r"""
+ template <
+   int K,
+   int BLOCK_M,
+   int BLOCK_N,
+   int BLOCK_K,
+   int SPLIT_K,
+   bool C_N_MAJOR,
+   int NUM_STAGES
+ >
+ __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
+ void kernel(
+   const __grid_constant__ CUtensorMap A_tmap,
+   const __grid_constant__ CUtensorMap B_tmap,
+   const char *SFA_ptr,
+   const char *SFB_ptr,
+   half *C_ptr,
+   float *buf_ptr,
+   int M, int N
+ ) {
+   const int tid = threadIdx.x;
+   const int bid_k = blockIdx.x;
+   const int bid = blockIdx.y;
   
- cuda_src = r"""
- #include <ATen/ATen.h>
- #include <cuda.h>
- #include <cuda_fp16.h>
- #include <cuda_runtime.h>
- #include <math.h>
- #include <stdint.h>
+   const int lane_id = tid % WARP_SIZE;
+   const int warp_id = tid / WARP_SIZE;
   
- __global__ void silu_mul_kernel(const half* g1, const half* g2, half* out, int64_t count, int64_t out_stride) {
-     int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
-     if (idx >= count) {
-         return;
+   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;
+   constexpr int SFA_size = 128 * BLOCK_K / 16;
+   constexpr int SFB_size = 128 * BLOCK_K / 16;
+   constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
+  
+   #pragma nv_diag_suppress static_var_with_dynamic_init
+   __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
+   const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
+   const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
+   const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
+  
+   constexpr int SFA_tmem = BLOCK_N;
+   constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
+  
+   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));
+   }
+   __syncthreads();
+  
+   constexpr int num_iters = K / BLOCK_K / SPLIT_K;
+  
+   if (warp_id == NUM_WARPS - 2 && elect_sync()) {
+     uint64_t cache_A, cache_B;
+     if (M > N) {
+       cache_A = EVICT_FIRST;
+       cache_B = EVICT_LAST;
+     } else {
+       cache_A = EVICT_LAST;
+       cache_B = EVICT_FIRST;
      }
-     float x = __half2float(g1[idx]);
-     float y = __half2float(g2[idx]);
-     float silu = x / (1.0f + expf(-x));
-     out[idx * out_stride] = __float2half(silu * y);
+  
+     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 B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_size;
+  
+       const int off_k = SPLIT_K == 1 ? iter_k * BLOCK_K : (iter_k * SPLIT_K + bid_k) * BLOCK_K;
+       tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
+       tma_3d_gmem2smem(B_smem, &B_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 *SFB_src = SFB_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(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
+  
+       asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
+                   :: "r"(mbar_addr), "r"(STAGE_SIZE) : "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()) {
+     constexpr int MMA_N = BLOCK_N;
+     constexpr int MMA_M = 128;
+     constexpr uint32_t i_desc = (1U << 7U)
+                               | (1U << 10U)
+                               | ((uint32_t)MMA_N >> 3U << 17U)
+                               | ((uint32_t)MMA_M >> 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 B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_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 SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
+  
+       for (int k = 0; k < BLOCK_K / MMA_K; k++) {
+         uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
+         uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
+         tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
+         tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
+       }
+  
+       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 b_desc = make_desc_AB(B_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_B_tmem = SFB_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, b_desc, i_desc, scale_A_tmem, scale_B_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) {
+     mbarrier_wait(mainloop_mbar_addr, 0);
+     asm volatile("tcgen05.fence::after_thread_sync;");
+  
+     auto epilogue_M_major = [&]() {
+       constexpr int WIDTH = (BLOCK_N < 64 ? BLOCK_N : 64);
+  
+       for (int n = 0; n < BLOCK_N / WIDTH; n++) {
+         float tmp[WIDTH];
+         if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp, warp_id * 32, n * WIDTH);
+         if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp, warp_id * 32, n * WIDTH);
+         if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp, warp_id * 32, n * WIDTH);
+         asm volatile("tcgen05.wait::ld.sync.aligned;");
+  
+         for (int i = 0; i < WIDTH; i++) {
+           const int row = off_n + n * WIDTH + i;
+           const int col = off_m + tid;
+  
+           if constexpr (SPLIT_K == 1)
+             C_ptr[row * M + col] = __float2half(tmp[i]);
+           else
+             atomicAdd(buf_ptr + row * M + col, tmp[i]);
+         }
+       }
+     };
+     auto epilogue_N_major = [&]() {
+       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, 0);
+         if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);
+         if constexpr (BLOCK_N == 32) tcgen05_ld_16x256bx4(tmp, warp_id * 32 + m * 16, 0);
+         asm volatile("tcgen05.wait::ld.sync.aligned;");
+  
+         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;
+  
+           if constexpr (SPLIT_K == 1) {
+             reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
+             reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
+           } else {
+             atomicAdd(reinterpret_cast<float2 *>(buf_ptr + (row + 0) * N + col), float2({tmp[i * 4 + 0], tmp[i * 4 + 1]}));
+             atomicAdd(reinterpret_cast<float2 *>(buf_ptr + (row + 8) * N + col), float2({tmp[i * 4 + 2], tmp[i * 4 + 3]}));
+           }
+         }
+       }
+     };
+  
+     if constexpr (C_N_MAJOR)
+       epilogue_N_major();
+     else
+       epilogue_M_major();
+  
+     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));
+   }
  }
   
- void silu_mul_cuda(const at::Half* g1, const at::Half* g2, at::Half* out, int64_t count, int64_t out_stride) {
-     int threads = 256;
-     int blocks = static_cast<int>((count + threads - 1) / threads);
-     auto g1_ptr = reinterpret_cast<const half*>(g1);
-     auto g2_ptr = reinterpret_cast<const half*>(g2);
-     auto out_ptr = reinterpret_cast<half*>(out);
-     silu_mul_kernel<<<blocks, threads>>>(g1_ptr, g2_ptr, out_ptr, count, out_stride);
+ template <
+   int K,
+   int BLOCK_M,
+   int BLOCK_N,
+   int BLOCK_K,
+   int SPLIT_K,
+   bool SWAP_AB,
+   bool C_N_MAJOR,
+   int NUM_STAGES
+ >
+ at::Tensor gemm_launch(
+   const at::Tensor& A,
+   const at::Tensor& B,
+   const at::Tensor& SFA,
+   const at::Tensor& SFB,
+         at::Tensor& C,
+         at::Tensor& buf
+ ) {
+   static_assert(BLOCK_K % 256 == 0);
+  
+   const int M = A.size(0);
+   const int N = B.size(0);
+  
+   auto A_ptr   = reinterpret_cast<const char *>(A.data_ptr());
+   auto B_ptr   = reinterpret_cast<const char *>(B.data_ptr());
+   auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
+   auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
+   auto C_ptr   = reinterpret_cast<half *>(C.data_ptr());
+   auto buf_ptr = buf.data_ptr<float>();
+  
+   int new_M = M;
+   int new_N = N;
+   if constexpr (SWAP_AB) {
+     std::swap(A_ptr, B_ptr);
+     std::swap(SFA_ptr, SFB_ptr);
+     std::swap(new_M, new_N);
+   }
+  
+   CUtensorMap A_tmap, B_tmap;
+   init_AB_tmap(&A_tmap, A_ptr, new_M, K, BLOCK_M, BLOCK_K);
+   init_AB_tmap(&B_tmap, B_ptr, new_N, K, BLOCK_N, BLOCK_K);
+  
+   dim3 grid(SPLIT_K, (new_M / BLOCK_M) * (new_N / BLOCK_N));
+   int tb_size = BLOCK_M + 2 * WARP_SIZE;
+   int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+   int SFAB_size = 128 * (BLOCK_K / 16) * 2;
+   int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
+  
+   auto this_kernel = kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, C_N_MAJOR != SWAP_AB, NUM_STAGES>;
+   if (smem_size > 48000)
+     cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+   this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, buf_ptr, new_M, new_N);
+  
+   if constexpr (SPLIT_K == 1)
+     return C_N_MAJOR ? C : C.view({N, M, 1}).transpose(0, 1);
+   else
+     return C_N_MAJOR ? buf : buf.view({N, M, 1}).transpose(0, 1);
  }
+  
+ at::Tensor gemm(
+   const at::Tensor& A,
+   const at::Tensor& B,
+   const at::Tensor& SFA,
+   const at::Tensor& SFB,
+         at::Tensor& C,
+         at::Tensor& buf
+ ) {
+   const int K = A.size(1) * 2;
+  
+ #define LAUNCH(K_, BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, SWAP_AB, C_N_MAJOR, NUM_STAGES) \
+   else if (K == K_) C = gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, SWAP_AB, C_N_MAJOR, NUM_STAGES>(A, B, SFA, SFB, C, buf);
+  
+   if (false) {}
+   LAUNCH(7168, 128, 128, 256, 1, true, true, 5)
+   LAUNCH(4096, 128, 64, 256, 1, true, false, 6)
+   LAUNCH(2048, 128, 64, 256, 1, true, false, 8)
+   LAUNCH(2304, 128, 64, 256, 1, true, false, 6)
+   LAUNCH(1536, 128, 64, 256, 1, true, false, 6)
+   LAUNCH(512, 128, 64, 256, 1, true, false, 6)
+   LAUNCH(256, 128, 64, 256, 1, true, false, 6)
+  
+ #undef LAUNCH
+  
+   return C;
+ }
+  
+ __global__ void silu_mul_kernel(const half* g1, const half* g2, half* out, int64_t count) {
+   int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+   if (idx >= count) {
+     return;
+   }
+   float x = __half2float(g1[idx]);
+   float y = __half2float(g2[idx]);
+   float silu = x / (1.0f + expf(-x));
+   out[idx] = __float2half(silu * y);
+ }
+  
+ at::Tensor silu_mul(
+   const at::Tensor& g1,
+   const at::Tensor& g2,
+         at::Tensor& out
+ ) {
+   int64_t count = g1.numel();
+   int threads = 256;
+   int blocks = static_cast<int>((count + threads - 1) / threads);
+   silu_mul_kernel<<<blocks, threads>>>(
+     reinterpret_cast<const half *>(g1.data_ptr()),
+     reinterpret_cast<const half *>(g2.data_ptr()),
+     reinterpret_cast<half *>(out.data_ptr()),
+     count
+   );
+   return out;
+ }
+  
+ TORCH_LIBRARY(nvfp4_ops, m) {
+   m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, Tensor(b!) buf) -> Tensor");
+   m.def("silu_mul(Tensor g1, Tensor g2, Tensor(a!) out) -> Tensor");
+   m.impl("gemm", &gemm);
+   m.impl("silu_mul", &silu_mul);
+ }
  """
   
   
- ext = load_inline(
-     name="nvfp4_dual_gemm_ext",
-     cpp_sources=cpp_src,
-     cuda_sources=cuda_src,
-     functions=None,
-     with_cuda=True,
-     extra_cflags=[
-         "-O3",
-         "-std=c++17",
-     ],
-     extra_cuda_cflags=[
-         "-O3",
-         "--use_fast_math",
-         "-lineinfo",
-     ],
-     verbose=False,
- )
+ _EXT_READY = False
   
   
+ def _load_ext():
+     global _EXT_READY
+     if _EXT_READY:
+         return
+     if hasattr(torch.ops, "nvfp4_ops") and hasattr(torch.ops.nvfp4_ops, "gemm"):
+         _EXT_READY = True
+         return
+     load_inline(
+         name="nvfp4_dual_gemm_ext",
+         cpp_sources="",
+         cuda_sources=_CUDA_SRC_COMMON + _CUDA_SRC_V4,
+         functions=None,
+         with_cuda=True,
+         verbose=False,
+         is_python_module=False,
+         extra_cflags=["-O3", "-std=c++17"],
+         extra_cuda_cflags=[
+             "-O3",
+             "-gencode=arch=compute_100a,code=sm_100a",
+             "--use_fast_math",
+             "--expt-relaxed-constexpr",
+             "--relocatable-device-code=false",
+             "-lineinfo",
+         ],
+         extra_ldflags=["-lcuda"],
+     )
+     _EXT_READY = True
+  
+  
+ def _from_permuted_scale(scale_p: torch.Tensor) -> torch.Tensor:
+     # 使用 permuted scale 生成 kernel 期望的线性布局
+     if scale_p.dim() != 6:
+         raise RuntimeError("scale_p 维度不符合预期")
+     scale_l = scale_p.select(5, 0)
+     permuted = scale_l.permute(2, 4, 0, 1, 3).contiguous()
+     return permuted.reshape(-1)
+  
+  
+ def _gemm(a: torch.Tensor, b: torch.Tensor, sfa: torch.Tensor, sfb: torch.Tensor) -> torch.Tensor:
+     m = a.size(0)
+     n = b.size(0)
+     c = torch.empty((m, n, 1), device=a.device, dtype=torch.float16)
+     buf = torch.empty((1,), device=a.device, dtype=torch.float32)
+     return torch.ops.nvfp4_ops.gemm(a, b, sfa, sfb, c, buf).contiguous()
+  
+  
  def custom_kernel(data):
-     a, b1, b2, sfa, sfb1, sfb2, sfa_p, sfb1_p, sfb2_p, c = data
-     return ext.fused_dual_gemm(a, b1, b2, sfa, sfb1, sfb2, c)
+     a, b1, b2, _sfa, _sfb1, _sfb2, sfa_p, sfb1_p, sfb2_p, c = data
+     _load_ext()
+     out = c.contiguous()
   
+     m = a.size(0)
+     n = b1.size(0)
+     k = a.size(1) * 2
+     l = out.size(2)
+     if l != 1:
+         raise RuntimeError("仅支持 L=1")
+     if k not in (7168, 4096, 2304, 2048, 1536, 512, 256):
+         raise RuntimeError("仅支持预设 K 集合")
   
+     a_l = a.select(2, 0)
+     b1_l = b1.select(2, 0)
+     b2_l = b2.select(2, 0)
+     scale_a = _from_permuted_scale(sfa_p)
+     scale_b1 = _from_permuted_scale(sfb1_p)
+     scale_b2 = _from_permuted_scale(sfb2_p)
+  
+     g1 = _gemm(a_l, b1_l, scale_a, scale_b1)
+     g2 = _gemm(a_l, b2_l, scale_a, scale_b2)
+  
+     out_slice = out.view(m, n)
+     g1_view = g1.view(m, n)
+     g2_view = g2.view(m, n)
+     torch.ops.nvfp4_ops.silu_mul(g1_view, g2_view, out_slice)
+     return out
+  
+  
  __all__ = ["custom_kernel"]