Skip to content
KernelIndex
Search⌘K
Sign in

submission 371828

shiyeegao · python · License unknown
Kernel source · 740 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 740 lines, June 9 Researcher Reciprocity License v1.0.

submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-371828?include=source"
interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp8_e4m3, nvfp4

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
NVFP4 dual GEMMsuite of 4 cases
NVIDIA B200
30.8µs
#145 of 161
2026-01-18
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:4e759404aa5fc2fb1b4a149679f55a2e52e3190ebc2a02dc95431316d0aec085
license declaredunknown
license concludedunknown
authorsshiyeegao
imported2026-08-15

Techniques

Extracted from the mirrored source by pattern, never inferred. Each row cites its line.

mbarrier__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
shared-memoryextern __shared__ __align__(1024) char smem_ptr[];
tcgen05asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
tma"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
vector-width = float2reinterpret_cast<float2 *>(C_ptr + (row + 0) * N + col)[0] = float2{tmp[i * 4 + 0], tmp[i * 4 + 1]};

Kernel source

submission.py740 lines



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 <torch/library.h>
#include <ATen/core/Tensor.h>

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

constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;
constexpr uint64_t EVICT_LAST  = 0x14F0000000000000ULL;

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

__device__ __forceinline__ 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"
    "}\n\t"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

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

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

__device__ __forceinline__ 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__ __forceinline__ void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
  asm volatile(
    "cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
    "[%0], [%1, {%2, %3, %4}], [%5], %6;"
    :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
    : "memory"
  );
}

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

__device__ __forceinline__ void tcgen05_mma_nvfp4(
  uint64_t a_desc,
  uint64_t b_desc,
  uint32_t i_desc,
  int scale_A_tmem,
  int scale_B_tmem,
  int enable_input_d
) {
  const int d_tmem = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %6, 0;\n\t"
    "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
    "}\n\t"
    :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
       "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
  );
}

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

template <const char *SHAPE_, const char *NUM_>
__device__ __forceinline__ 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_)
  );
}

__device__ __forceinline__ void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
  tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}

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 init_AB_tmap(
  CUtensorMap *tmap,
  const char *ptr,
  uint64_t global_h, uint64_t global_w,
  uint32_t shared_h, uint32_t shared_w
) {
  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]       = {256, global_h, global_w / 256};
  uint64_t globalStrides[rank-1] = {global_w / 2, 128};
  uint32_t boxDim[rank]          = {256, shared_h, shared_w / 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
  );
  ck_cu(err);
}

template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void gemm_f32_kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B_tmap,
  const char *SFA_ptr,
  const char *SFB_ptr,
  float *C_ptr,
  int M, int N
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.y;

  const int lane_id = tid & 31;
  const int warp_id = tid >> 5;

  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 - bid_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;

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
  constexpr int A_size = BLOCK_M * BLOCK_K / 2;
  constexpr int B_size = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);

  if (warp_id == 0 && elect_sync()) {
    for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
  }
  __syncthreads();

  constexpr int num_iters = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    const uint64_t cache_A = EVICT_LAST;
    const uint64_t cache_B = EVICT_FIRST;

    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const int off_k = iter_k * BLOCK_K;
      tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
      tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      const int rest_k = K / 16 / 4;
      const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
      tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);

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

    constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < PRELOAD; 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) & 1;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)MMA_N >> 3U << 17U) | ((uint32_t)MMA_M >> 7U << 27U);

    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) & 1;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      auto make_desc_AB = [](int addr) -> uint64_t {
        const int SBO = 8 * 128;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
      };
      auto make_desc_SF = [](int addr) -> uint64_t {
        const int SBO = 8 * 16;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
      };

      constexpr uint64_t SF_desc = make_desc_SF(0);
      const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);

      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
        tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
      }

      for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);

          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }

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

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

    for (int mm = 0; mm < 2; mm++) {
      float tmp[BLOCK_N / 2];
      tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);
      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 + mm * 16 + lane_id / 4;
        const int col = off_n + i * 8 + (lane_id & 3) * 2;
        reinterpret_cast<float2 *>(C_ptr + (row + 0) * N + col)[0] = float2{tmp[i * 4 + 0], tmp[i * 4 + 1]};
        reinterpret_cast<float2 *>(C_ptr + (row + 8) * N + col)[0] = float2{tmp[i * 4 + 2], tmp[i * 4 + 3]};
      }
    }

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

template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
void gemm_silu_mul_kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B_tmap,
  const char *SFA_ptr,
  const char *SFB_ptr,
  const float *G1_ptr,
  half *Out_ptr,
  int M, int N
) {
  const int tid = threadIdx.x;
  const int bid = blockIdx.y;

  const int lane_id = tid & 31;
  const int warp_id = tid >> 5;

  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 - bid_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;

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
  constexpr int A_size = BLOCK_M * BLOCK_K / 2;
  constexpr int B_size = BLOCK_N * BLOCK_K / 2;
  constexpr int SFA_size = 128 * BLOCK_K / 16;
  constexpr int SFB_size = 128 * BLOCK_K / 16;
  constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;

  constexpr int SFA_tmem = BLOCK_N;
  constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);

  if (warp_id == 0 && elect_sync()) {
    for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
    asm volatile("fence.mbarrier_init.release.cluster;");
  } else if (warp_id == 1) {
    asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
  }
  __syncthreads();

  constexpr int num_iters = K / BLOCK_K;

  if (warp_id == NUM_WARPS - 2 && elect_sync()) {
    const uint64_t cache_A = EVICT_LAST;
    const uint64_t cache_B = EVICT_FIRST;

    auto issue_tma = [&](int iter_k, int stage_id) {
      const int mbar_addr = tma_mbar_addr + stage_id * 8;
      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      const int off_k = iter_k * BLOCK_K;
      tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
      tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);

      const int rest_k = K / 16 / 4;
      const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
      const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
      tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
      tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);

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

    constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
    for (int iter_k = 0; iter_k < PRELOAD; 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) & 1;
      mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
      issue_tma(iter_k, stage_id);
    }
  } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
    constexpr int MMA_N = BLOCK_N;
    constexpr int MMA_M = 128;
    constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)MMA_N >> 3U << 17U) | ((uint32_t)MMA_M >> 7U << 27U);

    for (int iter_k = 0; iter_k < num_iters; iter_k++) {
      const int stage_id = iter_k % NUM_STAGES;
      const int tma_phase = (iter_k / NUM_STAGES) & 1;
      mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);

      const int A_smem = smem + stage_id * STAGE_SIZE;
      const int B_smem = A_smem + A_size;
      const int SFA_smem = B_smem + B_size;
      const int SFB_smem = SFA_smem + SFA_size;

      auto make_desc_AB = [](int addr) -> uint64_t {
        const int SBO = 8 * 128;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
      };
      auto make_desc_SF = [](int addr) -> uint64_t {
        const int SBO = 8 * 16;
        return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
      };

      constexpr uint64_t SF_desc = make_desc_SF(0);
      const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
      const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);

      for (int k = 0; k < BLOCK_K / MMA_K; k++) {
        uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
        uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
        tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
        tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
      }

      for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
        for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
          uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
          uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);

          const int k_sf = k1 * 4 + k2;
          const int scale_A_tmem = SFA_tmem + k_sf * 4;
          const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);

          const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
          tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
        }

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

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

    for (int mm = 0; mm < 2; mm++) {
      float tmp[BLOCK_N / 2];
      tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);
      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 + mm * 16 + lane_id / 4;
        const int col = off_n + i * 8 + (lane_id & 3) * 2;
        const float2 x0 = reinterpret_cast<const float2 *>(G1_ptr + (row + 0) * N + col)[0];
        const float2 x8 = reinterpret_cast<const float2 *>(G1_ptr + (row + 8) * N + col)[0];
        const float2 y0 = float2{tmp[i * 4 + 0], tmp[i * 4 + 1]};
        const float2 y8 = float2{tmp[i * 4 + 2], tmp[i * 4 + 3]};

        float2 o0;
        float2 o8;
        const float s00 = 1.0f / (1.0f + __expf(-x0.x));
        const float s01 = 1.0f / (1.0f + __expf(-x0.y));
        const float s80 = 1.0f / (1.0f + __expf(-x8.x));
        const float s81 = 1.0f / (1.0f + __expf(-x8.y));
        o0.x = (x0.x * s00) * y0.x;
        o0.y = (x0.y * s01) * y0.y;
        o8.x = (x8.x * s80) * y8.x;
        o8.y = (x8.y * s81) * y8.y;

        reinterpret_cast<half2 *>(Out_ptr + (row + 0) * N + col)[0] = __float22half2_rn(o0);
        reinterpret_cast<half2 *>(Out_ptr + (row + 8) * N + col)[0] = __float22half2_rn(o8);
      }
    }

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

