submission 382077
yue · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 691 lines, June 9 Researcher Reciprocity License v1.0.
submit_v2.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-382077?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:a22cd92b336bd7ff51d8ca480cc28a165deb63b0d5c3415d323ad13502ac2ee4
license declaredunknown
license concludedunknown
authorsyue
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
mbarrier
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
asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;"Kernel source
submit_v2.py691 lines
# Dual GEMM: C = silu(A @ B1) * (A @ B2)
# Optimized version with better epilogue unrolling
import torch
from task import input_t, output_t
from torch.utils.cpp_extension import load_inline
CUDA_SRC_COMMON = r"""
#include <cudaTypedefs.h>
#include <cuda_fp16.h>
#include <torch/library.h>
#include <ATen/core/Tensor.h>
constexpr int WARP_SIZE = 32;
constexpr int MMA_K = 64; // 32 bytes
constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;
constexpr uint64_t EVICT_LAST = 0x14F0000000000000;
__device__ inline
constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
__device__
uint32_t elect_sync() {
uint32_t pred = 0;
asm volatile(
"{\n\t"
".reg .pred %%px;\n\t"
"elect.sync _|%%px, %1;\n\t"
"@%%px mov.s32 %0, 1;\n\t"
"}"
: "+r"(pred)
: "r"(0xFFFFFFFF)
);
return pred;
}
__device__ inline
void mbarrier_init(int mbar_addr, int count) {
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}
__device__ inline
void mbarrier_init_pred(int mbar_addr, int count, int pred) {
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %2, 0;\n\t"
"@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t"
"}"
:: "r"(mbar_addr), "r"(count), "r"(pred)
);
}
__device__
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"
"}"
:: "r"(mbar_addr), "r"(phase), "r"(ticks)
);
}
__device__ inline
void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;"
:: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy));
}
__device__ inline
void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
asm volatile("cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
"[%0], [%1, {%2, %3, %4}], [%5], %6;"
:: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
: "memory");
}
__device__ inline
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__ inline
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"
"}"
:: "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 _32x32b[] = ".32x32b";
};
struct NUM {
static constexpr char x32[] = ".x32";
static constexpr char x64[] = ".x64";
static constexpr char x128[] = ".x128";
};
template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_32regs(float *tmp, int row, int col) {
asm volatile("tcgen05.ld.sync.aligned%33%34.b32 "
"{ %0, %1, %2, %3, %4, %5, %6, %7, "
" %8, %9, %10, %11, %12, %13, %14, %15, "
" %16, %17, %18, %19, %20, %21, %22, %23, "
" %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
: "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
"=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
"=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
"=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
: "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}
template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_64regs(float *tmp, int row, int col) {
asm volatile("tcgen05.ld.sync.aligned%65%66.b32 "
"{ %0, %1, %2, %3, %4, %5, %6, %7, "
" %8, %9, %10, %11, %12, %13, %14, %15, "
" %16, %17, %18, %19, %20, %21, %22, %23, "
" %24, %25, %26, %27, %28, %29, %30, %31, "
" %32, %33, %34, %35, %36, %37, %38, %39, "
" %40, %41, %42, %43, %44, %45, %46, %47, "
" %48, %49, %50, %51, %52, %53, %54, %55, "
" %56, %57, %58, %59, %60, %61, %62, %63}, [%64];"
: "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
"=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
"=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
"=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31]),
"=f"(tmp[32]), "=f"(tmp[33]), "=f"(tmp[34]), "=f"(tmp[35]), "=f"(tmp[36]), "=f"(tmp[37]), "=f"(tmp[38]), "=f"(tmp[39]),
"=f"(tmp[40]), "=f"(tmp[41]), "=f"(tmp[42]), "=f"(tmp[43]), "=f"(tmp[44]), "=f"(tmp[45]), "=f"(tmp[46]), "=f"(tmp[47]),
"=f"(tmp[48]), "=f"(tmp[49]), "=f"(tmp[50]), "=f"(tmp[51]), "=f"(tmp[52]), "=f"(tmp[53]), "=f"(tmp[54]), "=f"(tmp[55]),
"=f"(tmp[56]), "=f"(tmp[57]), "=f"(tmp[58]), "=f"(tmp[59]), "=f"(tmp[60]), "=f"(tmp[61]), "=f"(tmp[62]), "=f"(tmp[63])
: "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}
template <const char *SHAPE, const char *NUM>
__device__ inline
void tcgen05_ld_128regs(float *tmp, int row, int col) {
asm volatile("tcgen05.ld.sync.aligned%129%130.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, %65, %66, %67, %68, %69, %70, %71, "
" %72, %73, %74, %75, %76, %77, %78, %79, "
" %80, %81, %82, %83, %84, %85, %86, %87, "
" %88, %89, %90, %91, %92, %93, %94, %95, "
" %96, %97, %98, %99,%100,%101,%102,%103, "
"%104,%105,%106,%107,%108,%109,%110,%111, "
"%112,%113,%114,%115,%116,%117,%118,%119, "
"%120,%121,%122,%123,%124,%125,%126,%127}, [%128];"
: "=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]),
"=f"(tmp[64]), "=f"(tmp[65]), "=f"(tmp[66]), "=f"(tmp[67]), "=f"(tmp[68]), "=f"(tmp[69]), "=f"(tmp[70]), "=f"(tmp[71]),
"=f"(tmp[72]), "=f"(tmp[73]), "=f"(tmp[74]), "=f"(tmp[75]), "=f"(tmp[76]), "=f"(tmp[77]), "=f"(tmp[78]), "=f"(tmp[79]),
"=f"(tmp[80]), "=f"(tmp[81]), "=f"(tmp[82]), "=f"(tmp[83]), "=f"(tmp[84]), "=f"(tmp[85]), "=f"(tmp[86]), "=f"(tmp[87]),
"=f"(tmp[88]), "=f"(tmp[89]), "=f"(tmp[90]), "=f"(tmp[91]), "=f"(tmp[92]), "=f"(tmp[93]), "=f"(tmp[94]), "=f"(tmp[95]),
"=f"(tmp[96]), "=f"(tmp[97]), "=f"(tmp[98]), "=f"(tmp[99]), "=f"(tmp[100]),"=f"(tmp[101]),"=f"(tmp[102]),"=f"(tmp[103]),
"=f"(tmp[104]),"=f"(tmp[105]),"=f"(tmp[106]),"=f"(tmp[107]),"=f"(tmp[108]),"=f"(tmp[109]),"=f"(tmp[110]),"=f"(tmp[111]),
"=f"(tmp[112]),"=f"(tmp[113]),"=f"(tmp[114]),"=f"(tmp[115]),"=f"(tmp[116]),"=f"(tmp[117]),"=f"(tmp[118]),"=f"(tmp[119]),
"=f"(tmp[120]),"=f"(tmp[121]),"=f"(tmp[122]),"=f"(tmp[123]),"=f"(tmp[124]),"=f"(tmp[125]),"=f"(tmp[126]),"=f"(tmp[127])
: "r"((row << 16) | col), "C"(SHAPE), "C"(NUM));
}
__device__ inline void tcgen05_ld_32x32bx32(float *tmp, int row, int col) { tcgen05_ld_32regs<SHAPE::_32x32b, NUM::x32>(tmp, row, col); }
__device__ inline void tcgen05_ld_32x32bx64(float *tmp, int row, int col) { tcgen05_ld_64regs<SHAPE::_32x32b, NUM::x64>(tmp, row, col); }
__device__ inline void tcgen05_ld_32x32bx128(float *tmp, int row, int col) { tcgen05_ld_128regs<SHAPE::_32x32b, NUM::x128>(tmp, row, col); }
// Fast silu using PTX SFU - optimized for pipelining
__device__ inline
float silu(float x) {
float result;
asm volatile(
"{\n\t"
".reg .f32 neg_scaled, exp_val, denom, rcp_val;\n\t"
"mul.f32 neg_scaled, %1, 0FBFB8AA3B;\n\t"
"ex2.approx.ftz.f32 exp_val, neg_scaled;\n\t"
"add.f32 denom, exp_val, 0F3F800000;\n\t"
"rcp.approx.ftz.f32 rcp_val, denom;\n\t"
"mul.f32 %0, %1, rcp_val;\n\t"
"}"
: "=f"(result)
: "f"(x)
);
return result;
}
void check_cu(CUresult err) {
if (err == CUDA_SUCCESS) return;
const char *error_msg_ptr;
if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
error_msg_ptr = "unable to get error string";
TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
}
void init_AB_tmap(
CUtensorMap *tmap,
const char *ptr,
uint64_t global_height, uint64_t global_width,
uint32_t shared_height, uint32_t shared_width
) {
constexpr uint32_t rank = 3;
uint64_t globalDim[rank] = {256, global_height, global_width / 256};
uint64_t globalStrides[rank-1] = {global_width / 2, 128};
uint32_t boxDim[rank] = {256, shared_height, shared_width / 256};
uint32_t elementStrides[rank] = {1, 1, 1};
auto err = cuTensorMapEncodeTiled(
tmap,
CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
rank,
(void *)ptr,
globalDim,
globalStrides,
boxDim,
elementStrides,
CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
);
check_cu(err);
}
"""
CUDA_SRC_DUAL_GEMM = r"""
template <
int K,
int BLOCK_M,
int BLOCK_N,
int BLOCK_K,
int NUM_STAGES
>
__global__
__launch_bounds__(BLOCK_M + 2 * WARP_SIZE, 1)
void dual_gemm_kernel(
const __grid_constant__ CUtensorMap A_tmap,
const __grid_constant__ CUtensorMap B1_tmap,
const __grid_constant__ CUtensorMap B2_tmap,
const char *SFA_ptr,
const char *SFB1_ptr,
const char *SFB2_ptr,
half *C_ptr,
int M, int N
) {
const int tid = threadIdx.x;
const int bid = blockIdx.x;
const int lane_id = tid % WARP_SIZE;
const int warp_id = tid / WARP_SIZE;
const int grid_m = M / BLOCK_M;
const int grid_n = N / BLOCK_N;
const int bid_m = bid / grid_n;
const int bid_n = bid % grid_n;
const int off_m = bid_m * BLOCK_M;
const int off_n = bid_n * BLOCK_N;
constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2;
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 * 2 + SFA_size + SFB_size * 2;
#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 SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);
mbarrier_init_pred(tma_mbar_addr + lane_id * 8, 1, (warp_id == 0) & (lane_id < NUM_STAGES * 2 + 1));
if (warp_id == 1) {
asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 4));
}
__syncthreads();
const int num_iters = K / BLOCK_K;
if (warp_id == 4 && elect_sync()) {
constexpr uint64_t cache_A = EVICT_LAST;
constexpr uint64_t cache_B = EVICT_FIRST;
constexpr int rest_k = K / 16 / 4;
const char *SFA_base = SFA_ptr + (off_m / 128) * rest_k * 512;
const char *SFB1_base = SFB1_ptr + (off_n / 128) * rest_k * 512;
const char *SFB2_base = SFB2_ptr + (off_n / 128) * rest_k * 512;
auto issue_tma = [&](int iter_k, int stage_id) {
const int mbar_addr = tma_mbar_addr + stage_id * 8;
const int stage_smem = smem + stage_id * STAGE_SIZE;
const int A_smem = stage_smem;
const int B1_smem = stage_smem + A_size;
const int B2_smem = stage_smem + A_size + B_size;
const int SFA_smem = stage_smem + A_size + B_size * 2;
const int SFB1_smem = stage_smem + A_size + B_size * 2 + SFA_size;
const int SFB2_smem = stage_smem + A_size + B_size * 2 + SFA_size + SFB_size;
const int k_idx = iter_k * (BLOCK_K / 256);
tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, k_idx, mbar_addr, cache_A);
tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, k_idx, mbar_addr, cache_B);
tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, k_idx, mbar_addr, cache_B);
const char *SFA_src = SFA_base + iter_k * (BLOCK_K / 64) * 512;
const char *SFB1_src = SFB1_base + iter_k * (BLOCK_K / 64) * 512;
const char *SFB2_src = SFB2_base + iter_k * (BLOCK_K / 64) * 512;
tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
tma_gmem2smem(SFB1_smem, SFB1_src, SFB_size, mbar_addr, cache_B);
tma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_B);
asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
:: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
};
for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++)
issue_tma(iter_k, iter_k);
for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
const int stage_id = iter_k % NUM_STAGES;
const int mma_phase = (iter_k / NUM_STAGES - 1) % 2;
mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
issue_tma(iter_k, stage_id);
}
}
else if (warp_id == 5 && elect_sync()) {
constexpr uint32_t i_desc = (1U << 7U)
| (1U << 10U)
| ((uint32_t)BLOCK_N >> 3U << 17U)
| ((uint32_t)128 >> 7U << 27U);
for (int iter_k = 0; iter_k < num_iters; iter_k++) {
const int stage_id = iter_k % NUM_STAGES;
const int tma_phase = (iter_k / NUM_STAGES) % 2;
mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);
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);
constexpr int d_tmem_acc2 = BLOCK_N * 2;
constexpr uint64_t SF_stride = 512ULL >> 4ULL;
const int scale_A_base = SFA_tmem + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
const int scale_B1_base = SFB1_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
const int scale_B2_base = SFB2_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
const uint64_t a_desc_base = make_desc_AB(A_smem);
const uint64_t b1_desc_base = make_desc_AB(B1_smem);
const uint64_t b2_desc_base = make_desc_AB(B2_smem);
tcgen05_cp_nvfp4(SFA_tmem, SFA_desc);
tcgen05_cp_nvfp4(SFB1_tmem, SFB1_desc);
tcgen05_cp_nvfp4(SFB2_tmem, SFB2_desc);
constexpr int A_k256_stride = BLOCK_M * 128;
constexpr int B_k256_stride = BLOCK_N * 128;
constexpr int NUM_K_ITERS = BLOCK_K / MMA_K;
#pragma unroll
for (int k = 0; k < NUM_K_ITERS; k++) {
const int k1 = k >> 2;
const int k2 = k & 3;
if (k + 1 < NUM_K_ITERS) {
const int next_k = k + 1;
tcgen05_cp_nvfp4(SFA_tmem + (next_k << 2), SFA_desc + next_k * SF_stride);
tcgen05_cp_nvfp4(SFB1_tmem + (next_k << 2), SFB1_desc + next_k * SF_stride);
tcgen05_cp_nvfp4(SFB2_tmem + (next_k << 2), SFB2_desc + next_k * SF_stride);
}
const int ab_off = k1 * A_k256_stride + (k2 << 5);
const int bb_off = k1 * B_k256_stride + (k2 << 5);
const uint64_t a_desc = a_desc_base + (ab_off >> 4);
const uint64_t b1_desc = b1_desc_base + (bb_off >> 4);
const uint64_t b2_desc = b2_desc_base + (bb_off >> 4);
const int scale_A = scale_A_base + (k << 2);
const int scale_B1 = scale_B1_base + (k << 2);
const int scale_B2 = scale_B2_base + (k << 2);
const int enable_d = (k == 0) ? iter_k : 1;
tcgen05_mma_nvfp4(0, a_desc, b1_desc, i_desc, scale_A, scale_B1, enable_d);
tcgen05_mma_nvfp4(d_tmem_acc2, a_desc, b2_desc, i_desc, scale_A, scale_B2, enable_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) {
// Optimized epilogue with 4-way unrolling
mbarrier_wait(mainloop_mbar_addr, 0);
asm volatile("tcgen05.fence::after_thread_sync;");
constexpr int WIDTH = std::min(BLOCK_N, 64);
const int col = off_m + tid;
for (int n = 0; n < BLOCK_N / WIDTH; n++) {
float tmp1[WIDTH];
float tmp2[WIDTH];
if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp1, warp_id * 32, n * WIDTH);
if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp1, warp_id * 32, n * WIDTH);
if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp1, warp_id * 32, n * WIDTH);
if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);
if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);
if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);
asm volatile("tcgen05.wait::ld.sync.aligned;");
const int base_row = off_n + n * WIDTH;
half * __restrict__ C_row = C_ptr + base_row * M + col;
// 4-way unrolled loop for better ILP
#pragma unroll
for (int i = 0; i < WIDTH; i += 4) {
float r0 = silu(tmp1[i]) * tmp2[i];
float r1 = silu(tmp1[i+1]) * tmp2[i+1];
float r2 = silu(tmp1[i+2]) * tmp2[i+2];
float r3 = silu(tmp1[i+3]) * tmp2[i+3];
C_row[i * M] = __float2half(r0);
C_row[(i+1) * M] = __float2half(r1);
C_row[(i+2) * M] = __float2half(r2);
C_row[(i+3) * M] = __float2half(r3);
}
}
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));
}
}
template <
int K,
int BLOCK_M,
int BLOCK_N,
int BLOCK_K,
int NUM_STAGES
>
at::Tensor dual_gemm_launch(
const at::Tensor& A,
const at::Tensor& B1,
const at::Tensor& B2,
const at::Tensor& SFA,
const at::Tensor& SFB1,
const at::Tensor& SFB2,
at::Tensor& C
) {
static_assert(BLOCK_K % 256 == 0);
const int M = A.size(0);
const int N = B1.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 C_ptr = reinterpret_cast<half *>(C.data_ptr());
CUtensorMap A_tmap, B1_tmap, B2_tmap;
init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
init_AB_tmap(&B1_tmap, B1_ptr, N, K, BLOCK_N, BLOCK_K);
init_AB_tmap(&B2_tmap, B2_ptr, N, K, BLOCK_N, BLOCK_K);
int grid = (M / BLOCK_M) * (N / BLOCK_N);
int tb_size = BLOCK_M + 2 * WARP_SIZE;
int AB_size = (BLOCK_M + BLOCK_N * 2) * (BLOCK_K / 2);
int SFAB_size = 128 * (BLOCK_K / 16) * 3;
int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
auto this_kernel = dual_gemm_kernel<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
if (smem_size > 48'000)
cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
this_kernel<<<grid, tb_size, smem_size>>>(A_tmap, B1_tmap, B2_tmap, SFA_ptr, SFB1_ptr, SFB2_ptr, C_ptr, M, N);
return C.view({N, M, 1}).transpose(0, 1);
}
at::Tensor dual_gemm(
const at::Tensor& A,
const at::Tensor& B1,
const at::Tensor& B2,
const at::Tensor& SFA,
const at::Tensor& SFB1,
const at::Tensor& SFB2,
at::Tensor& C
) {
const int K = A.size(1) * 2;
const int M = A.size(0);
const int N = B1.size(0);
#define LAUNCH_EXACT(K_, M_, N_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \
else if (K == K_ && M == M_ && N == N_) C = dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);
#define LAUNCH(K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \
else if (K == K_) C = dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);
if (false) {}
LAUNCH_EXACT(7168, 256, 4096, 128, 64, 256, 5)
LAUNCH_EXACT(7168, 512, 4096, 128, 128, 256, 4)
LAUNCH_EXACT(4096, 256, 3072, 128, 64, 256, 5)
LAUNCH_EXACT(7168, 512, 3072, 128, 128, 256, 4)
LAUNCH(16384, 128, 64, 256, 5)
LAUNCH( 7168, 128, 64, 256, 5)
LAUNCH( 4096, 128, 64, 256, 5)
LAUNCH( 3072, 128, 64, 256, 5)
LAUNCH( 2048, 128, 64, 256, 5)
LAUNCH( 2304, 128, 64, 256, 5)
LAUNCH( 1536, 128, 64, 256, 5)
LAUNCH( 1024, 128, 64, 256, 5)
LAUNCH( 512, 128, 64, 256, 5)
LAUNCH( 256, 128, 64, 256, 5)
#undef LAUNCH_EXACT
#undef LAUNCH
return C;
}
TORCH_LIBRARY(my_module_dual_gemm_v2, m) {
m.def("dual_gemm(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");
m.impl("dual_gemm", &dual_gemm);
}
"""
load_inline(
"dual_gemm_cuda_v2",
cpp_sources="",
cuda_sources=CUDA_SRC_COMMON + CUDA_SRC_DUAL_GEMM,
verbose=True,
is_python_module=False,
no_implicit_headers=True,
extra_cuda_cflags=[
"-O3",
"-gencode=arch=compute_100a,code=sm_100a",
"--use_fast_math",
"--expt-relaxed-constexpr",
"--relocatable-device-code=false",
"-lineinfo",
"-Xptxas=-v",
],
extra_ldflags=["-lcuda"],
)
dual_gemm_cuda = torch.ops.my_module_dual_gemm_v2.dual_gemm
def custom_kernel(data: input_t) -> output_t:
return dual_gemm_cuda(data[0], data[1], data[2], data[6], data[7], data[8], data[9])
scrolls · 691 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 375735.
# Dual GEMM: C = silu(A @ B1) * (A @ B2)- #!POPCORN leaderboard modal_nvfp4_dual_gemm- #!POPCORN gpu B200+ # Optimized version with better epilogue unrollingimport torchfrom task import input_t, output_tfrom torch.utils.cpp_extension import load_inline--CUDA_SRC_COMMON = r"""#include <cudaTypedefs.h>#include <cuda_fp16.h>⋯ 5 unchanged linesconstexpr int MMA_K = 64; // 32 bytes- // https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;constexpr uint64_t EVICT_FIRST = 0x12F0000000000000;constexpr uint64_t EVICT_LAST = 0x14F0000000000000;- // x & 0x3'FFFFULL: keep only the lower 18 bit, 0x3'FFFF is 256KB - 1- // limit the addressable range to 256KB (max shared memory size)- // >> 4ULL: right shift 4 bits, divide by 16, because the shared memory is 16-byte aligned- // the tcgen05_mma_nvfp4 will automatically do the * 16 to get the actual shared memory byte address.- // It designed as this to save bits (save 4 bits, also force alignment)+__device__ inlineconstexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };- // https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cute/arch/cluster_sm90.hpp#L180__device__uint32_t elect_sync() {uint32_t pred = 0;⋯ 29 unchanged lines}- // https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cutlass/arch/barrier.h#L408__device__void mbarrier_wait(int mbar_addr, int phase) {uint32_t ticks = 0x989680;⋯ 13 unchanged lines__device__ inlinevoid tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {- // This is for SFA, SFB copy.- 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));+ 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__ inlinevoid tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {- // This is for transfer data for A, B, which is represent as 3d tensor.- // r is for 32-bit register for int values; l is for 64 bit register for pointers and uint64_t- // memory: tells compiler this instruction touches memory.- // cp.async.bulk: TMA copy- // tensor.3d: operating on 3d tensor- // shared::cta: destination is CTA-scoped shared memory; can use shared::cluster for all CTAs in the same cluster, all CTA get the same data in one load.- // global: source is global memory- // mbarrier::complete_tx::byte: use mbarrier to track completion, track bytes trasferred- // cta_group::1: cta group size is 1, only this CTA track completion- // L2::cache_hint: apply L2 cache policy hint- 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");+ 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");}- // The tcgen05.cp instruction is issued by warp 5 (the MMA warp in this code),- // but the actual TMEM storage is managed by the TMEM controller, which organizes- // data according to the warp structure of warps 0-3. (my understanding)__device__ inlinevoid tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {- // .32x128b corresponds to (32, 16) 8-bit scale -> 1 MMA for nvfp4.- // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-cp- // .32x128b require .warpx4, the data will be multicasted into all 4 warps.- // row 0 - 31 replicated to row 32 - 63 and row 64 - 95 and row 96 - 127,- // because each warp's tensor core can only access it's portion of tmem during mma.asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));}-- // setp.ne.b32 p, %6, 0;- // p = (enable_input_d != 0), p = false is overwrite mode,- // p = true is accumulate mode, when k = 0 is overwrite, when k > 0 is accumulate.- // block_scale: block wise scaling, not per element scaling- // .block16: 1 scale per 16 elements__device__ inlinevoid tcgen05_mma_nvfp4(int d_tmem,⋯ 16 unchanged lines}- // see https://docs.nvidia.com/cuda/inline-ptx-assembly/index.htmlstruct SHAPE {- static constexpr char _32x32b[] = ".32x32b"; // 32x1 tile for each warp+ static constexpr char _32x32b[] = ".32x32b";};⋯ 89 unchanged lines__device__ inline void tcgen05_ld_32x32bx128(float *tmp, int row, int col) { tcgen05_ld_128regs<SHAPE::_32x32b, NUM::x128>(tmp, row, col); }- // Fast silu using PTX SFU with FMA for better pipelining- // silu(x) = x * sigmoid(x) = x / (1 + exp(-x)) = x / (1 + 2^(-x*log2e))- // TODO: use emulation to not use SFU - ex2.approx which use SFU but only have 1 per TC- // reference FA4- // TODO: use f32x2 with x * y * sigmoid(x)+ // Fast silu using PTX SFU - optimized for pipelining__device__ inlinefloat silu(float x) {float result;asm volatile("{\n\t"".reg .f32 neg_scaled, exp_val, denom, rcp_val;\n\t"- "mul.f32 neg_scaled, %1, 0FBFB8AA3B;\n\t" // -x * log2(e) in one instruction- "ex2.approx.ftz.f32 exp_val, neg_scaled;\n\t" // exp(-x)- "add.f32 denom, exp_val, 0F3F800000;\n\t" // 1 + exp(-x)- "rcp.approx.ftz.f32 rcp_val, denom;\n\t" // 1/(1+exp(-x))- "mul.f32 %0, %1, rcp_val;\n\t" // x * sigmoid(x)+ "mul.f32 neg_scaled, %1, 0FBFB8AA3B;\n\t"+ "ex2.approx.ftz.f32 exp_val, neg_scaled;\n\t"+ "add.f32 denom, exp_val, 0F3F800000;\n\t"+ "rcp.approx.ftz.f32 rcp_val, denom;\n\t"+ "mul.f32 %0, %1, rcp_val;\n\t""}": "=f"(result): "f"(x)⋯ 11 unchanged lines}- void check_cuda(cudaError_t err) {- if (err == cudaSuccess) return;- TORCH_CHECK(false, cudaGetErrorString(err));- }--void init_AB_tmap(CUtensorMap *tmap,const char *ptr,⋯ 1 unchanged linesuint32_t shared_height, uint32_t shared_width) {constexpr uint32_t rank = 3;- // swizzle unit: 128 byte, which is 256 values for fp4uint64_t globalDim[rank] = {256, global_height, global_width / 256};- uint64_t globalStrides[rank-1] = {global_width / 2, 128}; // in bytes+ uint64_t globalStrides[rank-1] = {global_width / 2, 128};uint32_t boxDim[rank] = {256, shared_height, shared_width / 256};uint32_t elementStrides[rank] = {1, 1, 1};⋯ 17 unchanged lines"""--CUDA_SRC_DUAL_GEMM = r"""template <int K,⋯ 37 unchanged linesconst 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;- // always copy 128xBLOCK_K/16, because the tcgen05.cp that copy from smem to tmem- // has work work with fixed (32, 16) 8 bit scale, which is 4 tmem columns, each column has 32 rows- // and each column value has 4 8-bit scale factor as draw in https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x- // The scale factor need to layout in that way, so our SFA and SFB both need to be 128 rows,- // so each MMA work on (128, 4) scale factors which can be (32, 4, 4) which is (32, 16).- // so no matter what the block_m and block_n are we need the 128 * (BLOCK_K / 16), the BLOCK_K / 16- // includes the scale factors for all MMA in the block (e.g. BLOCK_K=256, MMA_K = 64, so it has shape- // [128, 4 (4 MMA), 32, 4, 4])constexpr int SFA_size = 128 * BLOCK_K / 16;constexpr int SFB_size = 128 * BLOCK_K / 16;- // For dual GEMM: A, B1, B2, SFA, SFB1, SFB2constexpr int STAGE_SIZE = A_size + B_size * 2 + SFA_size + SFB_size * 2;- // set up mbarriers and tmem- // we have NUM_STAGES mbars for TMA, NUM_STAGES mbars for MMA, 1 mbar for mainloop#pragma nv_diag_suppress static_var_with_dynamic_init__shared__ int64_t mbars[NUM_STAGES * 2 + 1];const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));- // byte offset: each barrier takes int64_t which is 8 bytes, so *8 here.const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;- // TMEM layout for dual GEMM:- // - Acc1: columns 0 to BLOCK_N-1 (d_tmem = 0)- // - SFA: columns BLOCK_N onwards- // - SFB1, SFB2: after SFA- // - Acc2: columns BLOCK_N*2 onwards (d_tmem = BLOCK_N*2)constexpr int SFA_tmem = BLOCK_N;- // The columns of SFA_tmem is 4 * 4, the first 4 is 4 tmem columns- // (128, 64 / 16) = (128, 4) per mma = (32, 4, 4) as in https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x- // the second 4 mmas per block_k.constexpr int SFB1_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);constexpr int SFB2_tmem = SFB1_tmem + 4 * (BLOCK_K / MMA_K);mbarrier_init_pred(tma_mbar_addr + lane_id * 8, 1, (warp_id == 0) & (lane_id < NUM_STAGES * 2 + 1));- // We actually not need BLOCK_N * 4 TMEM, but if we use BLOCK_N * 3 here it will error out,- // so somehow TMEM require it to be power of 2?if (warp_id == 1) {- // TMEM allocation MUST be power of 2! Tested: 384 fails, 448 fails, 512 works- // For BLOCK_N=64: need 192 cols (Acc2 ends at 191), so 256 works- // For BLOCK_N=128: need 384 cols (Acc2 ends at 383), so 512 neededasm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(smem), "r"(BLOCK_N * 4));}__syncthreads();⋯ 3 unchanged linesif (warp_id == 4 && elect_sync()) {- // TMA warp - cache eviction policy (optimized for M < N benchmark cases)- // A tile reuse = N/BLOCK_N (high when N > M) → keep in cache longest- // B tile reuse = M/BLOCK_M (low when M < N) → evict firstconstexpr uint64_t cache_A = EVICT_LAST;constexpr uint64_t cache_B = EVICT_FIRST;- // K / 16 / 4: 16 elements share same scale factor, 4 fp8 packed to 1 value- // it's the total number of (4 scale factors pack to one 32 bit) scale factors along K.constexpr int rest_k = K / 16 / 4;- // layout of SFA IS [M/128, rest_k, 32, 4, 4]- // so SFA_base:- // Move ptr to the current block's position for the SFA, SFB - not including- // the k iter offset yet, which will be added during the loop along k.const char *SFA_base = SFA_ptr + (off_m / 128) * rest_k * 512;- // SFB is [N/128, rest_k, 32, 4, 4]- // When block_n = 64, because we are always using 128 here, whcih is required by- // tcgen05, then block 0 and block 1 both load the entire 128 n rows, and block 0- // use the first half and block 1 use the second half.const char *SFB1_base = SFB1_ptr + (off_n / 128) * rest_k * 512;const char *SFB2_base = SFB2_ptr + (off_n / 128) * rest_k * 512;⋯ 9 unchanged linesconst int SFB2_smem = stage_smem + A_size + B_size * 2 + SFA_size + SFB_size;- // issue TMA for A, B1, B2- // iter_k * (BLOCK_K / 256), BLOCK_K = 256- const int k_idx = iter_k;- // the A_tmap has globalDim as (256, M, K/256), off_m is for row (M), k_idx is for col (K/256).- // 0 means start from the first value in the block.+ const int k_idx = iter_k * (BLOCK_K / 256);tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, k_idx, mbar_addr, cache_A);tma_3d_gmem2smem(B1_smem, &B1_tmap, 0, off_n, k_idx, mbar_addr, cache_B);tma_3d_gmem2smem(B2_smem, &B2_tmap, 0, off_n, k_idx, mbar_addr, cache_B);- // Scale factor TMA with pre-computed baseconst char *SFA_src = SFA_base + iter_k * (BLOCK_K / 64) * 512;const char *SFB1_src = SFB1_base + iter_k * (BLOCK_K / 64) * 512;const char *SFB2_src = SFB2_base + iter_k * (BLOCK_K / 64) * 512;⋯ 2 unchanged linestma_gmem2smem(SFB2_smem, SFB2_src, SFB_size, mbar_addr, cache_B);- // mbarrier: memory barrier- // .arrive: signal arrival at the barrier- // .expect_tx: expect a transaction - tells barrier how many bytes to wait for- // .release: ensure all prior writes are visible before arriving- // .cta: within the CTA barrier sync- // .shared::cta: barrier is in CTA-scoped shared memory- // .b64: barrier is 64 bit- // _ discard the return value- // This is a non blocking instruction that tells the barrier to expect bytes from- // TMA transaction and send arrival signal in the following mbarrier_wait.asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;":: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");};- // prefetch: issue NUM_STAGES TMA callfor (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;- // e.g. when NUM_STAGES == 5, when iter_k = 5 - 9, wait mma_phase = 0,- // when iter_k = 10 - 14, wait mma_phase = 1,- // when iter_k = 15 - 19, wait mma_phase = 0, ...const int mma_phase = (iter_k / NUM_STAGES - 1) % 2;mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);issue_tma(iter_k, stage_id);}}else if (warp_id == 5 && elect_sync()) {- // from https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instuction-desc-kind-mxf4-mxf4nvf4- // see E2M1 = 1, so we do (1U << 7U) and (1U << 10U)- // also according to the table, the block_n need >> 3U, and block_m need >>7U- // which is assuming N must be 8X, and M must be 128X, so those 3 or 7 bits are- // all zero, so we can save some bit for better compression.- // MMA instruction descriptor.- constexpr uint32_t i_desc = (1U << 7U) // atype=E2M1- | (1U << 10U) // btype=E2M1- | ((uint32_t)BLOCK_N >> 3U << 17U) // MMA_N- | ((uint32_t)128 >> 7U << 27U); // MMA_M+ constexpr uint32_t i_desc = (1U << 7U)+ | (1U << 10U)+ | ((uint32_t)BLOCK_N >> 3U << 17U)+ | ((uint32_t)128 >> 7U << 27U);for (int iter_k = 0; iter_k < num_iters; iter_k++) {⋯ 9 unchanged linesconst int SFB1_smem = SFA_smem + SFA_size;const int SFB2_smem = SFB1_smem + SFB_size;- // This matches the swizzle as specified in init_AB_tmap CU_TENSOR_MAP_SWIZZLE_128B- // which is how TMA saved data in SMEM.- // The bit shifting follows https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-shared-memory-descriptor.- // Need to set bits that are not 0.+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);};- // no swizzling for scale factorsauto make_desc_SF = [](int addr) -> uint64_t {const int SBO = 8 * 16;return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);⋯ 1 unchanged linesconstexpr uint64_t SF_desc = make_desc_SF(0);- // address need to make 4 bit right shift, same as in desc_encode.- // SFA_desc: describe the SFA in shared memory - it's address, swizzle mode, etcconst uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);const uint64_t SFB1_desc = SF_desc + ((uint64_t)SFB1_smem >> 4ULL);const uint64_t SFB2_desc = SF_desc + ((uint64_t)SFB2_smem >> 4ULL);constexpr int d_tmem_acc2 = BLOCK_N * 2;- // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x- // 32 * 4 * 4 = 512constexpr uint64_t SF_stride = 512ULL >> 4ULL;- // TMEM column addresses for scale factors.- // Scale factors are always loaded in 128-row chunks (tcgen05.cp requirement).- // Each (32, 4, 4) tile contains scale factors for 128 M/N rows.- // When BLOCK_M/N < 128, one tile covers multiple MMA tiles, so we offset to- // select the portion corresponding to this CTA's M/N range:- // - (bid % (128/BLOCK)): which sub-portion of the 128 rows this MMA tile uses- // - (BLOCK/32): TMEM columns per sub-portion (32 rows per TMEM column)- // E.g., BLOCK_N=64: bid_n=0 uses cols 0-1 (N=0-63), bid_n=1 uses cols 2-3 (N=64-127)- // When BLOCK_M/N=128: offset is always 0 (MMA tile uses entire scale factor tile)- // This is the offset calculating for the row (N) direction - when BLOCK_N = 64, for bid_n = 1, 3, 5, ...- // they will have 2 TMEM column offset in the TMEM, and the first 2 TMEM rows are wasted.- // Loaded: Full 128 N rows of scale factors (tcgen05.cp requirement)- // Used: Only 64 N rows (half), corresponding to this CTA's N range- // What's loaded into TMEM (per K iteration):┌─────────────┬─────────────┐│ Cols 0-1 │ Cols 2-3 ││ N = 0-63 │ N = 64-127 ││ (64 rows) │ (64 rows) │└─────────────┴─────────────┘ ↑ ↑ bid_n=0 bid_n=1 uses this uses this (half) (half)- // So there's some inefficiency when BLOCK_N < 128:- // You load 128 rows but only use BLOCK_N rows- // The unused portion wastes TMEM space and memory bandwidthconst int scale_A_base = SFA_tmem + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);const int scale_B1_base = SFB1_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);const int scale_B2_base = SFB2_tmem + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);⋯ 4 unchanged linesconst uint64_t b2_desc_base = make_desc_AB(B2_smem);- // Prefetch first scale factortcgen05_cp_nvfp4(SFA_tmem, SFA_desc);tcgen05_cp_nvfp4(SFB1_tmem, SFB1_desc);tcgen05_cp_nvfp4(SFB2_tmem, SFB2_desc);++ constexpr int A_k256_stride = BLOCK_M * 128;+ constexpr int B_k256_stride = BLOCK_N * 128;constexpr int NUM_K_ITERS = BLOCK_K / MMA_K;#pragma unrollfor (int k = 0; k < NUM_K_ITERS; k++) {- // Prefetch next scale (only if within bounds)+ const int k1 = k >> 2;+ const int k2 = k & 3;++if (k + 1 < NUM_K_ITERS) {const int next_k = k + 1;- // Always move 4 TMEM cols to next MMA no matter what the BLOCK_N is.tcgen05_cp_nvfp4(SFA_tmem + (next_k << 2), SFA_desc + next_k * SF_stride);tcgen05_cp_nvfp4(SFB1_tmem + (next_k << 2), SFB1_desc + next_k * SF_stride);tcgen05_cp_nvfp4(SFB2_tmem + (next_k << 2), SFB2_desc + next_k * SF_stride);}- // Offset within shared memory for this MMA_K=64 chunk- // Each MMA_K=64 nvfp4 = 32 bytes, so offset = k * 32 = k << 5- // we are assuming use BLOCK_K = 256 here. Otherwisae we need to do:- // ab_off = k1 * A_k256_stride + (k2 << 5), bb_off = k1 * B_k256_stride + (k2 << 5)- const int ab_off = k << 5;- const int bb_off = k << 5;+ const int ab_off = k1 * A_k256_stride + (k2 << 5);+ const int bb_off = k1 * B_k256_stride + (k2 << 5);const uint64_t a_desc = a_desc_base + (ab_off >> 4);const uint64_t b1_desc = b1_desc_base + (bb_off >> 4);const uint64_t b2_desc = b2_desc_base + (bb_off >> 4);- // jump 4 TMEM columns per MMA iteration.+const int scale_A = scale_A_base + (k << 2);const int scale_B1 = scale_B1_base + (k << 2);const int scale_B2 = scale_B2_base + (k << 2);const int enable_d = (k == 0) ? iter_k : 1;- // we don't need explicit wait on TMEM here, because the TC tracks dependencies on TMEM and stalls if data not ready.tcgen05_mma_nvfp4(0, a_desc, b1_desc, i_desc, scale_A, scale_B1, enable_d);tcgen05_mma_nvfp4(d_tmem_acc2, a_desc, b2_desc, i_desc, scale_A, scale_B2, enable_d);}- // signal the mma barrier of arrival, consumer is TMA warp+asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];":: "r"(mma_mbar_addr + stage_id * 8) : "memory");}⋯ 3 unchanged lines:: "r"(mainloop_mbar_addr) : "memory");}else if (tid < BLOCK_M) {- // epilogue warps - read both accumulators and apply silu(Acc1) * Acc2+ // Optimized epilogue with 4-way unrollingmbarrier_wait(mainloop_mbar_addr, 0);- // Fence before reading TMEMasm volatile("tcgen05.fence::after_thread_sync;");- // The WIDTH constant controls how many TMEM columns are read per tcgen05.ld instruction, and it's capped at 64 due to register pressure constraints.- constexpr int WIDTH = 64;+ constexpr int WIDTH = std::min(BLOCK_N, 64);+ const int col = off_m + tid;+for (int n = 0; n < BLOCK_N / WIDTH; n++) {float tmp1[WIDTH];float tmp2[WIDTH];- // Issue BOTH TMEM reads first to overlap latency+ if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp1, warp_id * 32, n * WIDTH);if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp1, warp_id * 32, n * WIDTH);+ if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp1, warp_id * 32, n * WIDTH);+ if constexpr (WIDTH == 128) tcgen05_ld_32x32bx128(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);if constexpr (WIDTH == 64) tcgen05_ld_32x32bx64(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);- //synchronization barrier that stalls the thread until all previously issued tcgen05.ld instructions have finished loading data into registers.+ if constexpr (WIDTH == 32) tcgen05_ld_32x32bx32(tmp2, warp_id * 32, n * WIDTH + BLOCK_N * 2);asm volatile("tcgen05.wait::ld.sync.aligned;");- const int col = off_m + tid;const int base_row = off_n + n * WIDTH;+ half * __restrict__ C_row = C_ptr + base_row * M + col;+ // 4-way unrolled loop for better ILP#pragma unroll- for (int i = 0; i < WIDTH; i += 2) {+ for (int i = 0; i < WIDTH; i += 4) {float r0 = silu(tmp1[i]) * tmp2[i];float r1 = silu(tmp1[i+1]) * tmp2[i+1];- // TODO: use f32x2 with x * y * sigmoid(x) for the convert- C_ptr[(base_row + i) * M + col] = __float2half(r0);- C_ptr[(base_row + i + 1) * M + col] = __float2half(r1);+ float r2 = silu(tmp1[i+2]) * tmp2[i+2];+ float r3 = silu(tmp1[i+3]) * tmp2[i+3];+ C_row[i * M] = __float2half(r0);+ C_row[(i+1) * M] = __float2half(r1);+ C_row[(i+2) * M] = __float2half(r2);+ C_row[(i+3) * M] = __float2half(r3);}}⋯ 12 unchanged linesint BLOCK_K,int NUM_STAGES>- at::Tensor cutlass_dual_gemm_launch(+ at::Tensor dual_gemm_launch(const at::Tensor& A,const at::Tensor& B1,const at::Tensor& B2,⋯ 56 unchanged lines#define LAUNCH_EXACT(K_, M_, N_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \- else if (K == K_ && M == M_ && N == N_) C = cutlass_dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);+ else if (K == K_ && M == M_ && N == N_) C = dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);#define LAUNCH(K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES) \- else if (K == K_) C = cutlass_dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);+ else if (K == K_) C = dual_gemm_launch<K_, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>(A, B1, B2, SFA, SFB1, SFB2, C);if (false) {}⋯ 21 unchanged lines}- TORCH_LIBRARY(my_module_dual_gemm, m) {+ TORCH_LIBRARY(my_module_dual_gemm_v2, m) {m.def("dual_gemm(Tensor A, Tensor B1, Tensor B2, Tensor SFA, Tensor SFB1, Tensor SFB2, Tensor(a!) C) -> Tensor");m.impl("dual_gemm", &dual_gemm);}"""--load_inline(- "dual_gemm_cuda_best_v1",+ "dual_gemm_cuda_v2",cpp_sources="",cuda_sources=CUDA_SRC_COMMON + CUDA_SRC_DUAL_GEMM,verbose=True,⋯ 7 unchanged lines"--relocatable-device-code=false","-lineinfo","-Xptxas=-v",- "-maxrregcount=192",],extra_ldflags=["-lcuda"],)- dual_gemm_cuda = torch.ops.my_module_dual_gemm.dual_gemm+ dual_gemm_cuda = torch.ops.my_module_dual_gemm_v2.dual_gemm--def custom_kernel(data: input_t) -> output_t:- # data contains: A, B1, B2, _, _, _, SFA, SFB1, SFB2, C- # Dual GEMM computes: C = silu(A @ B1) * (A @ B2)return dual_gemm_cuda(data[0], data[1], data[2], data[6], data[7], data[8], data[9])--
scrolls · 562 diff lines total
Best evidence level for this revision: reported
JSON