submission 411277
novo_force · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 1292 lines, June 9 Researcher Reciprocity License v1.0.
submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-group-gemm-411277?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:370b87c69f07b448aa8dc39a43c9b1405161a61f6165d8fe7389d0b6e6f0619d
license declaredunknown
license concludedunknown
authorsnovo_force
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));tile-k = 256
constexpr int BLOCK_K = 256;tile-m = 128
constexpr int BLOCK_M = 128;tma
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "vector-width = half2
reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});Kernel source
submission.py1292 lines
from __future__ import annotations
import os
from typing import List
import torch
from torch.utils.cpp_extension import load_inline
_FORCE_NO_GROUPED = False
_FORCE_BN64 = False
_FORCE_RAW_SF = False
_EXT_READY = False
_OPS_READY = False
_GEMM = None
_GEMM_GROUPED = None
_SCRATCH_A: dict = {}
_SCRATCH_A_G: dict = {}
def _load_ext() -> None:
global _EXT_READY, _OPS_READY, _GEMM, _GEMM_GROUPED
if _EXT_READY:
return
cuda_src = r"""
#include <cuda.h>
#include <cudaTypedefs.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <torch/extension.h>
#include <torch/library.h>
#include <ATen/ATen.h>
#include <cstdint>
#include <vector>
#include <array>
#ifndef NVFP4_GGEMM_CHECK
#define NVFP4_GGEMM_CHECK 0
#endif
#if NVFP4_GGEMM_CHECK
#define GG_CHECK(x, msg) TORCH_CHECK((x), msg)
#else
#define GG_CHECK(x, msg) ((void)0)
#endif
constexpr int WARP_SIZE = 32;
constexpr int MMA_K = 64;
constexpr uint64_t EVICT_FIRST = 0x12F0000000000000ULL;
constexpr uint64_t EVICT_LAST = 0x14F0000000000000ULL;
__device__ __forceinline__ constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3FFFFULL) >> 4ULL; }
__device__ __forceinline__ uint32_t elect_sync() {
uint32_t pred = 0;
asm volatile(
"{\n\t"
".reg .pred %%px;\n\t"
"elect.sync _|%%px, %1;\n\t"
"@%%px mov.s32 %0, 1;\n\t"
"}"
: "+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"
"}"
:: "r"(mbar_addr), "r"(phase), "r"(ticks)
);
}
__device__ __forceinline__ void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
asm volatile(
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
"[%0], [%1], %2, [%3], %4;"
:: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy)
);
}
__device__ __forceinline__ void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
asm volatile(
"cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
"[%0], [%1, {%2, %3, %4}], [%5], %6;"
:: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
: "memory"
);
}
__device__ __forceinline__ void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}
__device__ __forceinline__ void tcgen05_mma_nvfp4(
uint64_t a_desc,
uint64_t b_desc,
uint32_t i_desc,
int scale_A_tmem,
int scale_B_tmem,
int enable_input_d
) {
const int d_tmem = 0;
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %6, 0;\n\t"
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
"}"
:: "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_V, const char *NUM_V>
__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_V), "C"(NUM_V));
}
template <const char *SHAPE_V, const char *NUM_V>
__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_V), "C"(NUM_V));
}
__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 __forceinline__ void check_cu(CUresult err) {
if (err == CUDA_SUCCESS) return;
const char *msg = "unknown";
cuGetErrorString(err, &msg);
TORCH_CHECK(false, msg);
}
struct TmapKey {
uint64_t ptr;
uint64_t global_height;
uint64_t global_width;
uint32_t shared_height;
uint32_t shared_width;
int32_t dev;
};
static __forceinline__ bool tmap_key_eq(const TmapKey &a, const TmapKey &b) {
return a.ptr == b.ptr
&& a.global_height == b.global_height
&& a.global_width == b.global_width
&& a.shared_height == b.shared_height
&& a.shared_width == b.shared_width
&& a.dev == b.dev;
}
template <int CAP>
struct TmapCache {
std::array<TmapKey, CAP> keys;
std::array<CUtensorMap, CAP> vals;
std::array<uint8_t, CAP> used;
int head;
TmapCache() : used{}, head(0) {}
bool lookup(const TmapKey &k, CUtensorMap *out) {
#pragma unroll
for (int i = 0; i < CAP; i++) {
if (used[(size_t)i] && tmap_key_eq(keys[(size_t)i], k)) {
*out = vals[(size_t)i];
return true;
}
}
return false;
}
void insert(const TmapKey &k, const CUtensorMap &v) {
keys[(size_t)head] = k;
vals[(size_t)head] = v;
used[(size_t)head] = 1;
head++;
if (head >= CAP) head = 0;
}
};
static TmapCache<64> g_tmap_cache;
static __forceinline__ 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
) {
int dev = 0;
cudaGetDevice(&dev);
TmapKey key;
key.ptr = (uint64_t)ptr;
key.global_height = global_height;
key.global_width = global_width;
key.shared_height = shared_height;
key.shared_width = shared_width;
key.dev = (int32_t)dev;
if (g_tmap_cache.lookup(key, tmap)) return;
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);
g_tmap_cache.insert(key, *tmap);
}
struct __align__(16) GroupDesc {
CUtensorMap A_tmap;
CUtensorMap B_tmap;
uint64_t SFA_ptr;
uint64_t SFB_ptr;
uint64_t C_ptr;
int M;
int N;
int K;
};
__device__ __forceinline__ int active_threads_128(int M, int off_m) {
int rem = M - off_m;
if (rem <= 0) return 0;
int thr = (rem + 31) & ~31;
if (thr > 128) thr = 128;
return thr;
}
template <int BLOCK_N, int NUM_STAGES>
__global__ __launch_bounds__(128 + 2 * WARP_SIZE)
void kernel(
const __grid_constant__ CUtensorMap A_tmap,
const __grid_constant__ CUtensorMap B_tmap,
const char *SFA_ptr,
const char *SFB_ptr,
half *C_ptr,
int M, int N, int K
) {
constexpr int BLOCK_M = 128;
constexpr int BLOCK_K = 256;
const int tid = (int)threadIdx.x;
const int bid_n = (int)blockIdx.x;
const int bid_m = (int)blockIdx.y;
const int lane_id = tid & (WARP_SIZE - 1);
const int warp_id = tid >> 5;
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 = (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 = (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()) {
#pragma unroll
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();
const int num_iters = K / BLOCK_K;
if (warp_id == NUM_WARPS - 2 && elect_sync()) {
const int grid_m = (M + 127) >> 7;
int grid_n = 0;
if constexpr (BLOCK_N == 128) grid_n = (N + 127) >> 7;
else grid_n = (N + 63) >> 6;
const bool keep_A = (grid_n >= grid_m);
const uint64_t cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;
const uint64_t cache_B = keep_A ? EVICT_FIRST : 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 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");
};
const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
for (int iter_k = 0; iter_k < init_stage; 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);
auto make_desc_AB = [] __device__ (int addr) -> uint64_t {
const int SBO = 8 * 128;
return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
};
auto make_desc_SF = [] __device__ (int addr) -> uint64_t {
const int SBO = 8 * 16;
return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL);
};
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;
const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
const uint64_t SFB_desc = make_desc_SF(0) + ((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;
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) {
const int thr = active_threads_128(M, off_m);
if (tid >= thr) return;
mbarrier_wait(mainloop_mbar_addr, 0);
asm volatile("tcgen05.fence::after_thread_sync;");
#pragma unroll
for (int m = 0; m < 2; m++) {
const int trow = warp_id * 32 + m * 16;
if constexpr (BLOCK_N == 128) {
float tmp[64];
tcgen05_ld_16x256bx16(tmp, trow, 0);
asm volatile("tcgen05.wait::ld.sync.aligned;");
#pragma unroll
for (int i = 0; i < 16; i++) {
const int row = off_m + trow + (lane_id >> 2);
const int col = off_n + i * 8 + ((lane_id & 3) << 1);
if (row < M) {
reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
}
const int row2 = row + 8;
if (row2 < M) {
reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
}
}
} else {
const bool full_n = (off_n + BLOCK_N) <= N;
float tmp[BLOCK_N / 2];
tcgen05_ld_16x256bx8(tmp, trow, 0);
asm volatile("tcgen05.wait::ld.sync.aligned;");
#pragma unroll
for (int i = 0; i < BLOCK_N / 8; i++) {
const int row = off_m + trow + (lane_id >> 2);
const int col = off_n + i * 8 + ((lane_id & 3) << 1);
if (row < M) {
if (full_n || (col + 1) < N) {
reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
} else if (col < N) {
C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);
}
}
const int row2 = row + 8;
if (row2 < M) {
if (full_n || (col + 1) < N) {
reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
} else if (col < N) {
C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);
}
}
}
}
}
asm volatile("bar.sync 1, %0;" :: "r"(thr) : "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 BLOCK_N, int NUM_STAGES>
__global__ __launch_bounds__(128 + 2 * WARP_SIZE)
void kernel_grouped(const GroupDesc *descs) {
constexpr int BLOCK_M = 128;
constexpr int BLOCK_K = 256;
const int gid = (int)blockIdx.z;
const GroupDesc *desc = descs + gid;
const int M = desc->M;
const int N = desc->N;
const int K = desc->K;
const int grid_m = (M + 127) / 128;
const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;
const int bid_n = (int)blockIdx.x;
const int bid_m = (int)blockIdx.y;
if (bid_n >= grid_n || bid_m >= grid_m) return;
const CUtensorMap *A_tmap = &desc->A_tmap;
const CUtensorMap *B_tmap = &desc->B_tmap;
const char *SFA_ptr = (const char *)desc->SFA_ptr;
const char *SFB_ptr = (const char *)desc->SFB_ptr;
half *C_ptr = (half *)desc->C_ptr;
const int tid = (int)threadIdx.x;
const int lane_id = tid & (WARP_SIZE - 1);
const int warp_id = tid >> 5;
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 = (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 = (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()) {
#pragma unroll
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();
const int num_iters = K / BLOCK_K;
if (warp_id == NUM_WARPS - 2 && elect_sync()) {
const bool keep_A = (grid_n >= grid_m);
const uint64_t cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;
const uint64_t cache_B = keep_A ? EVICT_FIRST : 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 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");
};
const int init_stage = (num_iters < NUM_STAGES) ? num_iters : NUM_STAGES;
for (int iter_k = 0; iter_k < init_stage; 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);
auto make_desc_AB = [] __device__ (int addr) -> uint64_t {
const int SBO = 8 * 128;
return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
};
auto make_desc_SF = [] __device__ (int addr) -> uint64_t {
const int SBO = 8 * 16;
return desc_encode((uint64_t)addr) | (desc_encode((uint64_t)SBO) << 32ULL) | (1ULL << 46ULL);
};
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;
const uint64_t SFA_desc = make_desc_SF(0) + ((uint64_t)SFA_smem >> 4ULL);
const uint64_t SFB_desc = make_desc_SF(0) + ((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;
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) {
const int thr = active_threads_128(M, off_m);
if (tid >= thr) return;
mbarrier_wait(mainloop_mbar_addr, 0);
asm volatile("tcgen05.fence::after_thread_sync;");
#pragma unroll
for (int m = 0; m < 2; m++) {
const int trow = warp_id * 32 + m * 16;
if constexpr (BLOCK_N == 128) {
float tmp[64];
tcgen05_ld_16x256bx16(tmp, trow, 0);
asm volatile("tcgen05.wait::ld.sync.aligned;");
#pragma unroll
for (int i = 0; i < 16; i++) {
const int row = off_m + trow + (lane_id >> 2);
const int col = off_n + i * 8 + ((lane_id & 3) << 1);
if (row < M) {
reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
}
const int row2 = row + 8;
if (row2 < M) {
reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
}
}
} else {
const bool full_n = (off_n + BLOCK_N) <= N;
float tmp[BLOCK_N / 2];
tcgen05_ld_16x256bx8(tmp, trow, 0);
asm volatile("tcgen05.wait::ld.sync.aligned;");
#pragma unroll
for (int i = 0; i < BLOCK_N / 8; i++) {
const int row = off_m + trow + (lane_id >> 2);
const int col = off_n + i * 8 + ((lane_id & 3) << 1);
if (row < M) {
if (full_n || (col + 1) < N) {
reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
} else if (col < N) {
C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);
}
}
const int row2 = row + 8;
if (row2 < M) {
if (full_n || (col + 1) < N) {
reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
} else if (col < N) {
C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);
}
}
}
}
}
asm volatile("bar.sync 1, %0;" :: "r"(thr) : "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 BLOCK_N, int NUM_STAGES>
static __forceinline__ void gemm_launch(
const at::Tensor& A,
const at::Tensor& B,
const at::Tensor& SFA,
const at::Tensor& SFB,
at::Tensor& C,
int M, int N, int K
) {
const int Apad = (int)A.size(0);
const char *A_ptr = (const char *)A.data_ptr();
const char *B_ptr = (const char *)B.data_ptr();
const char *SFA_ptr = (const char *)SFA.data_ptr();
const char *SFB_ptr = (const char *)SFB.data_ptr();
half *C_ptr = (half *)C.data_ptr<at::Half>();
CUtensorMap A_tmap, B_tmap;
init_AB_tmap(&A_tmap, A_ptr, (uint64_t)Apad, (uint64_t)K, 128, 256);
init_AB_tmap(&B_tmap, B_ptr, (uint64_t)N, (uint64_t)K, (uint32_t)BLOCK_N, 256);
const int grid_m = (M + 127) / 128;
const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;
const int tb_size = 128 + 2 * WARP_SIZE;
const int A_size = 128 * 256 / 2;
const int B_size = BLOCK_N * 256 / 2;
const int SF_size = 128 * 256 / 16;
const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;
auto k = kernel<BLOCK_N, NUM_STAGES>;
static int last_dev = -1;
int dev = -1;
cudaGetDevice(&dev);
if (dev != last_dev) {
auto err = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
GG_CHECK(err == cudaSuccess, "cudaFuncSetAttribute failed");
last_dev = dev;
}
dim3 grid((unsigned)grid_n, (unsigned)grid_m, 1);
k<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N, K);
}
static __forceinline__ int pick_stage_from_K(int K) {
const int iters = K >> 8;
if (iters <= 8) return 2;
if (iters <= 16) return 3;
return 4;
}
static __forceinline__ int pick_stage_bn128_from_K(int K) {
const int iters = K >> 8;
if (iters <= 8) return 2;
if (iters >= 24) return 4;
return 3;
}
at::Tensor gemm(
const at::Tensor& A,
const at::Tensor& B,
const at::Tensor& SFA,
const at::Tensor& SFB,
at::Tensor& C,
int64_t M,
int64_t N,
int64_t K
) {
const int Mi = (int)M;
const int Ni = (int)N;
const int Ki = (int)K;
if (((Ni & 127) == 0) && (Ni >= 128)) {
const int stage = pick_stage_bn128_from_K(Ki);
if (stage == 2) gemm_launch<128, 2>(A, B, SFA, SFB, C, Mi, Ni, Ki);
else if (stage == 3) gemm_launch<128, 3>(A, B, SFA, SFB, C, Mi, Ni, Ki);
else gemm_launch<128, 4>(A, B, SFA, SFB, C, Mi, Ni, Ki);
} else {
const int stage = pick_stage_from_K(Ki);
if (stage == 2) gemm_launch<64, 2>(A, B, SFA, SFB, C, Mi, Ni, Ki);
else if (stage == 3) gemm_launch<64, 3>(A, B, SFA, SFB, C, Mi, Ni, Ki);
else gemm_launch<64, 4>(A, B, SFA, SFB, C, Mi, Ni, Ki);
}
GG_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");
return C;
}
struct GroupWorkspace {
at::Tensor descs_d;
int64_t cap_G;
int64_t dev;
uint64_t last_hash;
uint64_t last_hash2;
int64_t last_G;
int last_block_n;
int last_max_grid_m;
int last_max_grid_n;
int last_stage;
GroupWorkspace()
: cap_G(0),
dev(-1),
last_hash(0),
last_hash2(0),
last_G(0),
last_block_n(0),
last_max_grid_m(0),
last_max_grid_n(0),
last_stage(0) {}
};
static GroupWorkspace g_ws;
static __forceinline__ void ensure_ws(int64_t dev, int64_t G) {
if (g_ws.dev != dev || g_ws.cap_G < G || !g_ws.descs_d.defined()) {
g_ws.dev = dev;
g_ws.cap_G = G;
at::TensorOptions opt_u8 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kByte);
g_ws.descs_d = at::empty({G, (int64_t)sizeof(GroupDesc)}, opt_u8);
g_ws.last_hash = 0;
g_ws.last_hash2 = 0;
g_ws.last_G = 0;
}
}
static __forceinline__ uint64_t fnv1a_mix_u64(uint64_t h, uint64_t x) {
h ^= x;
h *= 1099511628211ULL;
return h;
}
template <int BLOCK_N, int NUM_STAGES>
static __forceinline__ void grouped_launch(
int64_t dev,
int max_grid_m,
int max_grid_n,
int64_t G
) {
const int tb_size = 128 + 2 * WARP_SIZE;
const int A_size = 128 * 256 / 2;
const int B_size = BLOCK_N * 256 / 2;
const int SF_size = 128 * 256 / 16;
const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;
auto k = kernel_grouped<BLOCK_N, NUM_STAGES>;
static int last_dev = -1;
if ((int)dev != last_dev) {
auto err = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
GG_CHECK(err == cudaSuccess, "cudaFuncSetAttribute failed");
last_dev = (int)dev;
}
dim3 grid((unsigned)max_grid_n, (unsigned)max_grid_m, (unsigned)G);
k<<<grid, tb_size, smem_size>>>((const GroupDesc *)g_ws.descs_d.data_ptr());
}
void gemm_grouped(
at::TensorList A_list,
at::TensorList B_list,
at::TensorList SFA_list,
at::TensorList SFB_list,
at::TensorList C_list,
bool force_bn64
) {
const int64_t G = (int64_t)A_list.size();
if (G <= 0) return;
const auto &A0 = A_list[0];
const int64_t dev = (int64_t)A0.get_device();
if (g_ws.dev != dev) cudaSetDevice((int)dev);
ensure_ws(dev, G);
bool use_bn128 = !force_bn64;
for (int i = 0; i < (int)G; i++) {
const auto &C = C_list[i];
const int N = (int)C.size(1);
if (N < 128 || ((N & 127) != 0)) { use_bn128 = false; break; }
}
const int block_n = use_bn128 ? 128 : 64;
uint64_t h1 = 1469598103934665603ULL;
uint64_t h2 = 1099511628211ULL ^ 0x9E3779B97F4A7C15ULL;
h1 = fnv1a_mix_u64(h1, (uint64_t)G);
h1 = fnv1a_mix_u64(h1, (uint64_t)block_n);
h2 = fnv1a_mix_u64(h2, (uint64_t)G ^ 0xD6E8FEB86659FD93ULL);
h2 = fnv1a_mix_u64(h2, (uint64_t)block_n ^ 0xA5A3564E27F3A7A9ULL);
int max_grid_m = 0;
int max_grid_n = 0;
int max_K = 0;
for (int i = 0; i < (int)G; i++) {
const auto &A = A_list[i];
const auto &B = B_list[i];
const auto &SFA = SFA_list[i];
const auto &SFB = SFB_list[i];
const auto &C = C_list[i];
const int M = (int)C.size(0);
const int N = (int)C.size(1);
const int K = (int)A.size(1) * 2;
const int Apad = (int)A.size(0);
const int grid_m = (M + 127) / 128;
const int grid_n = (N + block_n - 1) / block_n;
if (grid_m > max_grid_m) max_grid_m = grid_m;
if (grid_n > max_grid_n) max_grid_n = grid_n;
if (K > max_K) max_K = K;
const uint64_t A_ptr = (uint64_t)A.data_ptr();
const uint64_t B_ptr = (uint64_t)B.data_ptr();
const uint64_t SFA_ptr = (uint64_t)SFA.data_ptr();
const uint64_t SFB_ptr = (uint64_t)SFB.data_ptr();
const uint64_t C_ptr = (uint64_t)C.data_ptr<at::Half>();
const uint64_t shape_pack = ((uint64_t)K << 32) | (uint64_t)(uint32_t)Apad;
h1 = fnv1a_mix_u64(h1, A_ptr);
h1 = fnv1a_mix_u64(h1, B_ptr);
h1 = fnv1a_mix_u64(h1, SFA_ptr);
h1 = fnv1a_mix_u64(h1, SFB_ptr);
h1 = fnv1a_mix_u64(h1, C_ptr);
h1 = fnv1a_mix_u64(h1, (uint64_t)M);
h1 = fnv1a_mix_u64(h1, (uint64_t)N);
h1 = fnv1a_mix_u64(h1, shape_pack);
h2 = fnv1a_mix_u64(h2, A_ptr ^ 0x8B8B8B8B8B8B8B8BULL);
h2 = fnv1a_mix_u64(h2, B_ptr ^ 0xC3C3C3C3C3C3C3C3ULL);
h2 = fnv1a_mix_u64(h2, SFA_ptr ^ 0x1313131313131313ULL);
h2 = fnv1a_mix_u64(h2, SFB_ptr ^ 0x5D5D5D5D5D5D5D5DULL);
h2 = fnv1a_mix_u64(h2, C_ptr ^ 0xA7A7A7A7A7A7A7A7ULL);
h2 = fnv1a_mix_u64(h2, ((uint64_t)M << 32) ^ (uint64_t)N);
h2 = fnv1a_mix_u64(h2, shape_pack ^ 0x9E3779B97F4A7C15ULL);
}
int stage = 0;
if (use_bn128) stage = pick_stage_bn128_from_K(max_K);
else stage = pick_stage_from_K(max_K);
const bool hit = (g_ws.last_hash == h1)
&& (g_ws.last_hash2 == h2)
&& (g_ws.last_G == G)
&& (g_ws.last_block_n == block_n)
&& (g_ws.last_stage == stage)
&& g_ws.descs_d.defined();
if (!hit) {
std::vector<GroupDesc> descs((size_t)G);
for (int i = 0; i < (int)G; i++) {
const auto &A = A_list[i];
const auto &B = B_list[i];
const auto &SFA = SFA_list[i];
const auto &SFB = SFB_list[i];
const auto &C = C_list[i];
const int M = (int)C.size(0);
const int N = (int)C.size(1);
const int K = (int)A.size(1) * 2;
const int Apad = (int)A.size(0);
GroupDesc d;
init_AB_tmap(&d.A_tmap, (const char *)A.data_ptr(), (uint64_t)Apad, (uint64_t)K, 128, 256);
init_AB_tmap(&d.B_tmap, (const char *)B.data_ptr(), (uint64_t)N, (uint64_t)K, (uint32_t)block_n, 256);
d.SFA_ptr = (uint64_t)SFA.data_ptr();
d.SFB_ptr = (uint64_t)SFB.data_ptr();
d.C_ptr = (uint64_t)C.data_ptr<at::Half>();
d.M = M;
d.N = N;
d.K = K;
descs[(size_t)i] = d;
}
cudaError_t cperr = cudaMemcpy(
g_ws.descs_d.data_ptr(),
descs.data(),
(size_t)G * sizeof(GroupDesc),
cudaMemcpyHostToDevice
);
GG_CHECK(cperr == cudaSuccess, "memcpy fail");
g_ws.last_hash = h1;
g_ws.last_hash2 = h2;
g_ws.last_G = G;
g_ws.last_block_n = block_n;
g_ws.last_max_grid_m = max_grid_m;
g_ws.last_max_grid_n = max_grid_n;
g_ws.last_stage = stage;
} else {
max_grid_m = g_ws.last_max_grid_m;
max_grid_n = g_ws.last_max_grid_n;
}
if (use_bn128) {
if (stage == 2) grouped_launch<128, 2>(dev, max_grid_m, max_grid_n, G);
else if (stage == 3) grouped_launch<128, 3>(dev, max_grid_m, max_grid_n, G);
else grouped_launch<128, 4>(dev, max_grid_m, max_grid_n, G);
} else {
if (stage == 2) grouped_launch<64, 2>(dev, max_grid_m, max_grid_n, G);
else if (stage == 3) grouped_launch<64, 3>(dev, max_grid_m, max_grid_n, G);
else grouped_launch<64, 4>(dev, max_grid_m, max_grid_n, G);
}
GG_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");
}
TORCH_LIBRARY(nvfp4_group_gemm_opt, m) {
m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int M, int N, int K) -> Tensor");
m.impl("gemm", &gemm);
m.def("gemm_grouped(Tensor[] A, Tensor[] B, Tensor[] SFA, Tensor[] SFB, Tensor[] C, bool force_bn64) -> ()");
m.impl("gemm_grouped", &gemm_grouped);
}
"""
build_dir = os.path.join(os.path.dirname(__file__), ".build_nvfp4_group_gemm_opt")
os.makedirs(build_dir, exist_ok=True)
load_inline(
name="nvfp4_group_gemm_opt_ext",
cpp_sources="",
cuda_sources=cuda_src,
functions=None,
extra_cflags=["-O3"],
extra_cuda_cflags=[
"-O3",
"-gencode=arch=compute_100a,code=sm_100a",
"--use_fast_math",
"--expt-extended-lambda",
"--expt-relaxed-constexpr",
"--relocatable-device-code=false",
"-std=c++17",
],
extra_ldflags=["-lcuda"],
with_cuda=True,
is_python_module=False,
no_implicit_headers=True,
build_directory=build_dir,
verbose=False,
)
_EXT_READY = True
_OPS_READY = False
_GEMM = None
_GEMM_GROUPED = None
def _get_ops():
global _OPS_READY, _GEMM, _GEMM_GROUPED
if not _EXT_READY:
_load_ext()
if not _OPS_READY:
_GEMM = torch.ops.nvfp4_group_gemm_opt.gemm
_GEMM_GROUPED = torch.ops.nvfp4_group_gemm_opt.gemm_grouped
_OPS_READY = True
return _GEMM, _GEMM_GROUPED
def _as_u8(x: torch.Tensor) -> torch.Tensor:
if x.dtype == torch.uint8:
return x
if x.element_size() != 1:
raise RuntimeError("packed tensor must have 1-byte elements")
return x.view(torch.uint8)
def _reorder_scale_from_raw(scale_u8_2d: torch.Tensor, rows_pad: int) -> torch.Tensor:
if scale_u8_2d.dim() != 2:
raise RuntimeError("scale must be 2D")
rows = int(scale_u8_2d.size(0))
k16 = int(scale_u8_2d.size(1))
if (k16 & 3) != 0:
raise RuntimeError("K//16 must be multiple of 4")
if (rows_pad & 127) != 0:
raise RuntimeError("rows_pad must be multiple of 128")
blk_m = rows_pad // 128
blk_k = k16 // 4
buf = torch.empty((rows_pad, k16), device=scale_u8_2d.device, dtype=torch.uint8)
buf[:rows].copy_(scale_u8_2d)
v = buf.view(blk_m, 32, 4, blk_k, 4).permute(0, 3, 1, 2, 4).contiguous()
return v
def _get_scratch_a(device: torch.device, m_pad: int, k2: int) -> torch.Tensor:
key = (int(device.index), int(m_pad), int(k2))
buf = _SCRATCH_A.get(key)
if buf is None or (not buf.is_cuda) or buf.numel() != (m_pad * k2):
buf = torch.empty((m_pad, k2, 1), device=device, dtype=torch.uint8)
_SCRATCH_A[key] = buf
return buf
def _get_scratch_a_grouped(device: torch.device, group_i: int, m_pad: int, k2: int) -> torch.Tensor:
key = (int(device.index), int(group_i), int(m_pad), int(k2))
buf = _SCRATCH_A_G.get(key)
if buf is None or (not buf.is_cuda) or buf.numel() != (m_pad * k2):
buf = torch.empty((m_pad, k2, 1), device=device, dtype=torch.uint8)
_SCRATCH_A_G[key] = buf
return buf
def _must_sfp_layout(x: torch.Tensor) -> None:
if (not x.is_cuda) or (x.dim() != 6) or (x.element_size() != 1) or (int(x.storage_offset()) != 0):
raise RuntimeError("bad sfx_p layout")
if int(x.size(0)) != 32 or int(x.size(1)) != 4 or int(x.size(3)) != 4:
raise RuntimeError("bad sfx_p shape")
def custom_kernel(data):
abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
gemm, gemm_grouped = _get_ops()
g = len(problem_sizes)
if g == 0:
return []
dev0 = int(abc_tensors[0][0].device.index)
all_l1 = True
for i in range(g):
if int(problem_sizes[i][3]) != 1:
all_l1 = False
break
if int(abc_tensors[i][0].device.index) != dev0:
all_l1 = False
break
outs: List[torch.Tensor] = []
if all_l1 and (not _FORCE_NO_GROUPED):
a_list: List[torch.Tensor] = []
b_list: List[torch.Tensor] = []
sfa_list: List[torch.Tensor] = []
sfb_list: List[torch.Tensor] = []
c_list: List[torch.Tensor] = []
for i in range(g):
a, b, c = abc_tensors[i]
sfa, sfb = sfasfb_tensors[i]
sfa_p, sfb_p = sfasfb_reordered_tensors[i]
m, n, _k, _l = problem_sizes[i]
if (not a.is_contiguous()) or (not b.is_contiguous()) or (not c.is_contiguous()):
raise RuntimeError("grouped fast path requires contiguous a/b/c")
m_int = int(m)
n_int = int(n)
a_u8 = _as_u8(a)
b_u8 = _as_u8(b)
m_pad = ((m_int + 127) // 128) * 128
if a_u8.size(0) != m_pad:
a_pad = _get_scratch_a_grouped(a_u8.device, i, m_pad, int(a_u8.size(1)))
a_pad[: a_u8.size(0)].copy_(a_u8)
else:
a_pad = a_u8
if _FORCE_RAW_SF:
sfa2 = _as_u8(sfa[..., 0]).contiguous()
sfb2 = _as_u8(sfb[..., 0]).contiguous()
sfa_arg = _reorder_scale_from_raw(sfa2, m_pad)
sfb_arg = _reorder_scale_from_raw(sfb2, ((n_int + 127) // 128) * 128)
else:
_must_sfp_layout(sfa_p)
_must_sfp_layout(sfb_p)
sfa_arg = sfa_p
sfb_arg = sfb_p
a_list.append(a_pad)
b_list.append(b_u8)
sfa_list.append(sfa_arg)
sfb_list.append(sfb_arg)
c_list.append(c)
gemm_grouped(a_list, b_list, sfa_list, sfb_list, c_list, _FORCE_BN64)
return c_list
for i in range(g):
a, b, c = abc_tensors[i]
sfa, sfb = sfasfb_tensors[i]
sfa_p, sfb_p = sfasfb_reordered_tensors[i]
m, n, k, l = problem_sizes[i]
m_int = int(m)
n_int = int(n)
k_int = int(k)
l_int = int(l)
c_out = c
if not c_out.is_contiguous():
c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)
else:
c_tmp = c_out
if l_int == 1:
a_u8 = _as_u8(a).contiguous()
b_u8 = _as_u8(b).contiguous()
m_pad = ((m_int + 127) // 128) * 128
if a_u8.size(0) != m_pad:
a_pad = _get_scratch_a(a_u8.device, m_pad, int(a_u8.size(1)))
a_pad[: a_u8.size(0)].copy_(a_u8)
else:
a_pad = a_u8
if _FORCE_RAW_SF:
sfa2 = _as_u8(sfa[..., 0]).contiguous()
sfb2 = _as_u8(sfb[..., 0]).contiguous()
sfa_arg = _reorder_scale_from_raw(sfa2, m_pad)
sfb_arg = _reorder_scale_from_raw(sfb2, ((n_int + 127) // 128) * 128)
else:
_must_sfp_layout(sfa_p)
_must_sfp_layout(sfb_p)
sfa_arg = sfa_p
sfb_arg = sfb_p
gemm(a_pad, b_u8, sfa_arg, sfb_arg, c_tmp, m_int, n_int, k_int)
else:
for li in range(l_int):
a2 = _as_u8(a[..., li]).contiguous().unsqueeze(-1)
b2 = _as_u8(b[..., li]).contiguous().unsqueeze(-1)
m_pad = ((m_int + 127) // 128) * 128
if a2.size(0) != m_pad:
a_pad = _get_scratch_a(a2.device, m_pad, int(a2.size(1)))
a_pad[: a2.size(0)].copy_(a2)
else:
a_pad = a2
sfa2 = _as_u8(sfa[..., li]).contiguous()
sfb2 = _as_u8(sfb[..., li]).contiguous()
sfa_r = _reorder_scale_from_raw(sfa2, m_pad)
sfb_r = _reorder_scale_from_raw(sfb2, ((n_int + 127) // 128) * 128)
c2 = torch.empty((m_int, n_int, 1), device=c_tmp.device, dtype=torch.float16)
gemm(a_pad, b2, sfa_r, sfb_r, c2, m_int, n_int, k_int)
c_tmp[..., li].copy_(c2[..., 0])
if c_tmp is not c_out:
c_out.copy_(c_tmp)
outs.append(c_out)
return outs
__all__ = ["custom_kernel"]
scrolls · 1292 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 411192.
⋯ 9 unchanged lines- _FORCE_NO_GROUPED = False- _FORCE_BN64 = False- _FORCE_RAW_SF = False- _FORCE_BN128_STAGE4 = True- _FORCE_BN128_EPILOGUE_X16 = False++ _FORCE_NO_GROUPED = False+ _FORCE_BN64 = False+ _FORCE_RAW_SF = False+_EXT_READY = False_OPS_READY = False⋯ 23 unchanged lines#include <vector>#include <array>- // 默认关闭检查，极致压缩 host 热路径分支#ifndef NVFP4_GGEMM_CHECK#define NVFP4_GGEMM_CHECK 0#endif⋯ 30 unchanged linesasm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));}- #ifndef NVFP4_GGEMM_WAIT_BACKOFF- #define NVFP4_GGEMM_WAIT_BACKOFF 0- #endif-- __device__ __forceinline__ void mbarrier_wait_basic(int mbar_addr, int phase) {+ __device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {uint32_t ticks = 0x989680;asm volatile("{\n\t"⋯ 8 unchanged lines);}- __device__ __forceinline__ void mbarrier_wait_backoff(int mbar_addr, int phase) {- uint32_t ticks0 = 0x80;- uint32_t ticks1 = 0x400;- uint32_t ticks2 = 0x989680;- asm volatile(- "{\n\t"- ".reg .pred P1;\n\t"- ".reg .b32 cnt;\n\t"- "mov.b32 cnt, 0;\n\t"- "LAB_WAIT0:\n\t"- "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"- "@P1 bra.uni DONE;\n\t"- "add.s32 cnt, cnt, 1;\n\t"- "setp.lt.s32 P1, cnt, 8;\n\t"- "@P1 bra.uni LAB_WAIT0;\n\t"- "mov.b32 cnt, 0;\n\t"- "LAB_WAIT1:\n\t"- "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %3;\n\t"- "@P1 bra.uni DONE;\n\t"- "add.s32 cnt, cnt, 1;\n\t"- "setp.lt.s32 P1, cnt, 16;\n\t"- "@P1 bra.uni LAB_WAIT1;\n\t"- "LAB_WAIT2:\n\t"- "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %4;\n\t"- "@P1 bra.uni DONE;\n\t"- "bra.uni LAB_WAIT2;\n\t"- "DONE:\n\t"- "}"- :: "r"(mbar_addr), "r"(phase), "r"(ticks0), "r"(ticks1), "r"(ticks2)- );- }-- __device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {- #if NVFP4_GGEMM_WAIT_BACKOFF- mbarrier_wait_backoff(mbar_addr, phase);- #else- mbarrier_wait_basic(mbar_addr, phase);- #endif- }-__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 "⋯ 188 unchanged linesg_tmap_cache.insert(key, *tmap);}- // grouped 元数据：一次性搬到 device，kernel 只读 descs[gid]struct __align__(16) GroupDesc {CUtensorMap A_tmap;CUtensorMap B_tmap;⋯ 13 unchanged linesreturn thr;}- template <int BLOCK_N, int NUM_STAGES, bool EPILOGUE_X16>+ template <int BLOCK_N, int NUM_STAGES>__global__ __launch_bounds__(128 + 2 * WARP_SIZE)void kernel(const __grid_constant__ CUtensorMap A_tmap,⋯ 47 unchanged linesconst int num_iters = K / BLOCK_K;if (warp_id == NUM_WARPS - 2 && elect_sync()) {- uint64_t cache_A, cache_B;const int grid_m = (M + 127) >> 7;- const int grid_n = (N + BLOCK_N - 1) / BLOCK_N;- if (grid_n >= (grid_m << 2)) {- cache_A = EVICT_LAST;- cache_B = EVICT_FIRST;- } else {- const bool keep_A = (grid_n >= grid_m);- cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;- cache_B = keep_A ? EVICT_FIRST : EVICT_LAST;- }+ int grid_n = 0;+ if constexpr (BLOCK_N == 128) grid_n = (N + 127) >> 7;+ else grid_n = (N + 63) >> 6;+ const bool keep_A = (grid_n >= grid_m);+ const uint64_t cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;+ const uint64_t cache_B = keep_A ? EVICT_FIRST : EVICT_LAST;auto issue_tma = [&](int iter_k, int stage_id) {const int mbar_addr = tma_mbar_addr + stage_id * 8;⋯ 85 unchanged linesasm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];":: "r"(mainloop_mbar_addr) : "memory");} else if (tid < BLOCK_M) {- const bool full_m = (off_m + 128) <= M;- int thr = 128;- if (!full_m) {- thr = active_threads_128(M, off_m);- if (tid >= thr) return;- }+ const int thr = active_threads_128(M, off_m);+ if (tid >= thr) return;mbarrier_wait(mainloop_mbar_addr, 0);asm volatile("tcgen05.fence::after_thread_sync;");- const bool full_n = (off_n + BLOCK_N) <= N;+ #pragma unroll+ for (int m = 0; m < 2; m++) {+ const int trow = warp_id * 32 + m * 16;- if (full_m && full_n) {- #pragma unroll- for (int m = 0; m < 2; m++) {- if constexpr (BLOCK_N == 128) {- if constexpr (EPILOGUE_X16) {- float tmp[64];- tcgen05_ld_16x256bx16(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 16; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }- } else {- float tmp[32];-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 64);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int ii = i + 8;- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + ii * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }- }- } else {- float tmp[BLOCK_N / 2];- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 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 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ if constexpr (BLOCK_N == 128) {+ float tmp[64];+ tcgen05_ld_16x256bx16(tmp, trow, 0);+ asm volatile("tcgen05.wait::ld.sync.aligned;");+ #pragma unroll+ for (int i = 0; i < 16; i++) {+ const int row = off_m + trow + (lane_id >> 2);+ const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ if (row < M) {reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;+ }+ const int row2 = row + 8;+ if (row2 < M) {reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});}}- }- } else {- #pragma unroll- for (int m = 0; m < 2; m++) {- if constexpr (BLOCK_N == 128) {- if constexpr (EPILOGUE_X16) {- float tmp[64];- tcgen05_ld_16x256bx16(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 16; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ } else {+ const bool full_n = (off_n + BLOCK_N) <= N;+ float tmp[BLOCK_N / 2];+ tcgen05_ld_16x256bx8(tmp, trow, 0);+ asm volatile("tcgen05.wait::ld.sync.aligned;");- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }+ #pragma unroll+ for (int i = 0; i < BLOCK_N / 8; i++) {+ const int row = off_m + trow + (lane_id >> 2);+ const int col = off_n + i * 8 + ((lane_id & 3) << 1);- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }+ if (row < M) {+ if (full_n || (col + 1) < N) {+ reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});+ } else if (col < N) {+ C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);}- } else {- float tmp[32];-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }- }-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 64);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int ii = i + 8;- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + ii * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }- }}- } else {- float tmp[BLOCK_N / 2];- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 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 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }+ const int row2 = row + 8;+ if (row2 < M) {+ if (full_n || (col + 1) < N) {+ reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});+ } else if (col < N) {+ C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);}-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }}}}⋯ 4 unchanged lines}}- template <int BLOCK_N, int NUM_STAGES, bool EPILOGUE_X16>+ template <int BLOCK_N, int NUM_STAGES>__global__ __launch_bounds__(128 + 2 * WARP_SIZE)void kernel_grouped(const GroupDesc *descs) {constexpr int BLOCK_M = 128;⋯ 56 unchanged linesconst int num_iters = K / BLOCK_K;if (warp_id == NUM_WARPS - 2 && elect_sync()) {- uint64_t cache_A, cache_B;- if (grid_n >= (grid_m << 2)) {- cache_A = EVICT_LAST;- cache_B = EVICT_FIRST;- } else {- const bool keep_A = (grid_n >= grid_m);- cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;- cache_B = keep_A ? EVICT_FIRST : EVICT_LAST;- }+ const bool keep_A = (grid_n >= grid_m);+ const uint64_t cache_A = keep_A ? EVICT_LAST : EVICT_FIRST;+ const uint64_t cache_B = keep_A ? EVICT_FIRST : EVICT_LAST;auto issue_tma = [&](int iter_k, int stage_id) {const int mbar_addr = tma_mbar_addr + stage_id * 8;⋯ 85 unchanged linesasm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];":: "r"(mainloop_mbar_addr) : "memory");} else if (tid < BLOCK_M) {- const bool full_m = (off_m + 128) <= M;- int thr = 128;- if (!full_m) {- thr = active_threads_128(M, off_m);- if (tid >= thr) return;- }+ const int thr = active_threads_128(M, off_m);+ if (tid >= thr) return;mbarrier_wait(mainloop_mbar_addr, 0);asm volatile("tcgen05.fence::after_thread_sync;");- const bool full_n = (off_n + BLOCK_N) <= N;+ #pragma unroll+ for (int m = 0; m < 2; m++) {+ const int trow = warp_id * 32 + m * 16;- if (full_m && full_n) {- #pragma unroll- for (int m = 0; m < 2; m++) {- if constexpr (BLOCK_N == 128) {- if constexpr (EPILOGUE_X16) {- float tmp[64];- tcgen05_ld_16x256bx16(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 16; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }- } else {- float tmp[32];-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 64);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int ii = i + 8;- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + ii * 8 + ((lane_id & 3) << 1);- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- }- }- } else {- float tmp[BLOCK_N / 2];- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 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 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ if constexpr (BLOCK_N == 128) {+ float tmp[64];+ tcgen05_ld_16x256bx16(tmp, trow, 0);+ asm volatile("tcgen05.wait::ld.sync.aligned;");+ #pragma unroll+ for (int i = 0; i < 16; i++) {+ const int row = off_m + trow + (lane_id >> 2);+ const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ if (row < M) {reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- const int row2 = row + 8;+ }+ const int row2 = row + 8;+ if (row2 < M) {reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});}}- }- } else {- #pragma unroll- for (int m = 0; m < 2; m++) {- if constexpr (BLOCK_N == 128) {- if constexpr (EPILOGUE_X16) {- float tmp[64];- tcgen05_ld_16x256bx16(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 16; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);+ } else {+ const bool full_n = (off_n + BLOCK_N) <= N;+ float tmp[BLOCK_N / 2];+ tcgen05_ld_16x256bx8(tmp, trow, 0);+ asm volatile("tcgen05.wait::ld.sync.aligned;");- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }+ #pragma unroll+ for (int i = 0; i < BLOCK_N / 8; i++) {+ const int row = off_m + trow + (lane_id >> 2);+ const int col = off_n + i * 8 + ((lane_id & 3) << 1);- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }+ if (row < M) {+ if (full_n || (col + 1) < N) {+ reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});+ } else if (col < N) {+ C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);}- } else {- float tmp[32];-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }- }-- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 64);- asm volatile("tcgen05.wait::ld.sync.aligned;");- #pragma unroll- for (int i = 0; i < 8; i++) {- const int ii = i + 8;- const int row = off_m + warp_id * 32 + m * 16 + (lane_id >> 2);- const int col = off_n + ii * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }- }-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }- }}- } else {- float tmp[BLOCK_N / 2];- tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 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 + m * 16 + (lane_id >> 2);- const int col = off_n + i * 8 + ((lane_id & 3) << 1);-- if (row < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row * N + col)[0] = __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});- } else if (col < N) {- C_ptr[row * N + col] = __float2half_rn(tmp[i * 4 + 0]);- }+ const int row2 = row + 8;+ if (row2 < M) {+ if (full_n || (col + 1) < N) {+ reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});+ } else if (col < N) {+ C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);}-- const int row2 = row + 8;- if (row2 < M) {- if (full_n || (col + 1) < N) {- reinterpret_cast<half2 *>(C_ptr + row2 * N + col)[0] = __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});- } else if (col < N) {- C_ptr[row2 * N + col] = __float2half_rn(tmp[i * 4 + 2]);- }- }}}}⋯ 4 unchanged lines}}- template <int BLOCK_N, int NUM_STAGES, bool EPILOGUE_X16>+ template <int BLOCK_N, int NUM_STAGES>static __forceinline__ void gemm_launch(const at::Tensor& A,const at::Tensor& B,⋯ 22 unchanged linesconst int SF_size = 128 * 256 / 16;const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;- auto k = kernel<BLOCK_N, NUM_STAGES, EPILOGUE_X16>;- // 只在首次使用该设备时设置一次，避免每次调用都走一次 runtime API+ auto k = kernel<BLOCK_N, NUM_STAGES>;static int last_dev = -1;int dev = -1;cudaGetDevice(&dev);⋯ 8 unchanged lines}static __forceinline__ int pick_stage_from_K(int K) {- // K 一定是 256 的倍数const int iters = K >> 8;if (iters <= 8) return 2;if (iters <= 16) return 3;return 4;}+ static __forceinline__ int pick_stage_bn128_from_K(int K) {+ const int iters = K >> 8;+ if (iters <= 8) return 2;+ if (iters >= 24) return 4;+ return 3;+ }+at::Tensor gemm(const at::Tensor& A,const at::Tensor& B,⋯ 2 unchanged linesat::Tensor& C,int64_t M,int64_t N,- int64_t K,- bool enable_bn128_stage4,- bool epilogue_x16+ int64_t K) {const int Mi = (int)M;const int Ni = (int)N;const int Ki = (int)K;if (((Ni & 127) == 0) && (Ni >= 128)) {- const int iters = Ki >> 8;- int stage = (iters <= 8) ? 2 : 3;- if (enable_bn128_stage4 && iters >= 24) {- // 形状感知：stage=4 的 shared 成本很大，仅在 tile 总数较小（更像“少波次长 K”）时启用- const int grid_m = (Mi + 127) >> 7;- const int grid_n = Ni >> 7;- const int tiles = grid_m * grid_n;- if ((grid_m <= 2) && (tiles <= 128)) stage = 4;- }- if (epilogue_x16) {- if (stage == 2) gemm_launch<128, 2, true>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else if (stage == 3) gemm_launch<128, 3, true>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else gemm_launch<128, 4, true>(A, B, SFA, SFB, C, Mi, Ni, Ki);- } else {- if (stage == 2) gemm_launch<128, 2, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else if (stage == 3) gemm_launch<128, 3, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else gemm_launch<128, 4, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);- }+ const int stage = pick_stage_bn128_from_K(Ki);+ if (stage == 2) gemm_launch<128, 2>(A, B, SFA, SFB, C, Mi, Ni, Ki);+ else if (stage == 3) gemm_launch<128, 3>(A, B, SFA, SFB, C, Mi, Ni, Ki);+ else gemm_launch<128, 4>(A, B, SFA, SFB, C, Mi, Ni, Ki);} else {const int stage = pick_stage_from_K(Ki);- if (stage == 2) gemm_launch<64, 2, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else if (stage == 3) gemm_launch<64, 3, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);- else gemm_launch<64, 4, false>(A, B, SFA, SFB, C, Mi, Ni, Ki);+ if (stage == 2) gemm_launch<64, 2>(A, B, SFA, SFB, C, Mi, Ni, Ki);+ else if (stage == 3) gemm_launch<64, 3>(A, B, SFA, SFB, C, Mi, Ni, Ki);+ else gemm_launch<64, 4>(A, B, SFA, SFB, C, Mi, Ni, Ki);}GG_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");⋯ 4 unchanged linesat::Tensor descs_d;int64_t cap_G;int64_t dev;- uint64_t last_h0;- uint64_t last_h1;+ uint64_t last_hash;+ uint64_t last_hash2;int64_t last_G;int last_block_n;int last_max_grid_m;int last_max_grid_n;int last_stage;- uint64_t last_head[8];- uint64_t last_mid[8];- uint64_t last_tail[8];GroupWorkspace(): cap_G(0),dev(-1),- last_h0(0),- last_h1(0),+ last_hash(0),+ last_hash2(0),last_G(0),last_block_n(0),last_max_grid_m(0),last_max_grid_n(0),- last_stage(0),- last_head{},- last_mid{},- last_tail{} {}+ last_stage(0) {}};static GroupWorkspace g_ws;⋯ 4 unchanged linesg_ws.cap_G = G;at::TensorOptions opt_u8 = at::TensorOptions().device(at::kCUDA, (int)dev).dtype(at::kByte);g_ws.descs_d = at::empty({G, (int64_t)sizeof(GroupDesc)}, opt_u8);- g_ws.last_h0 = 0;- g_ws.last_h1 = 0;+ g_ws.last_hash = 0;+ g_ws.last_hash2 = 0;g_ws.last_G = 0;- for (int i = 0; i < 8; i++) { g_ws.last_head[i] = 0; g_ws.last_mid[i] = 0; g_ws.last_tail[i] = 0; }}}⋯ 3 unchanged linesreturn h;}- static __forceinline__ uint64_t lcg_mix_u64(uint64_t h, uint64_t x) {- return h * 6364136223846793005ULL + x + 1442695040888963407ULL;- }-- template <int BLOCK_N, int NUM_STAGES, bool EPILOGUE_X16>+ template <int BLOCK_N, int NUM_STAGES>static __forceinline__ void grouped_launch(int64_t dev,int max_grid_m,int max_grid_n,- int64_t G,- const GroupDesc *descs_ptr+ int64_t G) {const int tb_size = 128 + 2 * WARP_SIZE;const int A_size = 128 * 256 / 2;⋯ 1 unchanged linesconst int SF_size = 128 * 256 / 16;const int smem_size = (A_size + B_size + SF_size * 2) * NUM_STAGES;- auto k = kernel_grouped<BLOCK_N, NUM_STAGES, EPILOGUE_X16>;+ auto k = kernel_grouped<BLOCK_N, NUM_STAGES>;static int last_dev = -1;if ((int)dev != last_dev) {auto err = cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);⋯ 2 unchanged lines}dim3 grid((unsigned)max_grid_n, (unsigned)max_grid_m, (unsigned)G);- k<<<grid, tb_size, smem_size>>>(descs_ptr);+ k<<<grid, tb_size, smem_size>>>((const GroupDesc *)g_ws.descs_d.data_ptr());}void gemm_grouped(⋯ 2 unchanged linesat::TensorList SFA_list,at::TensorList SFB_list,at::TensorList C_list,- bool force_bn64,- bool enable_bn128_stage4,- bool epilogue_x16+ bool force_bn64) {const int64_t G = (int64_t)A_list.size();if (G <= 0) return;⋯ 11 unchanged lines}const int block_n = use_bn128 ? 128 : 64;- uint64_t head[8];- uint64_t mid[8];- uint64_t tail[8];- const int mid_i = (int)G >> 1;+ uint64_t h1 = 1469598103934665603ULL;+ uint64_t h2 = 1099511628211ULL ^ 0x9E3779B97F4A7C15ULL;+ h1 = fnv1a_mix_u64(h1, (uint64_t)G);+ h1 = fnv1a_mix_u64(h1, (uint64_t)block_n);+ h2 = fnv1a_mix_u64(h2, (uint64_t)G ^ 0xD6E8FEB86659FD93ULL);+ h2 = fnv1a_mix_u64(h2, (uint64_t)block_n ^ 0xA5A3564E27F3A7A9ULL);- uint64_t h0 = 1469598103934665603ULL;- uint64_t h1 = 0x9E3779B97F4A7C15ULL;- h0 = fnv1a_mix_u64(h0, (uint64_t)G);- h0 = fnv1a_mix_u64(h0, (uint64_t)block_n);- h1 = lcg_mix_u64(h1, (uint64_t)G);- h1 = lcg_mix_u64(h1, (uint64_t)block_n);-int max_grid_m = 0;- int min_grid_m = 0x7fffffff;int max_grid_n = 0;int max_K = 0;- int64_t sum_grid_mn = 0;for (int i = 0; i < (int)G; i++) {const auto &A = A_list[i];⋯ 9 unchanged linesconst int grid_m = (M + 127) / 128;const int grid_n = (N + block_n - 1) / block_n;- sum_grid_mn += (int64_t)grid_m * (int64_t)grid_n;if (grid_m > max_grid_m) max_grid_m = grid_m;- if (grid_m < min_grid_m) min_grid_m = grid_m;if (grid_n > max_grid_n) max_grid_n = grid_n;if (K > max_K) max_K = K;⋯ 2 unchanged linesconst uint64_t SFA_ptr = (uint64_t)SFA.data_ptr();const uint64_t SFB_ptr = (uint64_t)SFB.data_ptr();const uint64_t C_ptr = (uint64_t)C.data_ptr<at::Half>();- const uint64_t MNKA = ((uint64_t)K << 32) | (uint64_t)(uint32_t)Apad;+ const uint64_t shape_pack = ((uint64_t)K << 32) | (uint64_t)(uint32_t)Apad;- if (i == 0) {- head[0] = A_ptr;- head[1] = B_ptr;- head[2] = SFA_ptr;- head[3] = SFB_ptr;- head[4] = C_ptr;- head[5] = (uint64_t)M;- head[6] = (uint64_t)N;- head[7] = MNKA;- }- if (i == mid_i) {- mid[0] = A_ptr;- mid[1] = B_ptr;- mid[2] = SFA_ptr;- mid[3] = SFB_ptr;- mid[4] = C_ptr;- mid[5] = (uint64_t)M;- mid[6] = (uint64_t)N;- mid[7] = MNKA;- }- if (i + 1 == (int)G) {- tail[0] = A_ptr;- tail[1] = B_ptr;- tail[2] = SFA_ptr;- tail[3] = SFB_ptr;- tail[4] = C_ptr;- tail[5] = (uint64_t)M;- tail[6] = (uint64_t)N;- tail[7] = MNKA;- }+ h1 = fnv1a_mix_u64(h1, A_ptr);+ h1 = fnv1a_mix_u64(h1, B_ptr);+ h1 = fnv1a_mix_u64(h1, SFA_ptr);+ h1 = fnv1a_mix_u64(h1, SFB_ptr);+ h1 = fnv1a_mix_u64(h1, C_ptr);+ h1 = fnv1a_mix_u64(h1, (uint64_t)M);+ h1 = fnv1a_mix_u64(h1, (uint64_t)N);+ h1 = fnv1a_mix_u64(h1, shape_pack);- h0 = fnv1a_mix_u64(h0, A_ptr);- h0 = fnv1a_mix_u64(h0, B_ptr);- h0 = fnv1a_mix_u64(h0, SFA_ptr);- h0 = fnv1a_mix_u64(h0, SFB_ptr);- h0 = fnv1a_mix_u64(h0, C_ptr);- h0 = fnv1a_mix_u64(h0, (uint64_t)M);- h0 = fnv1a_mix_u64(h0, (uint64_t)N);- h0 = fnv1a_mix_u64(h0, MNKA);-- h1 = lcg_mix_u64(h1, A_ptr);- h1 = lcg_mix_u64(h1, B_ptr);- h1 = lcg_mix_u64(h1, SFA_ptr);- h1 = lcg_mix_u64(h1, SFB_ptr);- h1 = lcg_mix_u64(h1, C_ptr);- h1 = lcg_mix_u64(h1, (uint64_t)M);- h1 = lcg_mix_u64(h1, (uint64_t)N);- h1 = lcg_mix_u64(h1, MNKA);+ h2 = fnv1a_mix_u64(h2, A_ptr ^ 0x8B8B8B8B8B8B8B8BULL);+ h2 = fnv1a_mix_u64(h2, B_ptr ^ 0xC3C3C3C3C3C3C3C3ULL);+ h2 = fnv1a_mix_u64(h2, SFA_ptr ^ 0x1313131313131313ULL);+ h2 = fnv1a_mix_u64(h2, SFB_ptr ^ 0x5D5D5D5D5D5D5D5DULL);+ h2 = fnv1a_mix_u64(h2, C_ptr ^ 0xA7A7A7A7A7A7A7A7ULL);+ h2 = fnv1a_mix_u64(h2, ((uint64_t)M << 32) ^ (uint64_t)N);+ h2 = fnv1a_mix_u64(h2, shape_pack ^ 0x9E3779B97F4A7C15ULL);}int stage = 0;- if (use_bn128) {- const int iters = max_K >> 8;- stage = (iters <= 8) ? 2 : 3;- if (enable_bn128_stage4 && iters >= 24) {- const int tiles = max_grid_m * max_grid_n;- if ((max_grid_m <= 2) && (tiles <= 128)) stage = 4;- }- } else {- stage = pick_stage_from_K(max_K);- }- bool head_ok = true;- bool mid_ok = true;- bool tail_ok = true;- for (int j = 0; j < 8; j++) {- head_ok = head_ok && (g_ws.last_head[j] == head[j]);- mid_ok = mid_ok && (g_ws.last_mid[j] == mid[j]);- tail_ok = tail_ok && (g_ws.last_tail[j] == tail[j]);- }- const bool hit = (g_ws.last_h0 == h0)- && (g_ws.last_h1 == h1)+ if (use_bn128) stage = pick_stage_bn128_from_K(max_K);+ else stage = pick_stage_from_K(max_K);++ const bool hit = (g_ws.last_hash == h1)+ && (g_ws.last_hash2 == h2)&& (g_ws.last_G == G)&& (g_ws.last_block_n == block_n)&& (g_ws.last_stage == stage)- && head_ok- && mid_ok- && tail_ok&& g_ws.descs_d.defined();- // P1：默认单次 launch；仅当空 block 浪费足够大时才启用分桶双次 launch- int split_m = max_grid_m;- int G0 = (int)G;- int G1 = 0;- int max_grid_m0 = max_grid_m;- int max_grid_n0 = max_grid_n;- int max_grid_m1 = 0;- int max_grid_n1 = 0;-- bool enable_split = false;- if (min_grid_m < max_grid_m) {- const int64_t total_blocks = (int64_t)G * (int64_t)max_grid_m * (int64_t)max_grid_n;- const int64_t waste_blocks = total_blocks - sum_grid_mn;- // 启发式：必须满足“绝对浪费大”或“浪费占比高且规模足够”- if (G >= 4) {- const bool huge_waste = (waste_blocks >= 8192);- const bool high_ratio = (waste_blocks >= 2048) && (waste_blocks * 3 >= total_blocks);- enable_split = huge_waste || high_ratio;- }- }-- if (enable_split) {- split_m = (min_grid_m + max_grid_m) >> 1;- G0 = 0;- G1 = 0;- max_grid_m0 = 0;- max_grid_n0 = 0;- max_grid_m1 = 0;- max_grid_n1 = 0;- for (int i = 0; i < (int)G; i++) {- const auto &C = C_list[i];- const int M = (int)C.size(0);- const int N = (int)C.size(1);- const int grid_m = (M + 127) / 128;- const int grid_n = (N + block_n - 1) / block_n;- if (grid_m <= split_m) {- G0++;- if (grid_m > max_grid_m0) max_grid_m0 = grid_m;- if (grid_n > max_grid_n0) max_grid_n0 = grid_n;- } else {- G1++;- if (grid_m > max_grid_m1) max_grid_m1 = grid_m;- if (grid_n > max_grid_n1) max_grid_n1 = grid_n;- }- }- if (G0 <= 0 || G1 <= 0) {- G0 = (int)G;- G1 = 0;- split_m = max_grid_m;- max_grid_m0 = max_grid_m;- max_grid_n0 = max_grid_n;- }- }-if (!hit) {std::vector<GroupDesc> descs((size_t)G);- int w0 = 0;- int w1 = G0;for (int i = 0; i < (int)G; i++) {const auto &A = A_list[i];const auto &B = B_list[i];⋯ 15 unchanged linesd.M = M;d.N = N;d.K = K;- const int grid_m = (M + 127) / 128;- const int pos = (G1 == 0 || grid_m <= split_m) ? (w0++) : (w1++);- descs[(size_t)pos] = d;+ descs[(size_t)i] = d;}- // 单次 memcpy：避免多次 Host→GPU 调用与同步点cudaError_t cperr = cudaMemcpy(g_ws.descs_d.data_ptr(),descs.data(),⋯ 2 unchanged lines);GG_CHECK(cperr == cudaSuccess, "memcpy fail");- g_ws.last_h0 = h0;- g_ws.last_h1 = h1;+ g_ws.last_hash = h1;+ g_ws.last_hash2 = h2;g_ws.last_G = G;g_ws.last_block_n = block_n;g_ws.last_max_grid_m = max_grid_m;g_ws.last_max_grid_n = max_grid_n;g_ws.last_stage = stage;- for (int j = 0; j < 8; j++) { g_ws.last_head[j] = head[j]; g_ws.last_mid[j] = mid[j]; g_ws.last_tail[j] = tail[j]; }} else {max_grid_m = g_ws.last_max_grid_m;max_grid_n = g_ws.last_max_grid_n;}- const GroupDesc *descs_ptr = (const GroupDesc *)g_ws.descs_d.data_ptr();-if (use_bn128) {- if (epilogue_x16) {- if (stage == 2) {- grouped_launch<128, 2, true>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 2, true>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else if (stage == 3) {- grouped_launch<128, 3, true>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 3, true>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else {- grouped_launch<128, 4, true>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 4, true>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- }- } else {- if (stage == 2) {- grouped_launch<128, 2, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 2, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else if (stage == 3) {- grouped_launch<128, 3, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 3, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else {- grouped_launch<128, 4, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<128, 4, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- }- }+ if (stage == 2) grouped_launch<128, 2>(dev, max_grid_m, max_grid_n, G);+ else if (stage == 3) grouped_launch<128, 3>(dev, max_grid_m, max_grid_n, G);+ else grouped_launch<128, 4>(dev, max_grid_m, max_grid_n, G);} else {- if (stage == 2) {- grouped_launch<64, 2, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<64, 2, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else if (stage == 3) {- grouped_launch<64, 3, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<64, 3, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- } else {- grouped_launch<64, 4, false>(dev, max_grid_m0, max_grid_n0, (int64_t)G0, descs_ptr);- if (G1) grouped_launch<64, 4, false>(dev, max_grid_m1, max_grid_n1, (int64_t)G1, descs_ptr + G0);- }+ if (stage == 2) grouped_launch<64, 2>(dev, max_grid_m, max_grid_n, G);+ else if (stage == 3) grouped_launch<64, 3>(dev, max_grid_m, max_grid_n, G);+ else grouped_launch<64, 4>(dev, max_grid_m, max_grid_n, G);}GG_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");}TORCH_LIBRARY(nvfp4_group_gemm_opt, m) {- m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int M, int N, int K, bool enable_bn128_stage4, bool epilogue_x16) -> Tensor");+ m.def("gemm(Tensor A, Tensor B, Tensor SFA, Tensor SFB, Tensor(a!) C, int M, int N, int K) -> Tensor");m.impl("gemm", &gemm);- m.def("gemm_grouped(Tensor[] A, Tensor[] B, Tensor[] SFA, Tensor[] SFB, Tensor[] C, bool force_bn64, bool enable_bn128_stage4, bool epilogue_x16) -> ()");+ m.def("gemm_grouped(Tensor[] A, Tensor[] B, Tensor[] SFA, Tensor[] SFB, Tensor[] C, bool force_bn64) -> ()");m.impl("gemm_grouped", &gemm_grouped);}"""⋯ 1 unchanged linesbuild_dir = os.path.join(os.path.dirname(__file__), ".build_nvfp4_group_gemm_opt")os.makedirs(build_dir, exist_ok=True)- wait_backoff = 1 if os.environ.get("NVFP4_WAIT_BACKOFF", "0") == "1" else 0-load_inline(name="nvfp4_group_gemm_opt_ext",cpp_sources="",⋯ 8 unchanged lines"--expt-relaxed-constexpr","--relocatable-device-code=false","-std=c++17",- f"-DNVFP4_GGEMM_WAIT_BACKOFF={wait_backoff}",],extra_ldflags=["-lcuda"],with_cuda=True,⋯ 64 unchanged linesdef _must_sfp_layout(x: torch.Tensor) -> None:-if (not x.is_cuda) or (x.dim() != 6) or (x.element_size() != 1) or (int(x.storage_offset()) != 0):raise RuntimeError("bad sfx_p layout")if int(x.size(0)) != 32 or int(x.size(1)) != 4 or int(x.size(3)) != 4:⋯ 27 unchanged linessfa_list: List[torch.Tensor] = []sfb_list: List[torch.Tensor] = []c_list: List[torch.Tensor] = []- c_out_list: List[torch.Tensor] = []- c_tmp_list: List[torch.Tensor] = []for i in range(g):a, b, c = abc_tensors[i]⋯ 1 unchanged linessfa_p, sfb_p = sfasfb_reordered_tensors[i]m, n, _k, _l = problem_sizes[i]+ if (not a.is_contiguous()) or (not b.is_contiguous()) or (not c.is_contiguous()):+ raise RuntimeError("grouped fast path requires contiguous a/b/c")+m_int = int(m)n_int = int(n)-- if not a.is_contiguous():- a = a.contiguous()- if not b.is_contiguous():- b = b.contiguous()-- c_out = c- if not c_out.is_contiguous():- c_tmp = torch.empty_like(c_out, memory_format=torch.contiguous_format)- else:- c_tmp = c_out-a_u8 = _as_u8(a)b_u8 = _as_u8(b)⋯ 19 unchanged linesb_list.append(b_u8)sfa_list.append(sfa_arg)sfb_list.append(sfb_arg)- c_list.append(c_tmp)- c_out_list.append(c_out)- c_tmp_list.append(c_tmp)+ c_list.append(c)- gemm_grouped(- a_list,- b_list,- sfa_list,- sfb_list,- c_list,- _FORCE_BN64,- _FORCE_BN128_STAGE4,- _FORCE_BN128_EPILOGUE_X16,- )+ gemm_grouped(a_list, b_list, sfa_list, sfb_list, c_list, _FORCE_BN64)+ return c_list- for i in range(g):- c_out = c_out_list[i]- c_tmp = c_tmp_list[i]- if c_tmp is not c_out:- c_out.copy_(c_tmp)- outs.append(c_out)-- return outs-for i in range(g):a, b, c = abc_tensors[i]sfa, sfb = sfasfb_tensors[i]⋯ 33 unchanged linessfa_arg = sfa_psfb_arg = sfb_p- gemm(- a_pad,- b_u8,- sfa_arg,- sfb_arg,- c_tmp,- m_int,- n_int,- k_int,- _FORCE_BN128_STAGE4,- _FORCE_BN128_EPILOGUE_X16,- )+ gemm(a_pad, b_u8, sfa_arg, sfb_arg, c_tmp, m_int, n_int, k_int)else:for li in range(l_int):a2 = _as_u8(a[..., li]).contiguous().unsqueeze(-1)⋯ 12 unchanged linessfb_r = _reorder_scale_from_raw(sfb2, ((n_int + 127) // 128) * 128)⋯ diff truncated
scrolls · 1202 diff lines total
Best evidence level for this revision: reported
JSON