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

PV SS MMA with 8 K-tile accumulation

Browse Source
This commit is contained in:
biondizzle
2026-05-28 14:30:09 +00:00
parent d5b0941f2e
commit 727c509454
1 changed files with 16 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tests/unit/test_pv_ss_128.cu
View File

@@ -37,7 +37,7 @@ test_pv_ss_128()
for (int j = tid; j < 128; j += 128) {
int core_k = j / 8, lc = j % 8;
int dst_idx = core_k * 16 * 64 + 0 * 64 + 0 * 8 + lc;
sP[dst_idx] = f32_to_bf16(0.5f);
sP[dst_idx] = f32_to_bf16(0.5f); // All 128 positions = 0.5
}
__syncthreads();
@@ -57,14 +57,20 @@ test_pv_ss_128()
__syncthreads();
uint32_t tb = *sTmemBase;
// PV SS MMA: A = P K-tile 0 at sP+0, B = V at sV
// K-tile 0 of (128,128): g_k=[0,1], offset=0, size=2048 BF16
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sP), BLOCK_MN);
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
uint32_t idesc = make_idesc(BLOCK_MN, HD);
if (tid == 0) umma_ss_f16(tb, dp, dv, idesc, false);
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
// PV SS MMA: 8 K-tiles with accumulation
// K-tile kt of (128,128): g_k=[2*kt, 2*kt+1], offset = kt * 2048 BF16
// P row 0 = 0.5 for first 16 positions, 0.3 for positions 16-127
{
uint64_t dv = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sV), 16);
uint32_t idesc = make_idesc(BLOCK_MN, HD);
for (int kt = 0; kt < 8; kt++) {
bf16_t* sp = sP + kt * 2048;
uint64_t dp = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sp), BLOCK_MN);
if (tid == 0) umma_ss_f16(tb, dp, dv, idesc, kt > 0);
asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
__syncthreads();
}
}
// Read O from TMEM
if (wid == 0) {
@@ -82,7 +88,7 @@ test_pv_ss_128()
printf("O[0,0..15] (raw): ");
for (int d=0;d<HD;d++) printf("%.2f ", o_vals[d]);
// Expected: sum(P[0,0..15] * V[j,0..15]) = 16 * 0.5 * 1.0 = 8.0 (raw, no MMA scale)
printf("\nExpected ~8.0 (raw, MMA scale unknown)\n");
printf("\nExpected ~64.0 (8 K-tiles × 16 values × 0.5 × 1.0, MMA scale 1.0)\n");
}
}