diff --git a/dsv4/kernels/attention/fmha_sm100.cuh b/dsv4/kernels/attention/fmha_sm100.cuh index b90fc264..bd69d497 100644 --- a/dsv4/kernels/attention/fmha_sm100.cuh +++ b/dsv4/kernels/attention/fmha_sm100.cuh @@ -93,6 +93,58 @@ constexpr int TOTAL_SMEM(int hd) { return (sQ + sK + sV + sC) * sizeof(__nv_bfloat16); } +// ===================================================================== +// SMEM Descriptor construction for tcgen05.mma +// ===================================================================== +// The SMEM descriptor is a 64-bit value encoding the SMEM address, +// layout, and dimensions for the MMA operation. +// See cute/arch/mma_sm100_desc.hpp for the bitfield format. + +enum class SmemSwizzle : uint8_t { + NONE = 0, + SWIZZLE_128B_BASE32B = 1, + SWIZZLE_128B = 2, + SWIZZLE_64B = 4, +}; + +/** + * Construct a UMMA SMEM descriptor for BF16 K-major layout. + * + * For Q (A-matrix): K-major, shape (TILE_M, HEAD_DIM), ld = HEAD_DIM * sizeof(bf16) + * For K (B-matrix): K-major, shape (SMEM_TILE_K, HEAD_DIM), ld = HEAD_DIM * sizeof(bf16) + * + * @param smem_ptr SMEM pointer (must be 16B-aligned) + * @param ld_bytes Leading dimension in bytes (row stride) + * @param swizzle Swizzle mode + */ +__device__ __forceinline__ uint64_t make_umma_smem_desc_bf16( + const void* smem_ptr, + uint32_t ld_bytes, + SmemSwizzle swizzle = SmemSwizzle::SWIZZLE_128B +) { + uint64_t desc = 0; + + // start_address: bits [0,14), shifted left by 4 (16B-aligned) + uint32_t addr = static_cast(__cvta_generic_to_shared(smem_ptr)); + uint16_t start_addr = static_cast(addr >> 4); + + // leading_byte_offset: bits [16,30), shifted left by 4 + uint16_t leading_offset = static_cast(ld_bytes >> 4); + + // stride_byte_offset: bits [32,46), shifted left by 4 + // For 2D matrices, stride = ld * num_rows (for TMA, this is the 2nd dimension stride) + // But for MMA SMEM descriptors, stride is the offset between tiles (0 for single-tile) + uint16_t stride_offset = 0; + + // Pack into 64-bit descriptor + desc = static_cast(start_addr) | + (static_cast(leading_offset) << 16) | + (static_cast(stride_offset) << 32) | + (static_cast(static_cast(swizzle)) << 61); + + return desc; +} + // ===================================================================== // Device helpers // =====================================================================