CUTLASS CUTE 7 hpc-ops 之 group gemm & tma¶

__global__ void group_gemm_pertensor_fp8_kernel(...) {

    // ===== PRELOGUE =====
    int idx = threadIdx.x;
    bool is_producer = (idx >= 256);  // 128 threads for load
    bool is_consumer = (idx < 256);    // 256 threads for math

    extern __shared__ uint8_t shm_data[];
    auto* shm_a = (Tin*)shm_data;          // [kTileM, kTileK, kStage]
    auto* shm_b = shm_a + ...;              // [kTileN, kTileK, kStage]
    int* shm_tiles = (int*)(shm_data + ...); // For scheduler

    // Initialize mbarriers (producer-consumer sync)
    if (is_leader) {
    for (int s = 0; s < kStage; s++) {
        initialize_barrier(readable[s], 1);  // Consumer waits on this
        initialize_barrier(writable[s], num_mma_warpgroup);  // Producer waits on this
    }
    }
    // Load scheduler metadata to shared memory
    for (int i = idx; i < num_group; i += 384)
    shm_tiles[i] = tiles_ptr[i];
    __syncthreads();

    // ===== MAINLOOP: Producer-Consumer Pipeline =====
    int phase = 0;

    if (is_producer && is_leader_in_load) {
    // PRODUCER: Load data via TMA
    int s_write = 0;  // Current stage to write
    while (true) {
        // SCHEDULER: Get next tile (igroup, itile_m, itile_n)
        if (!get_next_tile(shm_tiles, iblock, ...)) break;

        for (int k = 0; k < ntile_k; k++) {
        // MBARRIER: Wait for consumer to release stage
        wait_barrier(writable[s_write], phase);

        // TMA load X and W tiles to shared memory
        tma_copy(shm_a[_, _, s_write], global_X[igroup, itile_m, k]);
        tma_copy(shm_b[_, _, s_write], global_W[igroup, itile_n, k]);

        // MBARRIER: Signal consumer data is ready
        set_barrier_transaction_bytes(readable[s_write], ...);

        // Circular stage buffer
        s_write = (s_write + 1) % kStage;
        if (s_write == 0) phase ^= 1;
        }
    }
    }

    if (is_consumer) {
    // CONSUMER: Compute via GMMA
    int s_read = 0;  // Current stage to read
    while (true) {
        // SCHEDULER: Same as producer, get next tile
        if (!get_next_tile(shm_tiles, iblock, ...)) break;

        for (int k = 0; k < ntile_k; k++) {
        // MBARRIER: Wait for producer to fill stage
        wait_barrier(readable[s_read], phase);

        // GMMA: Compute on shared memory data
        gemm(shm_a[_, _, s_read], shm_b[_, _, s_read], accum);

        // MBARRIER: Signal producer stage is consumed
        if (is_leader_in_warpgroup)
            arrive_barrier(writable[s_read]);

        // Circular stage buffer
        s_read = (s_read + 1) % kStage;
        if (s_read == 0) phase ^= 1;
        }

        // ===== EPILOGUE =====
        cast_to_bf16(accum, output, pertensor_scale);
        tma_store(output, global_Y[igroup, itile_m, itile_n]);
    }
    }
}

constexpr int kGroupPerThread = 8;
constexpr int kThreadPerBlock = 32;
kernels::update_grouped_tma<...>
    <<<num_group + 1, kThreadPerBlock, 0, stream>>>(...);

template <typename Tin, typename Tout, typename TmaX, typename TmaY,
            int kTileM, int kGroupPerThread, int kThreadPerBlock>