template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
static inline void launch_gemm_f32(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA,
  const at::Tensor& SFB,
  at::Tensor& C
) {
  const int M = (int)A.size(0);
  const int N = (int)B.size(0);

  auto A_ptr = reinterpret_cast<const char *>(A.data_ptr());
  auto B_ptr = reinterpret_cast<const char *>(B.data_ptr());
  auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
  auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
  auto C_ptr = reinterpret_cast<float *>(C.data_ptr());

  CUtensorMap A_tmap, B_tmap;
  init_AB_tmap(&A_tmap, A_ptr, (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
  init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);

  dim3 grid(1, (unsigned)((M / BLOCK_M) * (N / BLOCK_N)));
  const int tb_size = BLOCK_M + 2 * WARP_SIZE;
  const int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
  const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
  const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;

  auto kptr = gemm_f32_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
  if (smem_size > 48'000) cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  kptr<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N);
}

template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
static inline void launch_gemm_silu_mul(
  const at::Tensor& A,
  const at::Tensor& B,
  const at::Tensor& SFA,
  const at::Tensor& SFB,
  const at::Tensor& g1,
  at::Tensor& out
) {
  const int M = (int)A.size(0);
  const int N = (int)B.size(0);

  auto A_ptr = reinterpret_cast<const char *>(A.data_ptr());
  auto B_ptr = reinterpret_cast<const char *>(B.data_ptr());
  auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
  auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
  auto G1_ptr = reinterpret_cast<const float *>(g1.data_ptr());
  auto Out_ptr = reinterpret_cast<half *>(out.data_ptr());

  CUtensorMap A_tmap, B_tmap;
  init_AB_tmap(&A_tmap, A_ptr, (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
  init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);

  dim3 grid(1, (unsigned)((M / BLOCK_M) * (N / BLOCK_N)));
  const int tb_size = BLOCK_M + 2 * WARP_SIZE;
  const int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
  const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
  const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;

  auto kptr = gemm_silu_mul_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
  if (smem_size > 48'000) cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
  kptr<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, G1_ptr, Out_ptr, M, N);
}

__global__ void silu_mul_f32_vec2(const float* __restrict__ x, const float* __restrict__ y, half* __restrict__ out, int64_t n2) {
  const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
  if (idx >= n2) return;
  const float2 fx = reinterpret_cast<const float2*>(x)[idx];
  const float2 fy = reinterpret_cast<const float2*>(y)[idx];
  float2 o;
  const float sx0 = 1.0f / (1.0f + __expf(-fx.x));
  const float sx1 = 1.0f / (1.0f + __expf(-fx.y));
  o.x = (fx.x * sx0) * fy.x;
  o.y = (fx.y * sx1) * fy.y;
  reinterpret_cast<half2*>(out)[idx] = __float22half2_rn(o);
}

static inline void launch_silu_mul_f32(const at::Tensor& g1, const at::Tensor& g2, at::Tensor& out) {
  const int64_t n = out.numel();
  TORCH_CHECK((n & 1) == 0, "n");
  const int64_t n2 = n >> 1;
  const int threads = 256;
  const int blocks = (int)((n2 + threads - 1) / threads);
  silu_mul_f32_vec2<<<blocks, threads>>>(
    reinterpret_cast<const float*>(g1.data_ptr()),
    reinterpret_cast<const float*>(g2.data_ptr()),
    reinterpret_cast<half*>(out.data_ptr()),
    n2
  );
}

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,
  at::Tensor& g1,
  at::Tensor& g2
) {
  TORCH_CHECK(A.is_cuda() && B1.is_cuda() && B2.is_cuda(), "cuda");
  TORCH_CHECK(SFA.is_cuda() && SFB1.is_cuda() && SFB2.is_cuda(), "cuda");
  TORCH_CHECK(out.is_cuda() && g1.is_cuda() && g2.is_cuda(), "cuda");
  TORCH_CHECK(A.dim() == 3 && B1.dim() == 3 && B2.dim() == 3, "dim");
  TORCH_CHECK(out.dim() == 3 && g1.dim() == 3 && g2.dim() == 3, "dim");

  const int64_t M = A.size(0);
  const int64_t Kp = A.size(1);
  const int64_t L = A.size(2);
  const int64_t N = B1.size(0);
  TORCH_CHECK(L == 1, "l");
  TORCH_CHECK(B1.size(1) == Kp && B1.size(2) == L, "b1");
  TORCH_CHECK(B2.size(1) == Kp && B2.size(2) == L, "b2");
  TORCH_CHECK(out.size(0) == M && out.size(1) == N && out.size(2) == L, "out");
  TORCH_CHECK(g1.size(0) == M && g1.size(1) == N && g1.size(2) == L, "g1");
  TORCH_CHECK(g2.size(0) == M && g2.size(1) == N && g2.size(2) == L, "g2");

  TORCH_CHECK((M % 128) == 0, "m");
  TORCH_CHECK((N % 64) == 0, "n");

  const int K = (int)(Kp * 2);
  if (K == 7168) {
    launch_gemm_f32<7168, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    if (M == 256 && N == 4096) {
      launch_gemm_silu_mul<7168, 128, 64, 256, 8>(A, B2, SFA, SFB2, g1, out);
    } else {
      launch_gemm_f32<7168, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
      launch_silu_mul_f32(g1, g2, out);
    }
  } else if (K == 4096) {
    launch_gemm_f32<4096, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<4096, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else if (K == 2304) {
    launch_gemm_f32<2304, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<2304, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else if (K == 2048) {
    launch_gemm_f32<2048, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<2048, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else if (K == 1536) {
    launch_gemm_f32<1536, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<1536, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else if (K == 512) {
    launch_gemm_f32<512, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<512, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else if (K == 256) {
    launch_gemm_f32<256, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
    launch_gemm_f32<256, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
    launch_silu_mul_f32(g1, g2, out);
  } else {
    TORCH_CHECK(false, "k ", K);
  }

  return out;
}

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


_loaded = False


def _load():
    global _loaded
    if _loaded:
        return
    load_inline(
        name="nvfp4_dual_ext_tc_v3",
        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",
            "-lineinfo",
        ],
        extra_ldflags=["-lcuda"],
        verbose=False,
        is_python_module=False,
        no_implicit_headers=True,
    )
    _loaded = True


_buf_cache = {}


def _get_buf(tag, shape, device):
    key = (tag, shape, device)
    t = _buf_cache.get(key)
    if t is None or t.shape != shape or t.device != device:
        t = torch.empty(shape, device=device, dtype=torch.float32)
        _buf_cache[key] = t
    return t


def custom_kernel(data):
    _load()
    a, b1, b2, _sfa, _sfb1, _sfb2, sfa_p, sfb1_p, sfb2_p, c = data
    g1 = _get_buf(1, c.shape, a.device)
    g2 = _get_buf(2, c.shape, a.device)
    return torch.ops.nvfp4_dual_lib.fused(a, b1, b2, sfa_p, sfb1_p, sfb2_p, c, g1, g2)


__all__ = ["custom_kernel"]
scrolls · 740 lines total

Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0

Changes from previous submission

Against this author's previous submission submission 242245.

- #!POPCORN leaderboard nvfp4_dual_gemm
- #!POPCORN gpu NVIDIA
   
+  
+  
  import torch
  from torch.utils.cpp_extension import load_inline
   
  ⋯ 10 unchanged lines
  constexpr int WARP_SIZE = 32;
  constexpr int MMA_K = 64;
   
- constexpr uint64_t EVICT_NORMAL = 0x1000000000000000ULL;
- constexpr uint64_t EVICT_FIRST  = 0x12F0000000000000ULL;
- constexpr uint64_t EVICT_LAST   = 0x14F0000000000000ULL;
+ constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;
+ constexpr uint64_t EVICT_LAST  = 0x14F0000000000000ULL;
   
  __device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }
   
  ⋯ 52 unchanged lines
  }
   
  __device__ __forceinline__ void tcgen05_mma_nvfp4(
-   int d_tmem,
    uint64_t a_desc,
    uint64_t b_desc,
    uint32_t i_desc,
  ⋯ 1 unchanged lines
    int scale_B_tmem,
    int enable_input_d
  ) {
+   const int d_tmem = 0;
    asm volatile(
      "{\n\t"
      ".reg .pred p;\n\t"
  ⋯ 6 unchanged lines
  }
   
  struct SHAPE { static constexpr char _16x256b[] = ".16x256b"; };
- struct NUM { static constexpr char x4[] = ".x4"; static constexpr char x8[] = ".x8"; };
+ struct NUM { static constexpr char x8[] = ".x8"; };
   
  template <const char *SHAPE_, const char *NUM_>
- __device__ __forceinline__ void tcgen05_ld_16regs(float *tmp, int row, int col) {
-   asm volatile(
-     "tcgen05.ld.sync.aligned%17%18.b32 "
-     "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
-     "  %8,  %9, %10, %11, %12, %13, %14, %15}, [%16];"
-     : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
-       "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15])
-     : "r"((row << 16) | col), "C"(SHAPE_), "C"(NUM_)
-   );
- }
-  
- template <const char *SHAPE_, const char *NUM_>
  __device__ __forceinline__ void tcgen05_ld_32regs(float *tmp, int row, int col) {
    asm volatile(
      "tcgen05.ld.sync.aligned%33%34.b32 "
  ⋯ 9 unchanged lines
    );
  }
   
- __device__ __forceinline__ void tcgen05_ld_16x256bx4(float *tmp, int row, int col) {
-   tcgen05_ld_16regs<SHAPE::_16x256b, NUM::x4>(tmp, row, col);
- }
-  
  __device__ __forceinline__ void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
    tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
  }
  ⋯ 34 unchanged lines
    ck_cu(err);
  }
   
- struct TmapCacheKey {
-   const char *ptr;
-   uint64_t global_h;
-   uint64_t global_w;
-   uint32_t shared_h;
-   uint32_t shared_w;
- };
-  
- struct TmapCacheEntry {
-   TmapCacheKey key;
-   CUtensorMap tmap;
-   bool valid;
- };
-  
- static inline const CUtensorMap& get_or_init_tmap(TmapCacheEntry &cache, const char *ptr, uint64_t global_h, uint64_t global_w, uint32_t shared_h, uint32_t shared_w) {
-   if (cache.valid &&
-       cache.key.ptr == ptr &&
-       cache.key.global_h == global_h &&
-       cache.key.global_w == global_w &&
-       cache.key.shared_h == shared_h &&
-       cache.key.shared_w == shared_w) {
-     return cache.tmap;
-   }
-   init_AB_tmap(&cache.tmap, ptr, global_h, global_w, shared_h, shared_w);
-   cache.key.ptr = ptr;
-   cache.key.global_h = global_h;
-   cache.key.global_w = global_w;
-   cache.key.shared_h = shared_h;
-   cache.key.shared_w = shared_w;
-   cache.valid = true;
-   return cache.tmap;
- }
-  
  template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
- __global__ __launch_bounds__(BLOCK_M + 3 * WARP_SIZE)
- void dual_gemm_silu_mul_kernel(
+ __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
+ void gemm_f32_kernel(
    const __grid_constant__ CUtensorMap A_tmap,
-   const __grid_constant__ CUtensorMap B1_tmap,
-   const __grid_constant__ CUtensorMap B2_tmap,
+   const __grid_constant__ CUtensorMap B_tmap,
    const char *SFA_ptr,
-   const char *SFB1_ptr,
-   const char *SFB2_ptr,
-   half *Out_ptr,
+   const char *SFB_ptr,
+   float *C_ptr,
    int M, int N
  ) {
    const int tid = threadIdx.x;
  ⋯ 2 unchanged lines
    const int lane_id = tid & 31;
    const int warp_id = tid >> 5;
   
+   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 - bid_m * grid_n;
  ⋯ 1 unchanged lines
    const int off_m = bid_m * BLOCK_M;
    const int off_n = bid_n * BLOCK_N;
   
-   constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 3;
+   constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
   
    extern __shared__ __align__(1024) char smem_ptr[];
    const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
-  
-   constexpr int A_size   = BLOCK_M * BLOCK_K / 2;
-   constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
+   constexpr int A_size = BLOCK_M * BLOCK_K / 2;
+   constexpr int B_size = BLOCK_N * BLOCK_K / 2;
    constexpr int SFA_size = 128 * BLOCK_K / 16;
    constexpr int SFB_size = 128 * BLOCK_K / 16;
-   constexpr int STAGE_SIZE = A_size + 2 * B_size + SFA_size + 2 * SFB_size;
+   constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
   
    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  ⋯ 1 unchanged lines
    const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
    const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
   
-   constexpr int D1_tmem = 0;
-   constexpr int D2_tmem = BLOCK_N;
-   constexpr int SFA1_tmem  = 2 * BLOCK_N;
-   constexpr int SFB1_tmem  = SFA1_tmem + 4 * (BLOCK_K / MMA_K);
-   constexpr int SFA2_tmem  = SFB1_tmem + 4 * (BLOCK_K / MMA_K);
-   constexpr int SFB2_tmem  = SFA2_tmem + 4 * (BLOCK_K / MMA_K);
+   constexpr int SFA_tmem = BLOCK_N;
+   constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
   
    if (warp_id == 0 && elect_sync()) {
-     for (int i = 0; i < NUM_STAGES; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
-     for (int i = 0; i < NUM_STAGES; i++) mbarrier_init(mma_mbar_addr + i * 8, 2);
-     mbarrier_init(mainloop_mbar_addr, 2);
+     for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
      asm volatile("fence.mbarrier_init.release.cluster;");
    } else if (warp_id == 1) {
-     asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 4));
+     asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
    }
    __syncthreads();
   
    constexpr int num_iters = K / BLOCK_K;
   
-   if (warp_id == NUM_WARPS - 3 && elect_sync()) {
-     const uint64_t cache_A = EVICT_FIRST;
+   if (warp_id == NUM_WARPS - 2 && elect_sync()) {
+     const uint64_t cache_A = EVICT_LAST;
      const uint64_t cache_B = EVICT_FIRST;
   
-     constexpr int rest_k = K / 16 / 4;
-  
      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 + B_size;
-       const int SFA_smem = B2_smem + B_size;
-       const int SFB1_smem = SFA_smem + SFA_size;
-       const int SFB2_smem = SFB1_smem + SFB_size;
+       const int A_smem = smem + stage_id * STAGE_SIZE;
+       const int B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_size;
   
        const int off_k = 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);
+       tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
+       tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
   
-       const int sf_off = off_k / (16 * 4);
-       const char *SFA_src  = SFA_ptr  + ((off_m / 128) * rest_k + sf_off) * 512;
-       const char *SFB1_src = SFB1_ptr + ((off_n / 128) * rest_k + sf_off) * 512;
-       const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + sf_off) * 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);
+       const int rest_k = K / 16 / 4;
+       const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
+       const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
+       tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
+       tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
   
        asm volatile(
          "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
  ⋯ 4 unchanged lines
   
      constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
      for (int iter_k = 0; iter_k < PRELOAD; iter_k++) issue_tma(iter_k, iter_k);
- #pragma unroll 1
      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) & 1;
        mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
        issue_tma(iter_k, stage_id);
      }
-   } else if ((warp_id == NUM_WARPS - 2 || warp_id == NUM_WARPS - 1) && elect_sync()) {
+   } else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
      constexpr int MMA_N = BLOCK_N;
      constexpr int MMA_M = 128;
      constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)MMA_N >> 3U << 17U) | ((uint32_t)MMA_M >> 7U << 27U);
   
