submission 372551
shiyeegao · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 911 lines, June 9 Researcher Reciprocity License v1.0.
submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-372551?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
Reported · How evidence levels are derived →
Source and license
sourceavailable
revision digestsha256:83ac7c8cd5d571f743a2d32967b4e2eb7c767ae125c1258ddfb2dea7a03ecbe9
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-memory
extern __shared__ __align__(1024) char smem_ptr[];tcgen05
asm 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 = float2
const float2 x0 = float2{tmp_x[i * 4 + 0], tmp_x[i * 4 + 1]};Kernel source
submission.py911 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_NORMAL = 0x1000000000000000ULL;
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__ constexpr int i_align_up(int x, int a) {
return ((x + a - 1) / a) * a;
}
__device__ __forceinline__ float fast_sigmoid(float x) {
return 1.0f / (1.0f + __expf(-x));
}
__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__ int mbarrier_try_wait(int mbar_addr, int phase) {
uint32_t pred = 0;
uint32_t ticks = 0;
asm volatile(
"{\n\t"
".reg .pred P1;\n\t"
"mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%1], %2, %3;\n\t"
"@P1 mov.u32 %0, 1;\n\t"
"}\n\t"
: "=r"(pred)
: "r"(mbar_addr), "r"(phase), "r"(ticks)
);
return (int)pred;
}
__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(
int d_tmem,
uint64_t a_desc,
uint64_t b_desc,
uint32_t i_desc,
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"
"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";
static constexpr char x16[] = ".x16";
};
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_)
);
}
template <const char *SHAPE_, const char *NUM_>
__device__ __forceinline__ 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_)
);
}
__device__ __forceinline__ void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
}
__device__ __forceinline__ void tcgen05_ld_16x256bx16(float *tmp, int row, int col) {
tcgen05_ld_64regs<SHAPE::_16x256b, NUM::x16>(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 dual_gemm_silu_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 *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 + 2 * B_size + SFA_size + 2 * 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 K_STEPS = BLOCK_K / MMA_K;
constexpr int SF_COLS = 4 * K_STEPS;
constexpr int OUT1_tmem = BLOCK_N;
constexpr int SFA_tmem = 2 * BLOCK_N;
constexpr int SFB1_tmem = SFA_tmem + SF_COLS;
constexpr int SFB2_tmem = SFB1_tmem + SF_COLS;
constexpr int N_SLICE_MAX = (128 / BLOCK_N - 1) * (BLOCK_N / 32);
constexpr int TMEM_MAX = SFB2_tmem + SF_COLS - 1 + N_SLICE_MAX;
constexpr int TMEM_MIN = i_align_up(TMEM_MAX + 1, 32);
static_assert((TMEM_MIN & 31) == 0, "tmem");
static_assert(TMEM_MIN <= 256, "tmem");
constexpr int TMEM_COLS = 256;
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"(TMEM_COLS));
}
__syncthreads();
constexpr int num_iters = K / BLOCK_K;
if (warp_id == NUM_WARPS - 2 && elect_sync()) {
const uint64_t cache_A_data = EVICT_LAST;
const uint64_t cache_B_data = EVICT_FIRST;
const uint64_t cache_SF_A = EVICT_LAST;
const uint64_t cache_SF_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 + 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 off_k = iter_k * BLOCK_K;
tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A_data);
tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);
tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);
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_SF_A);
tma_gmem2smem(SFB1_smem, SFB1_src, SFB_size, mbar_addr, cache_SF_B);
tma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_SF_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 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;
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);
#pragma unroll
for (int k = 0; k < K_STEPS; k++) {
const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
const uint64_t sfb1_desc = SFB1_desc + (uint64_t)k * (512ULL >> 4ULL);
const 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 n_slice = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
#pragma unroll
for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
#pragma unroll
for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
const uint64_t b1_desc = make_desc_AB(B1_smem + k1 * BLOCK_N * 128 + k2 * 32);
const uint64_t b2_desc = make_desc_AB(B2_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 = SFB1_tmem + k_sf * 4 + n_slice;
const int scale_B2_tmem = SFB2_tmem + k_sf * 4 + n_slice;
const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
tcgen05_mma_nvfp4(0, a_desc, b1_desc, i_desc, scale_A_tmem, scale_B1_tmem, enable_input_d);
tcgen05_mma_nvfp4(OUT1_tmem, a_desc, b2_desc, i_desc, scale_A_tmem, scale_B2_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;");
#pragma unroll
for (int mm = 0; mm < 2; mm++) {
float tmp_x[BLOCK_N / 2];
float tmp_y[BLOCK_N / 2];
tcgen05_ld_16x256bx8(tmp_x, warp_id * 32 + mm * 16, 0);
tcgen05_ld_16x256bx8(tmp_y, warp_id * 32 + mm * 16, OUT1_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 + mm * 16 + lane_id / 4;
const int col = off_n + i * 8 + (lane_id & 3) * 2;
const float2 x0 = float2{tmp_x[i * 4 + 0], tmp_x[i * 4 + 1]};
const float2 x8 = float2{tmp_x[i * 4 + 2], tmp_x[i * 4 + 3]};
const float2 y0 = float2{tmp_y[i * 4 + 0], tmp_y[i * 4 + 1]};
const float2 y8 = float2{tmp_y[i * 4 + 2], tmp_y[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"(TMEM_COLS));
}
}
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_kernel_splitb(
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 *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_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;
constexpr int PHASE1_SIZE = B_size + SFB_size;
#pragma nv_diag_suppress static_var_with_dynamic_init
__shared__ int64_t mbars[NUM_STAGES * 4 + 1];
const int tma0_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
const int tma1_mbar_addr = tma0_mbar_addr + NUM_STAGES * 8;
const int b1_mbar_addr = tma1_mbar_addr + NUM_STAGES * 8;
const int done_mbar_addr = b1_mbar_addr + NUM_STAGES * 8;
const int mainloop_mbar_addr = done_mbar_addr + NUM_STAGES * 8;
constexpr int K_STEPS = BLOCK_K / MMA_K;
constexpr int SF_COLS = 4 * K_STEPS;
constexpr int OUT1_tmem = BLOCK_N;
constexpr int SFA_tmem = 2 * BLOCK_N;
constexpr int SFB1_tmem = SFA_tmem + SF_COLS;
constexpr int SFB2_tmem = SFB1_tmem + SF_COLS;
constexpr int N_SLICE_MAX = (128 / BLOCK_N - 1) * (BLOCK_N / 32);
constexpr int TMEM_MAX = SFB2_tmem + SF_COLS - 1 + N_SLICE_MAX;
constexpr int TMEM_MIN = i_align_up(TMEM_MAX + 1, 32);
static_assert((TMEM_MIN & 31) == 0, "tmem");
static_assert(TMEM_MIN <= 256, "tmem");
constexpr int TMEM_COLS = 256;
if (warp_id == 0 && elect_sync()) {
for (int i = 0; i < NUM_STAGES * 4 + 1; i++) mbarrier_init(tma0_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"(TMEM_COLS));
}
__syncthreads();
constexpr int num_iters = K / BLOCK_K;
if (warp_id == NUM_WARPS - 2 && elect_sync()) {
const uint64_t cache_A_data = EVICT_LAST;
const uint64_t cache_B_data = EVICT_FIRST;
const uint64_t cache_SF_A = EVICT_LAST;
const uint64_t cache_SF_B = EVICT_LAST;
auto issue_tma0 = [&](int iter_k, int stage_id) {
const int mbar_addr = tma0_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_data);
tma_3d_gmem2smem(B_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);
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;
tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_SF_A);
tma_gmem2smem(SFB_smem, SFB1_src, SFB_size, mbar_addr, cache_SF_B);
asm volatile(
"mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
:: "r"(mbar_addr), "r"(STAGE_SIZE)
: "memory"
);
};
auto issue_tma1 = [&](int iter_k, int stage_id) {
const int mbar_addr = tma1_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(B_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);
const int rest_k = K / 16 / 4;
const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
tma_gmem2smem(SFB_smem, SFB2_src, SFB_size, mbar_addr, cache_SF_B);
asm volatile(
"mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
:: "r"(mbar_addr), "r"(PHASE1_SIZE)
: "memory"
);
};
constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
for (int iter_k = 0; iter_k < PRELOAD; iter_k++) issue_tma0(iter_k, iter_k);
for (int iter_k = 0; iter_k < num_iters; iter_k++) {
const int stage_id = iter_k % NUM_STAGES;
const int phase = (iter_k / NUM_STAGES) & 1;
mbarrier_wait(b1_mbar_addr + stage_id * 8, phase);
issue_tma1(iter_k, stage_id);
mbarrier_wait(done_mbar_addr + stage_id * 8, phase);
const int next_iter = iter_k + NUM_STAGES;
if (next_iter < num_iters) issue_tma0(next_iter, 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);
constexpr int PHASE_GAP = 2;
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 int n_slice = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
for (int step = 0; step < num_iters + PHASE_GAP; step++) {
if (step >= PHASE_GAP) {
const int iter_k = step - PHASE_GAP;
const int stage_id = iter_k % NUM_STAGES;
const int tma_phase = (iter_k / NUM_STAGES) & 1;
mbarrier_wait(tma1_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;
const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
#pragma unroll
for (int k = 0; k < K_STEPS; k++) {
const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
const uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
tcgen05_cp_nvfp4(SFB2_tmem + k * 4, sfb_desc);
}
#pragma unroll
for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
#pragma unroll
for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
const 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 = SFB2_tmem + k_sf * 4 + n_slice;
const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
tcgen05_mma_nvfp4(OUT1_tmem, 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"(done_mbar_addr + stage_id * 8)
: "memory"
);
}
if (step < num_iters) {
const int iter_k = step;
const int stage_id = iter_k % NUM_STAGES;
const int tma_phase = (iter_k / NUM_STAGES) & 1;
mbarrier_wait(tma0_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;
const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
#pragma unroll
for (int k = 0; k < K_STEPS; k++) {
const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
const uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
tcgen05_cp_nvfp4(SFB1_tmem + k * 4, sfb_desc);
}
#pragma unroll
for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {
#pragma unroll
for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
const 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 = SFB1_tmem + k_sf * 4 + n_slice;
const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
tcgen05_mma_nvfp4(0, 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"(b1_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;");
#pragma unroll
for (int mm = 0; mm < 2; mm++) {
float tmp_x[BLOCK_N / 2];
float tmp_y[BLOCK_N / 2];
if constexpr (BLOCK_N == 128) {
tcgen05_ld_16x256bx16(tmp_x, warp_id * 32 + mm * 16, 0);
tcgen05_ld_16x256bx16(tmp_y, warp_id * 32 + mm * 16, OUT1_tmem);
} else {
tcgen05_ld_16x256bx8(tmp_x, warp_id * 32 + mm * 16, 0);
tcgen05_ld_16x256bx8(tmp_y, warp_id * 32 + mm * 16, OUT1_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 + mm * 16 + lane_id / 4;
const int col = off_n + i * 8 + (lane_id & 3) * 2;
const float2 x0 = float2{tmp_x[i * 4 + 0], tmp_x[i * 4 + 1]};
const float2 x8 = float2{tmp_x[i * 4 + 2], tmp_x[i * 4 + 3]};
const float2 y0 = float2{tmp_y[i * 4 + 0], tmp_y[i * 4 + 1]};
const float2 y8 = float2{tmp_y[i * 4 + 2], tmp_y[i * 4 + 3]};
float2 o0;
float2 o8;
const float s00 = fast_sigmoid(x0.x);
const float s01 = fast_sigmoid(x0.y);
const float s80 = fast_sigmoid(x8.x);
const float s81 = fast_sigmoid(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"(TMEM_COLS));
}
}
template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
static inline void launch_dual(
const at::Tensor& A,
const at::Tensor& B1,
const at::Tensor& B2,
const at::Tensor& SFA,
const at::Tensor& SFB1,
const at::Tensor& SFB2,
at::Tensor& out
) {
const int M = (int)A.size(0);
const int N = (int)B1.size(0);
const char *A_ptr = reinterpret_cast<const char *>(A.data_ptr());
const char *B1_ptr = reinterpret_cast<const char *>(B1.data_ptr());
const char *B2_ptr = reinterpret_cast<const char *>(B2.data_ptr());
const char *SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
const char *SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());
const char *SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());
half *Out_ptr = reinterpret_cast<half *>(out.data_ptr());
CUtensorMap A_tmap, B1_tmap, B2_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(&B1_tmap, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
init_AB_tmap(&B2_tmap, 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 smem_size = (BLOCK_M * BLOCK_K / 2 + 2 * (BLOCK_N * BLOCK_K / 2) + 128 * (BLOCK_K / 16) + 2 * (128 * (BLOCK_K / 16))) * NUM_STAGES;
auto kptr = dual_gemm_silu_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, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);
}
template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
static inline void launch_dual_splitb(
const at::Tensor& A,
const at::Tensor& B1,
const at::Tensor& B2,
const at::Tensor& SFA,
const at::Tensor& SFB1,
const at::Tensor& SFB2,
at::Tensor& out
) {
const int M = (int)A.size(0);
const int N = (int)B1.size(0);
const char *A_ptr = reinterpret_cast<const char *>(A.data_ptr());
const char *B1_ptr = reinterpret_cast<const char *>(B1.data_ptr());
const char *B2_ptr = reinterpret_cast<const char *>(B2.data_ptr());
const char *SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());
const char *SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());
const char *SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());
half *Out_ptr = reinterpret_cast<half *>(out.data_ptr());
CUtensorMap A_tmap, B1_tmap, B2_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(&B1_tmap, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);
init_AB_tmap(&B2_tmap, 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 smem_size = (BLOCK_M * BLOCK_K / 2 + (BLOCK_N * BLOCK_K / 2) + 2 * (128 * (BLOCK_K / 16))) * NUM_STAGES;
auto kptr = dual_gemm_silu_kernel_splitb<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, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);
}
at::Tensor fused(
const at::Tensor& A,
const at::Tensor& B1,
const at::Tensor& B2,
const at::Tensor& SFA,
const at::Tensor& SFB1,
const at::Tensor& SFB2,
at::Tensor& out
) {
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(), "cuda");
TORCH_CHECK(A.dim() == 3 && B1.dim() == 3 && B2.dim() == 3, "dim");
TORCH_CHECK(out.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((M % 128) == 0, "m");
TORCH_CHECK((N % 64) == 0, "n");
const int K = (int)(Kp * 2);
const bool is_rank =
(L == 1) &&
((M == 256 || M == 512)) &&
((N == 3072 || N == 4096)) &&
((K == 4096 || K == 7168));
if (is_rank) {
if (K == 7168) {
launch_dual<7168, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 4096) {
launch_dual<4096, 128, 64, 256, 5>(A, B1, B2, SFA, SFB1, SFB2, out);
} else {
TORCH_CHECK(false, "k ", K);
}
return out;
}
if (K == 7168) {
launch_dual<7168, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 4096) {
launch_dual<4096, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 2304) {
launch_dual<2304, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 2048) {
launch_dual<2048, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 1536) {
launch_dual<1536, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 512) {
launch_dual<512, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, out);
} else if (K == 256) {
launch_dual<256, 128, 64, 256, 3>(A, B1, B2, SFA, SFB1, SFB2, 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");
m.impl("fused", &fused);
}
"""
_LOADED = False
def _load() -> None:
global _LOADED
if _LOADED:
return
load_inline(
name="nvfp4_dual_ext_tc_v4",
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
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)
__all__ = ["custom_kernel"]
scrolls · 911 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 372114.
-import torchfrom torch.utils.cpp_extension import load_inline⋯ 10 unchanged linesconstexpr int WARP_SIZE = 32;constexpr int MMA_K = 64;- // cache hint（来自 CuTe/cutlass 的经验值）- constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;- constexpr uint64_t EVICT_LAST = 0x14F0000000000000ULL;+ constexpr uint64_t EVICT_NORMAL = 0x1000000000000000ULL;+ constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;+ constexpr uint64_t EVICT_LAST = 0x14F0000000000000ULL;- // ranked (K=7168,M=512,N=3072) 专用：最小变量实验（仅调参，不改主循环结构）- constexpr int STAGE_7168_512_3072_BN64 = 8;- constexpr int STAGE_7168_512_3072_BN128 = 3;- constexpr bool USE_BN128_7168_512_3072 = true;-__device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }+ __device__ __forceinline__ constexpr int i_align_up(int x, int a) {+ return ((x + a - 1) / a) * a;+ }++ __device__ __forceinline__ float fast_sigmoid(float x) {+ return 1.0f / (1.0f + __expf(-x));+ }+__device__ __forceinline__ uint32_t elect_sync() {uint32_t pred = 0;asm volatile(⋯ 27 unchanged lines);}- // 低阶 1D 搬运：scale 用（512B/块，延迟不敏感）+ __device__ __forceinline__ int mbarrier_try_wait(int mbar_addr, int phase) {+ uint32_t pred = 0;+ uint32_t ticks = 0;+ asm volatile(+ "{\n\t"+ ".reg .pred P1;\n\t"+ "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%1], %2, %3;\n\t"+ "@P1 mov.u32 %0, 1;\n\t"+ "}\n\t"+ : "=r"(pred)+ : "r"(mbar_addr), "r"(phase), "r"(ticks)+ );+ return (int)pred;+ }+__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 "⋯ 2 unchanged lines);}- // 3D tensor map 搬运：A/B 用（对齐与 swizzle 由 tensor map 保证）__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 "⋯ 8 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 linesint 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 x8[] = ".x8"; static constexpr char x16[] = ".x16"; };+ struct NUM {+ static constexpr char x8[] = ".x8";+ static constexpr char x16[] = ".x16";+ };template <const char *SHAPE_, const char *NUM_>__device__ __forceinline__ void tcgen05_ld_32regs(float *tmp, int row, int col) {⋯ 81 unchanged linestemplate <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>__global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)- void gemm_to_half_kernel(+ void dual_gemm_silu_kernel(const __grid_constant__ CUtensorMap A_tmap,- const __grid_constant__ CUtensorMap B_tmap,+ const __grid_constant__ CUtensorMap B1_tmap,+ const __grid_constant__ CUtensorMap B2_tmap,const char *SFA_ptr,- const char *SFB_ptr,- half *C_ptr,+ const char *SFB1_ptr,+ const char *SFB2_ptr,+ half *Out_ptr,int M, int N) {const int tid = threadIdx.x;⋯ 2 unchanged linesconst 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;⋯ 5 unchanged linesconstexpr 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;+ constexpr int STAGE_SIZE = A_size + 2 * B_size + SFA_size + 2 * SFB_size;#pragma nv_diag_suppress static_var_with_dynamic_init__shared__ int64_t mbars[NUM_STAGES * 2 + 1];⋯ 1 unchanged linesconst 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);+ constexpr int K_STEPS = BLOCK_K / MMA_K;+ constexpr int SF_COLS = 4 * K_STEPS;+ constexpr int OUT1_tmem = BLOCK_N;+ constexpr int SFA_tmem = 2 * BLOCK_N;+ constexpr int SFB1_tmem = SFA_tmem + SF_COLS;+ constexpr int SFB2_tmem = SFB1_tmem + SF_COLS;+ constexpr int N_SLICE_MAX = (128 / BLOCK_N - 1) * (BLOCK_N / 32);+ constexpr int TMEM_MAX = SFB2_tmem + SF_COLS - 1 + N_SLICE_MAX;+ constexpr int TMEM_MIN = i_align_up(TMEM_MAX + 1, 32);+ static_assert((TMEM_MIN & 31) == 0, "tmem");+ static_assert(TMEM_MIN <= 256, "tmem");+ constexpr int TMEM_COLS = 256;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));+ asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(TMEM_COLS));}__syncthreads();constexpr int num_iters = K / BLOCK_K;if (warp_id == NUM_WARPS - 2 && elect_sync()) {- uint64_t cache_A, cache_B;- if (M > N) { cache_A = EVICT_FIRST; cache_B = EVICT_LAST; }- else { cache_A = EVICT_LAST; cache_B = EVICT_FIRST; }+ const uint64_t cache_A_data = EVICT_LAST;+ const uint64_t cache_B_data = EVICT_FIRST;+ const uint64_t cache_SF_A = EVICT_LAST;+ const uint64_t cache_SF_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 B_smem = A_smem + A_size;- const int SFA_smem = B_smem + B_size;- const int SFB_smem = SFA_smem + SFA_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 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);+ tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A_data);+ tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);+ tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);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);+ 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_SF_A);+ tma_gmem2smem(SFB1_smem, SFB1_src, SFB_size, mbar_addr, cache_SF_B);+ tma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_SF_B);asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"⋯ 2 unchanged lines);};- for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++) issue_tma(iter_k, iter_k);+ 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;⋯ 11 unchanged linesmbarrier_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;+ 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;auto make_desc_AB = [](int addr) -> uint64_t {const int SBO = 8 * 128;⋯ 6 unchanged linesconstexpr 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 SFB1_desc = SF_desc + ((uint64_t)SFB1_smem >> 4ULL);+ const uint64_t SFB2_desc = SF_desc + ((uint64_t)SFB2_smem >> 4ULL);#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 sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);+ for (int k = 0; k < K_STEPS; k++) {+ const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);+ const uint64_t sfb1_desc = SFB1_desc + (uint64_t)k * (512ULL >> 4ULL);+ const uint64_t sfb2_desc = SFB2_desc + (uint64_t)k * (512ULL >> 4ULL);tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);- tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);+ tcgen05_cp_nvfp4(SFB1_tmem + k * 4, sfb1_desc);+ tcgen05_cp_nvfp4(SFB2_tmem + k * 4, sfb2_desc);}+ const int n_slice = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);+#pragma unroll- for (int k1 = 0; k1 < BLOCK_K / 256; k1++)+ for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {#pragma unrollfor (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 uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);+ const uint64_t b1_desc = make_desc_AB(B1_smem + k1 * BLOCK_N * 128 + k2 * 32);+ const uint64_t b2_desc = make_desc_AB(B2_smem + k1 * BLOCK_N * 128 + k2 * 32);const int k_sf = k1 * 4 + k2;- const int scale_A_tmem = SFA_tmem + k_sf * 4 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);- const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);+ const int scale_A_tmem = SFA_tmem + k_sf * 4;+ const int scale_B1_tmem = SFB1_tmem + k_sf * 4 + n_slice;+ const int scale_B2_tmem = SFB2_tmem + k_sf * 4 + n_slice;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);+ tcgen05_mma_nvfp4(0, a_desc, b1_desc, i_desc, scale_A_tmem, scale_B1_tmem, enable_input_d);+ tcgen05_mma_nvfp4(OUT1_tmem, a_desc, b2_desc, i_desc, scale_A_tmem, scale_B2_tmem, enable_input_d);}+ }asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"⋯ 13 unchanged lines#pragma unrollfor (int mm = 0; mm < 2; mm++) {- float tmp[BLOCK_N / 2];- if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);- if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp, warp_id * 32 + mm * 16, 0);+ float tmp_x[BLOCK_N / 2];+ float tmp_y[BLOCK_N / 2];+ tcgen05_ld_16x256bx8(tmp_x, warp_id * 32 + mm * 16, 0);+ tcgen05_ld_16x256bx8(tmp_y, warp_id * 32 + mm * 16, OUT1_tmem);asm volatile("tcgen05.wait::ld.sync.aligned;");#pragma unrollfor (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<half2 *>(C_ptr + (row + 0) * N + col)[0] =- __float22half2_rn(float2{tmp[i * 4 + 0], tmp[i * 4 + 1]});- reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] =- __float22half2_rn(float2{tmp[i * 4 + 2], tmp[i * 4 + 3]});++ const float2 x0 = float2{tmp_x[i * 4 + 0], tmp_x[i * 4 + 1]};+ const float2 x8 = float2{tmp_x[i * 4 + 2], tmp_x[i * 4 + 3]};+ const float2 y0 = float2{tmp_y[i * 4 + 0], tmp_y[i * 4 + 1]};+ const float2 y8 = float2{tmp_y[i * 4 + 2], tmp_y[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));+ if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(TMEM_COLS));}}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_to_f32_kernel(+ void dual_gemm_silu_kernel_splitb(const __grid_constant__ CUtensorMap A_tmap,- const __grid_constant__ CUtensorMap B_tmap,+ const __grid_constant__ CUtensorMap B1_tmap,+ const __grid_constant__ CUtensorMap B2_tmap,const char *SFA_ptr,- const char *SFB_ptr,- float *C_ptr,+ const char *SFB1_ptr,+ const char *SFB2_ptr,+ half *Out_ptr,int M, int N) {const int tid = threadIdx.x;⋯ 5 unchanged linesconst 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;⋯ 6 unchanged linesconstexpr 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;+ constexpr int PHASE1_SIZE = B_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;+ __shared__ int64_t mbars[NUM_STAGES * 4 + 1];+ const int tma0_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));+ const int tma1_mbar_addr = tma0_mbar_addr + NUM_STAGES * 8;+ const int b1_mbar_addr = tma1_mbar_addr + NUM_STAGES * 8;+ const int done_mbar_addr = b1_mbar_addr + NUM_STAGES * 8;+ const int mainloop_mbar_addr = done_mbar_addr + NUM_STAGES * 8;- constexpr int SFA_tmem = BLOCK_N;- constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);+ constexpr int K_STEPS = BLOCK_K / MMA_K;+ constexpr int SF_COLS = 4 * K_STEPS;+ constexpr int OUT1_tmem = BLOCK_N;+ constexpr int SFA_tmem = 2 * BLOCK_N;+ constexpr int SFB1_tmem = SFA_tmem + SF_COLS;+ constexpr int SFB2_tmem = SFB1_tmem + SF_COLS;+ constexpr int N_SLICE_MAX = (128 / BLOCK_N - 1) * (BLOCK_N / 32);+ constexpr int TMEM_MAX = SFB2_tmem + SF_COLS - 1 + N_SLICE_MAX;+ constexpr int TMEM_MIN = i_align_up(TMEM_MAX + 1, 32);+ static_assert((TMEM_MIN & 31) == 0, "tmem");+ static_assert(TMEM_MIN <= 256, "tmem");+ constexpr int TMEM_COLS = 256;if (warp_id == 0 && elect_sync()) {- for (int i = 0; i < NUM_STAGES * 2 + 1; i++) mbarrier_init(tma_mbar_addr + i * 8, 1);+ for (int i = 0; i < NUM_STAGES * 4 + 1; i++) mbarrier_init(tma0_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));+ asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(TMEM_COLS));}__syncthreads();constexpr int num_iters = K / BLOCK_K;if (warp_id == NUM_WARPS - 2 && elect_sync()) {- uint64_t cache_A, cache_B;- if (M > N) { cache_A = EVICT_FIRST; cache_B = EVICT_LAST; }- else { cache_A = EVICT_LAST; cache_B = EVICT_FIRST; }+ const uint64_t cache_A_data = EVICT_LAST;+ const uint64_t cache_B_data = EVICT_FIRST;+ const uint64_t cache_SF_A = EVICT_LAST;+ const uint64_t cache_SF_B = EVICT_LAST;- auto issue_tma = [&](int iter_k, int stage_id) {- const int mbar_addr = tma_mbar_addr + stage_id * 8;+ auto issue_tma0 = [&](int iter_k, int stage_id) {+ const int mbar_addr = tma0_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);+ tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A_data);+ tma_3d_gmem2smem(B_smem, &B1_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);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);+ const char *SFB1_src = SFB1_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;+ tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_SF_A);+ tma_gmem2smem(SFB_smem, SFB1_src, SFB_size, mbar_addr, cache_SF_B);asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"⋯ 2 unchanged lines);};- for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++) issue_tma(iter_k, iter_k);- for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {- const int stage_id = iter_k % NUM_STAGES;- const int mma_phase = (iter_k / NUM_STAGES - 1) & 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);-+ auto issue_tma1 = [&](int iter_k, int stage_id) {+ const int mbar_addr = tma1_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;- 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);-- #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 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);- }-- #pragma unroll- for (int k1 = 0; k1 < BLOCK_K / 256; k1++)- #pragma unroll- 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 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);- const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);-- const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;- tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);- }-- asm volatile(- "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"- :: "r"(mma_mbar_addr + stage_id * 8)- : "memory"- );- }-- asm volatile(- "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"- :: "r"(mainloop_mbar_addr)- : "memory"- );- } else if (tid < BLOCK_M) {- mbarrier_wait(mainloop_mbar_addr, 0);- asm volatile("tcgen05.fence::after_thread_sync;");-- #pragma unroll- for (int mm = 0; mm < 2; mm++) {- float tmp[BLOCK_N / 2];- if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);- if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(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 half *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_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()) {- uint64_t cache_A, cache_B;- if (M > N) { cache_A = EVICT_FIRST; cache_B = EVICT_LAST; }- else { cache_A = EVICT_LAST; cache_B = EVICT_FIRST; }-- 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);+ tma_3d_gmem2smem(B_smem, &B2_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B_data);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);+ const char *SFB2_src = SFB2_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;+ tma_gmem2smem(SFB_smem, SFB2_src, SFB_size, mbar_addr, cache_SF_B);asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"- :: "r"(mbar_addr), "r"(STAGE_SIZE)+ :: "r"(mbar_addr), "r"(PHASE1_SIZE): "memory");};- for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++) issue_tma(iter_k, iter_k);- for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {+ constexpr int PRELOAD = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;+ for (int iter_k = 0; iter_k < PRELOAD; iter_k++) issue_tma0(iter_k, iter_k);++ for (int iter_k = 0; 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);+ const int phase = (iter_k / NUM_STAGES) & 1;++ mbarrier_wait(b1_mbar_addr + stage_id * 8, phase);+ issue_tma1(iter_k, stage_id);++ mbarrier_wait(done_mbar_addr + stage_id * 8, phase);+ const int next_iter = iter_k + NUM_STAGES;+ if (next_iter < num_iters) issue_tma0(next_iter, 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);+ constexpr int PHASE_GAP = 2;- 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);+ 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 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 SF_desc = make_desc_SF(0);+ const int n_slice = (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);- 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);- };+ for (int step = 0; step < num_iters + PHASE_GAP; step++) {+ if (step >= PHASE_GAP) {+ const int iter_k = step - PHASE_GAP;+ const int stage_id = iter_k % NUM_STAGES;+ const int tma_phase = (iter_k / NUM_STAGES) & 1;+ mbarrier_wait(tma1_mbar_addr + stage_id * 8, tma_phase);- 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 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;- #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 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);- }+ const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);+ const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);+ #pragma unroll+ for (int k = 0; k < K_STEPS; k++) {+ const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);+ const uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);+ tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);+ tcgen05_cp_nvfp4(SFB2_tmem + k * 4, sfb_desc);+ }- #pragma unroll- for (int k1 = 0; k1 < BLOCK_K / 256; k1++)#pragma unroll- 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);+ for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {+ #pragma unroll+ for (int k2 = 0; k2 < 256 / MMA_K; k2++) {+ const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);+ const 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 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);- const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);+ const int k_sf = k1 * 4 + k2;+ const int scale_A_tmem = SFA_tmem + k_sf * 4;+ const int scale_B_tmem = SFB2_tmem + k_sf * 4 + n_slice;- 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);+ const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;+ tcgen05_mma_nvfp4(OUT1_tmem, 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;");-- #pragma unroll- for (int mm = 0; mm < 2; mm++) {- float tmp[BLOCK_N / 2];- if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);- if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(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 half2 hx0 = reinterpret_cast<const half2 *>(G1_ptr + (row + 0) * N + col)[0];- const half2 hx8 = reinterpret_cast<const half2 *>(G1_ptr + (row + 8) * N + col)[0];- const float2 x0 = __half22float2(hx0);- const float2 x8 = __half22float2(hx8);- 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, 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(+ "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"+ :: "r"(done_mbar_addr + stage_id * 8)+ : "memory"+ );}- }- 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));- }- }+ if (step < num_iters) {+ const int iter_k = step;+ const int stage_id = iter_k % NUM_STAGES;+ const int tma_phase = (iter_k / NUM_STAGES) & 1;+ mbarrier_wait(tma0_mbar_addr + stage_id * 8, tma_phase);- 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_f32g1_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 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 lane_id = tid & 31;- const int warp_id = tid >> 5;+ 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 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;+ #pragma unroll+ for (int k = 0; k < K_STEPS; k++) {+ const uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);+ const uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);+ tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);+ tcgen05_cp_nvfp4(SFB1_tmem + k * 4, sfb_desc);+ }- 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()) {- uint64_t cache_A, cache_B;- if (M > N) { cache_A = EVICT_FIRST; cache_B = EVICT_LAST; }- else { cache_A = EVICT_LAST; cache_B = EVICT_FIRST; }-- 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"- );- };-- for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++) issue_tma(iter_k, iter_k);- for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {- const int stage_id = iter_k % NUM_STAGES;- const int mma_phase = (iter_k / NUM_STAGES - 1) & 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);-- #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 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);- }-- #pragma unroll- for (int k1 = 0; k1 < BLOCK_K / 256; k1++)#pragma unroll- 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);+ for (int k1 = 0; k1 < BLOCK_K / 256; k1++) {+ #pragma unroll+ for (int k2 = 0; k2 < 256 / MMA_K; k2++) {+ const uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);+ const 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 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);- const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);+ const int k_sf = k1 * 4 + k2;+ const int scale_A_tmem = SFA_tmem + k_sf * 4;+ const int scale_B_tmem = SFB1_tmem + k_sf * 4 + n_slice;- 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);+ const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;+ tcgen05_mma_nvfp4(0, 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"(b1_mbar_addr + stage_id * 8)+ : "memory"+ );+ }}asm volatile(⋯ 7 unchanged lines#pragma unrollfor (int mm = 0; mm < 2; mm++) {- float tmp[BLOCK_N / 2];- if constexpr (BLOCK_N == 64) tcgen05_ld_16x256bx8(tmp, warp_id * 32 + mm * 16, 0);- if constexpr (BLOCK_N == 128) tcgen05_ld_16x256bx16(tmp, warp_id * 32 + mm * 16, 0);+ float tmp_x[BLOCK_N / 2];+ float tmp_y[BLOCK_N / 2];+ if constexpr (BLOCK_N == 128) {+ tcgen05_ld_16x256bx16(tmp_x, warp_id * 32 + mm * 16, 0);+ tcgen05_ld_16x256bx16(tmp_y, warp_id * 32 + mm * 16, OUT1_tmem);+ } else {+ tcgen05_ld_16x256bx8(tmp_x, warp_id * 32 + mm * 16, 0);+ tcgen05_ld_16x256bx8(tmp_y, warp_id * 32 + mm * 16, OUT1_tmem);+ }asm volatile("tcgen05.wait::ld.sync.aligned;");#pragma unroll⋯ 1 unchanged linesconst 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]};+ const float2 x0 = float2{tmp_x[i * 4 + 0], tmp_x[i * 4 + 1]};+ const float2 x8 = float2{tmp_x[i * 4 + 2], tmp_x[i * 4 + 3]};+ const float2 y0 = float2{tmp_y[i * 4 + 0], tmp_y[i * 4 + 1]};+ const float2 y8 = float2{tmp_y[i * 4 + 2], tmp_y[i * 4 + 3]};- float2 o0, 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));+ float2 o0;+ float2 o8;+ const float s00 = fast_sigmoid(x0.x);+ const float s01 = fast_sigmoid(x0.y);+ const float s80 = fast_sigmoid(x8.x);+ const float s81 = fast_sigmoid(x8.y);o0.x = (x0.x * s00) * y0.x;o0.y = (x0.y * s01) * y0.y;o8.x = (x8.x * s80) * y8.x;⋯ 5 unchanged lines}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));+ if (warp_id == 0) asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(0), "r"(TMEM_COLS));}}template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>- static inline void launch_gemm_to_half(+ static inline void launch_dual(const at::Tensor& A,- const at::Tensor& B,+ const at::Tensor& B1,+ const at::Tensor& B2,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<half *>(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_to_half_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_to_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_to_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,+ const at::Tensor& SFB1,+ const at::Tensor& SFB2,at::Tensor& out) {const int M = (int)A.size(0);- const int N = (int)B.size(0);+ const int N = (int)B1.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 half *>(g1.data_ptr());- auto Out_ptr = reinterpret_cast<half *>(out.data_ptr());+ const char *A_ptr = reinterpret_cast<const char *>(A.data_ptr());+ const char *B1_ptr = reinterpret_cast<const char *>(B1.data_ptr());+ const char *B2_ptr = reinterpret_cast<const char *>(B2.data_ptr());+ const char *SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());+ const char *SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());+ const char *SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());+ half *Out_ptr = reinterpret_cast<half *>(out.data_ptr());- CUtensorMap A_tmap, B_tmap;+ CUtensorMap A_tmap, B1_tmap, B2_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);+ init_AB_tmap(&B1_tmap, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);+ init_AB_tmap(&B2_tmap, 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 + BLOCK_N) * (BLOCK_K / 2);- const int SFAB_size = 128 * (BLOCK_K / 16) * 2;- const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;+ const int smem_size = (BLOCK_M * BLOCK_K / 2 + 2 * (BLOCK_N * BLOCK_K / 2) + 128 * (BLOCK_K / 16) + 2 * (128 * (BLOCK_K / 16))) * NUM_STAGES;- auto kptr = gemm_silu_mul_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;+ auto kptr = dual_gemm_silu_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);+ kptr<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_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_f32g1(+ static inline void launch_dual_splitb(const at::Tensor& A,- const at::Tensor& B,+ const at::Tensor& B1,+ const at::Tensor& B2,const at::Tensor& SFA,- const at::Tensor& SFB,- const at::Tensor& g1,+ const at::Tensor& SFB1,+ const at::Tensor& SFB2,at::Tensor& out) {const int M = (int)A.size(0);- const int N = (int)B.size(0);+ const int N = (int)B1.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());+ const char *A_ptr = reinterpret_cast<const char *>(A.data_ptr());+ const char *B1_ptr = reinterpret_cast<const char *>(B1.data_ptr());+ const char *B2_ptr = reinterpret_cast<const char *>(B2.data_ptr());+ const char *SFA_ptr = reinterpret_cast<const char *>(SFA.data_ptr());+ const char *SFB1_ptr = reinterpret_cast<const char *>(SFB1.data_ptr());+ const char *SFB2_ptr = reinterpret_cast<const char *>(SFB2.data_ptr());+ half *Out_ptr = reinterpret_cast<half *>(out.data_ptr());- CUtensorMap A_tmap, B_tmap;+ CUtensorMap A_tmap, B1_tmap, B2_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);+ init_AB_tmap(&B1_tmap, B1_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, (uint32_t)BLOCK_K);+ init_AB_tmap(&B2_tmap, 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 + BLOCK_N) * (BLOCK_K / 2);- const int SFAB_size = 128 * (BLOCK_K / 16) * 2;- const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;+ const int smem_size = (BLOCK_M * BLOCK_K / 2 + (BLOCK_N * BLOCK_K / 2) + 2 * (128 * (BLOCK_K / 16))) * NUM_STAGES;- auto kptr = gemm_silu_mul_f32g1_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;+ auto kptr = dual_gemm_silu_kernel_splitb<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);+ kptr<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, Out_ptr, M, N);}- static inline void fused_dispatch(⋯ diff truncated
scrolls · 1201 diff lines total
Best evidence level for this revision: reported
JSON