FMHA SM100: Add SMEM descriptor construction for tcgen05.mma

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<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+    uint16_t start_addr = static_cast<uint16_t>(addr >> 4);
+    
+    // leading_byte_offset: bits [16,30), shifted left by 4
+    uint16_t leading_offset = static_cast<uint16_t>(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<uint64_t>(start_addr) | 
+           (static_cast<uint64_t>(leading_offset) << 16) |
+           (static_cast<uint64_t>(stride_offset) << 32) |
+           (static_cast<uint64_t>(static_cast<uint8_t>(swizzle)) << 61);
+    
+    return desc;
+}
+
 // =====================================================================
 // Device helpers
 // =====================================================================