UMMA: fix descriptor + idesc — use gau-nernst tutorial values
- 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
This commit is contained in:
@@ -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<int ROWS, int COLS>
|
||||
__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<int ROWS>
|
||||
__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<int HD>
|
||||
__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<int ROWS, int COLS>
|
||||
__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<int SK, int HD>
|
||||
__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<int SK, int HD>
|
||||
__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<uint64_t>(smem_ptr >> 4) & 0x3FFF); // start_address
|
||||
desc |= (static_cast<uint64_t>(1) & 0x3FFF) << 16; // LBO = 1 (16B)
|
||||
desc |= (static_cast<uint64_t>(2) & 0x3FFF) << 32; // SBO = 2 (16B units)
|
||||
desc |= (static_cast<uint64_t>(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<uint64_t>(smem_ptr >> 4) & 0x3FFF); // start_address
|
||||
desc |= (static_cast<uint64_t>(1) & 0x3FFF) << 16; // LBO = 1 (16B)
|
||||
desc |= (static_cast<uint64_t>(block_k / 8) & 0x3FFF) << 32; // SBO
|
||||
desc |= (static_cast<uint64_t>(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<uint64_t>(smem_ptr >> 4) & 0x3FFF); // start_address
|
||||
desc |= (static_cast<uint64_t>(1) & 0x3FFF) << 16; // LBO = 1 (128B)
|
||||
desc |= (static_cast<uint64_t>(block_k / 8) & 0x3FFF) << 32; // SBO
|
||||
desc |= (static_cast<uint64_t>(1) << 46); // version = 1
|
||||
desc |= (static_cast<uint64_t>(1) << 52); // lbo_mode = 1
|
||||
desc |= (static_cast<uint64_t>(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<uint32_t>(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<uint32_t>(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
|
||||
|
||||
@@ -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 <cuda_runtime.h>
|
||||
@@ -13,27 +14,29 @@
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
// 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<<<grid, block, smem_size>>>(
|
||||
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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user