biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

auto: pre-test commit

Browse Source
This commit is contained in:
biondizzle
2026-05-28 15:48:15 +00:00
parent 80fd612132
commit 1cdb90462f
1 changed files with 10 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
tests/unit/test_tmem_layout_pv64.cu
View File

@@ -16,10 +16,11 @@
using namespace dsv4::kernels::attention;
static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
constexpr int HD = 64, SK = 128, BLOCK_MN = 128;
constexpr int MMA_K_BF16 = 16;
constexpr int TILE_SZ = BLOCK_MN * MMA_K_BF16;
constexpr int LOCAL_MMA_K = 16;
constexpr int TILE_SZ = BLOCK_MN * LOCAL_MMA_K;
constexpr int V_TILE_SZ = (HD / 8) * 2 * 64; // 1024
__global__ void __launch_bounds__(128)
@@ -40,19 +41,19 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
for (int kt = 0; kt < 4; kt++) {
bf16_t* sq = sQ0 + kt * TILE_SZ;
for (int i = tid; i < TILE_SZ; i += 128) sq[i] = 0;
for (int d = tid; d < MMA_K_BF16; d += 128) {
for (int d = tid; d < LOCAL_MMA_K; d += 128) {
int ck = d / 8, lc = d % 8;
sq[ck * 16 * 64 + lc] = q[kt * MMA_K_BF16 + d];
sq[ck * 16 * 64 + lc] = q[kt * LOCAL_MMA_K + d];
}
}
for (int kt = 0; kt < 4; kt++) {
bf16_t* sk = sK0 + kt * TILE_SZ;
for (int i = tid; i < TILE_SZ; i += 128) sk[i] = 0;
for (int r = 0; r < SK; r++) {
for (int d = tid; d < MMA_K_BF16; d += 128) {
for (int d = tid; d < LOCAL_MMA_K; d += 128) {
int ck = d / 8, lc = d % 8;
int tmn = r / 8, lr = r % 8;
sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * MMA_K_BF16 + d];
sk[ck * 16 * 64 + tmn * 64 + lr * 8 + lc] = k[r * HD + kt * LOCAL_MMA_K + d];
}
}
}
@@ -60,8 +61,8 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
bf16_t* sv = sV + kt * V_TILE_SZ;
for (int i = tid; i < V_TILE_SZ; i += 128) sv[i] = 0;
for (int d = tid; d < HD; d += 128) {
for (int lr = 0; lr < MMA_K_BF16; lr++) {
int r = kt * MMA_K_BF16 + lr;
for (int lr = 0; lr < LOCAL_MMA_K; lr++) {
int r = kt * LOCAL_MMA_K + lr;
int g_mn = d / 8, g_k = lr / 8;
int llr = d % 8, lc = lr % 8;
sv[g_k * 8 * 64 + g_mn * 64 + llr * 8 + lc] = v[d * SK + r];
@@ -131,7 +132,7 @@ test_tmem_layout_pv(const bf16_t* q, const bf16_t* k, const bf16_t* v,
if (tid < 16) {
int c = tid;
int ck = c / 8, lc = c % 8;
sPk[ck * 16 * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * MMA_K_BF16 + c]);
sPk[ck * 16 * 64 + 0 * 64 + 0 * 8 + lc] = f32_to_bf16(s_p_vals[kt * LOCAL_MMA_K + c]);
}
__syncthreads();