submission 209319

#!POPCORN leaderboard nvfp4_dual_gemm
#!POPCORN gpu NVIDIA

import os
from typing import Any

import torch
from torch.utils.cpp_extension import load_inline

try:
    from task import input_t, output_t  # type: ignore
except Exception:  # pragma: no cover
    input_t = Any  # type: ignore
    output_t = Any  # type: ignore


# Ensure we compile for Blackwell (B200 = sm_100a).
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "10.0a")


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

#include <cstdint>

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

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

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

__device__ inline int64_t globaltimer() {
  int64_t t;
  asm volatile("mov.u64 %0, %globaltimer;" : "=l"(t) :: "memory");
  return t;
}

// Minimal intra-kernel profiler (inspired by learn-cuda/02e_matmul_sm100/profiler.h).
// Layout: profile[(bid * NUM_WARPS + warp_id) * PROF_FIELDS + field]
enum ProfField : int {
  PROF_WAIT = 0,
  PROF_ISSUE = 1,
  PROF_OTHER = 2,
  PROF_TOTAL = 3,
};
constexpr int PROF_FIELDS = 4;

__device__ inline constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };

// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cute/arch/cluster_sm90.hpp#L180
__device__ uint32_t elect_sync() {
  uint32_t pred = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred %%px;\n\t"
    "elect.sync _|%%px, %1;\n\t"
    "@%%px mov.s32 %0, 1;\n\t"
    "}"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

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

// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cutlass/arch/barrier.h#L408
__device__ void mbarrier_wait(int mbar_addr, int phase) {
  uint32_t ticks = 0x989680;  // optional
  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");
}

template <int CTA_GROUP>
__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) {
  // .cta_group::2 allows mbar_addr and dst to be in different CTA's smem.
  asm volatile("cp.async.bulk.tensor.3d.cta_group::%7.shared::cluster.global.mbarrier::complete_tx::bytes.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), "n"(CTA_GROUP)
              : "memory");
}

__device__ inline void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  // .32x128b corresponds to (32, 16) 8-bit scale -> 1 MMA for nvfp4.
  // .warpx4 duplicates data across 32-lane groups.
  asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

template <int CTA_GROUP>
__device__ inline void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
  // .32x128b corresponds to (32, 16) 8-bit scale -> 1 MMA for nvfp4.
  // .warpx4 duplicates data across 32-lane groups.
  asm volatile("tcgen05.cp.cta_group::%2.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc), "n"(CTA_GROUP));
}

__device__ inline void tcgen05_mma_nvfp4(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"  // predicate register enable-input-d
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

template <int CTA_GROUP>
__device__ inline void tcgen05_mma_nvfp4(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"  // predicate register enable-input-d
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d), "n"(CTA_GROUP)
  );
}