-     const int col_off = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
-     const bool is_d1 = (warp_id == NUM_WARPS - 2);
-     const int D_tmem = is_d1 ? D1_tmem : D2_tmem;
-     const int SFA_tmem = is_d1 ? SFA1_tmem : SFA2_tmem;
-     const int SFB_tmem = is_d1 ? SFB1_tmem : SFB2_tmem;
-     const int scale_B_base = SFB_tmem + col_off;
-  
- #pragma unroll 1
      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) & 1;
        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 + B_size;
-       const int SFA_smem = B2_smem + B_size;
-       const int SFB1_smem = SFA_smem + SFA_size;
-       const int SFB2_smem = SFB1_smem + SFB_size;
+       const int A_smem = smem + stage_id * STAGE_SIZE;
+       const int B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_size;
   
-       constexpr uint64_t AB_meta = (desc_encode((uint64_t)(8 * 128)) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
-       constexpr uint64_t SF_meta = (desc_encode((uint64_t)(8 * 16)) << 32ULL) | (1ULL << 46ULL);
+       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);
+       };
   
-       const uint64_t A_desc0  = AB_meta | desc_encode((uint64_t)A_smem);
-       const uint64_t B1_desc0 = AB_meta | desc_encode((uint64_t)B1_smem);
-       const uint64_t B2_desc0 = AB_meta | desc_encode((uint64_t)B2_smem);
+       constexpr uint64_t SF_desc = make_desc_SF(0);
+       const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
+       const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
   
