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

fix: align SMEM buffers to 16 bytes for UMMA descriptors

Browse Source
This commit is contained in:
biondizzle
2026-05-28 08:02:53 +00:00
parent 3549a2388b
commit c5f7a9a15c
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
dsv4/kernels/attention/fmha_qk_verify.cuh
View File

@@ -11,6 +11,7 @@
#include "fmha_common.cuh"
#include "fmha_umma_desc.cuh"
#include <cstdint>
namespace dsv4::kernels::attention {
@@ -29,10 +30,13 @@ fmha_qk_verify(
const bf16_t* kb = k + batch*bstride_kv;
// SMEM: sQ (128×HD BF16) + sK (128×HD BF16) + tmem_base (4B)
// SMEM: sTmemBase (4B, aligned to 16B) + sQ (128×HD BF16) + sK (128×HD BF16)
// CRITICAL: sQ and sK must be 16-byte aligned for UMMA descriptors
extern __shared__ char sbuf[];
uint32_t* sTmemBase = (uint32_t*)sbuf;
bf16_t* sQ = (bf16_t*)(sbuf + 4);
bf16_t* sK = (bf16_t*)(sbuf + 4 + 128 * HD * sizeof(bf16_t));
// Align sQ to 16 bytes
bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~15);
bf16_t* sK = sQ + 128 * HD;
// Load Q: (1, HD) padded to (128, HD) with zeros
for (int i = tid; i < 128 * HD; i += NTHREADS) sQ[i] = 0;