// Blackwell operand-collector variants: reuse A across back-to-back MMAs (B1 then B2).
__device__ inline void tcgen05_mma_nvfp4_collector_a_fill(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16.collector::a::fill [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::fill [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

template <int CTA_GROUP>
__device__ inline void tcgen05_mma_nvfp4_collector_a_fill(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.block16.collector::a::fill [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::fill [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d), "n"(CTA_GROUP)
  );
}

__device__ inline void tcgen05_mma_nvfp4_collector_a_use(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16.collector::a::use [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::use [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

template <int CTA_GROUP>
__device__ inline void tcgen05_mma_nvfp4_collector_a_use(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.block16.collector::a::use [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::use [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d), "n"(CTA_GROUP)
  );
}

__device__ inline void tcgen05_mma_nvfp4_collector_a_lastuse(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16.collector::a::lastuse [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::lastuse [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

template <int CTA_GROUP>
__device__ inline void tcgen05_mma_nvfp4_collector_a_lastuse(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int d_tmem,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.block16.collector::a::lastuse [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#else
    "tcgen05.mma.cta_group::%7.kind::mxf4nvf4.block_scale.scale_vec::4X.collector::a::lastuse [%0], %1, %2, %3, [%4], [%5], p;\n\t"
#endif
    "}"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d), "n"(CTA_GROUP)
  );
}

// tcgen05.ld helpers (only what we need for BLOCK_N=64 path)
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));
}

struct SHAPE {
  static constexpr char _16x256b[] = ".16x256b";
};

struct NUM {
  static constexpr char x8[] = ".x8";
  static constexpr char x16[] = ".x16";
};

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

// Safe/exact silu using expf: silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
__device__ inline float sigmoid_safe(float x) {
  return 1.0f / (1.0f + expf(-x));
}

__device__ inline float silu_safe(float x) {
  return x * sigmoid_safe(x);
}

void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *error_msg_ptr;
  if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS) error_msg_ptr = "unable to get error string";
  TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
}

void check_cuda(cudaError_t err) {
  if (err == cudaSuccess) return;
  TORCH_CHECK(false, cudaGetErrorString(err));
}

void init_AB_tmap(
  CUtensorMap *tmap,
  const char *ptr,
  uint64_t global_height, uint64_t global_width,
  uint32_t shared_height, uint32_t shared_width
) {
  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
  uint64_t globalStrides[rank-1] = {global_width / 2, 128};  // in bytes
  uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
  uint32_t elementStrides[rank]  = {1, 1, 1};

  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
    rank,
    (void *)ptr,
    globalDim,
    globalStrides,
    boxDim,
    elementStrides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
}

void init_SF_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;
  // Pack scale factors as uint16 so the fastest-moving dim is 256 (i.e., 512 bytes).
  // This matches common SM100 TMA constraints (dim0=256) while preserving byte layout.
  uint64_t globalDim[rank]       = {256, global_height, global_width};
  uint64_t globalStrides[rank-1] = {global_width * 512, 512};  // in bytes
  uint32_t boxDim[rank]          = {256, shared_height, shared_width};
  uint32_t elementStrides[rank]  = {1, 1, 1};

  auto err = cuTensorMapEncodeTiled(
    tmap,
    CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT16,
    rank,
    (void *)ptr,
    globalDim,
    globalStrides,
    boxDim,
    elementStrides,
    CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
    CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
    CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
}
"""


CUDA_SRC = r"""
template <
  int K,
  int BLOCK_M,
  int BLOCK_N,
  int BLOCK_K,
  int NUM_STAGES,
  bool DO_PROFILE
>
__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B1_tmap,
  const __grid_constant__ CUtensorMap B2_tmap,
  const char *SFA_ptr,
  const char *SFB1_ptr,
  const char *SFB2_ptr,
  half *C_ptr,
  int M, int N,
  int64_t *profile_ptr
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.x;

  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  const int grid_m = M / BLOCK_M;
  const int grid_n = N / BLOCK_N;
  const int bid_m = bid / grid_n;
  const int bid_n = bid % grid_n;

  const int off_m = bid_m * BLOCK_M;
  const int off_n = bid_n * BLOCK_N;

  constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;

  // Optional profiling: each warp leader accumulates cycles into 4 fields.
  int64_t prof_wait = 0;
  int64_t prof_issue = 0;
  int64_t prof_other = 0;
  int64_t prof_total_start = 0;
  if constexpr (DO_PROFILE) {
    if (lane_id == 0) prof_total_start = globaltimer();
  }

  // set up smem
  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 B1_size  = BLOCK_N * BLOCK_K / 2;
  constexpr int B2_size  = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;  // always copy 128xBLOCK_K/16
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B1_size + B2_size + SFA_size + SFB_size + SFB_size;

  // set up mbarriers and tmem
  // we have NUM_STAGES mbars for TMA
  //         NUM_STAGES mbars for MMA
  //                  1 mbar  for mainloop
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  // tmem layout:
  // - ACC1: [0, BLOCK_N)
  // - ACC2: [BLOCK_N, 2*BLOCK_N)
  // - SFA : starts at 2*BLOCK_N
  constexpr int ACC1_tmem = 0;
  constexpr int ACC2_tmem = BLOCK_N;
  constexpr int SFA_tmem  = 2 * BLOCK_N;
  constexpr int SF_COLS_PER_K = 4 * (BLOCK_K / MMA_K);  // 4 cols per MMA_K segment
  constexpr int SFB1_tmem = SFA_tmem + SF_COLS_PER_K;
  constexpr int SFB2_tmem = SFB1_tmem + SF_COLS_PER_K;
  constexpr int ALLOC_COLS = BLOCK_N * 4;  // conservative (2 acc + 3 SF tiles + padding)

  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;");  // visible to async proxy
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(ALLOC_COLS));
  }
  __syncthreads();

  constexpr int num_iters = K / BLOCK_K;

  // warp-specialization
  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    // TMA warp
    uint64_t cache_A, cache_B;
    if (M > N) {
      cache_A = EVICT_FIRST;
      cache_B = EVICT_LAST;
    } else {
      // Keep both A and B hot when M<=N; the benchmark working sets fit in B200's L2.
      cache_A = EVICT_LAST;
      cache_B = EVICT_LAST;
    }

    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 B1_smem = A_smem  + A_size;
      const int B2_smem = B1_smem + B1_size;
      const int SFA_smem  = B2_smem + B2_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_size;

      const int off_k = iter_k * BLOCK_K;
      tma_3d_gmem2smem(A_smem,  &A_tmap,  0, off_m, off_k / 256, mbar_addr, cache_A);
      tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
      tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      // layout of SFA is [M/128, rest_k, 32, 4, 4]
      //           SFB is [N/128, rest_k, 32, 4, 4]
      const int rest_k = K / 16 / 4;
      const char *SFA_src  = SFA_ptr  + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB1_src = SFB1_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      tma_gmem2smem(SFA_smem,  SFA_src,  SFA_size,  mbar_addr, cache_A);
      tma_gmem2smem(SFB1_smem, SFB1_src, SFB_size, mbar_addr, cache_B);
      tma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_B);

      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 < num_iters; 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;
      if constexpr (DO_PROFILE) {
        int64_t t0 = globaltimer();
        mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
        if (lane_id == 0) prof_wait += globaltimer() - t0;
      } else {
        mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      }
      int64_t t1 = 0;
      if constexpr (DO_PROFILE) {
        if (lane_id == 0) t1 = globaltimer();
      }
      issue_tma(iter_k, stage_id);
      if constexpr (DO_PROFILE) {
        if (lane_id == 0) prof_issue += globaltimer() - t1;
      }
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    // MMA warp
    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
                              | (1U << 10U)  // btype=E2M1
                              | ((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;
      if constexpr (DO_PROFILE) {
        int64_t t0 = globaltimer();
        mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);
        if (lane_id == 0) prof_wait += globaltimer() - t0;
      } else {
        mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);
      }

      const int A_smem  = smem + stage_id * STAGE_SIZE;
      const int B1_smem = A_smem  + A_size;
      const int B2_smem = B1_smem + B1_size;
      const int SFA_smem  = B2_smem + B2_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_size;

      int64_t t1 = 0;
      if constexpr (DO_PROFILE) {
        if (lane_id == 0) t1 = globaltimer();
      }

      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 SFB1_desc = SF_desc + ((uint64_t)SFB1_smem >> 4ULL);
      const uint64_t SFB2_desc = SF_desc + ((uint64_t)SFB2_smem >> 4ULL);

      // Load all SF upfront (batch loading - proven faster for 1-SM kernel)
      #pragma unroll
      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc  = SFA_desc  + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb1_desc = SFB1_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb2_desc = SFB2_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem  + k * 4, sfa_desc);
        tcgen05_cp_nvfp4(SFB1_tmem + k * 4, sfb1_desc);
        tcgen05_cp_nvfp4(SFB2_tmem + k * 4, sfb2_desc);
      }

      const int b_half = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);  // 0 or 2 for BLOCK_N=64
      // BLOCK_K is fixed to 256 for this kernel: 4 MMA_K steps per stage.
      uint64_t a_desc  = make_desc_AB(A_smem);
      uint64_t b1_desc = make_desc_AB(B1_smem);
      uint64_t b2_desc = make_desc_AB(B2_smem);

      int scale_A_tmem  = SFA_tmem;
      int scale_B1_tmem = SFB1_tmem + b_half;
      int scale_B2_tmem = SFB2_tmem + b_half;

      // k2 = 0: allow zero-init on iter_k==0 via enable_input_d=0.
      tcgen05_mma_nvfp4_collector_a_fill(a_desc, b1_desc, i_desc, ACC1_tmem, scale_A_tmem, scale_B1_tmem, iter_k);
      tcgen05_mma_nvfp4_collector_a_lastuse(a_desc, b2_desc, i_desc, ACC2_tmem, scale_A_tmem, scale_B2_tmem, iter_k);

      #pragma unroll
      for (int k2 = 1; k2 < 256 / MMA_K; k2++) {
        a_desc += 2;
        b1_desc += 2;
        b2_desc += 2;
        scale_A_tmem += 4;
        scale_B1_tmem += 4;
        scale_B2_tmem += 4;
        tcgen05_mma_nvfp4_collector_a_fill(a_desc, b1_desc, i_desc, ACC1_tmem, scale_A_tmem, scale_B1_tmem, 1);
        tcgen05_mma_nvfp4_collector_a_lastuse(a_desc, b2_desc, i_desc, ACC2_tmem, scale_A_tmem, scale_B2_tmem, 1);
      }

      asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                  :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
      if constexpr (DO_PROFILE) {
        if (lane_id == 0) prof_issue += globaltimer() - t1;
      }
    }

    asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                :: "r"(mainloop_mbar_addr) : "memory");
  } else if (tid < BLOCK_M) {
    // epilogue warps
    if constexpr (DO_PROFILE) {
      int64_t t0 = globaltimer();
      mbarrier_wait(mainloop_mbar_addr, 0);
      if (lane_id == 0) prof_wait += globaltimer() - t0;
    } else {
      mbarrier_wait(mainloop_mbar_addr, 0);
    }
    asm volatile("tcgen05.fence::after_thread_sync;");

    int64_t t1 = 0;
    if constexpr (DO_PROFILE) {
      if (lane_id == 0) t1 = globaltimer();
    }

    // C is N-major
    for (int m = 0; m < 32 / 16; m++) {
      float tmp1[BLOCK_N / 2];
      float tmp2[BLOCK_N / 2];
      if constexpr (BLOCK_N == 128) {
        tcgen05_ld_16x256bx16(tmp1, warp_id * 32 + m * 16, 0);
        tcgen05_ld_16x256bx16(tmp2, warp_id * 32 + m * 16, ACC2_tmem);
      } else {
        tcgen05_ld_16x256bx8(tmp1, warp_id * 32 + m * 16, 0);
        tcgen05_ld_16x256bx8(tmp2, warp_id * 32 + m * 16, ACC2_tmem);
      }
      asm volatile("tcgen05.wait::ld.sync.aligned;");

      #pragma unroll
      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;

        float x0 = tmp1[i * 4 + 0];
        float x1 = tmp1[i * 4 + 1];
        float x2 = tmp1[i * 4 + 2];
        float x3 = tmp1[i * 4 + 3];
        float y0 = tmp2[i * 4 + 0];
        float y1 = tmp2[i * 4 + 1];
        float y2 = tmp2[i * 4 + 2];
        float y3 = tmp2[i * 4 + 3];

        // Safe silu: silu(x) * y = (x * y) / (1 + exp(-x)).
        // ILP: issue exp() first; then reciprocals; then final muls.
        float xy0 = x0 * y0;
        float xy1 = x1 * y1;
        float xy2 = x2 * y2;
        float xy3 = x3 * y3;

        float e0 = expf(-x0);
        float e1 = expf(-x1);
        float e2 = expf(-x2);
        float e3 = expf(-x3);

        float sig0 = 1.0f / (1.0f + e0);
        float sig1 = 1.0f / (1.0f + e1);
        float sig2 = 1.0f / (1.0f + e2);
        float sig3 = 1.0f / (1.0f + e3);

        float v00 = xy0 * sig0;
        float v01 = xy1 * sig1;
        float v10 = xy2 * sig2;
        float v11 = xy3 * sig3;

        reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] = __float22half2_rn({v00, v01});
        reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] = __float22half2_rn({v10, v11});
      }
    }

    asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
    if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(ALLOC_COLS));

    if constexpr (DO_PROFILE) {
      if (lane_id == 0) {
        prof_other += globaltimer() - t1;
      }
    }
  }

  if constexpr (DO_PROFILE) {
    if (profile_ptr && lane_id == 0) {
      int64_t *out = profile_ptr + (int64_t)(bid * NUM_WARPS + warp_id) * PROF_FIELDS;
      out[PROF_WAIT] = prof_wait;
      out[PROF_ISSUE] = prof_issue;
      out[PROF_OTHER] = prof_other;
      out[PROF_TOTAL] = globaltimer() - prof_total_start;
    }
  }
}