-       const uint64_t SFA_desc0  = SF_meta | desc_encode((uint64_t)SFA_smem);
-       const uint64_t SFB1_desc0 = SF_meta | desc_encode((uint64_t)SFB1_smem);
-       const uint64_t SFB2_desc0 = SF_meta | desc_encode((uint64_t)SFB2_smem);
-  
        for (int k = 0; k < BLOCK_K / MMA_K; k++) {
-         const uint64_t off = (uint64_t)k * (512ULL >> 4ULL);
-         tcgen05_cp_nvfp4(SFA_tmem + k * 4, SFA_desc0 + off);
-         if (is_d1) {
-           tcgen05_cp_nvfp4(SFB_tmem + k * 4, SFB1_desc0 + off);
-         } else {
-           tcgen05_cp_nvfp4(SFB_tmem + k * 4, SFB2_desc0 + off);
-         }
+         uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
+         uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
+         tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
+         tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
        }
   
        for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
          for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
-           const uint64_t ab_off_a  = (uint64_t)(k1 * BLOCK_M * 128 + k2 * 32) >> 4ULL;
-           const uint64_t ab_off_b  = (uint64_t)(k1 * BLOCK_N * 128 + k2 * 32) >> 4ULL;
-           const uint64_t a_desc  = A_desc0  + ab_off_a;
-           const uint64_t b_desc  = (is_d1 ? B1_desc0 : B2_desc0) + ab_off_b;
+           uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
+           uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);
   
            const int k_sf = k1 * 4 + k2;
-           const int scale_A_tmem  = SFA_tmem + k_sf * 4;
-           const int scale_B_tmem  = scale_B_base + k_sf * 4;
+           const int scale_A_tmem = SFA_tmem + k_sf * 4;
+           const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
   
            const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
