diff --git a/dsv4/kernels/attention/fmha_qk_verify.cuh b/dsv4/kernels/attention/fmha_qk_verify.cuh index ba0aa21b..d754969f 100644 --- a/dsv4/kernels/attention/fmha_qk_verify.cuh +++ b/dsv4/kernels/attention/fmha_qk_verify.cuh @@ -1,11 +1,12 @@ /** - * DSV4 FMHA Phase 3 — Tensor-core QK GEMM verification. + * DSV4 FMHA — QK GEMM verification with canonical UMMA SMEM layout. * - * STEP 1 ONLY: Verify tcgen05.mma SS produces correct QK output. - * Load Q and K to SMEM, run tcgen05.mma, read S from TMEM, compare - * against CPU reference Q @ K^T. + * STEP 1: Verify tcgen05.mma SS produces correct QK output + * using the proper SWIZZLE_128B canonical SMEM layout. * - * This is a diagnostic kernel, not the production FMHA. + * The SMEM data must be written in the swizzled layout that the UMMA + * hardware expects. We use umma_smem_write_mn_sw128 and + * umma_smem_write_k_sw128 to write data in the correct format. */ #pragma once @@ -29,25 +30,35 @@ fmha_qk_verify( const bf16_t* qh = q + batch*bstride_q + head*HD; 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 + // SMEM: sQ (128×HD BF16 swizzled) + sK (128×HD BF16 swizzled) + tmem_base + // Size: 4 + 128*HD*2 + 128*HD*2 = 4 + 512*HD bytes + // For SW128 layout, the actual SMEM needed is the same as row-major + // because the swizzle is just a permutation of the same data. extern __shared__ char sbuf[]; uint32_t* sTmemBase = (uint32_t*)sbuf; - // Align sQ to 16 bytes - bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 15) & ~15); - bf16_t* sK = sQ + 128 * HD; + // Align to 128 bytes for UMMA descriptor + bf16_t* sQ = (bf16_t*)(((uintptr_t)(sbuf + 4) + 127) & ~127); + bf16_t* sK = sQ + 1024 * 8; // MN_SW128 atom = 1024*8 BF16 = 8192 BF16 per atom - // Load Q: (1, HD) padded to (128, HD) with zeros - for (int i = tid; i < 128 * HD; i += NTHREADS) sQ[i] = 0; - for (int d = tid; d < HD; d += NTHREADS) sQ[d] = qh[d]; // row 0 + // Load Q into swizzled SMEM: (1, HD) padded to (128, HD) + // Write row 0 with actual Q data, rows 1-127 are zero + for (int d = tid; d < HD; d += NTHREADS) { + umma_smem_write_mn_sw128(sQ, 0, d, HD, qh[d]); + } + // Zero rows 1-127 (only need to zero the elements we'll read) + for (int row = 1 + tid / HD; row < 128; row += NTHREADS / HD) { + for (int d = tid % HD; d < HD; d += HD) { + umma_smem_write_mn_sw128(sQ, row, d, HD, 0); + } + } - // Load K: (min(128, s_k), HD) padded to (128, HD) with zeros + // Load K into swizzled SMEM: (min(128, s_k), HD) padded to (128, HD) int kv_len = min(128, s_k); - for (int i = tid; i < 128 * HD; i += NTHREADS) { - int r = i / HD, c = i % HD; - if (r < kv_len) sK[i] = kb[r * HD + c]; - else sK[i] = 0; + for (int r = 0; r < 128; r++) { + for (int d = tid; d < HD; d += NTHREADS) { + bf16_t val = (r < kv_len) ? kb[r * HD + d] : 0; + umma_smem_write_k_sw128(sK, r, d, HD, val); + } } __syncthreads(); @@ -72,20 +83,11 @@ fmha_qk_verify( // ================================================================ // QK GEMM: S = Q @ K^T via tcgen05.mma SS // ================================================================ - // Q is (128, HD) in SMEM, MN-major - // K is (128, HD) in SMEM, transposed via K-major UMMA descriptor - // S is (128, 128) in TMEM uint32_t sQ_smem = __cvta_generic_to_shared(sQ); uint32_t sK_smem = __cvta_generic_to_shared(sK); - if (tid == 0) { - printf("[qk] sQ_smem=0x%x sK_smem=0x%x sQ_align=%d sK_align=%d\n", - sQ_smem, sK_smem, sQ_smem % 16, sK_smem % 16); - } - __syncthreads(); - - uint64_t desc_q = make_umma_desc_bf16(sQ_smem, 128, HD, HD * 2, UmmaMajor::MN); - uint64_t desc_k = make_umma_desc_bf16(sK_smem, 128, HD, HD * 2, UmmaMajor::K); + uint64_t desc_q = make_umma_desc(sQ_smem, 128, HD, UmmaMajor::MN); + uint64_t desc_k = make_umma_desc(sK_smem, 128, HD, UmmaMajor::K); if (tid == 0) { printf("[qk] desc_q=0x%016llx desc_k=0x%016llx\n", @@ -93,7 +95,7 @@ fmha_qk_verify( } __syncthreads(); - // MMA is called by ONE lane (elect_one_sync pattern) + // MMA: called by ONE lane (elect_one_sync pattern) if (wid == 0 && lane == 0) { umma_ss_f16(tmem_s, desc_q, desc_k, /*accumulate=*/false); } @@ -103,31 +105,18 @@ fmha_qk_verify( } __syncthreads(); - // ================================================================ // Read S from TMEM and output row 0 - // ================================================================ - // Row 0, lane 0's register 0 in each column if (wid == 0) { for (int col = lane; col < 128; col += WARP) { uint32_t u0, u1, u2, u3; tmem_load(tmem_s + col, u0, u1, u2, u3); - // Lane i's u0 = row i*4+0, u1 = row i*4+1, etc. - // Lane 0's u0 = row 0 (which is our Q row) - if (lane < 128) { - float val = u32_to_f32(u0); // row lane*4+0 - // Write to global output (row 0 values only) - if (lane * 4 < kv_len) { - s_out[lane * 4 + 0] = val * scale; - } - if (lane * 4 + 1 < kv_len) { - s_out[lane * 4 + 1] = u32_to_f32(u1) * scale; - } - if (lane * 4 + 2 < kv_len) { - s_out[lane * 4 + 2] = u32_to_f32(u2) * scale; - } - if (lane * 4 + 3 < kv_len) { - s_out[lane * 4 + 3] = u32_to_f32(u3) * scale; - } + // Lane i's u0 = row i*4+0, etc. + if (lane < 32) { + // Write row 0 (lane 0's u0) + if (lane * 4 + 0 < kv_len) s_out[lane * 4 + 0] = u32_to_f32(u0) * scale; + if (lane * 4 + 1 < kv_len) s_out[lane * 4 + 1] = u32_to_f32(u1) * scale; + if (lane * 4 + 2 < kv_len) s_out[lane * 4 + 2] = u32_to_f32(u2) * scale; + if (lane * 4 + 3 < kv_len) s_out[lane * 4 + 3] = u32_to_f32(u3) * scale; } } } diff --git a/dsv4/kernels/attention/fmha_umma_desc.cuh b/dsv4/kernels/attention/fmha_umma_desc.cuh index 176d8d82..774dbde5 100644 --- a/dsv4/kernels/attention/fmha_umma_desc.cuh +++ b/dsv4/kernels/attention/fmha_umma_desc.cuh @@ -1,47 +1,57 @@ /** - * DSV4 FMHA — UMMA SMEM descriptor construction for tcgen05.mma. + * DSV4 FMHA — Canonical UMMA SMEM layout for tcgen05.mma. * * ================================================================== - * UMMA SMEM DESCRIPTOR BITFIELD (from cute/arch/mma_sm100_desc.hpp) + * CANONICAL SMEM LAYOUT FOR UMMA ON BLACKWELL SM100 * ================================================================== * - * The descriptor is a 64-bit value with this layout: + * The tcgen05.mma instruction requires matrix operands in SMEM to be + * laid out in a specific canonical format with swizzle. The UMMA + * descriptor encodes the strides in this canonical layout. * - * Bits [0,14) start_address — SMEM byte address >> 4 (16B aligned) - * Bits [14,16) unused - * Bits [16,30) leading_byte_offset — byte offset between rows, >> 4 (16B aligned) - * Bits [30,32) unused - * Bits [32,46) stride_byte_offset — byte offset for stride dimension, >> 4 - * Bits [46,48) version — 1 for Blackwell - * Bit [48) unused - * Bits [49,52) base_offset - * Bit [52,53) lbo_mode — leading byte offset mode (0=legacy) - * Bits [53,56) unused - * Bits [56,61) unused - * Bits [61,64) layout_type — 0=NONE, 1=128B_BASE32B, 2=128B, 4=64B, 6=32B + * For BF16 operands with SWIZZLE_128B (the default for FMHA): * - * The upper 32 bits encode dimensions and format: - * Bits [64,72) dim_m or dim_n — depends on major - * Bits [72,80) dim_k or dim_m — depends on major - * Bits [80,88) dim_n or dim_k — depends on major - * Bits [88,96) lot_size - * Bits [96,104) lot_offset - * Bits [104,112) unused - * Bits [112,120) unused - * Bits [120,128) format — F16F32Format (0=F16, 1=BF16, 2=TF32) + * MN-major A (Q): Swizzle<3,4,3> applied to (128, n_uint128) layout + * K-major B (K): Swizzle<3,4,3> applied to (k_uint128, n_uint128) layout + * + * The swizzle pattern XORs address bits to avoid bank conflicts. + * The hardware expects data in the SWIZZLED order, not row-major. * * ================================================================== - * KEY INSIGHT: layout_type=0 (SWIZZLE_NONE) IS VALID + * SWIZZLE_128B PATTERN FOR BF16 * ================================================================== - * The CUTLASS source lists SWIZZLE_NONE as a valid layout_type (value 0). - * For our simple row-major SMEM layout, SWIZZLE_NONE should work. - * The swizzle types (128B, 64B, 32B) are for bank-conflict-free access - * patterns, but the UMMA hardware can operate on unswizzled data too. + * + * Swizzle<3,4,3> XORs bits [6:4] of the uint128_t offset with bits [9:7]. + * This permutes 128-bit "rows" within a 1024-bit (8 × 128-bit) block. + * + * For a contiguous row of n_uint128 128-bit values: + * row_offset = row * row_stride_uint128 + * swizzled_offset = row_offset ^ ((row_offset & 0x70) >> 3) + * + * But actually, the swizzle in CUTLASS is applied differently: + * base = offset & ~0x3F (clear bottom 6 bits of 128B = 8 uint128_t) + * swizzle_bits = (offset >> 4) ^ (offset >> 7) (3,4,3 pattern) + * swizzled = base | (swizzle_bits & 0x7) + * + * This is getting complex. Let me just implement the swizzle as a + * device function and build the SMEM write/read functions. * * ================================================================== - * tcgen05.mma is called by ONE lane per warp (elect_one_sync). - * Unlike TMEM ld/st, MMA is NOT warp-collective. + * UMMA DESCRIPTOR FOR SWIZZLE_128B BF16 * ================================================================== + * + * For MN-major A with SWIZZLE_128B: + * layout_type = 2 (SWIZZLE_128B) + * start_address = smem_ptr >> 4 + * leading_byte_offset = (row_stride_bytes) >> 4 (in uint128_t units) + * stride_byte_offset = same as leading for simple layout + * version = 1 + * + * The descriptor's stride fields are in uint128_t units (16B granularity), + * describing the canonical layout strides AFTER swizzle is applied. + * + * For K-major B, the layout is the same but the K-dimension is the + * "leading" dimension and the MN-dimension is the "stride" dimension. */ #pragma once @@ -49,127 +59,272 @@ namespace dsv4::kernels::attention { -enum class UmmaMajor { MN, K }; -enum class UmmaLayout { NONE = 0, B128_BASE32B = 1, B128 = 2, B64 = 4, B32 = 6 }; +// ================================================================== +// Swizzle<3,4,3> for 128B-swizzled SMEM layout +// ================================================================== +// The swizzle XORs bits to permute 128-bit rows within 1024-bit blocks. +// This avoids bank conflicts when the MMA engine reads SMEM. +// +// Swizzle<3,4,3> means: XOR 3 bits starting at bit 4 with 3 bits +// starting at bit 7. The result is a permutation of 8 consecutive +// 128-bit rows within each 1024-bit (128-byte = 8 × 16-byte) block. +// +// Input: 128-bit row index within a 128B block (0-7) +// Output: permuted 128-bit row index within the block +// ================================================================== -/** - * Construct a UMMA SMEM descriptor for a BF16 matrix in SMEM. - * - * @param smem_ptr SMEM pointer (from __cvta_generic_to_shared), 32-bit - * @param dim_m M dimension - * @param dim_n N dimension - * @param row_stride_bytes Byte stride between consecutive rows (must be 16B aligned) - * @param major MN-major (row-major A) or K-major (transposed B) - * @param layout Swizzle type (NONE for simple row-major) - */ -__device__ __forceinline__ uint64_t make_umma_desc_bf16( +__device__ __forceinline__ int swizzle_128b(int offset_128b) { + // Swizzle<3,4,3>: XOR bits [6:4] with bits [9:7] + // offset_128b is the uint128_t index + // Within a 128B block (8 uint128_t), the swizzle permutes the 8 rows + int block_base = offset_128b & ~7; // Clear bottom 3 bits (8-aligned) + int row_in_block = offset_128b & 7; // 0-7 within the block + + // Swizzle<3,4,3> on the row index: + // XOR the 3 bits of row_in_block with the 3 bits of the block index + // But wait — the swizzle is on the BYTE address, not the row index. + // Let me be more precise. + // + // The Swizzle<3,4,3> operates on the 128-bit vector index. + // For a contiguous layout, vector index = row * stride + col. + // The swizzle XORs bits [6:4] (3 bits at position 4) with bits [9:7] (3 bits at position 7). + // Since each vector is 16 bytes, bit 4 corresponds to 16*2^4 = 256 bytes offset, + // and bit 7 corresponds to 16*2^7 = 2048 bytes offset. + // + // Actually, the swizzle is applied to the LAYOUT, not the address. + // The CuTe Swizzle<3,4,3> with a contiguous layout means: + // For each element at offset i: + // swizzled_i = i ^ ((i >> 3) & 0x7) << 4) + // + // Hmm, this isn't right either. Let me look at the CuTe implementation. + + // Simple approach: the swizzle permutes the 8 rows within each + // 128B block. The permutation is: + // row 0 → row 0 + // row 1 → row 1 + // row 2 → row 2 + // row 3 → row 3 + // row 4 → row 4 + // row 5 → row 5 + // row 6 → row 6 + // row 7 → row 7 + // Wait, Swizzle<3,4,3> with a stride-1 layout on 8 elements: + // The XOR is: (i & 0x70) ^ ((i & 0x0E) << 3) ... no. + // + // Let me just compute it from the CuTe definition. + // Swizzle<3,4,3> means: baz=3, b4=4, b3=3 + // The swizzle XORs (base >> b3) with (base >> (b3+b4)), taking baz bits. + // So: ((offset >> 3) ^ (offset >> 7)) & 0x7 + // This is applied to the uint128_t offset within the full layout. + + int swizzled_row = row_in_block ^ (((offset_128b >> 3) ^ (offset_128b >> 7)) & 0x7); + // Actually, the swizzle is on the ELEMENT offset, not the row-in-block. + // Let me redo this properly. + + // For Swizzle<3,4,3> applied to offset i (in uint128_t units): + // swizzled = i ^ (((i >> 3) ^ (i >> 7)) & 0x7) << 4) + // Wait, that shifts by 4 bits, but we're in uint128_t units. + + // The CuTe Swizzle<3,4,3> definition: + // template + // struct Swizzle { + // CUTE_HOST_DEVICE constexpr + // uint64_t operator()(uint64_t offset) const { + // return offset ^ ((offset >> b3) ^ (offset >> (b3+b4))) & ((1 << baz) - 1)) << b3; + // } + // }; + // For baz=3, b4=4, b3=3: + // swizzled = offset ^ (((offset >> 3) ^ (offset >> 7)) & 0x7) << 3 + + // This is in ELEMENT units (BF16), not uint128_t units. + // But we need to apply it to the uint128_t offset. + // The atom layout for BF16 MN_SW128 has shape (1024, 8) with stride (1, 1024). + // In uint128_t: (128, 8) with stride (1, 128). + // The swizzle is applied to the 1D offset = row + col * 128 (in uint128_t units). + + // Actually, I realize the swizzle operates on the 1D address (in the element + // space), not on the 2D coordinates. The CuTe layout maps (row, col) to a 1D + // offset, then the swizzle permutes that offset. + + // For MN_SW128 BF16: + // Layout: (1024, 8) with stride (1, 1024) + // 1D offset = m + n * 1024 (in BF16 elements) + // Swizzle: offset ^ (((offset >> 3) ^ (offset >> 7)) & 0x7) << 3 + + // To write BF16 element at (row, col) in the swizzled SMEM: + // 1. Compute 1D offset = row + col * 1024 (BF16 element index) + // 2. Apply swizzle: swizzled = offset ^ (((offset >> 3) ^ (offset >> 7)) & 0x7) << 3 + // 3. SMEM byte address = smem_base + swizzled * 2 (2 bytes per BF16) + + // To convert to uint128_t units: + // uint128_offset = swizzled / 8 + // But the swizzle operates on BF16 elements, not uint128_t. + + // This is the key: the swizzle is on the ELEMENT offset, not the vector offset. + // So I need to compute the element-level swizzled offset, then convert to bytes. + + // For our (128, HD) matrix with MN_SW128: + // HD=64 BF16 → n = 8 uint128_t per row + // But the layout atom has n = 8 elements in the N dimension (in uint128_t: 1) + // Wait, the atom shape is (1024, 8) in BF16 = (128, 1) in uint128_t. + // But HD=64 means the K dimension is 64 BF16 = 4 uint128_t. + + // I think the confusion is that the atom shape (1024, 8) means 1024 rows × 8 columns + // in BF16. For our matrix (128, 64), we tile the atom: + // M: 128 rows, atom covers 1024 → 1 tile (atom is bigger than matrix) + // N: 64 columns, atom covers 8 → 8 tiles + // So the full layout is: tile_to_shape(atom, (128, 64)) + // = repeat the 8-column atom 8 times along the N dimension + + // OK let me just compute this numerically for a few elements to verify. + + return block_base + (row_in_block & 7); // placeholder +} + +// ================================================================== +// UMMA SMEM write: write a BF16 element to swizzled SMEM +// ================================================================== +// Given a (row, col) position in the logical matrix, compute the +// swizzled SMEM address and write the BF16 value. +// ================================================================== + +__device__ __forceinline__ void umma_smem_write_mn_sw128( + bf16_t* smem_base, int row, int col, int hd, bf16_t val +) { + // MN_SW128 BF16 layout: atom shape (1024, 8), stride (1, 1024) + // For (128, HD) matrix: + // 1D element offset = row + col * 1024 (strided by 1024 in the N dimension) + // Wait, that's wrong. The N-dim stride is 1024 because the atom has 1024 rows. + // But our matrix only has 128 rows. The extra 896 rows are padding. + + // Actually, the full tiled layout for (128, HD) with atom (1024, 8) stride (1, 1024): + // logical_offset(m, n) = (m % 1024) + (m / 1024) * 1024 * 8 + (n % 8) * 1024 + (n / 8) * 8 + // Wait, tile_to_shape repeats the atom pattern. + + // For MN_SW128, the K-dim stride is 1024 BF16 elements per atom column. + // With our 128-row matrix (which fits in 1 atom column of 1024 rows): + // logical_offset(m, n) = m + (n / 8) * 1024 * 8 + (n % 8) * 1024 + // = m + (n / 8) * 8192 + (n % 8) * 1024 + + // Then swizzle: swizzled_offset = logical_offset ^ (((logical_offset >> 3) ^ (logical_offset >> 7)) & 0x7) << 3 + + int logical = row + (col / 8) * 8192 + (col % 8) * 1024; + int swizzled = logical ^ ((((logical >> 3) ^ (logical >> 7)) & 0x7) << 3); + + smem_base[swizzled] = val; +} + +// ================================================================== +// UMMA SMEM read: read a BF16 element from swizzled SMEM +// ================================================================== + +__device__ __forceinline__ bf16_t umma_smem_read_mn_sw128( + const bf16_t* smem_base, int row, int col, int hd +) { + int logical = row + (col / 8) * 8192 + (col % 8) * 1024; + int swizzled = logical ^ ((((logical >> 3) ^ (logical >> 7)) & 0x7) << 3); + return smem_base[swizzled]; +} + +// ================================================================== +// K-major SMEM write/read for SW128 +// ================================================================== +// For K-major B (K^T): the matrix is stored with K as the major dimension. +// The atom for K_SW128 is the same shape but with K as the leading dim. +// Layout: (k_elements, n_atoms) with K-major strides. +// ================================================================== + +__device__ __forceinline__ void umma_smem_write_k_sw128( + bf16_t* smem_base, int row, int col, int hd, bf16_t val +) { + // K_SW128 BF16: atom shape (8, 1024), stride (1, 8) + // For (128, HD) matrix in K-major: + // row = K index (0..127), col = MN index (0..HD-1) + // logical_offset = (k % 8) + (k / 8) * 8 * 128 + (mn % 8) * 8 + (mn / 8) * 1024 + // Hmm, this needs to be worked out properly. + + // For K_SW128 BF16 atom: Shape<(8, 1024)>, Stride<(1, 8)> + // In uint128_t: Shape<(1, 128)>, Stride<(1, 1)> + // tile_to_shape for (128, HD): + // K: 128 BF16 = 16 uint128_t → 16 tiles of the K atom (1 uint128_t each) + // MN: HD BF16 = 8 uint128_t → 8 tiles of the MN atom (128 uint128_t each) + + // Actually for K_SW128: + // atom: (8, 1024) BF16, stride (1, 8) + // For (128, 64) matrix: + // logical_offset(k, mn) = (k % 8) + (mn % 1024) * 8 + (k / 8) * 8 * 1024 + (mn / 1024) * 8 + // = (k % 8) + mn * 8 + (k / 8) * 8192 + + // Wait, this doesn't seem right. The K_SW128 atom for BF16 is: + // Shape<(8, 1024)> in elements, Stride<(1, 8)> + // This means: 8 contiguous elements along K, then stride 8 to the next group + // 1024 groups along MN + + // For our (128, 64) matrix in K-major: + // k ranges 0..127, mn ranges 0..63 + // Atom covers 8 K and 1024 MN. Need 16 K-tiles and 1 MN-tile. + // logical_offset(k, mn) = (k % 8) + (mn % 1024) * 8 + (k / 8) * 8 * 1024 + // = (k % 8) + mn * 8 + (k / 8) * 8192 + + int logical = (row % 8) + col * 8 + (row / 8) * 8192; + int swizzled = logical ^ ((((logical >> 3) ^ (logical >> 7)) & 0x7) << 3); + smem_base[swizzled] = val; +} + +// ================================================================== +// UMMA SMEM descriptor construction (correct format) +// ================================================================== + +enum class UmmaMajor { MN, K }; +enum class UmmaLayout { NONE = 0, MN_SW128 = 2, K_SW128 = 2 }; + +__device__ __forceinline__ uint64_t make_umma_desc( uint32_t smem_ptr, - int dim_m, int dim_n, int row_stride_bytes, - UmmaMajor major, UmmaLayout layout = UmmaLayout::NONE + int dim_m, int dim_n, + UmmaMajor major, UmmaLayout layout = UmmaLayout::MN_SW128 ) { uint64_t desc = 0; - // start_address: bits [0,14) = smem_ptr >> 4 + // start_address: bits [0,14) = smem_ptr >> 4 (in 16B units) desc |= (static_cast(smem_ptr >> 4) & 0x3FFF); - // leading_byte_offset: bits [16,30) = row_stride_bytes >> 4 - desc |= (static_cast(row_stride_bytes >> 4) & 0x3FFF) << 16; + // For SW128 layout: + // The leading_byte_offset and stride_byte_offset describe the canonical + // layout in uint128_t units. + // + // MN_SW128: atom (1024, 8) stride (1, 1024) in BF16 elements + // = (128, 1) stride (1, 128) in uint128_t units + // leading_byte_offset (K-dim stride) = 1024 BF16 * 2 / 16 = 128 uint128_t + // stride_byte_offset (inter-tile MN stride) = same for 1 MN tile + // + // K_SW128: atom (8, 1024) stride (1, 8) in BF16 elements + // = (1, 128) stride (1, 1) in uint128_t units + // leading_byte_offset (MN-dim stride) = 8 BF16 * 2 / 16 = 1 uint128_t + // stride_byte_offset (inter-tile K stride) = 8192 BF16 * 2 / 16 = 1024 uint128_t - // stride_byte_offset: bits [32,46) = same as leading for simple layout - desc |= (static_cast(row_stride_bytes >> 4) & 0x3FFF) << 32; - - // version: bits [46,48) = 1 (Blackwell) - desc |= (static_cast(1) << 46); - - // base_offset: bits [49,52) = 0 - // lbo_mode: bit [52,53) = 0 (legacy) - - // layout_type: bits [61,64) - desc |= (static_cast(static_cast(layout) & 0x7) << 61); - - // Upper 32 bits (bits [64,128)) - uint32_t hi = 0; - - // Format: bits [120,128) = 1 (BF16) - hi |= (1 << 24); // F16F32Format::BF16 = 1, at byte offset 15 from bit 64 - - // Dimensions depend on major if (major == UmmaMajor::MN) { - // For MN-major A: dim_m, dim_n, dim_k at [64,72), [72,80), [80,88) - // Wait, this depends on the exact CUTLASS layout. Let me look at the - // actual C++ code for how it sets these. - // From the SS MMA: the idescE encodes the dimensions. - // CUTLASS sets these based on the tile shape and major mode. - hi |= (dim_n & 0xFF); // byte 0 - hi |= ((dim_m & 0xFF) << 8); // byte 1 - // dim_k at byte 2 - hi |= (0 << 16); // lot_size at byte 3 + // MN_SW128: + // leading_byte_offset = 128 (uint128_t units) = 128 * 16 = 2048 bytes + // stride_byte_offset = 128 (uint128_t units) — same for simple case + int lbo = 128; // 1024 BF16 * 2 bytes / 16 = 128 uint128_t + int sbo = 128; + desc |= (static_cast(lbo) & 0x3FFF) << 16; + desc |= (static_cast(sbo) & 0x3FFF) << 32; } else { - // K-major B: dimensions in different order - hi |= (dim_m & 0xFF); // byte 0 - hi |= ((dim_n & 0xFF) << 8); // byte 1 - hi |= (0 << 16); + // K_SW128: + // leading_byte_offset = 1 (uint128_t units) = 8 BF16 per MN group + // stride_byte_offset = 1024 (uint128_t units) + int lbo = 1; + int sbo = 1024; + desc |= (static_cast(lbo) & 0x3FFF) << 16; + desc |= (static_cast(sbo) & 0x3FFF) << 32; } - desc |= (static_cast(hi) << 64); + // version: bits [46,48) = 1 + desc |= (static_cast(1) << 46); - // Wait, uint64_t only has 64 bits. The "upper 32 bits" are actually - // the bits [32,64) of the 64-bit value. Let me re-read the descriptor. - // - // Looking at the CUTLASS source again: - // The descriptor is a uint64_t. The struct has lo (uint32_t) and hi (uint32_t). - // The bitfield members span all 64 bits: - // [0,14) start_address, [14,16) unused, [16,30) leading_byte_offset, - // [30,32) unused, [32,46) stride_byte_offset, [46,48) version, - // [48,49) unused, [49,52) base_offset, [52,53) lbo_mode, [53,56) unused, - // [56,61) unused, [61,64) layout_type - // - // So the entire 64-bit descriptor is consumed by these fields. - // The "idescE" passed to tcgen05.mma is a SEPARATE 32-bit value, - // not part of the 64-bit descriptor! - // - // From the CUTLASS MMA code: - // uint32_t idescE = static_cast(tensor_a >> 32); - // So idescE IS the upper 32 bits of desc_a! - // - // But the bitfield already uses bits [32,64) for stride_byte_offset + version! - // How can idescE also be at bits [32,64)? - // - // AH: the bitfield interpretation and the "idescE" interpretation are - // DIFFERENT views of the same 64-bit value. The hardware reads the - // descriptor as raw bits, not as the C bitfield. The C bitfield - // is just a convenient way to set the bits. - // - // But the MMA PTX takes TWO arguments: desc (64-bit) and idescE (32-bit). - // The idescE is NOT just desc >> 32. It's a SEPARATE encoding. - // - // Let me re-read the CUTLASS MMA code more carefully... - - // Actually, from the CUTLASS MMA implementation: - // uint32_t idescE = static_cast(tensor_a >> 32); - // This IS desc >> 32. The hardware uses the full 64-bit descriptor, - // but the MMA PTX instruction takes it as two separate arguments: - // - desc_a as a 64-bit "l" constraint - // - idescE as a 32-bit "r" constraint (= desc_a >> 32) - // - // So the upper 32 bits of the descriptor serve double duty: - // - As part of the 64-bit descriptor (stride_byte_offset, version, layout_type) - // - As the idescE parameter to MMA (which the hardware interprets separately) - // - // This means my descriptor construction is correct up to bit 64. - // The "dimensions" I was trying to add at bits [64,128) don't exist - // in the 64-bit descriptor. The dimensions are encoded in the - // stride_byte_offset, leading_byte_offset, etc. - // - // Actually wait, let me look at the CUTLASS code again. The bitfield - // has stride_byte_offset_ at [32,46) and version_ at [46,48). That's - // only 48 bits used. What about [48,64)? - // - // The answer: the upper 16 bits [48,64) contain base_offset, lbo_mode, - // and layout_type. The full 64-bit descriptor is accounted for. - // - // The idescE is just desc >> 32, which the MMA hardware reinterprets - // as a separate parameter. The hardware knows how to decode it. + // layout_type: bits [61,64) = 2 (SWIZZLE_128B) + desc |= (static_cast(2) << 61); return desc; } @@ -178,24 +333,12 @@ __device__ __forceinline__ uint64_t make_umma_desc_bf16( // tcgen05.mma PTX wrappers // ================================================================== -/** - * QK GEMM: S←S (both operands in SMEM, result in TMEM). - * - * Called by ONE lane per warp (elect_one_sync pattern). - */ __device__ void umma_ss_f16( - uint32_t tmem_c, - uint64_t desc_a, uint64_t desc_b, + uint32_t tmem_c, uint64_t desc_a, uint64_t desc_b, bool accumulate = false ) { - // idescE = upper 32 bits of descriptor A uint32_t idescE = static_cast(desc_a >> 32); - - // scaleC: 0.0f for init, 1.0f for accumulate - // The PTX uses: setp.ne.b32 pred, %scaleC, 0 uint32_t scaleC_bits = accumulate ? 0x3F800000u : 0u; - - // Mask: all zeros = apply to all N-dim tiles uint32_t mask0 = 0, mask1 = 0, mask2 = 0, mask3 = 0; asm volatile( @@ -212,21 +355,12 @@ __device__ void umma_ss_f16( ); } -/** - * PV GEMM: T←S (A in TMEM, B in SMEM, result in TMEM). - * - * Called by ONE lane per warp (elect_one_sync). - */ __device__ void umma_ts_f16( - uint32_t tmem_c, - uint32_t tmem_a, - uint64_t desc_b, + uint32_t tmem_c, uint32_t tmem_a, uint64_t desc_b, 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( @@ -234,7 +368,6 @@ __device__ void umma_ts_f16( ".reg .pred p;\n\t" "setp.ne.b32 p, %4, 0;\n\t" "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), diff --git a/tests/unit/test_qk_mma.cu b/tests/unit/test_qk_mma.cu index 0e78d2bf..ce5ed323 100644 --- a/tests/unit/test_qk_mma.cu +++ b/tests/unit/test_qk_mma.cu @@ -64,8 +64,14 @@ int main() { cudaMemcpy(dk, hkb, B*s_k*HD*2, cudaMemcpyHostToDevice); cudaMemset(ds_out, 0, s_k*4); - // SMEM = 4 (tmem_base) + 128*HD*2 (sQ) + 128*HD*2 (sK) = 4 + 256*HD*2 - int smem = 4 + 128 * HD * 2 * 2 + 1024; // +1024 slack + // SMEM = 4 (tmem_base) + 128B align + 128*HD*2*2 (sQ + sK, swizzled) + slack + // Swizzled layout may use more space than row-major due to atom padding + // MN_SW128 atom = 1024*8 = 8192 BF16 per atom. For (128, HD) = 128*HD BF16. + // With 1 MN tile and ceil(HD/8) K tiles, total = 8192 * ceil(HD/8) BF16 + int atoms_n = (HD + 7) / 8; // number of K-tiles (atom has 8 BF16 in K dim) + int smem_q = 1024 * 8 * atoms_n; // BF16 elements for Q (padded) + int smem_k = 1024 * 8 * atoms_n; // BF16 elements for K (padded) + int smem = 4 + 128 + smem_q * 2 + smem_k * 2 + 1024; // bytes dim3 grid(1, H, B); dim3 block(NTHREADS);