template <
  int K,
  int BLOCK_M,
  int BLOCK_N,
  int BLOCK_K,
  int NUM_STAGES
>
__global__
__cluster_dims__(2, 1, 1)
__launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void kernel_2sm(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B1_tmap,
  const __grid_constant__ CUtensorMap B2_tmap,
  const __grid_constant__ CUtensorMap SFA_tmap,
  const __grid_constant__ CUtensorMap SFB1_tmap,
  const __grid_constant__ CUtensorMap SFB2_tmap,
  half *C_ptr,
  int M, int N
) {
  constexpr int CTA_GROUP = 2;
  static_assert((BLOCK_N % CTA_GROUP) == 0, "BLOCK_N must be divisible by CTA_GROUP");

  const int tid = threadIdx.x;
  const int bid = blockIdx.x;

  const int lane_id = tid % WARP_SIZE;
  const int warp_id = tid / WARP_SIZE;

  int cta_rank;
  asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank));

  const int grid_n = N / BLOCK_N;
  constexpr int GROUP_M = CTA_GROUP;
  const int bid_m = bid / (grid_n * GROUP_M) * GROUP_M + (bid % GROUP_M);
  const int bid_n = (bid / GROUP_M) % 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;

  // set up smem
  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_N_LOCAL = BLOCK_N / CTA_GROUP;
  constexpr int B1_size  = B_N_LOCAL * BLOCK_K / 2;
  constexpr int B2_size  = B_N_LOCAL * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;  // always copy 128xBLOCK_K/16
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B1_size + B2_size + SFA_size + SFB_size + SFB_size;

  // set up mbarriers and tmem
  // we have NUM_STAGES mbars for TMA
  //         NUM_STAGES mbars for MMA
  //                  1 mbar  for mainloop
  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  // tmem layout:
  // - ACC1: [0, BLOCK_N)
  // - ACC2: [BLOCK_N, 2*BLOCK_N)
  // - SFA : starts at 2*BLOCK_N
  constexpr int ACC1_tmem = 0;
  constexpr int ACC2_tmem = BLOCK_N;
  constexpr int SFA_tmem  = 2 * BLOCK_N;
  constexpr int SF_COLS_PER_K = 4 * (BLOCK_K / MMA_K);  // 4 cols per MMA_K segment
  constexpr int SFB1_tmem = SFA_tmem + SF_COLS_PER_K;
  constexpr int SFB2_tmem = SFB1_tmem + SF_COLS_PER_K;
  constexpr int ALLOC_COLS = BLOCK_N * 4;  // conservative (2 acc + 3 SF tiles + padding)

  if (warp_id == 0 && elect_sync()) {
    for (int i = 0; i < NUM_STAGES; i++) {
      mbarrier_init(tma_mbar_addr + i * 8, CTA_GROUP);  // both CTAs report TMA to CTA0 only
      mbarrier_init(mma_mbar_addr + i * 8, 1);          // CTA0 reports MMA to BOTH CTAs (multicast)
    }
    mbarrier_init(mainloop_mbar_addr, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");  // visible to async proxy
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(ALLOC_COLS));
  }
  // visible to all threads in a cluster
  asm volatile("barrier.cluster.arrive.release.aligned;");
  asm volatile("barrier.cluster.wait.acquire.aligned;");

  constexpr int num_iters = K / BLOCK_K;

  // warp-specialization
  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    // TMA warp
    uint64_t cache_A, cache_B;
    if (M > N) {
      cache_A = EVICT_FIRST;
      cache_B = EVICT_LAST;
    } else {
      // Keep both A and B hot when M<=N; the benchmark working sets fit in B200's L2.
      cache_A = EVICT_LAST;
      cache_B = EVICT_LAST;
    }

    auto issue_tma = [&](int iter_k, int stage_id) {
      int mbar_addr = tma_mbar_addr + stage_id * 8;
      mbar_addr &= 0xFEFFFFFF;  // CTA0 barrier

      const int A_smem  = smem + stage_id * STAGE_SIZE;
      const int B1_smem = A_smem  + A_size;
      const int B2_smem = B1_smem + B1_size;
      const int SFA_smem  = B2_smem + B2_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_size;

      const int off_k = iter_k * BLOCK_K;
      tma_3d_gmem2smem<CTA_GROUP>(A_smem,  &A_tmap,  0, off_m, off_k / 256, mbar_addr, cache_A);
      tma_3d_gmem2smem<CTA_GROUP>(B1_smem, &B1_tmap, 0, off_n + cta_rank * B_N_LOCAL, off_k / 256, mbar_addr, cache_B);
      tma_3d_gmem2smem<CTA_GROUP>(B2_smem, &B2_tmap, 0, off_n + cta_rank * B_N_LOCAL, off_k / 256, mbar_addr, cache_B);

      // SFA: rest_m = M/128, rest_k = K/64. Tile loads 4 rest_k chunks (BLOCK_K=256 => 4*64).
      const int sfa_y = off_m / 128;
      const int sfb_y = off_n / 128;
      const int sf_z  = off_k / 64;
      tma_3d_gmem2smem<CTA_GROUP>(SFA_smem,  &SFA_tmap,  0, sfa_y, sf_z, mbar_addr, cache_A);
      tma_3d_gmem2smem<CTA_GROUP>(SFB1_smem, &SFB1_tmap, 0, sfb_y, sf_z, mbar_addr, cache_B);
      tma_3d_gmem2smem<CTA_GROUP>(SFB2_smem, &SFB2_tmap, 0, sfb_y, sf_z, mbar_addr, cache_B);

      // Cutlass 2-SM TMA style: only one CTA sets the expected transaction bytes for
      // the shared barrier (which receives completion bytes from both CTAs).
      if (cta_rank == 0) {
        asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;"
                    :: "r"(mbar_addr), "r"(STAGE_SIZE * CTA_GROUP) : "memory");
      } else {
        asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" :: "r"(mbar_addr) : "memory");
      }
    };

    for (int iter_k = 0; iter_k < NUM_STAGES && iter_k < num_iters; 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 ((cta_rank == 0) && (warp_id == NUM_WARPS - 1) && elect_sync()) {
    // MMA warp (CTA0 only)
    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128 * CTA_GROUP;
    constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
                              | (1U << 10U)  // btype=E2M1
                              | ((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);
      asm volatile("tcgen05.fence::after_thread_sync;");

      const int A_smem  = smem + stage_id * STAGE_SIZE;
      const int B1_smem = A_smem  + A_size;
      const int B2_smem = B1_smem + B1_size;
      const int SFA_smem  = B2_smem + B2_size;
      const int SFB1_smem = SFA_smem + SFA_size;
      const int SFB2_smem = SFB1_smem + SFB_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 SFB1_desc = SF_desc + ((uint64_t)SFB1_smem >> 4ULL);
      const uint64_t SFB2_desc = SF_desc + ((uint64_t)SFB2_smem >> 4ULL);

      constexpr int NUM_SF = BLOCK_K / MMA_K;
      constexpr int SF_PREFETCH = 1;
      auto cp_sf = [&](int k_sf) {
        uint64_t sfa_desc  = SFA_desc  + (uint64_t)k_sf * (512ULL >> 4ULL);
        uint64_t sfb1_desc = SFB1_desc + (uint64_t)k_sf * (512ULL >> 4ULL);
        uint64_t sfb2_desc = SFB2_desc + (uint64_t)k_sf * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4<CTA_GROUP>(SFA_tmem  + k_sf * 4, sfa_desc);
        tcgen05_cp_nvfp4<CTA_GROUP>(SFB1_tmem + k_sf * 4, sfb1_desc);
        tcgen05_cp_nvfp4<CTA_GROUP>(SFB2_tmem + k_sf * 4, sfb2_desc);
      };

      #pragma unroll
      for (int k_sf = 0; k_sf < SF_PREFETCH; k_sf++) {
        if (k_sf < NUM_SF) cp_sf(k_sf);
      }

      const int b_half = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);  // 0 or 2 for BLOCK_N=64

      // BLOCK_K is fixed to 256 for this kernel.
      uint64_t a_desc  = make_desc_AB(A_smem);
      uint64_t b1_desc = make_desc_AB(B1_smem);
      uint64_t b2_desc = make_desc_AB(B2_smem);

      int scale_A_tmem  = SFA_tmem;
      int scale_B1_tmem = SFB1_tmem + b_half;
      int scale_B2_tmem = SFB2_tmem + b_half;

      // k2 = 0
      tcgen05_mma_nvfp4_collector_a_fill<CTA_GROUP>(a_desc, b1_desc, i_desc, ACC1_tmem, scale_A_tmem, scale_B1_tmem, iter_k);
      tcgen05_mma_nvfp4_collector_a_lastuse<CTA_GROUP>(a_desc, b2_desc, i_desc, ACC2_tmem, scale_A_tmem, scale_B2_tmem, iter_k);
      {
        int k_pf = SF_PREFETCH;
        if (k_pf < NUM_SF) cp_sf(k_pf);
      }

      #pragma unroll
      for (int k2 = 1; k2 < 256 / MMA_K; k2++) {
        a_desc += 2;
        b1_desc += 2;
        b2_desc += 2;
        scale_A_tmem += 4;
        scale_B1_tmem += 4;
        scale_B2_tmem += 4;
        tcgen05_mma_nvfp4_collector_a_fill<CTA_GROUP>(a_desc, b1_desc, i_desc, ACC1_tmem, scale_A_tmem, scale_B1_tmem, 1);
        tcgen05_mma_nvfp4_collector_a_lastuse<CTA_GROUP>(a_desc, b2_desc, i_desc, ACC2_tmem, scale_A_tmem, scale_B2_tmem, 1);

        int k_pf = k2 + SF_PREFETCH;
        if (k_pf < NUM_SF) cp_sf(k_pf);
      }

      // Signal both CTAs' local MMA mbarriers (multicast within the cluster).
      constexpr int16_t cta_mask = (1 << CTA_GROUP) - 1;
      asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
                  :: "r"(mma_mbar_addr + stage_id * 8), "h"(cta_mask) : "memory");
    }

    constexpr int16_t cta_mask = (1 << CTA_GROUP) - 1;
    asm volatile("tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
                :: "r"(mainloop_mbar_addr), "h"(cta_mask) : "memory");
	  } else if (tid < BLOCK_M) {
	    // epilogue warps
	    mbarrier_wait(mainloop_mbar_addr, 0);
	    asm volatile("tcgen05.fence::after_thread_sync;");

	    // C is N-major
	    // NOTE: kernel_2sm uses MMA_M=256 (layout A). Prefer 32x32b loads and 16B stores (int4).
	    const int trow_base = cta_rank * 128 + warp_id * 32;
	    const int row = off_m + warp_id * 32 + lane_id;

	    #pragma unroll
	    for (int n = 0; n < BLOCK_N / 8; n++) {
	      float acc1[8];
	      float acc2[8];
	      int addr1 = (trow_base << 16) | (ACC1_tmem + n * 8);
	      int addr2 = (trow_base << 16) | (ACC2_tmem + n * 8);

	      asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 "
	                  "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
	                  : "=f"(acc1[0]), "=f"(acc1[1]), "=f"(acc1[2]), "=f"(acc1[3]),
	                    "=f"(acc1[4]), "=f"(acc1[5]), "=f"(acc1[6]), "=f"(acc1[7])
	                  : "r"(addr1));
		      asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 "
		                  "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
		                  : "=f"(acc2[0]), "=f"(acc2[1]), "=f"(acc2[2]), "=f"(acc2[3]),
		                    "=f"(acc2[4]), "=f"(acc2[5]), "=f"(acc2[6]), "=f"(acc2[7])
		                  : "r"(addr2));
		      asm volatile("tcgen05.wait::ld.sync.aligned;");

		      half2 out_h2[4];
		      #pragma unroll
		      for (int j = 0; j < 4; j++) {
		        float x0 = acc1[j * 2 + 0];
		        float x1 = acc1[j * 2 + 1];
		        float y0 = acc2[j * 2 + 0];
		        float y1 = acc2[j * 2 + 1];

		        // Safe silu: silu(x) * y = (x * y) / (1 + exp(-x)).
		        float xy0 = x0 * y0;
		        float xy1 = x1 * y1;

		        float e0 = expf(-x0);
		        float e1 = expf(-x1);

		        float sig0 = 1.0f / (1.0f + e0);
		        float sig1 = 1.0f / (1.0f + e1);

		        float v0 = xy0 * sig0;
		        float v1 = xy1 * sig1;

		        out_h2[j] = __float22half2_rn({v0, v1});
		      }

		      half *out_ptr = C_ptr + row * N + (off_n + n * 8);
		      reinterpret_cast<int4 *>(out_ptr)[0] = reinterpret_cast<int4 *>(out_h2)[0];
		    }

	    asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
	    if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(ALLOC_COLS));
	  }
}

