diff --git a/dsv4/kernels/attention/fmha_6warp_tma.cuh b/dsv4/kernels/attention/fmha_6warp_tma.cuh index 5579fa67..5658421b 100644 --- a/dsv4/kernels/attention/fmha_6warp_tma.cuh +++ b/dsv4/kernels/attention/fmha_6warp_tma.cuh @@ -99,13 +99,10 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { float* sRowSum = (float*)(sbuf + off); off += MAX_ROWS * sizeof(float); off = (off + 127) & ~(size_t)127; - bf16_t* sQ_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t); + bf16_t* sTmaBuf = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t); // TMA staging buffer (128, 16) row-major off = (off + 127) & ~(size_t)127; - bf16_t* sK_tma = (bf16_t*)(sbuf + off); off += TMA_TILE_BF16 * sizeof(bf16_t); - - off = (off + 127) & ~(size_t)127; - bf16_t* sQ = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); + bf16_t* sQ = (bf16_t*)(sbuf + off); off += 128 * HD * sizeof(bf16_t); // full Q (128, HD) canonical off = (off + 127) & ~(size_t)127; bf16_t* sK = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); @@ -113,9 +110,6 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { off = (off + 127) & ~(size_t)127; bf16_t* sPk = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); - off = (off + 127) & ~(size_t)127; - bf16_t* sV_tma = (bf16_t*)(sbuf + off); off += 16 * 128 * sizeof(bf16_t); - off = (off + 127) & ~(size_t)127; bf16_t* sV = (bf16_t*)(sbuf + off); off += TILE_SZ * sizeof(bf16_t); @@ -143,31 +137,52 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { const bool my_row_active = my_warp_active && (my_row < T); // ================================================================== - // QK GEMM → S in TMEM (loop over K sub-tiles) + // Load full Q into SMEM (128, HD) canonical via TMA // ================================================================== - // mbarrier phase parity — tracks which phase we're waiting on. - // The mbarrier was initialized with count=1 before the loop. - // After each TMA + arrive.expect_tx + wait, we flip the phase. - // We do NOT re-init the mbarrier inside the loop. int phase = 0; - for (int kt = 0; kt < NKT_QK; kt++) { - // --- TMA load Q sub-tile --- + // Zero Q canonical buffer first + if (is_load_warp) { + for (int i = lane; i < 128 * HD; i += 32) sQ[i] = 0; + } + __syncthreads(); + + for (int qkt = 0; qkt < NKT_QK; qkt++) { if (is_load_warp && lane == 0) { - uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sQ_tma); - tma_load_2d(smem_dst, (uint64_t)tma_q, mbar_addr, kt * MMA_K_BF16, 0); + uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf); + tma_load_2d(smem_dst, (uint64_t)tma_q, mbar_addr, qkt * MMA_K_BF16, 0); tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); } tma_mbarrier_wait(mbar_addr, phase); phase ^= 1; __syncthreads(); - if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sQ, sQ_tma); + // Write (128, 16) row-major TMA buffer into the right position in (128, HD) canonical Q + // The qkt-th (128, 16) sub-tile in canonical = columns [qkt*16, qkt*16+16) + // canonical offset for core_k = qkt*16/8 = qkt*2, same core_mn layout + if (is_load_warp) { + constexpr int CORES_MN = 128 / 8; // 16 + constexpr int CORES_K_SUB = 16 / 8; // 2 + constexpr int SUB_TOTAL = 128 * 16; + for (int i = lane; i < SUB_TOTAL; i += 32) { + int r = i / 16, c = i % 16; + int core_mn = r / 8, local_r = r % 8; + int core_k_sub = c / 8, local_c = c % 8; + int core_k_full = qkt * 2 + core_k_sub; + int dst_idx = core_k_full * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c; + sQ[dst_idx] = sTmaBuf[i]; + } + } __syncthreads(); + } + // ================================================================== + // QK GEMM → S in TMEM (loop over K sub-tiles) + // ================================================================== + for (int kt = 0; kt < NKT_QK; kt++) { // --- TMA load K sub-tile --- if (is_load_warp && lane == 0) { - uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sK_tma); + uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf); tma_load_2d(smem_dst, (uint64_t)tma_k, mbar_addr, kt * MMA_K_BF16, 0); tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); } @@ -175,13 +190,15 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { phase ^= 1; __syncthreads(); - if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sK, sK_tma); + if (is_load_warp) write_smem_canonical<128, MMA_K_BF16, 32>(sK, sTmaBuf); __syncthreads(); - // MMA: sQ × sK → TMEM + // MMA: Q sub-tile × K sub-tile → TMEM + // Q's kt-th sub-tile starts at offset kt * 128 * 32 bytes in canonical SMEM if (is_mma_warp) { uint32_t idesc = make_idesc(128, 128); - uint64_t dq = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sQ), 128); + uint32_t sq_kt = (uint32_t)__cvta_generic_to_shared(sQ) + kt * 128 * 32; + uint64_t dq = make_umma_desc_kmajor_none(sq_kt, 128); uint64_t dk = make_umma_desc_kmajor_none(__cvta_generic_to_shared(sK), 128); if (tid == 128) umma_ss_f16(tb, dq, dk, idesc, kt > 0); asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); @@ -262,9 +279,8 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { } __syncthreads(); - // --- TMA load V sub-tile --- if (is_load_warp && lane == 0) { - uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sV_tma); + uint32_t smem_dst = (uint32_t)__cvta_generic_to_shared(sTmaBuf); // V is (HD, s_k). TMA 2D: coord {col_start, d_base} tma_load_2d(smem_dst, (uint64_t)tma_v, mbar_addr, col_start, d_base); tma_mbarrier_arrive_expect_tx(mbar_addr, TMA_TILE_BYTES); @@ -273,7 +289,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { phase ^= 1; __syncthreads(); - // Transpose sV_tma (16, 128) → sV (128, 16) canonical + // Transpose sTmaBuf (16, 128) → sV (128, 16) canonical if (is_load_warp) { constexpr int SV_CORES_MN = 128 / 8; for (int i = lane; i < TILE_SZ; i += 32) sV[i] = 0; @@ -282,7 +298,7 @@ fmha_6warp_tma_kernel(FmhaMultiRowTmaParams params) { int core_mn = r / 8, local_r = r % 8; int core_k = d / 8, local_c = d % 8; int dst_idx = core_k * SV_CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c; - sV[dst_idx] = sV_tma[i]; + sV[dst_idx] = sTmaBuf[i]; } } __syncthreads(); diff --git a/tests/unit/test_fmha_tma.cu b/tests/unit/test_fmha_tma.cu index 3faab79a..73229362 100644 --- a/tests/unit/test_fmha_tma.cu +++ b/tests/unit/test_fmha_tma.cu @@ -40,19 +40,15 @@ static int compute_smem_tma() { off += MAX_T * sizeof(float); // sRowMax off += MAX_T * sizeof(float); // sRowSum off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sQ_tma + off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sTmaBuf (TMA staging) off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK_tma + off += 128 * HD * sizeof(bf16_t); // sQ full (128, HD) canonical off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sQ canonical + off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK canonical (128, 16) off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sK canonical + off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sPk canonical (128, 16) off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sPk canonical - off = (off + 127) & ~(size_t)127; - off += 16 * 128 * sizeof(bf16_t); // sV_tma - off = (off + 127) & ~(size_t)127; - off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sV canonical + off += 128 * MMA_K_BF16 * sizeof(bf16_t); // sV canonical (128, 16) return (int)off; }