submission 339941

import hashlib as _hashlib
import threading as _threading

import torch as _torch
from torch.utils.cpp_extension import load_inline as _load_inline


_SRC_PARTS = (
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) 做强特化；不做通用性与回退机制。
// 计算（中文）：C = silu(A@B1) * (A@B2)，A/B1/B2 为 NVFP4 block-scaled，scale 使用预重排版本。
// 分发（中文）：ranked 形状走专用配置（更深 stages/更激进 epilogue），其它 shape 仅保证正确性；域外直接报错。

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 k_dual_nvfp4(
  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) {
        static_assert(BK == 256);
        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;
        #pragma unroll
        for (int k2 = 0; k2 < BK / MMA_K64; ++k2) {
          tc_scale_cp(td_sfa,  sfa_desc);
          tc_scale_cp(td_sfb1, sfb1_desc);
          tc_scale_cp(td_sfb2, sfb2_desc);
          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;
          sfa_desc  += (512ULL >> 4ULL);
          sfb1_desc += (512ULL >> 4ULL);
          sfb2_desc += (512ULL >> 4ULL);
          td_sfa  += 4U;
          td_sfb1 += 4U;
          td_sfb2 += 4U;
          sa += 4U;
          sb1 += 4U;
          sb2 += 4U;
        }
      } 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 = k_dual_nvfp4<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_L2_128B,
       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_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 pack_mnk(int m, int n, int k) {
  return ((uint64_t)(uint32_t)m << 32) | ((uint64_t)(uint32_t)n << 16) | (uint64_t)(uint32_t)k;
}

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

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 L = (int)A.size(2);
  const int N = (int)B1.size(0);
  const int K = Kp * 2;

  TORCH_CHECK(L == 1, "l ", L);
  TORCH_CHECK((int)B1.size(2) == 1, "l ", (int)B1.size(2));
  TORCH_CHECK((int)B2.size(2) == 1, "l ", (int)B2.size(2));
  TORCH_CHECK((int)out.size(2) == 1, "l ", (int)out.size(2));

  const uint64_t key = pack_mnk(M, N, K);
  switch (key) {
    case KEY_MNK(256, 4096, 7168):
      launch_cfg<7168, 128,  64, 256, 5, 1, 1, 0, 64, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KEY_MNK(512, 4096, 7168):
      launch_cfg<7168, 128, 128, 256, 4, 3, 1, 0, 64, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KEY_MNK(256, 3072, 4096):
      launch_cfg<4096, 128,  64, 256, 5, 3, 0, 0, 64, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KEY_MNK(512, 3072, 7168):
      launch_cfg<7168, 128, 128, 256, 4, 3, 1, 0, 64, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
      return out;
    case KEY_MNK(512, 3072, 4096):
      launch_cfg<4096, 128, 128, 256, 4, 3, 0, 0, 64, 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，无回退）
  switch (K) {
    case 7168:
      if (use_bn128) launch_auto<7168, 128, 128, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n128);
      else           launch_auto<7168, 128, 64,  256, 5, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 4096:
      if (use_bn128) launch_auto<4096, 128, 128, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n128);
      else           launch_auto<4096, 128, 64,  256, 5, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 2304:
      launch_auto<2304, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 2048:
      launch_auto<2048, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 1536:
      launch_auto<1536, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 512:
      launch_auto<512, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    case 256:
      launch_auto<256, 128, 64, 256, 4, 0, 32, 0>(A, B1, B2, SFA, SFB1, SFB2, out, cta_m, cta_n64);
      break;
    default:
      TORCH_CHECK(false, "k ", K);
  }
  return out;
}

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


def _name_pair(src: str) -> tuple[str, str]:
    tag = _hashlib.sha256(src.encode("utf-8")).hexdigest()[:12]
    return f"b200_nvfp4_dual_{tag}", f"b200_nvfp4_dual_ext_{tag}"


_RAW_SRC = "".join(_SRC_PARTS)
_LIB, _EXT = _name_pair(_RAW_SRC + "|v2")
_CUDA_SRC = _RAW_SRC.replace("__LIB__", _LIB)

_LOAD_GUARD = _threading.Lock()
_LOADED = False

def _ensure_loaded() -> None:
    global _LOADED
    if _LOADED:
        return
    with _LOAD_GUARD:
        if _LOADED:
            return
        _load_inline(
            name=_EXT,
            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,
        )
        _LOADED = True


def custom_kernel(data):
    _ensure_loaded()
    a, b1, b2, _sfa, _sfb1, _sfb2, sfa_p, sfb1_p, sfb2_p, c = data
    return getattr(_torch.ops, _LIB).fused(a, b1, b2, sfa_p, sfb1_p, sfb2_p, c)


__all__ = ["custom_kernel"]