at::Tensor dual_gemm(
  const at::Tensor& A,
  const at::Tensor& B1,
  const at::Tensor& B2,
  const at::Tensor& SFA,
  const at::Tensor& SFB1,
  const at::Tensor& SFB2,
        at::Tensor& C
) {
  TORCH_CHECK(A.is_cuda() && B1.is_cuda() && B2.is_cuda() && SFA.is_cuda() && SFB1.is_cuda() && SFB2.is_cuda() && C.is_cuda(), "all tensors must be CUDA");
  TORCH_CHECK(A.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2, "A must be torch.float4_e2m1fn_x2");
  TORCH_CHECK(B1.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2 && B2.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2, "B must be torch.float4_e2m1fn_x2");
  TORCH_CHECK(C.dtype() == at::kHalf, "C must be fp16");

  const int M = A.size(0);
  const int N = B1.size(0);
  const int K = A.size(1) * 2;
  TORCH_CHECK(B1.size(1) * 2 == K && B2.size(1) * 2 == K, "K mismatch");
  TORCH_CHECK(B2.size(0) == N, "N mismatch");
  TORCH_CHECK(A.size(2) == 1 && B1.size(2) == 1 && B2.size(2) == 1 && C.size(2) == 1, "L must be 1 for this kernel");
  TORCH_CHECK(M % 128 == 0 && (N % 64) == 0 && (K % 256) == 0, "shape constraints not met");

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

  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_K = 256;
  constexpr int NUM_STAGES_1SM = 4;
  constexpr int NUM_STAGES_2SM_BN128 = 5;
  // BN64: keep 5. (Stages=3 was a large regression on the long-K case; see dev.md.)
  constexpr int NUM_STAGES_2SM_BN64 = 5;

  // Heuristic: prefer wider N tiles only when it doesn't starve the GPU (B200 has 148 SMs).
  int block_n = 64;
  // Empirical for the benchmark set: (M=512, N=3072, K=7168) strongly prefers 128-wide tiles.
  if (M >= 512 && N == 3072) block_n = 128;
  if ((N % 128) == 0) {
    int grid64 = (M / BLOCK_M) * (N / 64);
    int grid128 = (M / BLOCK_M) * (N / 128);
    // Use 128-wide tiles only when we still have enough CTAs to fill most SMs.
    if (grid128 >= 120 && grid128 * 2 >= grid64) block_n = 128;
  }

  // 2-SM (2-CTA cluster) path:
  // - Primary: BLOCK_N=128 and M>=512 (good cluster fill).
  // - Experiment: BLOCK_N=64 for (M=256,N=4096) to reduce duplicated B traffic across the 2 M-tiles
  //   while keeping enough clusters to avoid severe underfill.
  const bool use_2sm = (block_n == 128 && M >= 512) ||
                       (block_n == 64 && M == 256 && (N == 4096 || N == 3072));
  if (use_2sm) {
    CUtensorMap A_tmap;
    init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
    CUtensorMap B1_tmap, B2_tmap;
    init_AB_tmap(&B1_tmap, B1_ptr, N, K, block_n / 2, BLOCK_K);
    init_AB_tmap(&B2_tmap, B2_ptr, N, K, block_n / 2, BLOCK_K);

    const int rest_m = M / 128;
    const int rest_n = N / 128;
    const int rest_k = K / 64;
    CUtensorMap SFA_tmap, SFB1_tmap, SFB2_tmap;
    init_SF_tmap(&SFA_tmap,  SFA_ptr,  rest_m, rest_k, 1, BLOCK_K / 64);
    init_SF_tmap(&SFB1_tmap, SFB1_ptr, rest_n, rest_k, 1, BLOCK_K / 64);
    init_SF_tmap(&SFB2_tmap, SFB2_ptr, rest_n, rest_k, 1, BLOCK_K / 64);

    dim3 grid((M / BLOCK_M) * (N / block_n));
    int tb_size = BLOCK_M + 2 * WARP_SIZE;
    int AB_bytes = (BLOCK_M + block_n) * (BLOCK_K / 2);
    int SF_bytes = 128 * (BLOCK_K / 16) * 3;
    int num_stages_2sm = (block_n == 128) ? NUM_STAGES_2SM_BN128 : NUM_STAGES_2SM_BN64;
    int smem_size = (AB_bytes + SF_bytes) * num_stages_2sm;

    if (K == 4096) {
      if (block_n == 128) {
        auto this_kernel = kernel_2sm<4096, BLOCK_M, 128, BLOCK_K, NUM_STAGES_2SM_BN128>;
        if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
        cudaFuncSetAttribute(this_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
        this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap, C_ptr, M, N);
      } else {
        auto this_kernel = kernel_2sm<4096, BLOCK_M, 64, BLOCK_K, NUM_STAGES_2SM_BN64>;
        if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
        cudaFuncSetAttribute(this_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
        this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap, C_ptr, M, N);
      }
    } else if (K == 7168) {
      if (block_n == 128) {
        auto this_kernel = kernel_2sm<7168, BLOCK_M, 128, BLOCK_K, NUM_STAGES_2SM_BN128>;
        if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
        cudaFuncSetAttribute(this_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
        this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap, C_ptr, M, N);
      } else {
        auto this_kernel = kernel_2sm<7168, BLOCK_M, 64, BLOCK_K, NUM_STAGES_2SM_BN64>;
        if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
        cudaFuncSetAttribute(this_kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
        this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_tmap, SFB1_tmap, SFB2_tmap, C_ptr, M, N);
      }
    } else {
      TORCH_CHECK(false, "unsupported K=", K, " (expected 4096 or 7168)");
    }
    return C;
  }

  // 1-SM (single CTA) path.
  CUtensorMap A_tmap;
  init_AB_tmap(&A_tmap,  A_ptr,  M, K, BLOCK_M, BLOCK_K);
  CUtensorMap B1_tmap, B2_tmap;
  init_AB_tmap(&B1_tmap, B1_ptr, N, K, block_n, BLOCK_K);
  init_AB_tmap(&B2_tmap, B2_ptr, N, K, block_n, BLOCK_K);

  dim3 grid((M / BLOCK_M) * (N / block_n));
  int tb_size = BLOCK_M + 2 * WARP_SIZE;
  int AB_bytes = (BLOCK_M + 2 * block_n) * (BLOCK_K / 2);
  int SF_bytes = 128 * (BLOCK_K / 16) * 3;

  if (K == 4096) {
    if (block_n == 128) {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES_1SM;
      auto this_kernel = kernel<4096, BLOCK_M, 128, BLOCK_K, NUM_STAGES_1SM, false>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, nullptr);
    } else {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES_1SM;
      auto this_kernel = kernel<4096, BLOCK_M, 64, BLOCK_K, NUM_STAGES_1SM, false>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, nullptr);
    }
  } else if (K == 7168) {
    if (block_n == 128) {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES_1SM;
      auto this_kernel = kernel<7168, BLOCK_M, 128, BLOCK_K, NUM_STAGES_1SM, false>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, nullptr);
    } else {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES_1SM;
      auto this_kernel = kernel<7168, BLOCK_M, 64, BLOCK_K, NUM_STAGES_1SM, false>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, nullptr);
    }
  } else {
    TORCH_CHECK(false, "unsupported K=", K, " (expected 4096 or 7168)");
  }

  return C;
}