-           tcgen05_mma_nvfp4(D_tmem, a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
+           tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
          }
   
        asm volatile(
  ⋯ 13 unchanged lines
      asm volatile("tcgen05.fence::after_thread_sync;");
   
      for (int mm = 0; mm < 2; mm++) {
-       #pragma unroll
-       for (int nbase = 0; nbase < BLOCK_N; nbase += 32) {
-         float xbuf[16];
-         float ybuf[16];
+       float tmp[BLOCK_N / 2];
+       tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);
+       asm volatile("tcgen05.wait::ld.sync.aligned;");
   
-         tcgen05_ld_16x256bx4(xbuf, warp_id * 32 + mm * 16, nbase);
-         tcgen05_ld_16x256bx4(ybuf, warp_id * 32 + mm * 16, BLOCK_N + nbase);
-         asm volatile("tcgen05.wait::ld.sync.aligned;");
-  
-         #pragma unroll
-         for (int i = 0; i < 4; i++) {
-           const int row = off_m + warp_id * 32 + mm * 16 + lane_id / 4;
-           const int col = off_n + nbase + i * 8 + (lane_id & 3) * 2;
-  
-           const float2 x0 = float2{xbuf[i * 4 + 0], xbuf[i * 4 + 1]};
-           const float2 x8 = float2{xbuf[i * 4 + 2], xbuf[i * 4 + 3]};
-           const float2 y0 = float2{ybuf[i * 4 + 0], ybuf[i * 4 + 1]};
-           const float2 y8 = float2{ybuf[i * 4 + 2], ybuf[i * 4 + 3]};
-  
-           float2 o0;
-           float2 o8;
-           const float s00 = 1.0f / (1.0f + __expf(-x0.x));
-           const float s01 = 1.0f / (1.0f + __expf(-x0.y));
-           const float s80 = 1.0f / (1.0f + __expf(-x8.x));
-           const float s81 = 1.0f / (1.0f + __expf(-x8.y));
-           o0.x = (x0.x * s00) * y0.x;
-           o0.y = (x0.y * s01) * y0.y;
-           o8.x = (x8.x * s80) * y8.x;
-           o8.y = (x8.y * s81) * y8.y;
-  
-           reinterpret_cast<half2 *>(Out_ptr + (row + 0) * N + col)[0] = __float22half2_rn(o0);
-           reinterpret_cast<half2 *>(Out_ptr + (row + 8) * N + col)[0] = __float22half2_rn(o8);
-         }
+       #pragma unroll
+       for (int i = 0; i < BLOCK_N / 8; i++) {
+         const int row = off_m + warp_id * 32 + mm * 16 + lane_id / 4;
+         const int col = off_n + i * 8 + (lane_id & 3) * 2;
+         reinterpret_cast<float2 *>(C_ptr + (row + 0) * N + col)[0] = float2{tmp[i * 4 + 0], tmp[i * 4 + 1]};
+         reinterpret_cast<float2 *>(C_ptr + (row + 8) * N + col)[0] = float2{tmp[i * 4 + 2], tmp[i * 4 + 3]};
        }
      }
   
      asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
-     if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 4));
+     if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
    }
  }
   
  template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
  __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
- void dual_gemm_silu_mul_kernel_1mma(
+ void gemm_silu_mul_kernel(
    const __grid_constant__ CUtensorMap A_tmap,
-   const __grid_constant__ CUtensorMap B1_tmap,
-   const __grid_constant__ CUtensorMap B2_tmap,
+   const __grid_constant__ CUtensorMap B_tmap,
    const char *SFA_ptr,
-   const char *SFB1_ptr,
-   const char *SFB2_ptr,
+   const char *SFB_ptr,
+   const float *G1_ptr,
    half *Out_ptr,
    int M, int N
  ) {
  ⋯ 3 unchanged lines
    const int lane_id = tid & 31;
    const int warp_id = tid >> 5;
   
+   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 - bid_m * grid_n;
  ⋯ 5 unchanged lines
   
    extern __shared__ __align__(1024) char smem_ptr[];
    const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
-  
-   constexpr int A_size   = BLOCK_M * BLOCK_K / 2;
-   constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
+   constexpr int A_size = BLOCK_M * BLOCK_K / 2;
+   constexpr int B_size = BLOCK_N * BLOCK_K / 2;
    constexpr int SFA_size = 128 * BLOCK_K / 16;
    constexpr int SFB_size = 128 * BLOCK_K / 16;
-   constexpr int STAGE_SIZE = A_size + 2 * B_size + SFA_size + 2 * SFB_size;
+   constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
   
    #pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
  ⋯ 1 unchanged lines
    const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
    const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
   
-   constexpr int D1_tmem = 0;
-   constexpr int D2_tmem = BLOCK_N;
-   constexpr int SFA_tmem  = 2 * BLOCK_N;
-   constexpr int SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
-   constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);
+   constexpr int SFA_tmem = BLOCK_N;
+   constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
   
    if (warp_id == 0 && elect_sync()) {
-     for (int i = 0; i < NUM_STAGES; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
-     for (int i = 0; i < NUM_STAGES; i++) mbarrier_init(mma_mbar_addr + i * 8, 1);
-     mbarrier_init(mainloop_mbar_addr, 1);
+     for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);
      asm volatile("fence.mbarrier_init.release.cluster;");
    } else if (warp_id == 1) {
-     asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 4));
+     asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 2));
    }
    __syncthreads();
   
    constexpr int num_iters = K / BLOCK_K;
   
    if (warp_id == NUM_WARPS - 2 && elect_sync()) {
-     const uint64_t cache_A = EVICT_FIRST;
+     const uint64_t cache_A = EVICT_LAST;
      const uint64_t cache_B = EVICT_FIRST;
   
-     constexpr int rest_k = K / 16 / 4;
-  
      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 + B_size;
-       const int SFA_smem = B2_smem + B_size;
-       const int SFB1_smem = SFA_smem + SFA_size;
-       const int SFB2_smem = SFB1_smem + SFB_size;
+       const int A_smem = smem + stage_id * STAGE_SIZE;
+       const int B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_size;
   
        const int off_k = 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);
+       tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
+       tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
   
-       const int sf_off = off_k / (16 * 4);
-       const char *SFA_src  = SFA_ptr  + ((off_m / 128) * rest_k + sf_off) * 512;
-       const char *SFB1_src = SFB1_ptr + ((off_n / 128) * rest_k + sf_off) * 512;
-       const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + sf_off) * 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);
+       const int rest_k = K / 16 / 4;
+       const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
+       const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
+       tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
+       tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
   
        asm volatile(
          "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
  ⋯ 4 unchanged lines
   
      constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
      for (int iter_k = 0; iter_k < PRELOAD; iter_k++) issue_tma(iter_k, iter_k);
- #pragma unroll 1
      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) & 1;
  ⋯ 5 unchanged lines
      constexpr int MMA_M = 128;
      constexpr uint32_t i_desc = (1U << 7U) | (1U << 10U) | ((uint32_t)MMA_N >> 3U << 17U) | ((uint32_t)MMA_M >> 7U << 27U);
   
-     const int col_off = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
-     const int scale_B1_base = SFB1_tmem + col_off;
-     const int scale_B2_base = SFB2_tmem + col_off;
-  
- #pragma unroll 1
      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) & 1;
        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 + B_size;
