submission 296470

import torch
from torch.utils.cpp_extension import load_inline

_CUDA_SRC = r"""
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <ATen/core/Tensor.h>
#include <torch/library.h>

// 说明（中文）：仅面向 B200(sm_100a) 与评测固定形状做极致特化；不做任何回退。
// 分发器（中文）：计分形状启用 FUSE_CP_MMA（scale cp 与 mma 交织）；其余形状保持稳定路径。
// 备注（中文）：EPILOGUE_SEG 保留为实验开关，当前计分区使用 SEG=32。

constexpr int WARP_SZ = 32;
constexpr int MMA_K64 = 64;
constexpr uint64_t L2_EVICT_FIRST = 0x12F0000000000000ULL;
constexpr uint64_t L2_EVICT_LAST  = 0x14F0000000000000ULL;

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

__device__ __forceinline__ uint32_t elect_one() {
  uint32_t pred = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred %%p;\n\t"
    "elect.sync _|%%p, %1;\n\t"
    "@%%p mov.s32 %0, 1;\n\t"
    "}\n\t"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

__device__ __forceinline__ uint64_t l2_policy_first() { return L2_EVICT_FIRST; }
__device__ __forceinline__ uint64_t l2_policy_last() { return L2_EVICT_LAST; }

__device__ __forceinline__ void mbar_init_shared(int addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(addr), "r"(count));
}

__device__ __forceinline__ void mbar_wait_parity(int addr, int phase) {
  uint32_t ticks = 0x989680;
  asm volatile(
    "{\n\t"
    ".reg .pred P;\n\t"
    "L_WAIT:\n\t"
    "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P, [%0], %1, %2;\n\t"
    "@P bra.uni L_DONE;\n\t"
    "bra.uni L_WAIT;\n\t"
    "L_DONE:\n\t"
    "}\n\t"
    :: "r"(addr), "r"(phase), "r"(ticks)
  );
}

__device__ __forceinline__ void tma_g2s_bytes(int dst, const void* src, int bytes, int mbar, uint64_t cache) {
  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"(bytes), "r"(mbar), "l"(cache)
  );
}

__device__ __forceinline__ void tma_g2s_3d(int dst, const void* tmap, int x, int y, int z, int mbar, uint64_t cache) {
  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), "r"(x), "r"(y), "r"(z), "r"(mbar), "l"(cache)
    : "memory"
  );
}

__device__ __forceinline__ void tc_scale_cp(uint32_t taddr, uint64_t sdesc) {
  asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(sdesc));
}

__device__ __forceinline__ void tc_mma_a_fill(
  uint32_t daddr,
  uint64_t adesc,
  uint64_t bdesc,
  uint32_t idesc,
  uint32_t scale_a,
  uint32_t scale_b,
  int enable_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.collector::a::fill [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}\n\t"
    :: "r"(daddr), "l"(adesc), "l"(bdesc), "r"(idesc), "r"(scale_a), "r"(scale_b), "r"(enable_d)
  );
}

__device__ __forceinline__ void tc_mma_a_last(
  uint32_t daddr,
  uint64_t adesc,
  uint64_t bdesc,
  uint32_t idesc,
  uint32_t scale_a,
  uint32_t scale_b,
  int enable_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.collector::a::lastuse [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}\n\t"
    :: "r"(daddr), "l"(adesc), "l"(bdesc), "r"(idesc), "r"(scale_a), "r"(scale_b), "r"(enable_d)
  );
}

struct _TC_SH { static constexpr char _16x256b[] = ".16x256b"; };
struct _TC_NM { static constexpr char x4[] = ".x4"; static constexpr char x8[] = ".x8"; };

template <const char* SH, const char* NM>
__device__ __forceinline__ void tc_ld16(float* out, uint32_t addr) {
  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"(out[ 0]), "=f"(out[ 1]), "=f"(out[ 2]), "=f"(out[ 3]), "=f"(out[ 4]), "=f"(out[ 5]), "=f"(out[ 6]), "=f"(out[ 7]),
      "=f"(out[ 8]), "=f"(out[ 9]), "=f"(out[10]), "=f"(out[11]), "=f"(out[12]), "=f"(out[13]), "=f"(out[14]), "=f"(out[15])
    : "r"(addr), "C"(SH), "C"(NM)
  );
}

template <const char* SH, const char* NM>
__device__ __forceinline__ void tc_ld32(float* out, uint32_t addr) {
  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"(out[ 0]), "=f"(out[ 1]), "=f"(out[ 2]), "=f"(out[ 3]), "=f"(out[ 4]), "=f"(out[ 5]), "=f"(out[ 6]), "=f"(out[ 7]),
      "=f"(out[ 8]), "=f"(out[ 9]), "=f"(out[10]), "=f"(out[11]), "=f"(out[12]), "=f"(out[13]), "=f"(out[14]), "=f"(out[15]),
      "=f"(out[16]), "=f"(out[17]), "=f"(out[18]), "=f"(out[19]), "=f"(out[20]), "=f"(out[21]), "=f"(out[22]), "=f"(out[23]),
      "=f"(out[24]), "=f"(out[25]), "=f"(out[26]), "=f"(out[27]), "=f"(out[28]), "=f"(out[29]), "=f"(out[30]), "=f"(out[31])
    : "r"(addr), "C"(SH), "C"(NM)
  );
}

__device__ __forceinline__ void tc_ld_16x256bx8(float* out, uint32_t addr) { tc_ld32<_TC_SH::_16x256b, _TC_NM::x8>(out, addr); }
__device__ __forceinline__ void tc_ld_16x256bx4(float* out, uint32_t addr) { tc_ld16<_TC_SH::_16x256b, _TC_NM::x4>(out, addr); }

static inline void ck_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char* msg = nullptr;
  if (cuGetErrorString(err, &msg) != CUDA_SUCCESS) msg = "cu err";
  TORCH_CHECK(false, msg);
}

static inline void ck_cuda(cudaError_t err) {
  if (err == cudaSuccess) return;
  const char* msg = cudaGetErrorString(err);
  TORCH_CHECK(false, msg ? msg : "cuda err");
}

static inline int sm_count_cached() {
  static int sm = -1;
  if (sm > 0) return sm;
  int dev = 0;
  ck_cuda(cudaGetDevice(&dev));
  cudaDeviceProp prop;
  ck_cuda(cudaGetDeviceProperties(&prop, dev));
  sm = prop.multiProcessorCount;
  return sm;
}

static inline void encode_tmap(
  CUtensorMap* tmap,
  const char* ptr,
  uint64_t h,
  uint64_t w,
  uint32_t sh,
  uint32_t sw,
  CUtensorMapL2promotion promo
) {
  constexpr uint32_t rank = 3;
  uint64_t gdim[rank]      = {256, h, w / 256};
  uint64_t gstride[rank-1] = {w / 2, 128};
  uint32_t bdim[rank]      = {256, sh, sw / 256};
  uint32_t estride[rank]   = {1, 1, 1};
  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
    rank,
    (void*)ptr,
    gdim,
    gstride,
    bdim,
    estride,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    promo,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  ck_cu(err);
}

struct CacheSel {
  uint64_t a_data;
  uint64_t a_sf;
  uint64_t b_data;
  uint64_t b_sf;
};

template <int K, int POLICY>
__device__ __forceinline__ CacheSel cache_sel() {
  const uint64_t p_first = l2_policy_first();
  const uint64_t p_last  = l2_policy_last();

  CacheSel r;
  if constexpr (K == 4096 || K == 7168) {
    if constexpr (POLICY <= 2) {
      if constexpr (POLICY == 1) {
        r.a_data = p_last;  r.a_sf = p_last;
        r.b_data = p_first; r.b_sf = p_first;
      } else if constexpr (POLICY == 2) {
        r.a_data = p_first; r.a_sf = p_first;
        r.b_data = p_last;  r.b_sf = p_last;
      } else {
        r.a_data = p_first; r.a_sf = p_last;
        r.b_data = p_first; r.b_sf = p_first;
      }
    } else {
      static_assert(POLICY <= (3 + 15));
      constexpr int mask = POLICY - 3;
      r.a_data = (mask & 0x1) ? p_last : p_first;
      r.a_sf   = (mask & 0x2) ? p_last : p_first;
      r.b_data = (mask & 0x4) ? p_last : p_first;
      r.b_sf   = (mask & 0x8) ? p_last : p_first;
    }
  } else {
    r.a_data = p_first; r.a_sf = p_first;
    r.b_data = p_last;  r.b_sf = p_last;
  }
  return r;
}

template <int K, int BM, int BN, int BK, int STAGES, int POLICY, int CTA_N_MAJOR, int FAST_SILU, int EPILOGUE_SEG, int FUSE_CP_MMA>
__global__ __launch_bounds__(BM + 2 * WARP_SZ, 1)
void kernel_dual_fused(
  const __grid_constant__ CUtensorMap A_t,
  const __grid_constant__ CUtensorMap B1_t,
  const __grid_constant__ CUtensorMap B2_t,
  const char* __restrict__ SFA,
  const char* __restrict__ SFB1,
  const char* __restrict__ SFB2,
  half* __restrict__ OUT,
  int M,
  int N
) {
  const int tid = (int)threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;

  int bid_m;
  int bid_n;
  if constexpr (K == 7168) {
    if constexpr (CTA_N_MAJOR) { bid_n = (int)blockIdx.x; bid_m = (int)blockIdx.y; }
    else { bid_m = (int)blockIdx.x; bid_n = (int)blockIdx.y; }
  } else {
    bid_n = (int)blockIdx.x;
    bid_m = (int)blockIdx.y;
  }

  const int off_m = bid_m * BM;
  const int off_n = bid_n * BN;

  constexpr int WARP_CNT = BM / WARP_SZ + 2;

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

  constexpr int A_BYTES   = BM * BK / 2;
  constexpr int B_BYTES   = BN * BK / 2;
  constexpr int SFA_BYTES = 128 * BK / 16;
  constexpr int SFB_BYTES = 128 * BK / 16;
  constexpr int STAGE_BYTES = A_BYTES + 2 * B_BYTES + SFA_BYTES + 2 * SFB_BYTES;

  constexpr int TMEM_NEED = 2 * BN + 12 * (BK / MMA_K64);
  constexpr int TMEM_COLS = (TMEM_NEED <= 256) ? 256 : 512;
  static_assert(TMEM_NEED <= 512);

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[STAGES * 2 + 1];

  const int tma_mbar  = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar  = tma_mbar + STAGES * 8;
  const int main_mbar = mma_mbar + STAGES * 8;

  if (warp == 0 && elect_one()) {
    #pragma unroll
    for (int i = 0; i < STAGES * 2 + 1; ++i) mbar_init_shared(tma_mbar + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  }

  if (warp == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(TMEM_COLS));
  }
  __syncthreads();

  constexpr uint32_t tmem_base = 0;
  constexpr uint32_t out1_col = 0;
  constexpr uint32_t out2_col = (uint32_t)BN;
  constexpr uint32_t sfa_col  = (uint32_t)(2 * BN);
  constexpr uint32_t sfb1_col = (uint32_t)(2 * BN + 4 * (BK / MMA_K64));
  constexpr uint32_t sfb2_col = (uint32_t)(2 * BN + 8 * (BK / MMA_K64));

  constexpr int iters = K / BK;

  if (warp == WARP_CNT - 1 && elect_one()) {
    const CacheSel pol = cache_sel<K, POLICY>();
    constexpr int REST_K = K / 16 / 4;
    constexpr int SF_STEP = (BK / (16 * 4)) * 512;
    constexpr int Z_STEP = (BK / 256);

    const int off_m128 = off_m >> 7;
    const int off_n128 = off_n >> 7;
    const char* sfa_src  = SFA  + off_m128 * REST_K * 512;
    const char* sfb1_src = SFB1 + off_n128 * REST_K * 512;
    const char* sfb2_src = SFB2 + off_n128 * REST_K * 512;

    int stage = 0;
    int wraps = 0;
    int stage_base = smem;
    int z = 0;
    for (int iter = 0; iter < iters; ++iter) {
      if (iter >= STAGES) {
        mbar_wait_parity(mma_mbar + stage * 8, (wraps - 1) & 1);
      }

      const int mbar = tma_mbar + stage * 8;
      const int a_s   = stage_base;
      const int b1_s  = a_s + A_BYTES;
      const int b2_s  = b1_s + B_BYTES;
      const int sfa_s = b2_s + B_BYTES;
      const int sfb1_s = sfa_s + SFA_BYTES;
      const int sfb2_s = sfb1_s + SFB_BYTES;

      tma_g2s_3d(a_s,  &A_t,  0, off_m, z, mbar, pol.a_data);
      tma_g2s_3d(b1_s, &B1_t, 0, off_n, z, mbar, pol.b_data);
      tma_g2s_3d(b2_s, &B2_t, 0, off_n, z, mbar, pol.b_data);

      tma_g2s_bytes(sfa_s,  sfa_src,  SFA_BYTES, mbar, pol.a_sf);
      tma_g2s_bytes(sfb1_s, sfb1_src, SFB_BYTES, mbar, pol.b_sf);
      tma_g2s_bytes(sfb2_s, sfb2_src, SFB_BYTES, mbar, pol.b_sf);

      asm volatile(
        "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
        :: "r"(mbar), "r"(STAGE_BYTES)
        : "memory"
      );

      z += Z_STEP;
      sfa_src += SF_STEP;
      sfb1_src += SF_STEP;
      sfb2_src += SF_STEP;

      ++stage;
      stage_base += STAGE_BYTES;
      if (stage == STAGES) {
        stage = 0;
        stage_base = smem;
        ++wraps;
      }
    }
  } else if (warp == WARP_CNT - 2 && elect_one()) {
    constexpr int MMA_N = BN;
    constexpr int MMA_M = 128;
    constexpr uint32_t idesc =
      (1U << 7U) | (1U << 10U) | ((uint32_t)MMA_N >> 3U << 17U) | ((uint32_t)MMA_M >> 7U << 27U);

    auto desc_ab = [](int addr) -> uint64_t {
      const int sbo = 8 * 128;
      return desc_pack(addr) | (desc_pack(sbo) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
    };
    auto desc_sf = [](int addr) -> uint64_t {
      const int sbo = 8 * 16;
      return desc_pack(addr) | (desc_pack(sbo) << 32ULL) | (1ULL << 46ULL);
    };
    const uint64_t ab0 = desc_ab(0);
    const uint64_t sf0 = desc_sf(0);
    constexpr uint32_t SB_PARTS = (uint32_t)(128 / BN);
    static_assert((128 % BN) == 0);
    const uint32_t sb_off =
      (SB_PARTS == 1) ? 0U : ((uint32_t)bid_n & (SB_PARTS - 1U)) * (uint32_t)(BN / 32);

    int stage = 0;
    int phase = 0;
    int stage_base = smem;
    for (int iter = 0; iter < iters; ++iter) {
      mbar_wait_parity(tma_mbar + stage * 8, phase);

      const int a_s   = stage_base;
      const int b1_s  = a_s + A_BYTES;
      const int b2_s  = b1_s + B_BYTES;
      const int sfa_s = b2_s + B_BYTES;
      const int sfb1_s = sfa_s + SFA_BYTES;
      const int sfb2_s = sfb1_s + SFB_BYTES;

      const uint64_t sfa_desc0  = sf0 + ((uint64_t)sfa_s >> 4ULL);
      const uint64_t sfb1_desc0 = sf0 + ((uint64_t)sfb1_s >> 4ULL);
      const uint64_t sfb2_desc0 = sf0 + ((uint64_t)sfb2_s >> 4ULL);

      const uint64_t a_base  = ab0 + ((uint64_t)a_s  >> 4ULL);
      const uint64_t b1_base = ab0 + ((uint64_t)b1_s >> 4ULL);
      const uint64_t b2_base = ab0 + ((uint64_t)b2_s >> 4ULL);
      if constexpr (FUSE_CP_MMA == 1) {
        static_assert(BK == 256);
        constexpr uint64_t SF_DESC_STEP = (512ULL >> 4ULL);
        uint32_t td_sfa  = tmem_base + sfa_col;
        uint32_t td_sfb1 = tmem_base + sfb1_col;
        uint32_t td_sfb2 = tmem_base + sfb2_col;
        uint32_t sa = tmem_base + sfa_col;
        uint32_t sb1 = tmem_base + sfb1_col + sb_off;
        uint32_t sb2 = tmem_base + sfb2_col + sb_off;
        uint64_t sfa_desc  = sfa_desc0;
        uint64_t sfb1_desc = sfb1_desc0;
        uint64_t sfb2_desc = sfb2_desc0;
        uint64_t a_desc = a_base;
        uint64_t b1_desc = b1_base;
        uint64_t b2_desc = b2_base;
        tc_scale_cp(td_sfa,  sfa_desc);
        tc_scale_cp(td_sfb1, sfb1_desc);
        tc_scale_cp(td_sfb2, sfb2_desc);
        #pragma unroll
        for (int k2 = 0; k2 < BK / MMA_K64; ++k2) {
          const int en = (k2 == 0) ? iter : 1;
          tc_mma_a_fill(tmem_base + out1_col, a_desc, b1_desc, idesc, sa, sb1, en);
          if (k2 + 1 < BK / MMA_K64) {
            tc_scale_cp(td_sfa + 4U,  sfa_desc  + SF_DESC_STEP);
            tc_scale_cp(td_sfb1 + 4U, sfb1_desc + SF_DESC_STEP);
            tc_scale_cp(td_sfb2 + 4U, sfb2_desc + SF_DESC_STEP);
          }
          tc_mma_a_last(tmem_base + out2_col, a_desc, b2_desc, idesc, sa, sb2, en);
          a_desc += 2ULL;
          b1_desc += 2ULL;
          b2_desc += 2ULL;
          sfa_desc  += SF_DESC_STEP;
          sfb1_desc += SF_DESC_STEP;
          sfb2_desc += SF_DESC_STEP;
          td_sfa  += 4U;
          td_sfb1 += 4U;
          td_sfb2 += 4U;
          sa += 4U;
          sb1 += 4U;
          sb2 += 4U;
        }
      } else if constexpr (FUSE_CP_MMA == 2) {
        static_assert(BK == 256);
        constexpr uint64_t SF_DESC_STEP = (512ULL >> 4ULL);

        const uint32_t td_sfa0  = tmem_base + sfa_col;
        const uint32_t td_sfb10 = tmem_base + sfb1_col;
        const uint32_t td_sfb20 = tmem_base + sfb2_col;

        const uint32_t sa0 = tmem_base + sfa_col;
        const uint32_t sb1_0 = tmem_base + sfb1_col + sb_off;
        const uint32_t sb2_0 = tmem_base + sfb2_col + sb_off;

        const uint64_t sfa_d0  = sfa_desc0;
        const uint64_t sfb1_d0 = sfb1_desc0;
        const uint64_t sfb2_d0 = sfb2_desc0;

        // 计分 BN=64：加大 cp 预取距离（k2=1 先行装载；k2=2/3 采用 k2+2 预取），试图用更多 mma slack 覆盖 cp 延迟。
        tc_scale_cp(td_sfa0,  sfa_d0);
        tc_scale_cp(td_sfb10, sfb1_d0);
        tc_scale_cp(td_sfb20, sfb2_d0);
        tc_scale_cp(td_sfa0 + 4U,  sfa_d0  + SF_DESC_STEP);
        tc_scale_cp(td_sfb10 + 4U, sfb1_d0 + SF_DESC_STEP);
        tc_scale_cp(td_sfb20 + 4U, sfb2_d0 + SF_DESC_STEP);

        const int en0 = iter;
        tc_mma_a_fill(tmem_base + out1_col, a_base, b1_base, idesc, sa0, sb1_0, en0);
        tc_scale_cp(td_sfa0 + 8U,  sfa_d0  + 2ULL * SF_DESC_STEP);
        tc_scale_cp(td_sfb10 + 8U, sfb1_d0 + 2ULL * SF_DESC_STEP);
        tc_scale_cp(td_sfb20 + 8U, sfb2_d0 + 2ULL * SF_DESC_STEP);
        tc_mma_a_last(tmem_base + out2_col, a_base, b2_base, idesc, sa0, sb2_0, en0);

        tc_mma_a_fill(tmem_base + out1_col, a_base + 2ULL, b1_base + 2ULL, idesc, sa0 + 4U, sb1_0 + 4U, 1);
        tc_scale_cp(td_sfa0 + 12U,  sfa_d0  + 3ULL * SF_DESC_STEP);
        tc_scale_cp(td_sfb10 + 12U, sfb1_d0 + 3ULL * SF_DESC_STEP);
        tc_scale_cp(td_sfb20 + 12U, sfb2_d0 + 3ULL * SF_DESC_STEP);
        tc_mma_a_last(tmem_base + out2_col, a_base + 2ULL, b2_base + 2ULL, idesc, sa0 + 4U, sb2_0 + 4U, 1);

        tc_mma_a_fill(tmem_base + out1_col, a_base + 4ULL, b1_base + 4ULL, idesc, sa0 + 8U, sb1_0 + 8U, 1);
        tc_mma_a_last(tmem_base + out2_col, a_base + 4ULL, b2_base + 4ULL, idesc, sa0 + 8U, sb2_0 + 8U, 1);

        tc_mma_a_fill(tmem_base + out1_col, a_base + 6ULL, b1_base + 6ULL, idesc, sa0 + 12U, sb1_0 + 12U, 1);
        tc_mma_a_last(tmem_base + out2_col, a_base + 6ULL, b2_base + 6ULL, idesc, sa0 + 12U, sb2_0 + 12U, 1);
      } else {
        uint32_t td_sfa  = tmem_base + sfa_col;
        uint32_t td_sfb1 = tmem_base + sfb1_col;
        uint32_t td_sfb2 = tmem_base + sfb2_col;
        uint64_t sfa_desc  = sfa_desc0;
        uint64_t sfb1_desc = sfb1_desc0;
        uint64_t sfb2_desc = sfb2_desc0;
        #pragma unroll
        for (int kk = 0; kk < BK / MMA_K64; ++kk) {
          tc_scale_cp(td_sfa,  sfa_desc);
          tc_scale_cp(td_sfb1, sfb1_desc);
          tc_scale_cp(td_sfb2, sfb2_desc);
          sfa_desc  += (512ULL >> 4ULL);
          sfb1_desc += (512ULL >> 4ULL);
          sfb2_desc += (512ULL >> 4ULL);
          td_sfa  += 4U;
          td_sfb1 += 4U;
          td_sfb2 += 4U;
        }

        uint32_t sa = tmem_base + sfa_col;
        uint32_t sb1 = tmem_base + sfb1_col + sb_off;
        uint32_t sb2 = tmem_base + sfb2_col + sb_off;

        if constexpr (BK == 256) {
          uint64_t a_desc = a_base;
          uint64_t b1_desc = b1_base;
          uint64_t b2_desc = b2_base;
          #pragma unroll
          for (int k2 = 0; k2 < BK / MMA_K64; ++k2) {
            const int en = (k2 == 0) ? iter : 1;
            tc_mma_a_fill(tmem_base + out1_col, a_desc, b1_desc, idesc, sa, sb1, en);
            tc_mma_a_last(tmem_base + out2_col, a_desc, b2_desc, idesc, sa, sb2, en);
            a_desc += 2ULL;
            b1_desc += 2ULL;
            b2_desc += 2ULL;
            sa += 4U;
            sb1 += 4U;
            sb2 += 4U;
          }
        } else {
          constexpr uint64_t A_STEP = ((uint64_t)BM * 128ULL) >> 4ULL;
          constexpr uint64_t B_STEP = ((uint64_t)BN * 128ULL) >> 4ULL;
          #pragma unroll
          for (int k2 = 0; k2 < BK / MMA_K64; ++k2) {
            const int en = (k2 == 0) ? iter : 1;
            const int k1 = k2 >> 2;
            const int kk = k2 & 3;
            const uint64_t a_desc = a_base  + (uint64_t)k1 * A_STEP + (uint64_t)kk * 2ULL;
            const uint64_t b1_desc = b1_base + (uint64_t)k1 * B_STEP + (uint64_t)kk * 2ULL;
            const uint64_t b2_desc = b2_base + (uint64_t)k1 * B_STEP + (uint64_t)kk * 2ULL;
            tc_mma_a_fill(tmem_base + out1_col, a_desc, b1_desc, idesc, sa, sb1, en);
            tc_mma_a_last(tmem_base + out2_col, a_desc, b2_desc, idesc, sa, sb2, en);
            sa += 4U;
            sb1 += 4U;
            sb2 += 4U;
          }
        }
      }

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

      ++stage;
      stage_base += STAGE_BYTES;
      if (stage == STAGES) {
        stage = 0;
        stage_base = smem;
        phase ^= 1;
      }
    }

    asm volatile(
      "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
      :: "r"(main_mbar)
      : "memory"
    );
  } else if (tid < BM) {
    mbar_wait_parity(main_mbar, 0);
    asm volatile("tcgen05.fence::after_thread_sync;");

    constexpr float LOG2E = 1.4426950408889634f;
    const int lane_row = lane >> 2;
    const int lane_h2  = lane & 3;

    constexpr int SEG = EPILOGUE_SEG;
    static_assert((BN % SEG) == 0);
    constexpr int SEGS = BN / SEG;

    #pragma unroll
    for (int mm = 0; mm < 2; ++mm) {
      const int row0 = warp * 32 + mm * 16;
      const uint32_t addr_x0 = tmem_base + (uint32_t)((row0 << 16) | (int)out1_col);
      const uint32_t addr_y0 = tmem_base + (uint32_t)((row0 << 16) | (int)out2_col);
      const int row = off_m + row0 + lane_row;
      half2* out_row0 = reinterpret_cast<half2*>(OUT + row * N);
      half2* out_row8 = reinterpret_cast<half2*>(OUT + (row + 8) * N);

      #pragma unroll
      for (int seg = 0; seg < SEGS; ++seg) {
        float x[SEG / 2];
        float y[SEG / 2];
        const uint32_t col_off = (uint32_t)(seg * SEG);
        if constexpr (SEG == 32) {
          tc_ld_16x256bx4(x, addr_x0 + col_off);
          tc_ld_16x256bx4(y, addr_y0 + col_off);
        } else if constexpr (SEG == 64) {
          tc_ld_16x256bx8(x, addr_x0 + col_off);
          tc_ld_16x256bx8(y, addr_y0 + col_off);
        } else {
          static_assert(SEG == 32 || SEG == 64);
        }
        const int col_base_h2 = (off_n >> 1) + seg * (SEG >> 1);
        asm volatile("tcgen05.wait::ld.sync.aligned;");

        #pragma unroll
        for (int i = 0; i < SEG / 8; ++i) {
          const int out_col_base = col_base_h2 + i * 4;

          const float x00 = x[i * 4 + 0];
          const float x01 = x[i * 4 + 1];
          const float x80 = x[i * 4 + 2];
          const float x81 = x[i * 4 + 3];
          const float y00 = y[i * 4 + 0];
          const float y01 = y[i * 4 + 1];
          const float y80 = y[i * 4 + 2];
          const float y81 = y[i * 4 + 3];

          float s00, s01, s80, s81;
          if constexpr (FAST_SILU) {
            float t00, t01, t80, t81;
            asm("ex2.approx.f32 %0, %1;" : "=f"(t00) : "f"((-x00) * LOG2E));
            asm("ex2.approx.f32 %0, %1;" : "=f"(t01) : "f"((-x01) * LOG2E));
            asm("ex2.approx.f32 %0, %1;" : "=f"(t80) : "f"((-x80) * LOG2E));
            asm("ex2.approx.f32 %0, %1;" : "=f"(t81) : "f"((-x81) * LOG2E));
            const float d00 = 1.0f + t00;
            const float d01 = 1.0f + t01;
            const float d80 = 1.0f + t80;
            const float d81 = 1.0f + t81;
            asm("rcp.approx.f32 %0, %1;" : "=f"(s00) : "f"(d00));
            asm("rcp.approx.f32 %0, %1;" : "=f"(s01) : "f"(d01));
            asm("rcp.approx.f32 %0, %1;" : "=f"(s80) : "f"(d80));
            asm("rcp.approx.f32 %0, %1;" : "=f"(s81) : "f"(d81));
          } else {
            s00 = __fdividef(1.0f, 1.0f + exp2f((-x00) * LOG2E));
            s01 = __fdividef(1.0f, 1.0f + exp2f((-x01) * LOG2E));
            s80 = __fdividef(1.0f, 1.0f + exp2f((-x80) * LOG2E));
            s81 = __fdividef(1.0f, 1.0f + exp2f((-x81) * LOG2E));
          }

          float2 o0;
          float2 o8;
          o0.x = (x00 * s00) * y00;
          o0.y = (x01 * s01) * y01;
          o8.x = (x80 * s80) * y80;
          o8.y = (x81 * s81) * y81;

          const int out_col_h2 = out_col_base + lane_h2;
          out_row0[out_col_h2] = __float22half2_rn(o0);
          out_row8[out_col_h2] = __float22half2_rn(o8);
        }
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(BM) : "memory");
    if (warp == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(TMEM_COLS));
  }
}

template <int K, int BM, int BN, int BK, int STAGES, int POLICY, int CTA_N_MAJOR, int FAST_SILU, int EPILOGUE_SEG, int FUSE_CP_MMA>
static inline void launch_cfg(
  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& out
) {
  auto call = [](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& out,
                 CUtensorMapL2promotion promo_a,
                 CUtensorMapL2promotion promo_b) {
    const int M = (int)A.size(0);
    const int N = (int)B1.size(0);

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

    struct TmapCache {
      const char* a_ptr;
      const char* b1_ptr;
      const char* b2_ptr;
      int m;
      int n;
      CUtensorMap a_t;
      CUtensorMap b1_t;
      CUtensorMap b2_t;
      bool valid;
    };

    #pragma nv_diag_suppress static_var_with_dynamic_init
    static TmapCache cache = {nullptr, nullptr, nullptr, 0, 0, {}, {}, {}, false};

    CUtensorMap A_t, B1_t, B2_t;
    if (cache.valid && cache.a_ptr == A_ptr && cache.b1_ptr == B1_ptr && cache.b2_ptr == B2_ptr && cache.m == M && cache.n == N) {
      A_t = cache.a_t;
      B1_t = cache.b1_t;
      B2_t = cache.b2_t;
    } else {
      encode_tmap(&A_t, A_ptr,  (uint64_t)M, (uint64_t)K, (uint32_t)BM, (uint32_t)BK, promo_a);
      encode_tmap(&B1_t, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BN, (uint32_t)BK, promo_b);
      encode_tmap(&B2_t, B2_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BN, (uint32_t)BK, promo_b);
      cache.a_ptr = A_ptr;
      cache.b1_ptr = B1_ptr;
      cache.b2_ptr = B2_ptr;
      cache.m = M;
      cache.n = N;
      cache.a_t = A_t;
      cache.b1_t = B1_t;
      cache.b2_t = B2_t;
      cache.valid = true;
    }

    dim3 grid;
    if constexpr (K == 7168) {
      if constexpr (CTA_N_MAJOR) grid = dim3((unsigned)(N / BN), (unsigned)(M / BM));
      else grid = dim3((unsigned)(M / BM), (unsigned)(N / BN));
    } else {
      grid = dim3((unsigned)(N / BN), (unsigned)(M / BM));
    }

    const int tb = BM + 2 * WARP_SZ;
    constexpr int A_BYTES   = BM * BK / 2;
    constexpr int B_BYTES   = BN * BK / 2;
    constexpr int SFA_BYTES = 128 * BK / 16;
    constexpr int SFB_BYTES = 128 * BK / 16;
    constexpr int smem_bytes = (A_BYTES + 2 * B_BYTES + SFA_BYTES + 2 * SFB_BYTES) * STAGES;
    constexpr int kMaxSmemBytes = 227 * 1024;
    TORCH_CHECK(smem_bytes <= kMaxSmemBytes, "smem ", smem_bytes);

    auto kptr = kernel_dual_fused<K, BM, BN, BK, STAGES, POLICY, CTA_N_MAJOR, FAST_SILU, EPILOGUE_SEG, FUSE_CP_MMA>;
    if constexpr (smem_bytes > 48000) {
      static bool attr_set = false;
      if (!attr_set) {
        ck_cuda(cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes));
        attr_set = true;
      }
    }
    kptr<<<grid, tb, smem_bytes>>>(A_t, B1_t, B2_t, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);
  };

  call(A, B1, B2, SFA, SFB1, SFB2, out,
       CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
       CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE);
}

template <int K, int BM, int BN, int BK, int STAGES, int FAST_SILU, int EPILOGUE_SEG, int FUSE_CP_MMA>
static inline void launch_policy_auto(
  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& out,
  int cta_m,
  int cta_n
) {
  if constexpr (K == 4096 || K == 7168) {
    if (cta_n >= (cta_m << 3))      launch_cfg<K, BM, BN, BK, STAGES, 1, 0, FAST_SILU, EPILOGUE_SEG, FUSE_CP_MMA>(A, B1, B2, SFA, SFB1, SFB2, out);
    else if (cta_m >= (cta_n << 3)) launch_cfg<K, BM, BN, BK, STAGES, 2, 0, FAST_SILU, EPILOGUE_SEG, FUSE_CP_MMA>(A, B1, B2, SFA, SFB1, SFB2, out);
    else                            launch_cfg<K, BM, BN, BK, STAGES, 0, 0, FAST_SILU, EPILOGUE_SEG, FUSE_CP_MMA>(A, B1, B2, SFA, SFB1, SFB2, out);
  } else {
    launch_cfg<K, BM, BN, BK, STAGES, 0, 0, FAST_SILU, EPILOGUE_SEG, FUSE_CP_MMA>(A, B1, B2, SFA, SFB1, SFB2, out);
  }
}

static inline uint64_t shape_key_u64(int m, int n, int k) {
  return ((uint64_t)(uint32_t)k << 32) | ((uint64_t)(uint32_t)m << 16) | (uint64_t)(uint32_t)n;
}

#define KKEY(M, N, K) ((((uint64_t)(K)) << 32) | (((uint64_t)(M)) << 16) | ((uint64_t)(N)))

at::Tensor fused(
  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& out
) {
  const int M = (int)A.size(0);
  const int Kp = (int)A.size(1);
  const int N = (int)B1.size(0);
  const int K = Kp * 2;

  const uint64_t key = shape_key_u64(M, N, K);
  switch (key) {
    // 计分区：固定 shape 强特化（启用 FUSE_CP_MMA；EPILOGUE_SEG=32）
    case KKEY(256, 4096, 7168):
      launch_cfg<7168, 128, 64, 256, 5, 1, 1, 0, 32, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KKEY(512, 4096, 7168):
      launch_cfg<7168, 128, 128, 256, 4, 3, 1, 0, 32, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KKEY(256, 3072, 4096):
      launch_cfg<4096, 128, 64, 256, 5, 3, 0, 0, 32, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KKEY(512, 3072, 4096):
      launch_cfg<4096, 128, 128, 256, 4, 3, 0, 0, 32, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KKEY(512, 3072, 7168):
      launch_cfg<7168, 128, 128, 256, 4, 3, 1, 0, 32, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    default:
      break;
  }

  const int sm = sm_count_cached();
  const int cta_m = M / 128;
  const int cta_n128 = N / 128;
  const int cta_128 = cta_m * cta_n128;
  const int cta128_threshold = (sm > 96) ? 96 : sm;
  const bool use_bn128 = ((N & 127) == 0) && (cta_128 >= cta128_threshold);
  const int cta_n64 = N / 64;

  // 正确性区：其余形状仅需通过测试（保守 epilogue SEG=32 + x4；仍然使用自定义 CUDA kernel，无回退）
  if (K == 7168) {
    if (use_bn128) launch_policy_auto<7168, 128, 128, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n128);
    else           launch_policy_auto<7168, 128, 64,  256, 5, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 4096) {
    if (use_bn128) launch_policy_auto<4096, 128, 128, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n128);
    else           launch_policy_auto<4096, 128, 64,  256, 5, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 2304) {
    launch_policy_auto<2304, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 2048) {
    launch_policy_auto<2048, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 1536) {
    launch_policy_auto<1536, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 512) {
    launch_policy_auto<512, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else if (K == 256) {
    launch_policy_auto<256, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
  } else {
    TORCH_CHECK(false, "k ", K);
  }
  return out;
}

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

class _Op:
    __slots__ = ("_fn",)

    def __init__(self) -> None:
        self._fn = None

    def _compile(self) -> None:
        load_inline(
            name="nvfp4_dual_ext_r220_src",
            cpp_sources="",
            cuda_sources=_CUDA_SRC,
            functions=None,
            with_cuda=True,
            extra_cuda_cflags=[
                "-O3",
                "-gencode=arch=compute_100a,code=sm_100a",
                "--use_fast_math",
                "--expt-relaxed-constexpr",
                "--relocatable-device-code=false",
            ],
            extra_ldflags=["-lcuda"],
            verbose=False,
            is_python_module=False,
            no_implicit_headers=True,
        )
        self._fn = getattr(getattr(torch.ops, "nvfp4_dual_lib_r220"), "fused")

    def __call__(self, data):
        fn = self._fn
        if fn is None:
            self._compile()
            fn = self._fn
        return fn(data[0], data[1], data[2], data[6], data[7], data[8], data[9])


_OP = _Op()


def custom_kernel(data):
    return _OP(data)


__all__ = ["custom_kernel"]