at::Tensor dual_gemm_profile(
  const at::Tensor& A,
  const at::Tensor& B1,
  const at::Tensor& B2,
  const at::Tensor& SFA,
  const at::Tensor& SFB1,
  const at::Tensor& SFB2,
        at::Tensor& C,
        at::Tensor& profile
) {
  TORCH_CHECK(profile.is_cuda() && profile.scalar_type() == c10::ScalarType::Long, "profile must be a CUDA int64 tensor");
  TORCH_CHECK(profile.is_contiguous(), "profile must be contiguous");
  auto profile_ptr = reinterpret_cast<int64_t *>(profile.data_ptr());

  TORCH_CHECK(A.is_cuda() && B1.is_cuda() && B2.is_cuda() && SFA.is_cuda() && SFB1.is_cuda() && SFB2.is_cuda() && C.is_cuda(), "all tensors must be CUDA");
  TORCH_CHECK(A.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2, "A must be torch.float4_e2m1fn_x2");
  TORCH_CHECK(B1.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2 && B2.scalar_type() == c10::ScalarType::Float4_e2m1fn_x2, "B must be torch.float4_e2m1fn_x2");
  TORCH_CHECK(C.dtype() == at::kHalf, "C must be fp16");

  const int M = A.size(0);
  const int N = B1.size(0);
  const int K = A.size(1) * 2;
  TORCH_CHECK(B1.size(1) * 2 == K && B2.size(1) * 2 == K, "K mismatch");
  TORCH_CHECK(B2.size(0) == N, "N mismatch");
  TORCH_CHECK(A.size(2) == 1 && B1.size(2) == 1 && B2.size(2) == 1 && C.size(2) == 1, "L must be 1 for this kernel");
  TORCH_CHECK(M % 128 == 0 && (N % 64) == 0 && (K % 256) == 0, "shape constraints not met");

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

  constexpr int BLOCK_M = 128;
  constexpr int BLOCK_K = 256;
  constexpr int NUM_STAGES = 4;

  int block_n = 64;
  if (M >= 512 && N == 3072) block_n = 128;
  if ((N % 128) == 0) {
    int grid64 = (M / BLOCK_M) * (N / 64);
    int grid128 = (M / BLOCK_M) * (N / 128);
    if (grid128 >= 120 && grid128 * 2 >= grid64) block_n = 128;
  }

  CUtensorMap A_tmap, B1_tmap, B2_tmap;
  init_AB_tmap(&A_tmap,  A_ptr,  M, K, BLOCK_M, BLOCK_K);
  init_AB_tmap(&B1_tmap, B1_ptr, N, K, block_n, BLOCK_K);
  init_AB_tmap(&B2_tmap, B2_ptr, N, K, block_n, BLOCK_K);

  dim3 grid((M / BLOCK_M) * (N / block_n));
  int tb_size = BLOCK_M + 2 * WARP_SIZE;
  int AB_bytes = (BLOCK_M + 2 * block_n) * (BLOCK_K / 2);
  int SF_bytes = 128 * (BLOCK_K / 16) * 3;

  // Expect profile to have space for grid.x * NUM_WARPS * PROF_FIELDS entries.
  TORCH_CHECK(profile.numel() >= (int64_t)grid.x * (int64_t)(BLOCK_M / WARP_SIZE + 2) * (int64_t)PROF_FIELDS, "profile buffer too small");

  if (K == 4096) {
    if (block_n == 128) {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES;
      auto this_kernel = kernel<4096, BLOCK_M, 128, BLOCK_K, NUM_STAGES, true>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, profile_ptr);
    } else {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES;
      auto this_kernel = kernel<4096, BLOCK_M, 64, BLOCK_K, NUM_STAGES, true>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, profile_ptr);
    }
  } else if (K == 7168) {
    if (block_n == 128) {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES;
      auto this_kernel = kernel<7168, BLOCK_M, 128, BLOCK_K, NUM_STAGES, true>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, profile_ptr);
    } else {
      int smem_size = (AB_bytes + SF_bytes) * NUM_STAGES;
      auto this_kernel = kernel<7168, BLOCK_M, 64, BLOCK_K, NUM_STAGES, true>;
      if (smem_size > 48'000) cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
      this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N, profile_ptr);
    }
  } else {
    TORCH_CHECK(false, "unsupported K=", K, " (expected 4096 or 7168)");
  }

  return C;
}

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