-       const int SFA_smem = B2_smem + B_size;
-       const int SFB1_smem = SFA_smem + SFA_size;
-       const int SFB2_smem = SFB1_smem + SFB_size;
+       const int A_smem = smem + stage_id * STAGE_SIZE;
+       const int B_smem = A_smem + A_size;
+       const int SFA_smem = B_smem + B_size;
+       const int SFB_smem = SFA_smem + SFA_size;
   
-       constexpr uint64_t AB_meta = (desc_encode((uint64_t)(8 * 128)) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
-       constexpr uint64_t SF_meta = (desc_encode((uint64_t)(8 * 16)) << 32ULL) | (1ULL << 46ULL);
+       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);
+       };
   
-       const uint64_t A_desc0  = AB_meta | desc_encode((uint64_t)A_smem);
-       const uint64_t B1_desc0 = AB_meta | desc_encode((uint64_t)B1_smem);
-       const uint64_t B2_desc0 = AB_meta | desc_encode((uint64_t)B2_smem);
+       constexpr uint64_t SF_desc = make_desc_SF(0);
+       const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
+       const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
   
-       const uint64_t SFA_desc0  = SF_meta | desc_encode((uint64_t)SFA_smem);
-       const uint64_t SFB1_desc0 = SF_meta | desc_encode((uint64_t)SFB1_smem);
-       const uint64_t SFB2_desc0 = SF_meta | desc_encode((uint64_t)SFB2_smem);
-  
        for (int k = 0; k < BLOCK_K / MMA_K; k++) {
-         const uint64_t off = (uint64_t)k * (512ULL >> 4ULL);
-         tcgen05_cp_nvfp4(SFA_tmem + k * 4,  SFA_desc0 + off);
-         tcgen05_cp_nvfp4(SFB1_tmem + k * 4, SFB1_desc0 + off);
-         tcgen05_cp_nvfp4(SFB2_tmem + k * 4, SFB2_desc0 + off);
+         uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
+         uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
+         tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
+         tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
        }
   
        for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
          for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
-           const uint64_t ab_off_a  = (uint64_t)(k1 * BLOCK_M * 128 + k2 * 32) >> 4ULL;
-           const uint64_t ab_off_b  = (uint64_t)(k1 * BLOCK_N * 128 + k2 * 32) >> 4ULL;
-           const uint64_t a_desc  = A_desc0  + ab_off_a;
-           const uint64_t b1_desc = B1_desc0 + ab_off_b;
-           const uint64_t b2_desc = B2_desc0 + ab_off_b;
+           uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
+           uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);
   
            const int k_sf = k1 * 4 + k2;
-           const int scale_A_tmem  = SFA_tmem + k_sf * 4;
-           const int scale_B1_tmem = scale_B1_base + k_sf * 4;
-           const int scale_B2_tmem = scale_B2_base + k_sf * 4;
+           const int scale_A_tmem = SFA_tmem + k_sf * 4;
+           const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
   
            const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
-           tcgen05_mma_nvfp4(D1_tmem, a_desc, b1_desc, i_desc, scale_A_tmem, scale_B1_tmem, enable_input_d);
-           tcgen05_mma_nvfp4(D2_tmem, a_desc, b2_desc, i_desc, scale_A_tmem, scale_B2_tmem, enable_input_d);
+           tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
          }
   
        asm volatile(
  ⋯ 13 unchanged lines
      asm volatile("tcgen05.fence::after_thread_sync;");
   
      for (int mm = 0; mm < 2; mm++) {
+       float tmp[BLOCK_N / 2];
+       tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);
+       asm volatile("tcgen05.wait::ld.sync.aligned;");
+  
        #pragma unroll
-       for (int nbase = 0; nbase < BLOCK_N; nbase += 32) {
-         float xbuf[16];
-         float ybuf[16];
+       for (int i = 0; i < BLOCK_N / 8; i++) {
+         const int row = off_m + warp_id * 32 + mm * 16 + lane_id / 4;
+         const int col = off_n + i * 8 + (lane_id & 3) * 2;
+         const float2 x0 = reinterpret_cast<const float2 *>(G1_ptr + (row + 0) * N + col)[0];
+         const float2 x8 = reinterpret_cast<const float2 *>(G1_ptr + (row + 8) * N + col)[0];
+         const float2 y0 = float2{tmp[i * 4 + 0], tmp[i * 4 + 1]};
+         const float2 y8 = float2{tmp[i * 4 + 2], tmp[i * 4 + 3]};
   
-         tcgen05_ld_16x256bx4(xbuf, warp_id * 32 + mm * 16, nbase);
-         tcgen05_ld_16x256bx4(ybuf, warp_id * 32 + mm * 16, BLOCK_N + nbase);
-         asm volatile("tcgen05.wait::ld.sync.aligned;");
+         float2 o0;
+         float2 o8;
+         const float s00 = 1.0f / (1.0f + __expf(-x0.x));
+         const float s01 = 1.0f / (1.0f + __expf(-x0.y));
+         const float s80 = 1.0f / (1.0f + __expf(-x8.x));
+         const float s81 = 1.0f / (1.0f + __expf(-x8.y));
+         o0.x = (x0.x * s00) * y0.x;
+         o0.y = (x0.y * s01) * y0.y;
+         o8.x = (x8.x * s80) * y8.x;
+         o8.y = (x8.y * s81) * y8.y;
   
-         #pragma unroll
-         for (int i = 0; i < 4; i++) {
-           const int row = off_m + warp_id * 32 + mm * 16 + lane_id / 4;
-           const int col = off_n + nbase + i * 8 + (lane_id & 3) * 2;
-  
-           const float2 x0 = float2{xbuf[i * 4 + 0], xbuf[i * 4 + 1]};
-           const float2 x8 = float2{xbuf[i * 4 + 2], xbuf[i * 4 + 3]};
-           const float2 y0 = float2{ybuf[i * 4 + 0], ybuf[i * 4 + 1]};
-           const float2 y8 = float2{ybuf[i * 4 + 2], ybuf[i * 4 + 3]};
-  
-           float2 o0;
-           float2 o8;
-           const float s00 = 1.0f / (1.0f + __expf(-x0.x));
-           const float s01 = 1.0f / (1.0f + __expf(-x0.y));
-           const float s80 = 1.0f / (1.0f + __expf(-x8.x));
-           const float s81 = 1.0f / (1.0f + __expf(-x8.y));
-           o0.x = (x0.x * s00) * y0.x;
-           o0.y = (x0.y * s01) * y0.y;
-           o8.x = (x8.x * s80) * y8.x;
-           o8.y = (x8.y * s81) * y8.y;
-  
-           reinterpret_cast<half2 *>(Out_ptr + (row + 0) * N + col)[0] = __float22half2_rn(o0);
-           reinterpret_cast<half2 *>(Out_ptr + (row + 8) * N + col)[0] = __float22half2_rn(o8);
-         }
+         reinterpret_cast<half2 *>(Out_ptr + (row + 0) * N + col)[0] = __float22half2_rn(o0);
+         reinterpret_cast<half2 *>(Out_ptr + (row + 8) * N + col)[0] = __float22half2_rn(o8);
        }
      }
   
      asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
-     if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 4));
+     if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(BLOCK_N * 2));
    }
  }
   
  template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
