diff --git a/dsv4/kernels/attention/fmha_qk_verify.cuh b/dsv4/kernels/attention/fmha_qk_verify.cuh index ea432ebc..ab885610 100644 --- a/dsv4/kernels/attention/fmha_qk_verify.cuh +++ b/dsv4/kernels/attention/fmha_qk_verify.cuh @@ -11,6 +11,7 @@ #include "fmha_common.cuh" #include "fmha_umma_desc.cuh" +#include 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;