def _build_extension():
    return load_inline(name="nvfp4_dual_gemm_work_sf_prefetch_codex",
        cpp_sources="",
        cuda_sources=CUDA_SRC_COMMON + CUDA_SRC,
        functions=[],
        extra_cuda_cflags=[
            "-O3",
            "-std=c++17",
            "-gencode=arch=compute_100a,code=sm_100a",
            "--use_fast_math",
            "--expt-relaxed-constexpr",
            "--relocatable-device-code=false",
            "-lineinfo",
        ],
        extra_ldflags=["-lcuda"],
        verbose=False,
        with_cuda=True,
        is_python_module=False,
    )


_EXT_BUILT = False


def _ensure_extension():
    global _EXT_BUILT
    if _EXT_BUILT:
        return
    _build_extension()
    _EXT_BUILT = True


def custom_kernel(data: input_t) -> output_t:
    a, b1, b2, sfa_ref, sfb1_ref, sfb2_ref, sfa_p, sfb1_p, sfb2_p, c = data
    m, k_half, l = a.shape
    n, k_half_b, l_b = b1.shape
    if l_b != l:
        raise ValueError("a/b1 L mismatch")
    if b2.shape != b1.shape:
        raise ValueError("b1/b2 shape mismatch")

    k = int(k_half) * 2
    if int(k_half_b) * 2 != k:
        raise ValueError("a/b1 K mismatch")

    # Fast path: our tcgen05/TMA kernel is specialized for the benchmark shapes.
    if l == 1 and k in (4096, 7168) and (m % 128 == 0) and (n % 64 == 0):
        _ensure_extension()
        torch.ops.nvfp4_dual_gemm.dual_gemm(a, b1, b2, sfa_p, sfb1_p, sfb2_p, c)
        return c

    # Fallback: correctness for arbitrary K (divisible by 256) / L.
    def ceil_div(x: int, y: int) -> int:
        return (x + y - 1) // y

    def to_blocked(x: torch.Tensor) -> torch.Tensor:
        rows, cols = x.shape
        n_row_blocks = ceil_div(rows, 128)
        n_col_blocks = ceil_div(cols, 4)
        blocks = x.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
        rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
        return rearranged.flatten().contiguous()

    for li in range(l):
        scale_a = to_blocked(sfa_ref[:, :, li])
        scale_b1 = to_blocked(sfb1_ref[:, :, li])
        scale_b2 = to_blocked(sfb2_ref[:, :, li])
        acc1 = torch._scaled_mm(
            a[:, :, li],
            b1[:, :, li].transpose(0, 1),
            scale_a,
            scale_b1,
            bias=None,
            out_dtype=torch.float32,
        )
        acc2 = torch._scaled_mm(
            a[:, :, li],
            b2[:, :, li].transpose(0, 1),
            scale_a,
            scale_b2,
            bias=None,
            out_dtype=torch.float32,
        )
        c[:, :, li] = (torch.nn.functional.silu(acc1) * acc2).to(torch.float16)
    return c