- static inline void launch_dual_gemm_silu_mul(
+ static inline void launch_gemm_f32(
    const at::Tensor& A,
-   const at::Tensor& B1,
-   const at::Tensor& B2,
+   const at::Tensor& B,
    const at::Tensor& SFA,
-   const at::Tensor& SFB1,
-   const at::Tensor& SFB2,
-   at::Tensor& out
+   const at::Tensor& SFB,
+   at::Tensor& C
  ) {
-   static_assert((BLOCK_K % 256) == 0);
    const int M = (int)A.size(0);
-   const int N = (int)B1.size(0);
+   const int N = (int)B.size(0);
   
-   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 Out_ptr = reinterpret_cast<half *>(out.data_ptr());
+   auto A_ptr = reinterpret_cast<const char *>(A.data_ptr());
+   auto B_ptr = reinterpret_cast<const char *>(B.data_ptr());
+   auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
+   auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
+   auto C_ptr = reinterpret_cast<float *>(C.data_ptr());
   
-   static TmapCacheEntry cacheA  = {{nullptr, 0, 0, 0, 0}, {}, false};
-   static TmapCacheEntry cacheB1 = {{nullptr, 0, 0, 0, 0}, {}, false};
-   static TmapCacheEntry cacheB2 = {{nullptr, 0, 0, 0, 0}, {}, false};
+   CUtensorMap A_tmap, B_tmap;
+   init_AB_tmap(&A_tmap, A_ptr, (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
+   init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
   
-   const CUtensorMap &A_tmap  = get_or_init_tmap(cacheA,  A_ptr,  (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
-   const CUtensorMap &B1_tmap = get_or_init_tmap(cacheB1, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
-   const CUtensorMap &B2_tmap = get_or_init_tmap(cacheB2, B2_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
-  
    dim3 grid(1, (unsigned)((M / BLOCK_M) * (N / BLOCK_N)));
-   const int tb_size = BLOCK_M + 3 * WARP_SIZE;
-   const int AB_size = (BLOCK_M + 2 * BLOCK_N) * (BLOCK_K / 2);
-   const int SF_size = 128 * (BLOCK_K / 16) * 3;
-   const int smem_size = (AB_size + SF_size) * NUM_STAGES;
+   const int tb_size = BLOCK_M + 2 * WARP_SIZE;
+   const int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+   const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
+   const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
   
-   auto kptr = dual_gemm_silu_mul_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-   if (smem_size > 48'000) {
-     static bool once = false;
-     if (!once) {
-       cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-       once = true;
-     }
-   }
-   kptr<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);
+   auto kptr = gemm_f32_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+   if (smem_size > 48'000) cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+   kptr<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N);
  }
   
  template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
- static inline void launch_dual_gemm_silu_mul_1mma(
+ static inline void launch_gemm_silu_mul(
    const at::Tensor& A,
-   const at::Tensor& B1,
-   const at::Tensor& B2,
+   const at::Tensor& B,
    const at::Tensor& SFA,
-   const at::Tensor& SFB1,
-   const at::Tensor& SFB2,
+   const at::Tensor& SFB,
+   const at::Tensor& g1,
    at::Tensor& out
  ) {
-   static_assert((BLOCK_K % 256) == 0);
    const int M = (int)A.size(0);
-   const int N = (int)B1.size(0);
+   const int N = (int)B.size(0);
   
-   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 A_ptr = reinterpret_cast<const char *>(A.data_ptr());
+   auto B_ptr = reinterpret_cast<const char *>(B.data_ptr());
+   auto SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
+   auto SFB_ptr = reinterpret_cast<const char *>(SFB.data_ptr());
+   auto G1_ptr = reinterpret_cast<const float *>(g1.data_ptr());
    auto Out_ptr = reinterpret_cast<half *>(out.data_ptr());
   
-   static TmapCacheEntry cacheA  = {{nullptr, 0, 0, 0, 0}, {}, false};
-   static TmapCacheEntry cacheB1 = {{nullptr, 0, 0, 0, 0}, {}, false};
-   static TmapCacheEntry cacheB2 = {{nullptr, 0, 0, 0, 0}, {}, false};
+   CUtensorMap A_tmap, B_tmap;
+   init_AB_tmap(&A_tmap, A_ptr, (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
+   init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
   
-   const CUtensorMap &A_tmap  = get_or_init_tmap(cacheA,  A_ptr,  (uint64_t)M, (uint64_t)K, (uint32_t)BLOCK_M, (uint32_t)BLOCK_K);
-   const CUtensorMap &B1_tmap = get_or_init_tmap(cacheB1, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
-   const CUtensorMap &B2_tmap = get_or_init_tmap(cacheB2, B2_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
-  
    dim3 grid(1, (unsigned)((M / BLOCK_M) * (N / BLOCK_N)));
    const int tb_size = BLOCK_M + 2 * WARP_SIZE;
-   const int AB_size = (BLOCK_M + 2 * BLOCK_N) * (BLOCK_K / 2);
-   const int SF_size = 128 * (BLOCK_K / 16) * 3;
-   const int smem_size = (AB_size + SF_size) * NUM_STAGES;
+   const int AB_size = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
+   const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
+   const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
   
-   auto kptr = dual_gemm_silu_mul_kernel_1mma<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-   if (smem_size > 48'000) {
-     static bool once = false;
-     if (!once) {
-       cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-       once = true;
-     }
-   }
-   kptr<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);
+   auto kptr = gemm_silu_mul_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
+   if (smem_size > 48'000) cudaFuncSetAttribute(kptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+   kptr<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, G1_ptr, Out_ptr, M, N);
  }
   
+ __global__ void silu_mul_f32_vec2(const float* __restrict__ x, const float* __restrict__ y, half* __restrict__ out, int64_t n2) {
+   const int64_t idx = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+   if (idx >= n2) return;
+   const float2 fx = reinterpret_cast<const float2*>(x)[idx];
+   const float2 fy = reinterpret_cast<const float2*>(y)[idx];
+   float2 o;
+   const float sx0 = 1.0f / (1.0f + __expf(-fx.x));
+   const float sx1 = 1.0f / (1.0f + __expf(-fx.y));
+   o.x = (fx.x * sx0) * fy.x;
+   o.y = (fx.y * sx1) * fy.y;
+   reinterpret_cast<half2*>(out)[idx] = __float22half2_rn(o);
+ }
+  
+ static inline void launch_silu_mul_f32(const at::Tensor& g1, const at::Tensor& g2, at::Tensor& out) {
+   const int64_t n = out.numel();
+   TORCH_CHECK((n & 1) == 0, "n");
+   const int64_t n2 = n >> 1;
+   const int threads = 256;
+   const int blocks = (int)((n2 + threads - 1) / threads);
+   silu_mul_f32_vec2<<<blocks, threads>>>(
+     reinterpret_cast<const float*>(g1.data_ptr()),
+     reinterpret_cast<const float*>(g2.data_ptr()),
+     reinterpret_cast<half*>(out.data_ptr()),
+     n2
+   );
+ }
+  
  at::Tensor fused(
    const at::Tensor& A,
    const at::Tensor& B1,
  ⋯ 1 unchanged lines
    const at::Tensor& SFA,
    const at::Tensor& SFB1,
    const at::Tensor& SFB2,
-   at::Tensor& out
+   at::Tensor& out,
+   at::Tensor& g1,
+   at::Tensor& g2
  ) {
+   TORCH_CHECK(A.is_cuda() && B1.is_cuda() && B2.is_cuda(), "cuda");
+   TORCH_CHECK(SFA.is_cuda() && SFB1.is_cuda() && SFB2.is_cuda(), "cuda");
+   TORCH_CHECK(out.is_cuda() && g1.is_cuda() && g2.is_cuda(), "cuda");
+   TORCH_CHECK(A.dim() == 3 && B1.dim() == 3 && B2.dim() == 3, "dim");
+   TORCH_CHECK(out.dim() == 3 && g1.dim() == 3 && g2.dim() == 3, "dim");
+  
    const int64_t M = A.size(0);
    const int64_t Kp = A.size(1);
+   const int64_t L = A.size(2);
    const int64_t N = B1.size(0);
+   TORCH_CHECK(L == 1, "l");
+   TORCH_CHECK(B1.size(1) == Kp && B1.size(2) == L, "b1");
+   TORCH_CHECK(B2.size(1) == Kp && B2.size(2) == L, "b2");
+   TORCH_CHECK(out.size(0) == M && out.size(1) == N && out.size(2) == L, "out");
+   TORCH_CHECK(g1.size(0) == M && g1.size(1) == N && g1.size(2) == L, "g1");
+   TORCH_CHECK(g2.size(0) == M && g2.size(1) == N && g2.size(2) == L, "g2");
   
+   TORCH_CHECK((M % 128) == 0, "m");
+   TORCH_CHECK((N % 64) == 0, "n");
+  
    const int K = (int)(Kp * 2);
-   const bool use_bn128 = ((N & 127) == 0) && (M >= 512);
-   if (use_bn128) {
-     if (K == 7168) {
-       launch_dual_gemm_silu_mul_1mma<7168, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 4096) {
-       launch_dual_gemm_silu_mul_1mma<4096, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 2304) {
-       launch_dual_gemm_silu_mul_1mma<2304, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 2048) {
-       launch_dual_gemm_silu_mul_1mma<2048, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 1536) {
-       launch_dual_gemm_silu_mul_1mma<1536, 128, 128, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 512) {
-       launch_dual_gemm_silu_mul_1mma<512, 128, 128, 256, 2>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 256) {
-       launch_dual_gemm_silu_mul_1mma<256, 128, 128, 256, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
+   if (K == 7168) {
+     launch_gemm_f32<7168, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     if (M == 256 && N == 4096) {
+       launch_gemm_silu_mul<7168, 128, 64, 256, 8>(A, B2, SFA, SFB2, g1, out);
      } else {
-       TORCH_CHECK(false, "k ", K);
+       launch_gemm_f32<7168, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+       launch_silu_mul_f32(g1, g2, out);
      }
+   } else if (K == 4096) {
+     launch_gemm_f32<4096, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<4096, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
+   } else if (K == 2304) {
+     launch_gemm_f32<2304, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<2304, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
+   } else if (K == 2048) {
+     launch_gemm_f32<2048, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<2048, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
+   } else if (K == 1536) {
+     launch_gemm_f32<1536, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<1536, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
+   } else if (K == 512) {
+     launch_gemm_f32<512, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<512, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
+   } else if (K == 256) {
+     launch_gemm_f32<256, 128, 64, 256, 8>(A, B1, SFA, SFB1, g1);
+     launch_gemm_f32<256, 128, 64, 256, 8>(A, B2, SFA, SFB2, g2);
+     launch_silu_mul_f32(g1, g2, out);
    } else {
-     if (K == 7168) {
-       launch_dual_gemm_silu_mul<7168, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 4096) {
-       launch_dual_gemm_silu_mul<4096, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 2304) {
-       launch_dual_gemm_silu_mul<2304, 128, 64, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 2048) {
-       launch_dual_gemm_silu_mul<2048, 128, 64, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 1536) {
-       launch_dual_gemm_silu_mul<1536, 128, 64, 256, 4>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 512) {
-       launch_dual_gemm_silu_mul<512, 128, 64, 256, 2>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else if (K == 256) {
-       launch_dual_gemm_silu_mul<256, 128, 64, 256, 1>(A, B1, B2, SFA, SFB1, SFB2, out);
-     } else {
-       TORCH_CHECK(false, "k ", K);
-     }
+     TORCH_CHECK(false, "k ", K);
    }
+  
    return out;
  }
   
  TORCH_LIBRARY(nvfp4_dual_lib, m) {
-   m.def("fused(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) out) -> Tensor");
+   m.def("fused(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) out, Tensor(b!) g1, Tensor(c!) g2) -> Tensor");
    m.impl("fused", &fused);
  }
  """
  ⋯ 7 unchanged lines
      if _loaded:
          return
      load_inline(
-         name="nvfp4_dual_ext_tc_1mma_mix_v2",
+         name="nvfp4_dual_ext_tc_v3",
          cpp_sources="",
          cuda_sources=_CUDA_SRC,
          functions=None,
  ⋯ 4 unchanged lines
              "--use_fast_math",
              "--expt-relaxed-constexpr",
              "--relocatable-device-code=false",
+             "-lineinfo",
          ],
          extra_ldflags=["-lcuda"],
          verbose=False,
  ⋯ 3 unchanged lines
      _loaded = True
   
   
+ _buf_cache = {}
+  
+  
+ def _get_buf(tag, shape, device):
+     key = (tag, shape, device)
+     t = _buf_cache.get(key)
+     if t is None or t.shape != shape or t.device != device:
+         t = torch.empty(shape, device=device, dtype=torch.float32)
+         _buf_cache[key] = t
+     return t
+  
+  
  def custom_kernel(data):
      _load()
      a, b1, b2, _sfa, _sfb1, _sfb2, sfa_p, sfb1_p, sfb2_p, c = data
-     return torch.ops.nvfp4_dual_lib.fused(a, b1, b2, sfa_p, sfb1_p, sfb2_p, c)
+     g1 = _get_buf(1, c.shape, a.device)
+     g2 = _get_buf(2, c.shape, a.device)
+     return torch.ops.nvfp4_dual_lib.fused(a, b1, b2, sfa_p, sfb1_p, sfb2_p, c, g1, g2)
   
   
  __all__ = ["custom_kernel"]
scrolls · 959 diff lines total

Best evidence level for this revision: reported

JSON