From 0f6907b0013ec7862a03bdb64f0d772a5a764065 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 28 May 2026 09:18:45 +0000 Subject: [PATCH] =?UTF-8?q?UMMA:=20fix=20descriptor=20+=20idesc=20?= =?UTF-8?q?=E2=80=94=20use=20gau-nernst=20tutorial=20values?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - LBO = BLOCK_MN * 16 (bytes), SBO = 128 (bytes) for K-major NONE - Canonical SMEM layout: column-major interleaving of core matrices - idesc is SEPARATE 32-bit value (was using desc_a>>32 = WRONG) - idesc encodes dtype/atype/btype/MMA_M/MMA_N - This was the root cause of 'misaligned address' errors --- dsv4/kernels/attention/fmha_umma_desc.cuh | 378 ++++++++++++---------- tests/unit/test_umma_qk.cu | 156 ++++----- 2 files changed, 262 insertions(+), 272 deletions(-) diff --git a/dsv4/kernels/attention/fmha_umma_desc.cuh b/dsv4/kernels/attention/fmha_umma_desc.cuh index af5d1b01..a5465c09 100644 --- a/dsv4/kernels/attention/fmha_umma_desc.cuh +++ b/dsv4/kernels/attention/fmha_umma_desc.cuh @@ -5,29 +5,40 @@ * SMEM LAYOUT FOR tcgen05.mma (from PTX spec + gau-nernst tutorial) * ================================================================== * - * tcgen05.mma expects SMEM data in a specific "canonical" layout. - * The data is organized as a hierarchy of "core matrices": + * tcgen05.mma expects SMEM data in a specific "canonical" layout + * described by the descriptor's LBO and SBO strides. * - * 8 rows × 8 columns per core matrix (for BF16, 128 bytes each) + * For K-major, SWIZZLE_NONE, BF16: + * LBO = BLOCK_MN * 16 (in bytes, NOT in 16B units) + * → desc field = BLOCK_MN (after >> 4) + * This is the stride from the "first column" to the "second column" + * of the 8×2 core matrix tile. Each column is 16 bytes wide and + * BLOCK_MN rows tall. * - * For K-major layout (the MMA expects K-major), the K dimension - * is the "contiguous" dimension. Within each 8×8 tile, K varies - * fastest (row-major within the tile). + * SBO = 8 * 16 = 128 (in bytes, NOT in 16B units) + * → desc field = 8 (after >> 4) + * This is the stride between 8-row groups. * - * The tile at (block_mn, block_k) is at SMEM offset: - * (block_mn * (BLOCK_K / 8) + block_k) * 64 (in BF16 elements) + * The hardware walks the SMEM as: + * For 8-row group g (0..BLOCK_MN/8-1): + * For column c (0 or 1): (2 columns per MMA, 16B each) + * Read 8×16B at address: start + g * SBO + c * LBO * - * Within each tile, element (local_mn, local_k) is at: - * local_mn * 8 + local_k + * This means the SMEM layout is: + * [core(0,0), core(1,0), ..., core(BLOCK_MN/8-1,0), + * core(0,1), core(1,1), ..., core(BLOCK_MN/8-1,1)] + * + * Where core(g, c) is an 8×8 BF16 tile (128 bytes) representing + * rows [8g, 8g+8) and columns [8c, 8c+8) of the logical matrix. * * ================================================================== * DESCRIPTOR BIT LAYOUT (64 bits) * ================================================================== * [0,14) start_address SMEM address >> 4 (in 16-byte units) * [14,16) (unused) - * [16,30) leading_byte_offset (LBO) + * [16,30) leading_byte_offset (LBO) in 16-byte units * [30,32) (unused) - * [32,46) stride_byte_offset (SBO) + * [32,46) stride_byte_offset (SBO) in 16-byte units * [46,48) version = 1 * [48,49) (unused) * [49,52) base_offset @@ -35,64 +46,44 @@ * [53,61) (unused) * [61,64) layout_type 0=NONE, 1=SW128, 2=SW64, 3=SW32 * - * ================================================================== - * DESCRIPTOR FOR K-MAJOR, NONE, BF16 - * ================================================================== - * For a (BLOCK_MN, BLOCK_K) BF16 matrix in K-major core-matrix layout: - * start_address = smem_ptr >> 4 - * LBO = 1 (16 bytes, in 16B units, lbo_mode=0) - * SBO = BLOCK_K / 8 (in 16B units) - * layout_type = 0 + * For SWIZZLE_NONE: + * LBO = BLOCK_MN * 16 (bytes) → field = BLOCK_MN (16B units) + * SBO = 128 (bytes) → field = 8 (16B units) + * lbo_mode = 0, layout_type = 0 * - * The hardware walks the SMEM as: - * For each 8-row group g (0..BLOCK_MN/8-1): - * For each k_group k (0..BLOCK_K/8-1): - * Read 8×8 BF16 tile at address: - * start + g * SBO * 16 + k * LBO * 16 - * (in bytes, relative to start_address << 4) + * For SWIZZLE_128B: + * SBO = 8 * 128 = 1024 (bytes) → field = 64 (16B units) + * lbo_mode = 1, layout_type = 1 + * (LBO not explicitly set — hardware assumes LBO=1 in 128B units) * * ================================================================== * MMA K-TILING * ================================================================== * tcgen05.mma.kind::f16 operates on K=16 BF16 per MMA call. - * For head_dim > 16, we tile along K with accumulate. - * Each K-tile needs its own descriptor pointing to the right - * 16-column slice of the matrix in SMEM. + * The descriptor describes a single MMA K-slice (16 BF16 = 32 bytes). + * For BLOCK_K > 16, we call MMA multiple times, each with a different + * descriptor pointing to the next K-slice in SMEM. * - * For a (128, 64) matrix with K-tile size 16: - * K-tile 0: columns 0-15 - * K-tile 1: columns 16-31 - * K-tile 2: columns 32-47 - * K-tile 3: columns 48-63 + * For NO swizzle, each 16B slice is BLOCK_MN * 16B bytes in SMEM. + * The k-th MMA uses descriptor with start_address pointing to + * the k-th (BLOCK_MN, 16B) slice. + * But wait — MMA needs 2 columns (32B). The descriptor covers 2 slices. + * The first slice is at start, the second at start + LBO. + * So for k-th MMA (k even), the start is at the k/2-th slice base. * - * Each K-tile descriptor has BLOCK_K=16, so: - * LBO = 1 - * SBO = 2 - * start_address = base + (k_tile * 16 / 8) * core_mn_stride + * Actually from gau-nernst: for NO swizzle, each MMA call gets a + * descriptor pointing to the start of the (BLOCK_MN, 32B) region. + * The 32B = 2 × 16B columns. LBO = BLOCK_MN * 16B gives the stride + * between the 2 columns. * - * The K-tile offset in the core-matrix layout: - * Each core matrix in K dimension covers 8 BF16 columns. - * K-tile k covers columns [16k, 16k+16) = 2 core-matrix columns. - * Offset = k * 2 * 64 = k * 128 (in BF16 elements) - * - * But this assumes the core-matrix tiles are arranged with K-tiles - * contiguous within each 8-row group. The full layout is: - * tile(g, c) at offset (g * (HD/8) + c) * 64 - * K-tile k uses columns c = 2k, 2k+1 - * So K-tile k's data is at offsets: - * (g * (HD/8) + 2k) * 64 and (g * (HD/8) + 2k + 1) * 64 - * These are NOT contiguous (separated by other core matrices in K). - * - * SOLUTION: Store the matrix as SEPARATE K-tiles in SMEM, each in - * core-matrix layout. This way each K-tile is contiguous and gets - * its own descriptor with start_address pointing to its base. - * - * SMEM layout for (128, HD) with K-tile size 16: - * ktile[0]: (128, 16) core-matrix layout, 4 KB - * ktile[1]: (128, 16) core-matrix layout, 4 KB - * ... - * ktile[HD/16-1]: (128, 16) core-matrix layout, 4 KB - * Total: (HD/16) * 4 KB + * The k-th MMA reads columns [16k, 16k+16) in SMEM. + * In the core-matrix layout, columns are BLOCK_MN * 16B bytes apart + * (each column is BLOCK_MN/8 core matrices × 128B = BLOCK_MN * 16B). + * Wait, each core matrix is 128B. BLOCK_MN/8 core matrices per column = BLOCK_MN/8 * 128 = BLOCK_MN * 16B. ✓ + * So the k-th 16B column starts at offset k * BLOCK_MN * 16B. + * MMA needs columns 2k and 2k+1, so the descriptor for the k-th MMA + * starts at column 2k: offset = 2k * BLOCK_MN * 16B = k * BLOCK_MN * 32B. + * This matches gau-nernst: A_smem + k * BLOCK_M * 32. */ #pragma once @@ -102,70 +93,125 @@ namespace dsv4::kernels::attention { // MMA K-tile size for BF16: 16 elements (32 bytes) -constexpr int MMA_K_TILE = 16; +constexpr int MMA_K_BF16 = 16; // ================================================================== -// SMEM layout: row-major → K-major core-matrix layout +// SMEM layout: row-major → canonical K-major core-matrix layout +// ================================================================== +// The canonical layout for NO swizzle interleaves core matrices: +// [core(0,0), core(1,0), ..., core(MN/8-1,0), core(0,1), ...] +// where core(g, c) = 8×8 BF16 tile at rows [8g, 8g+8), cols [8c, 8c+8) +// +// Each core matrix is 128 bytes (8 rows × 8 BF16). +// The offset of core(g, c) in BF16 elements: +// c * (BLOCK_MN/8) * 64 + g * 64 + local_r * 8 + local_c +// = c * BLOCK_MN * 8 + g * 64 + local_r * 8 + local_c +// (since BLOCK_MN/8 * 64 = BLOCK_MN * 8) +// +// And in bytes: (c * BLOCK_MN * 8 + g * 64 + ...) * 2 +// = c * BLOCK_MN * 16 + g * 128 + ... +// +// This matches: g * SBO + c * LBO + local_r * 16 + local_c * 2 +// where SBO = 128, LBO = BLOCK_MN * 16 ✓ // ================================================================== /** - * Write a (ROWS, COLS) BF16 matrix from row-major to K-major - * core-matrix layout in SMEM. + * Write a (ROWS, COLS) BF16 matrix from row-major to canonical + * K-major core-matrix layout for tcgen05.mma (NO swizzle). * - * ROWS and COLS must be multiples of 8. - * Each 8×8 BF16 tile (64 elements, 128 bytes) is contiguous. - * Tiles are laid out in MN-major order: tile(row/8, col/8). + * ROWS must be a multiple of 8 (ideally 128 for full MMA tile). + * COLS must be a multiple of 8. */ template -__device__ void write_smem_kmajor(bf16_t* __restrict__ dst, - const bf16_t* __restrict__ src) { - constexpr int TOTAL = ROWS * COLS; +__device__ void write_smem_canonical(bf16_t* __restrict__ dst, + const bf16_t* __restrict__ src) { + constexpr int CORES_MN = ROWS / 8; constexpr int CORES_K = COLS / 8; + constexpr int TOTAL = ROWS * COLS; for (int i = threadIdx.x; i < TOTAL; i += NTHREADS) { int r = i / COLS; int c = i % COLS; - int tile_mn = r / 8; - int tile_k = c / 8; + int core_mn = r / 8; + int core_k = c / 8; int local_r = r % 8; int local_c = c % 8; - int dst_idx = (tile_mn * CORES_K + tile_k) * 64 + local_r * 8 + local_c; + // Canonical layout: column-major in (core_mn, core_k) + // All core_mn values for core_k=0 first, then core_k=1, etc. + int dst_idx = core_k * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c; dst[dst_idx] = src[i]; } } /** - * Write a (ROWS, MMA_K_TILE) K-tile from row-major to core-matrix layout. - * Specialization for the single K-tile case (COLS = MMA_K_TILE = 16). + * Write Q (1, HD) to SMEM in canonical layout, padded to (128, HD). + * Only row 0 has actual data; rows 1-127 are zero. */ -template -__device__ void write_smem_ktile(bf16_t* __restrict__ dst, - const bf16_t* __restrict__ src, - int k_tile, int hd) { - // src is (ROWS, hd) row-major. We extract columns [k_tile*16, k_tile*16+16). - // dst is (ROWS, 16) in core-matrix layout. - constexpr int KT = MMA_K_TILE; // 16 - constexpr int CORES_K = KT / 8; // 2 - int col_start = k_tile * KT; - for (int i = threadIdx.x; i < ROWS * KT; i += NTHREADS) { - int r = i / KT; - int c = i % KT; - int tile_mn = r / 8; - int tile_k = c / 8; - int local_r = r % 8; +template +__device__ void write_q_to_smem(bf16_t* __restrict__ dst, + const bf16_t* __restrict__ q) { + constexpr int CORES_MN = 128 / 8; // 16 + constexpr int CORES_K = HD / 8; + // Zero all first + for (int i = threadIdx.x; i < 128 * HD; i += NTHREADS) dst[i] = 0; + // Write row 0 only: row 0 is in core_mn=0, local_r=0 + for (int c = threadIdx.x; c < HD; c += NTHREADS) { + int core_k = c / 8; int local_c = c % 8; - int dst_idx = (tile_mn * CORES_K + tile_k) * 64 + local_r * 8 + local_c; - dst[dst_idx] = src[r * hd + col_start + c]; + int dst_idx = core_k * CORES_MN * 64 + 0 * 64 + 0 * 8 + local_c; + dst[dst_idx] = q[c]; } } /** - * Zero a (ROWS, COLS) region in core-matrix layout. + * Write K (SK, HD) to SMEM in canonical layout, padded to (128, HD). */ -template -__device__ void zero_smem(bf16_t* __restrict__ dst) { - constexpr int TOTAL = ROWS * COLS; - for (int i = threadIdx.x; i < TOTAL; i += NTHREADS) { - dst[i] = 0; +template +__device__ void write_k_to_smem(bf16_t* __restrict__ dst, + const bf16_t* __restrict__ k) { + constexpr int CORES_MN = 128 / 8; // 16 + constexpr int CORES_K = HD / 8; + // Zero all first + for (int i = threadIdx.x; i < 128 * HD; i += NTHREADS) dst[i] = 0; + // Write actual rows + for (int i = threadIdx.x; i < SK * HD; i += NTHREADS) { + int r = i / HD; + int c = i % HD; + int core_mn = r / 8; + int core_k = c / 8; + int local_r = r % 8; + int local_c = c % 8; + int dst_idx = core_k * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c; + dst[dst_idx] = k[i]; + } +} + +/** + * Write V (HD, SK) to SMEM in canonical layout, padded to (HD, 128). + * V is stored as (HD, SK) in GMEM. For PV GEMM, V is the B matrix + * with shape (BLOCK_N=128, BLOCK_K=HD) K-major in SMEM. + * But V in GMEM is (HD, SK) — we need to transpose it. + * V's column c (position c in the sequence) maps to row c in B. + * B's row c has columns d=0..HD-1, which is V[d, c] = V[d * SK + c]. + * So B[r, d] = V[d * SK + r] for r < SK, d < HD. + */ +template +__device__ void write_v_to_smem(bf16_t* __restrict__ dst, + const bf16_t* __restrict__ v) { + // dst is (128, HD) in canonical layout + constexpr int CORES_MN = 128 / 8; + constexpr int CORES_K = HD / 8; + for (int i = threadIdx.x; i < 128 * HD; i += NTHREADS) dst[i] = 0; + // B = (128, HD) where B[r, d] = v[d * SK + r] for r < SK + for (int d = threadIdx.x; d < HD; d += NTHREADS) { + for (int r = 0; r < SK; r += NTHREADS) { + if (r >= SK) continue; + int core_mn = r / 8; + int core_k = d / 8; + int local_r = r % 8; + int local_c = d % 8; + int dst_idx = core_k * CORES_MN * 64 + core_mn * 64 + local_r * 8 + local_c; + dst[dst_idx] = v[d * SK + r]; + } } } @@ -174,78 +220,61 @@ __device__ void zero_smem(bf16_t* __restrict__ dst) { // ================================================================== /** - * Construct a 64-bit SMEM descriptor for tcgen05.mma with - * K-major, SWIZZLE_NONE, BF16, for a single K-tile of width 16. - * - * The matrix is (BLOCK_MN, 16) in core-matrix layout. - * LBO = 1 (16B stride between K-groups within 8-row block) - * SBO = 2 (16/8 = 2, stride between 8-row blocks, in 16B units) + * desc_encode: convert byte value to 16B-unit descriptor field. + * Same as value >> 4 (divide by 16). */ -__device__ __forceinline__ uint64_t make_umma_desc_kmajor_none_ktile( - uint32_t smem_ptr -) { - uint64_t desc = 0; - desc |= (static_cast(smem_ptr >> 4) & 0x3FFF); // start_address - desc |= (static_cast(1) & 0x3FFF) << 16; // LBO = 1 (16B) - desc |= (static_cast(2) & 0x3FFF) << 32; // SBO = 2 (16B units) - desc |= (static_cast(1) << 46); // version = 1 - // lbo_mode = 0 (16B units), layout_type = 0 (NONE) - return desc; +__device__ __forceinline__ uint64_t desc_encode(uint64_t byte_val) { + return byte_val >> 4; } /** - * Construct descriptor for K-major, SWIZZLE_NONE, BF16 with - * arbitrary BLOCK_K (must be multiple of 16 for MMA K-tiling). + * Construct a 64-bit SMEM descriptor for tcgen05.mma with + * K-major, SWIZZLE_NONE, BF16. * - * For a K-tile of width BLOCK_K: - * LBO = 1 (16B) - * SBO = BLOCK_K / 8 (in 16B units) + * From gau-nernst's tutorial: + * LBO = BLOCK_MN * 16 (bytes) → desc field = BLOCK_MN + * SBO = 128 (bytes) → desc field = 8 + * + * @param smem_addr SMEM address (from __cvta_generic_to_shared) + * @param block_mn The MN dimension of the MMA tile (e.g., 128 for M=128) */ __device__ __forceinline__ uint64_t make_umma_desc_kmajor_none( - uint32_t smem_ptr, int block_k + uint32_t smem_addr, int block_mn ) { + const uint64_t LBO = block_mn * 16; // bytes + const uint64_t SBO = 128; // bytes (8 * 16B) uint64_t desc = 0; - desc |= (static_cast(smem_ptr >> 4) & 0x3FFF); // start_address - desc |= (static_cast(1) & 0x3FFF) << 16; // LBO = 1 (16B) - desc |= (static_cast(block_k / 8) & 0x3FFF) << 32; // SBO - desc |= (static_cast(1) << 46); // version = 1 + desc |= desc_encode(smem_addr) & 0x3FFF; // start_address + desc |= (desc_encode(LBO) & 0x3FFF) << 16; // LBO + desc |= (desc_encode(SBO) & 0x3FFF) << 32; // SBO + desc |= 1ULL << 46; // version = 1 return desc; } /** * Construct descriptor for K-major, SWIZZLE_128B, BF16. - * SMEM must be 128-byte aligned. - * LBO = 1 (128B units, lbo_mode=1) - * SBO = BLOCK_K / 8 (in 128B units) - * layout_type = 1 (SW128) + * + * From gau-nernst's tutorial: + * SBO = 8 * 128 = 1024 (bytes) → desc field = 64 + * layout_type = 2 (SW128 in PTX encoding — wait, gau-nernst uses 2ULL << 61) + * lbo_mode = 1 (128B units) + * + * @param smem_addr SMEM address (from __cvta_generic_to_shared), must be 128B-aligned + * @param block_mn The MN dimension of the MMA tile */ __device__ __forceinline__ uint64_t make_umma_desc_kmajor_sw128( - uint32_t smem_ptr, int block_k + uint32_t smem_addr, int block_mn ) { + const uint64_t SBO = 8 * 128; // 1024 bytes uint64_t desc = 0; - desc |= (static_cast(smem_ptr >> 4) & 0x3FFF); // start_address - desc |= (static_cast(1) & 0x3FFF) << 16; // LBO = 1 (128B) - desc |= (static_cast(block_k / 8) & 0x3FFF) << 32; // SBO - desc |= (static_cast(1) << 46); // version = 1 - desc |= (static_cast(1) << 52); // lbo_mode = 1 - desc |= (static_cast(1) << 61); // layout_type = 1 (SW128) + desc |= desc_encode(smem_addr) & 0x3FFF; + desc |= (desc_encode(SBO) & 0x3FFF) << 32; + desc |= 1ULL << 46; // version + desc |= 1ULL << 52; // lbo_mode = 1 (128B units) + desc |= 2ULL << 61; // layout_type = 2 (SW128) return desc; } -// ================================================================== -// SW128 swizzle: write SMEM data in swizzled order -// ================================================================== -// The SW128 swizzle pattern (Swizzle<3,4,3>) XORs bits of the -// SMEM address to permute data across memory banks. -// For core-matrix K-major layout with 128B swizzle: -// Each 128B block is swizzled. The swizzle pattern is: -// swizzled_offset = offset ^ (((offset >> 4) ^ (offset >> 8)) & 7) << 4 -// where offset is in bytes, relative to the start of the tile. -// -// For the initial bringup, we use SWIZZLE_NONE to avoid the -// swizzle complexity. SW128 will be added after the MMA is verified. -// ================================================================== - // ================================================================== // tcgen05.mma PTX wrappers // ================================================================== @@ -253,26 +282,13 @@ __device__ __forceinline__ uint64_t make_umma_desc_kmajor_sw128( /** * tcgen05.mma SS: both operands from SMEM, result to TMEM. * A (SMEM) × B (SMEM) → D (TMEM). - * - * IMPORTANT: Only ONE thread per CTA calls this. The MMA is - * single-threaded. All other threads should be idle or doing - * independent work. - * - * The MMA is asynchronous — after calling, you MUST wait for - * completion before reading TMEM (use tmem_fence_store + syncthreads). - * - * @param tmem_c TMEM column address for the accumulator - * @param desc_a 64-bit SMEM descriptor for A matrix (K-major) - * @param desc_b 64-bit SMEM descriptor for B matrix (K-major) - * @param accumulate true = D += A*B, false = D = A*B (zero accumulator) + * Only ONE thread per CTA calls this. */ __device__ void umma_ss_f16( uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, - bool accumulate = false + uint32_t i_desc, bool accumulate = false ) { - uint32_t idescE = static_cast(desc_a >> 32); uint32_t scaleC_bits = accumulate ? 0x3F800000u : 0u; - uint32_t mask0 = 0, mask1 = 0, mask2 = 0, mask3 = 0; asm volatile( "{\n\t" ".reg .pred p;\n\t" @@ -280,23 +296,20 @@ __device__ void umma_ss_f16( "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, {%5, %6, %7, %8}, p;\n\t" "}" :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), - "r"(idescE), "r"(scaleC_bits), - "r"(mask0), "r"(mask1), "r"(mask2), "r"(mask3) + "r"(i_desc), "r"(scaleC_bits), + "r"(0), "r"(0), "r"(0), "r"(0) ); } /** * tcgen05.mma TS: A from TMEM, B from SMEM, result to TMEM. - * A (TMEM) × B (SMEM) → D (TMEM). - * Used for PV GEMM where P is in TMEM and V is in SMEM. + * Used for PV GEMM. */ __device__ void umma_ts_f16( uint32_t tmem_c, uint32_t tmem_a, uint64_t desc_b, - bool accumulate = true + uint32_t i_desc, bool accumulate = true ) { - uint32_t idescE = static_cast(desc_b >> 32); uint32_t scaleC_bits = accumulate ? 0x3F800000u : 0u; - uint32_t mask0 = 0, mask1 = 0, mask2 = 0, mask3 = 0; asm volatile( "{\n\t" ".reg .pred p;\n\t" @@ -304,9 +317,26 @@ __device__ void umma_ts_f16( "tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p;\n\t" "}" :: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), - "r"(idescE), "r"(scaleC_bits), - "r"(mask0), "r"(mask1), "r"(mask2), "r"(mask3) + "r"(i_desc), "r"(scaleC_bits), + "r"(0), "r"(0), "r"(0), "r"(0) ); } +/** + * Construct the instruction descriptor (idesc) for tcgen05.mma.kind::f16. + * From gau-nernst: + * bit 4: dtype=FP32 (1) + * bit 7: atype=BF16 (1) + * bit 10: btype=BF16 (1) + * bits 17-23: MMA_N (BLOCK_N >> 3) + * bits 24-28: MMA_M (BLOCK_M >> 4) + */ +__device__ __forceinline__ uint32_t make_idesc(int block_m, int block_n) { + return (1U << 4) // dtype=FP32 + | (1U << 7) // atype=BF16 + | (1U << 10) // btype=BF16 + | ((uint32_t)(block_n >> 3) << 17) // MMA_N + | ((uint32_t)(block_m >> 4) << 24); // MMA_M +} + } // namespace dsv4::kernels::attention diff --git a/tests/unit/test_umma_qk.cu b/tests/unit/test_umma_qk.cu index 31877d36..cf674729 100644 --- a/tests/unit/test_umma_qk.cu +++ b/tests/unit/test_umma_qk.cu @@ -2,9 +2,10 @@ * Standalone CUDA test for UMMA QK GEMM (tcgen05.mma SS, BF16). * * Tests that tcgen05.mma produces correct QK attention scores. - * Uses K-major, SWIZZLE_NONE, BF16 descriptors with core-matrix SMEM layout. + * Uses K-major, SWIZZLE_NONE, BF16 descriptors with canonical SMEM layout. + * Based on gau-nernst's tcgen05 tutorial. * - * First test: HD=16, SK=128 (single K-tile, single MMA call) + * Test: HD=16, SK=128 (single K-tile, single MMA call) */ #include @@ -13,27 +14,29 @@ #include #include -// Include our FMHA headers #include "dsv4/kernels/attention/fmha_common.cuh" #include "dsv4/kernels/attention/fmha_umma_desc.cuh" using namespace dsv4::kernels::attention; -// Host-side BF16 conversion (no PTX on CPU) +// Host-side BF16 conversion static bf16_t f32_to_bf16_host(float f) { - uint32_t u; - memcpy(&u, &f, 4); + uint32_t u; memcpy(&u, &f, 4); return (uint16_t)(u >> 16); } // ================================================================== -// Test kernel: UMMA QK GEMM for HD=16, SK=128 (single K-tile) +// Test kernel: UMMA QK GEMM for HD=16, SK=128 // ================================================================== +// A = Q: (128, 16) padded, K-major core-matrix layout +// B = K: (128, 16) in K-major core-matrix layout +// MMA computes S = A × B^T = Q × K^T → (128, 128) in TMEM + __global__ void __launch_bounds__(NTHREADS) test_umma_qk_hd16( const bf16_t* __restrict__ q, const bf16_t* __restrict__ k, float* __restrict__ s_out, // output: S[0, 0..127] from TMEM - float* __restrict__ s_scalar, // scalar reference: Q @ K^T in SMEM + float* __restrict__ s_scalar, // scalar reference float scale ) { const int tid = threadIdx.x; @@ -42,23 +45,24 @@ test_umma_qk_hd16( // ================================================================ // SMEM layout // ================================================================ - // sQ_ktile: (128, 16) BF16 in K-major core-matrix layout = 4 KB - // sK_ktile: (128, 16) BF16 in K-major core-matrix layout = 4 KB + // sQ: (128, 16) in canonical layout = 4096 bytes + // sK: (128, 16) in canonical layout = 4096 bytes // sTmemBase: 4 bytes - // sQ_row: (16) floats for scalar reference + // sQ_row: 16 floats for scalar ref extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - bf16_t* sQ_ktile = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); - bf16_t* sK_ktile = sQ_ktile + 128 * 16; - float* sQ_row = (float*)(sK_ktile + 128 * 16); + // Align to 16 bytes for UMMA + bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~(uintptr_t)15); + bf16_t* sK = sQ + 128 * 16; // 4096 bytes after Q + float* sQ_row = (float*)(sK + 128 * 16); - // Load Q to sQ_row (float for scalar reference) + // Load Q to sQ_row for scalar reference for (int d = tid; d < 16; d += NTHREADS) { sQ_row[d] = bf16_to_f32(q[d]); } // ================================================================ - // TMEM allocation: 128 columns for S (128, 128) + // TMEM allocation // ================================================================ if (wid == 0) { uint32_t smem_ptr = __cvta_generic_to_shared(sTmemBase); @@ -67,7 +71,7 @@ test_umma_qk_hd16( __syncthreads(); uint32_t tmem_base = *sTmemBase; - // Zero TMEM S accumulator + // Zero TMEM if (wid == 0) { for (int col = 0; col < 128; col++) { tmem_store(tmem_base + col, 0, 0, 0, 0); @@ -77,82 +81,55 @@ test_umma_qk_hd16( __syncthreads(); // ================================================================ - // Load Q and K into SMEM in core-matrix layout + // Load Q and K into SMEM in canonical layout // ================================================================ - // Zero both buffers first - for (int i = tid; i < 128 * 16; i += NTHREADS) { - sQ_ktile[i] = 0; - sK_ktile[i] = 0; - } - __syncthreads(); - - // Write Q (1, 16) padded to (128, 16) in core-matrix layout - // Core-matrix: each 8x8 BF16 tile at (tile_mn, tile_k) offset (tile_mn * 2 + tile_k) * 64 - // Q row 0, cols 0-15: tiles (0,0) and (0,1) - // Tile (0,0): positions 0-7 = Q[0..7], Tile (0,1): positions 0-7 = Q[8..15] - for (int d = tid; d < 16; d += NTHREADS) { - int tile_k = d / 8; - int local_c = d % 8; - int dst_idx = (0 * 2 + tile_k) * 64 + 0 * 8 + local_c; // row 0 (local_r=0) - sQ_ktile[dst_idx] = q[d]; - } - - // Write K (128, 16) in core-matrix layout - // K is (128, 16) row-major in GMEM. In core-matrix: - // tile (r/8, c/8) at offset ((r/8) * 2 + (c/8)) * 64 + (r%8)*8 + (c%8) - for (int i = tid; i < 128 * 16; i += NTHREADS) { - int r = i / 16; - int c = i % 16; - int tile_mn = r / 8; - int tile_k = c / 8; - int local_r = r % 8; - int local_c = c % 8; - int dst_idx = (tile_mn * 2 + tile_k) * 64 + local_r * 8 + local_c; - sK_ktile[dst_idx] = k[i]; - } + write_q_to_smem<16>(sQ, q); + write_k_to_smem<128, 16>(sK, k); __syncthreads(); // ================================================================ - // Construct UMMA descriptors and call MMA + // Construct descriptors and instruction descriptor // ================================================================ - // Both A (Q) and B (K) are (128, 16) in K-major core-matrix layout. - // For K-major NONE with BLOCK_K=16: LBO=1, SBO=2 - uint32_t sQ_smem = __cvta_generic_to_shared(sQ_ktile); - uint32_t sK_smem = __cvta_generic_to_shared(sK_ktile); + uint32_t sQ_smem = __cvta_generic_to_shared(sQ); + uint32_t sK_smem = __cvta_generic_to_shared(sK); - uint64_t desc_q = make_umma_desc_kmajor_none_ktile(sQ_smem); - uint64_t desc_k = make_umma_desc_kmajor_none_ktile(sK_smem); + // K-major NONE: LBO = BLOCK_MN * 16, SBO = 128 + uint64_t desc_q = make_umma_desc_kmajor_none(sQ_smem, 128); + uint64_t desc_k = make_umma_desc_kmajor_none(sK_smem, 128); - // Single-threaded MMA launch + // Instruction descriptor: dtype=FP32, atype=BF16, btype=BF16, M=128, N=128 + uint32_t idesc = make_idesc(128, 128); + + // ================================================================ + // Call tcgen05.mma SS + // ================================================================ if (tid == 0) { - umma_ss_f16(tmem_base, desc_q, desc_k, /*accumulate=*/false); + umma_ss_f16(tmem_base, desc_q, desc_k, idesc, /*accumulate=*/false); } __syncwarp(); - // Wait for MMA to complete + // Wait for MMA if (wid == 0 && lane == 0) { tmem_fence_store(); } __syncthreads(); // ================================================================ - // Read S from TMEM and write to output + // Read S from TMEM // ================================================================ - // S is (128, 128) FP32 in TMEM. We care about row 0 (the query row). - // TMEM column col stores 128 FP32. Lane 0's u0 = S[0, col]. if (wid == 0) { for (int col = 0; col < 128; col++) { uint32_t u0, u1, u2, u3; tmem_load(tmem_base + col, u0, u1, u2, u3); if (lane == 0) { - s_out[col] = u32_to_f32(u0); // S[0, col] (un-normalized) + s_out[col] = u32_to_f32(u0); // S[0, col] } } } __syncthreads(); // ================================================================ - // Scalar reference: compute Q @ K^T in SMEM (row 0 only) + // Scalar reference // ================================================================ if (tid == 0) { for (int c = 0; c < 128; c++) { @@ -172,54 +149,43 @@ test_umma_qk_hd16( } // ================================================================== -// Host test harness +// Host // ================================================================== int main() { printf("=== UMMA QK GEMM Test (HD=16, SK=128) ===\n"); - const int HD = 16; - const int SK = 128; + const int HD = 16, SK = 128; const float SCALE = 1.0f / sqrtf((float)HD); - // Allocate host memory bf16_t* h_q = (bf16_t*)malloc(HD * sizeof(bf16_t)); bf16_t* h_k = (bf16_t*)malloc(SK * HD * sizeof(bf16_t)); float* h_s_out = (float*)malloc(SK * sizeof(float)); float* h_s_scalar = (float*)malloc(SK * sizeof(float)); - // Initialize Q and K with random BF16 values srand(42); - for (int d = 0; d < HD; d++) { - float val = (float)(rand() % 100) / 100.0f - 0.5f; - h_q[d] = f32_to_bf16_host(val); - } - for (int i = 0; i < SK * HD; i++) { - float val = (float)(rand() % 100) / 100.0f - 0.5f; - h_k[i] = f32_to_bf16_host(val); - } + for (int d = 0; d < HD; d++) + h_q[d] = f32_to_bf16_host((float)(rand() % 100) / 100.0f - 0.5f); + for (int i = 0; i < SK * HD; i++) + h_k[i] = f32_to_bf16_host((float)(rand() % 100) / 100.0f - 0.5f); - // Allocate device memory - bf16_t* d_q; cudaMalloc(&d_q, HD * sizeof(bf16_t)); - bf16_t* d_k; cudaMalloc(&d_k, SK * HD * sizeof(bf16_t)); - float* d_s_out; cudaMalloc(&d_s_out, SK * sizeof(float)); - float* d_s_scalar; cudaMalloc(&d_s_scalar, SK * sizeof(float)); + bf16_t *d_q, *d_k; + float *d_s_out, *d_s_scalar; + cudaMalloc(&d_q, HD * sizeof(bf16_t)); + cudaMalloc(&d_k, SK * HD * sizeof(bf16_t)); + cudaMalloc(&d_s_out, SK * sizeof(float)); + cudaMalloc(&d_s_scalar, SK * sizeof(float)); cudaMemcpy(d_q, h_q, HD * sizeof(bf16_t), cudaMemcpyHostToDevice); cudaMemcpy(d_k, h_k, SK * HD * sizeof(bf16_t), cudaMemcpyHostToDevice); cudaMemset(d_s_out, 0, SK * sizeof(float)); cudaMemset(d_s_scalar, 0, SK * sizeof(float)); - // Compute SMEM size - // sTmemBase: 4 + alignment 16 + sQ: 128*16*2 + sK: 128*16*2 + sQ_row: 16*4 + padding int smem_size = 4 + 16 + 128*16*2 + 128*16*2 + 16*4 + 256; smem_size = (smem_size + 127) & ~127; printf("SMEM size: %d bytes\n", smem_size); - // Launch - dim3 grid(1, 1, 1); - dim3 block(NTHREADS); - test_umma_qk_hd16<<>>( + test_umma_qk_hd16<<<1, NTHREADS, smem_size>>>( d_q, d_k, d_s_out, d_s_scalar, SCALE); cudaError_t err = cudaDeviceSynchronize(); @@ -228,31 +194,25 @@ int main() { return 1; } - // Copy results back cudaMemcpy(h_s_out, d_s_out, SK * sizeof(float), cudaMemcpyDeviceToHost); cudaMemcpy(h_s_scalar, d_s_scalar, SK * sizeof(float), cudaMemcpyDeviceToHost); - // Compare - printf("\nS[0,0..7] (TMEM/MMA): "); + printf("\nS[0,0..7] (MMA): "); for (int c = 0; c < 8; c++) printf("%.4f ", h_s_out[c]); - printf("\nS[0,0..7] (scalar ref): "); + printf("\nS[0,0..7] (scalar): "); for (int c = 0; c < 8; c++) printf("%.4f ", h_s_scalar[c]); printf("\n"); float max_diff = 0.0f, max_val = 0.0f; for (int c = 0; c < SK; c++) { - float diff = fabsf(h_s_out[c] - h_s_scalar[c]); - max_diff = fmaxf(max_diff, diff); + max_diff = fmaxf(max_diff, fabsf(h_s_out[c] - h_s_scalar[c])); max_val = fmaxf(max_val, fabsf(h_s_scalar[c])); } float rel_err = (max_val > 0) ? max_diff / max_val : max_diff; - printf("Max absolute diff: %.6f\n", max_diff); - printf("Max relative error: %.6f\n", rel_err); + printf("Max abs diff: %.6f, Max rel err: %.6f\n", max_diff, rel_err); printf("Test %s\n", rel_err < 0.01f ? "PASSED" : "FAILED"); - // Cleanup cudaFree(d_q); cudaFree(d_k); cudaFree(d_s_out); cudaFree(d_s_scalar); free(h_q); free(h_k); free(h_s_out); free(h_s_scalar); - return (rel_err < 0.01f) ? 0 : 1; }