__global__ void update_grouped_tma(
    const vec_t<TmaDescriptor, 2> td_xy,  // template TMA descriptors (X, Y) from host
    TmaDescriptor *tma_xy,                // [num_group * 2]: X desc at 2*i, Y desc at 2*i+1
    const Tin *x_ptr,                     // X (activation) data
    const Tout *y_ptr,                    // Y (output) data
    const int *seqlens_ptr,               // [num_group]: seqlen per group
    const int *cu_seqlens_ptr,            // [num_group + 1]: cumulative seqlen
    int *tiles_ptr,                       // [num_group]: tile M count per group
    int *cu_tiles_ptr,                    // [num_group + 1]: cumulative tile M count
    int num_group, int m, int n, int k) {

    int idx = threadIdx.x;
    int igroup = blockIdx.x;

    if (igroup == num_group) {
    // ---- Case 1: Compute tile statistics (last block) ----
    int tiles[kGroupPerThread];

    // Step 1: Each thread computes tile counts for its assigned groups
    for (int i = 0; i < kGroupPerThread; i++) {
        int g = idx * kGroupPerThread + i;
        if (g < num_group) {
        tiles[i] = (seqlens_ptr[g] + kTileM - 1) / kTileM;
        tiles_ptr[g] = tiles[i];
        } else {
        tiles[i] = 0;
        }
    }

    // Step 2: Exclusive scan to compute cumulative tile M counts
    using BlockScan = cub::BlockScan<int, kThreadPerBlock>;
    __shared__ typename BlockScan::TempStorage temp_storage;
    int block_aggregate;
    BlockScan(temp_storage).ExclusiveSum(tiles, tiles, block_aggregate);

    // Step 3: Write cumulative results to global memory
    for (int i = 0; i < kGroupPerThread; i++) {
        int g = idx * kGroupPerThread + i;
        if (g < num_group) {
        cu_tiles_ptr[g] = tiles[i];
        }
    }
    if (idx == 0) {
        cu_tiles_ptr[num_group] = block_aggregate;  // total tile M
    }

    } else {
    // ---- Case 2: Update TMA descriptors for group igroup ----
    __shared__ TmaDescriptor smem_tma_desc[2];

    int num_seq = seqlens_ptr[igroup];
    int cu_seqlen = cu_seqlens_ptr[igroup];

    // Step 1: Copy template descriptors to shared memory
    if (idx < 2) {
        smem_tma_desc[idx] = td_xy[idx];
    }
    __syncwarp();

    // Step 2: Thread 0 — update X (activation) TMA descriptor
    if (idx == 0) {
        auto gX = make_tensor(
            make_gmem_ptr(x_ptr + cu_seqlen * k),
            make_shape(num_seq, k),
            make_stride(k, Int<1>{}));              // stride unchanged: same k for all groups
        update_tma_gtensor<TmaX>(smem_tma_desc[0], gX);
    }

    // Step 3: Thread 1 — update Y (output) TMA descriptor
    if (idx == 1) {
        auto gY = make_tensor(
            make_gmem_ptr(y_ptr + cu_seqlen * n),
            make_shape(n, num_seq),
            make_stride(Int<1>{}, n));              // stride unchanged: same n for all groups
        update_tma_gtensor<TmaY>(smem_tma_desc[1], gY);
    }

    // Step 4: Commit smem writes, then copy descriptors from smem to gmem
    for (int i = 0; i < 2; i++) {
        __syncwarp();
        if (elect_one_sync()) {
        tma_desc_commit_group();
        tma_desc_wait_group();
        }
        tma_descriptor_cp_fence_release(tma_xy + igroup * 2 + i, smem_tma_desc[i]);
    }
    }
}

copy(tma_origin.with(new_tma_descriptor, mbarrier), gmem_tensor, smem_tensor);

// itile_m is not a global idx, it is relative to the group
cute::copy(tma_a.with(td_x, readable[ismem_write]), tAg(_, itile_m, itile_k), tAs(_, 0, 0, ismem_write));

using SLayoutW = decltype(tile_to_shape(SLayoutWAtom{}, 
                            make_shape(Int<kTileN>{}, Int<kTileK>{}, Int<kStage>{})));
// w is 3-dim tensor, but copy box is 2-dim
auto tma_w = make_tma_copy(SM90_TMA_LOAD{}, w, take<0, 2>(SLayoutW{}));

// Line 88: define BlockScan type
using BlockScan = cub::BlockScan<int, kThreadPerBlock>;
// - template param 1: int - data type for scan
// - template param 2: kThreadPerBlock = 32 - threads per block

// Line 89: allocate shared memory
__shared__ typename BlockScan::TempStorage temp_storage;
// - TempStorage is a struct defined internally by cub
// - shared memory needed to coordinate thread communication
// - size computed automatically by cub

// Line 90: used to return total sum
int block_aggregate;

// Line 91: execute Exclusive Sum Scan
BlockScan(temp_storage).ExclusiveSum(tiles, tiles, block_aggregate);
// Parameter description:
// - tiles (input): data array contributed by each thread
// - tiles (output): scan result (modified in-place)
// - block_aggregate: returns total sum of the block

if (cute::elect_one_sync()) {
    cute::tma_desc_commit_group();
    cute::tma_desc_wait_group();
}

tma_descriptor_cp_fence_release(tma_xy + i, smem_tma_desc[i]);
// "Release": ensure all prior writes are visible

tma_descriptor_fence_acquire(td_xy + i);
// "Acquire": ensure subsequent reads see the complete descriptor