DeepGEMM CUDA 内核实现深度分析

deep_gemm/
├── common/          # 公共工具层（类型、数学、调度、TMA、架构工具）
├── mma/             # MMA 指令抽象（SM90 WGMMA / SM100 UMMA）
├── ptx/             # PTX 内联汇编封装（WGMMA、TCGen05、TMA、LD/ST）
├── impls/           # 内核实现（GEMM、MQA、MegaMoE）
├── scheduler/       # 块调度器（GEMM、MegaMoE、Paged MQA）
├── epilogue/        # 后处理（变换、SM100 TMEM→SMEM→GMEM 写回）
├── layout/          # 内存布局（MegaMoE workspace、对称缓冲）
├── comm/            # 通信原语（grid sync、NVLink barrier）
└── csrc/jit_kernels/# JIT C++ 端（启发式配置选择）

enum class MmaKind { BF16 = 0, MXFP8FP4 = 1 };
enum class GemmType { Normal, MGroupedContiguous, MGroupedMasked, KGroupedContiguous, Batched, MGroupedContiguousWithPsumLayout };
enum class KernelType { Kernel1D1D = 0, Kernel1D2D = 1, KernelNoSF = 2 };

flowchart TB
    subgraph "CTA 线程分组"
        TMA["TMA Warp-Group<br>(数据加载)"]
        MMA["Math Warp-Groups<br>(矩阵计算)"]
        EPI["Epilogue Warp-Groups<br>(结果写回)"]
    end

    subgraph "流水线阶段"
        PREFETCH["TMA Descriptor<br>Prefetch"]
        LOAD["TMA Load<br>GMEM → SMEM"]
        COMPUTE["MMA Compute<br>SMEM → Registers/TMEM"]
        STORE["TMA Store<br>SMEM → GMEM"]
    end

    SCHED["Scheduler<br>(Block Assignment)"]

    SCHED --> TMA
    PREFETCH --> LOAD
    LOAD -->|"full_barrier"| MMA
    MMA -->|"empty_barrier"| LOAD
    MMA -->|"tmem_full_barrier"| EPI
    EPI -->|"tmem_empty_barrier"| MMA
    EPI --> STORE

flowchart LR
    subgraph "Shared Memory (≤232KB)"
        D["D (Output)"]
        A["A Stages<br>(ping-pong)"]
        B["B Stages<br>(ping-pong)"]
        SFA["SFA Stages"]
        SFB["SFB Stages"]
        BARR["Barriers<br>(full/empty)"]
        TMEM["TMEM Ptr<br>(SM100)"]
    end

    D --> A --> B --> SFA --> SFB --> BARR --> TMEM

smem = TensorMap(optional) + D + [A × kNumStages] + [B × kNumStages] + [SFA × kNumStages] + [SFB × kNumStages] + Barriers

flowchart TB
    subgraph "SM90 (Hopper)"
        direction TB
        S90_MMA["WGMMA<br>wgmma.mma_async<br>Accumulator: 寄存器"]
        S90_SF["Float32 SF<br>per-128-channel"]
        S90_MULTICAST["TMA Multicast<br>动态 2-CTA"]
        S90_EPI["直接寄存器写回<br>STSM → SMEM → TMA Store"]
    end

    subgraph "SM100 (Blackwell)"
        direction TB
        S100_MMA["UMMA (tcgen05.mma)<br>Accumulator: TMEM (Tensor Memory)"]
        S100_SF["UE8M0 SF<br>Block-Scaled MXFP8/FP4"]
        S100_MULTICAST["2-CTA Mode<br>静态双 SM 协作"]
        S100_EPI["TMEM Load → STSM<br>→ SMEM → TMA Store"]
    end

// 1. 发射前 fence
warpgroup_fence_operand(accum[i]);  // 防止编译器重排
warpgroup_arrive();                   // wgmma.fence.sync.aligned

// 2. 发射 K 维度循环
for (k = 0; k < BLOCK_K / WGMMA::K; k++) {
    desc_a = make_smem_desc(smem_a + ...);
    desc_b = make_smem_desc(smem_b + ...);
    WGMMA::wgmma(desc_a, desc_b, accum, k);  // 异步发射
}

// 3. 提交并等待
warpgroup_commit_batch();     // wgmma.commit_group
warpgroup_fence_operand(accum[i]);
warpgroup_wait<0>();          // wgmma.wait_group 0

// 1. TMEM 分配（一次性）
Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem);

// 2. SF 通过 UTCCP 加载到 TMEM
cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4);

// 3. UMMA 发射（结果累加到 TMEM）
mma_t::fma(a_desc, b_desc, tmem_c_offset, scale_c, runtime_instr_desc, tmem_sfa, tmem_sfb);

// 4. 提交到 mbarrier
empty_barrier_arrive(is_last_k_block);  // tcgen05.commit

// 5. Epilogue: TMEM → 寄存器 → STSM → SMEM → TMA Store
SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values...);
st_shared(smem_ptr, values...);
SM90_TMA_STORE_2D::copy(&tensor_map_cd, smem_ptr, n_idx, m_idx);

// TMA copy 函数签名
template <uint32_t BLOCK_INNER, uint32_t BLOCK_OUTER,
          uint32_t kSwizzleMode, typename dtype_t, bool kIs3DTMA>
void tma_copy(void const* desc_ptr, ClusterTransactionBarrier* barrier_ptr,
              dtype_t* smem_ptr, inner_idx, outer_idx, num_tma_multicast, batch_idx);

template <uint32_t BLOCK_INNER, uint32_t kSwizzleMode, typename dtype_t>
constexpr uint32_t get_inner_block_atom_size() {
    return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t);
}

flowchart LR
    subgraph "SM90 TMA Multicast"
        GMEM1["GMEM"] -->|"single TMA"| L2
        L2 -->|"multicast"| CTA1["CTA 0"]
        L2 -->|"multicast"| CTA2["CTA 1"]
    end

    subgraph "SM100 2-CTA Mode"
        GMEM2["GMEM"] -->|"2SM TMA"| SM1["SM 0"]
        GMEM2 -->|"2SM TMA"| SM2["SM 1"]
        SM1 <-->|"cluster_sync"| SM2
    end

// SM90: 使用 TMA reduce-add 实现累加写回
cute::SM90_TMA_REDUCE_ADD_2D::copy(&tensor_map_cd, smem_d, n_idx, m_idx);
cute::tma_store_arrive();

// SM100: 支持 2D/3D TMA store
using cute_tma_t = cute::conditional_t<kWithAccumulation,
    cute::SM90_TMA_REDUCE_ADD_2D, cute::SM90_TMA_STORE_2D>;
cute_tma_t::copy(&tensor_map_cd, smem_cd, n_idx, m_idx);

// FP8: 64×N×32, K=32, SS_TN (Shared-Shared, TN layout)
MMA_64xNx32_F32E4M3E4M3_SS_TN()

// BF16: 64×N×16, K=16, SS (支持 K/MN major)
MMA_64xNx16_F32BF16BF16_SS<MajorA, MajorB>()

// TF32: 64×N×8, K=8, RS (Register-Shared)
MMA_64xNx8_F32TF32TF32_RS_TN()

// 非量化: tcgen05.mma.cta_group::1.kind::f16
// MXFP8/FP4: tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale
// FP8: tcgen05.mma.cta_group::1.kind::f8f6f4
// WS (Warp Specialization): tcgen05.mma.ws.cta_group::1.kind::f16

UMMA_M = 64  → UMMA_N ∈ [8, 256], step 8
UMMA_M = 128 → UMMA_N ∈ [16, 256], step 16
UMMA_M = 256 → UMMA_N ∈ [16, 256], step 16
UMMA_K = 32 (固定)

// K-major 布局（A 矩阵典型）
// Atom: 8 × (kSwizzleMode bytes on K)
// SBO = num_non_contiguous × BLOCK_K × sizeof(dtype)
// LBO = 0 (K 维度仅 1 个 atom)
make_smem_desc(base_ptr, layout_type, LBO=0, SBO);

// MN-major 布局（B 矩阵典型）
// Atom: (kSwizzleMode bytes on MN) × 8
// SBO = num_non_contiguous × BLOCK_MN_ATOM × sizeof(dtype)
// LBO = BLOCK_K × BLOCK_MN_ATOM × sizeof(dtype)
make_smem_desc(base_ptr, layout_type, LBO, SBO);

flowchart TB
    subgraph "SM90: 寄存器累加"
        WGMMA_IN["WGMMA 输出"] --> REG["float accum[N]<br>驻留寄存器"]
        REG --> SCALE["Scale 乘法<br>final += sa * sb * accum"]
        REG --> STSM["STSM → SMEM"]
    end

    subgraph "SM100: TMEM 累加"
        UMMA_IN["UMMA 输出"] --> TMEM["TMEM 累加器<br>2KB/SM, 512 cols"]
        TMEM --> TLOAD["TMEM Load<br>→ 寄存器"]
        TLOAD --> STSM2["STSM → SMEM"]
        STSM2 --> TMASTORE["TMA Store → GMEM"]
    end

enum class KernelType {
    Kernel1D1D = 0,   // A/B 都有 per-channel scaling
    Kernel1D2D = 1,   // A per-channel, B per-tile scaling
    KernelNoSF = 2    // 无 scaling factor（BF16）
};

A: [BLOCK_M, BLOCK_K] → SFA: [BLOCK_M, BLOCK_K/128] (per-row)
B: [BLOCK_N, BLOCK_K] → SFB: [BLOCK_N, BLOCK_K/128] (per-col)

// 1. WGMMA 计算 int/raw 结果
WGMMA::wgmma(desc_a, desc_b, accum, k);

// 2. MMA 后应用 scaling
final_accum[i*4+0] += scale_a_0 * scale_b_0 * accum[i*4+0];
final_accum[i*4+1] += scale_a_0 * scale_b_1 * accum[i*4+1];
final_accum[i*4+2] += scale_a_1 * scale_b_0 * accum[i*4+2];
final_accum[i*4+3] += scale_a_1 * scale_b_1 * accum[i*4+3];

// UE8M0 SF 计算：仅保留指数部分
void get_e4m3_sf_and_sf_inv(const float2& amax, float2& sf, float2& sf_inv) {
    const float2 finfo_factor = {1.0 / 448.0, 1.0 / 448.0};
    const auto scaled = __fmul2_rn(amax, finfo_factor);
    const auto exp_x = fast_log2_ceil(scaled.x);
    sf.x = fast_pow2(exp_x);      // 2^ceil(log2(amax/448))
    sf_inv.x = fast_pow2(-exp_x); // 2^(-exp)
}

// UTCCP 4x32dp128bit: 每次搬运 128 个 SF 元素到 TMEM
using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_1cta;  // 1-CTA
using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta;  // 2-CTA

// 需要 warp-level transpose 使 SF 布局匹配 UTCCP 要求
utccp_required_smem_warp_transpose(smem_ptr);
cute_utccp_t::copy(sf_desc, tmem_col_offset);

enum class GemmType {
    Normal,                              // 普通 GEMM
    MGroupedContiguous,                  // M 维度连续分组（指针跳转）
    MGroupedMasked,                      // M 维度 masked 分组（cumsum 边界）
    KGroupedContiguous,                  // K 维度连续分组（变长 K）
    Batched,                             // Batched GEMM
    MGroupedContiguousWithPsumLayout     // 带 partial sum 的连续分组
};

flowchart TB
    START["get_next_block()"] --> IDX["next_block_idx =<br>++iter * kNumSMs + blockIdx.x"]
    IDX --> CHECK{"GemmType?"}

    CHECK -->|"Normal"| SWIZZLE["swizzled_block_idx<br>L2 局部性优化"]
    CHECK -->|"MGroupedContiguous"| GROUP_OFFSET["grouped_layout[m_idx]<br>获取组偏移"]
    CHECK -->|"MGroupedMasked"| MASKED["cumsum 边界检查<br>跳过无效组"]
    CHECK -->|"KGroupedContiguous"| KGROUP["current_k_cumsum<br>current_sf_k_cumsum"]
    CHECK -->|"Batched"| BATCH["current_group_idx<br>= block_idx / num_blocks"]

    SWIZZLE --> RETURN["返回 m/n_block_idx"]
    GROUP_OFFSET --> RETURN
    MASKED --> RETURN
    KGROUP --> RETURN
    BATCH --> RETURN

// 将 block_idx 重映射为 (group, in_group) 形式
const auto num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup;
const auto group_idx = block_idx / num_blocks_per_group;
auto first_block_idx = group_idx * kNum1DBlocksPerGroup;
auto in_group_idx = block_idx % num_blocks_per_group;

// kNum1DBlocksPerGroup ∈ {8, 16}，选择使 SM 利用率最高的值

// TMA 描述符动态更新（tensormap.replace）
ptx::tensor_map_replace_global_addr_in_smem(smem_tensor_map_a, gmem_a_ptr + current_k_cumsum * shape_m);
ptx::tensor_map_replace_global_inner_dim_stride_in_smem(smem_tensor_map_a, current_shape_k, current_shape_k);
ptx::tensor_map_release_gpu();
ptx::tensor_map_acquire_gpu(gmem_tensor_map_a);

flowchart TB
    subgraph "Mega MoE Kernel (SM100)"
        direction TB

        DISPATCH["Dispatch Phase<br>EP AlltoAll 通信<br>(FP8 tokens)"]

        subgraph "GEMM Phase 1 (L1)"
            L1_MMA["FP8×FP4 GEMM<br>Hidden → Intermediate"]
            L1_SWIGLU["SwiGLU Activation"]
            L1_OUT["FP8 Output"]
        end

        subgraph "GEMM Phase 2 (L2)"
            L2_MMA["FP8×FP4 GEMM<br>Intermediate → Hidden"]
            L2_ADD["Residual Add"]
            L2_OUT["BF16 Output"]
        end

        COMBINE["Combine Phase<br>EP AlltoAll 通信<br>(BF16 tokens)"]
    end

    DISPATCH --> L1_MMA --> L1_SWIGLU --> L1_OUT --> L2_MMA --> L2_ADD --> L2_OUT --> COMBINE

template <uint32_t kNumRanks>
struct SymBuffer {
    int64_t base;                    // 本地基地址
    int64_t offsets[kNumMaxRanks];   // 各 rank 相对偏移
    uint32_t rank_idx;               // 当前 rank

    // 将本地指针映射到目标 rank 的地址空间
    template <typename ptr_t>
    ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {
        int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);
        return *reinterpret_cast<ptr_t*>(&mapped_ptr);
    }
};

flowchart TB
    START["for_each_block()"] --> FETCH["fetch_expert_recv_count()<br>等待所有 rank 的 token 计数"]
    FETCH --> WAVE["Expert Wave 遍历<br>kNumExpertsPerWave 个专家/波"]

    WAVE --> L1["L1 Phase<br>L1_SHAPE_N = 2 × Intermediate<br>L1_SHAPE_K = Hidden"]
    L1 --> L1_ASSIGN["fetch_next_l1_block()<br>分配 L1 GEMM 块"]

    L1_ASSIGN --> L2["L2 Phase<br>L2_SHAPE_N = Hidden<br>L2_SHAPE_K = Intermediate"]
    L2 --> L2_ASSIGN["fetch_next_l2_block()<br>分配 L2 GEMM 块"]

    L2_ASSIGN --> NEXT_EXPERT["advance_expert_idx()"]
    NEXT_EXPERT --> WAVE

// Dispatch warps: 负责 EP 通信
if (warp_idx < kNumDispatchWarps) {
    // 1. 从远程 rank 拉取 token
    // 2. 写入本地 smem/gmem
    // 3. 更新 expert recv count
}

// Non-epilogue warps: 负责 MMA 发射
else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) {
    // 1. 等待 token 到达（l1_arrival_count）
    // 2. 发射 UMMA 指令
}

// Epilogue warps: 负责结果写回 + Combine
else {
    // 1. TMEM → SMEM → GMEM
    // 2. 触发 Combine 通信
}

// Grid Sync: 全局 CTA 同步（基于原子操作）
template <uint32_t kNumSMs, uint32_t kGridSyncIndex, typename sync_scope_t>
void grid_sync(const Workspace& workspace, sm_idx, thread_idx, sync_scope) {
    // SM 0 发起，其他 SM 递增计数器
    // 使用 release/acquire 语义保证内存可见性
    atomic_add_rel(count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1);
    // 等待 tag 翻转
    while ((ld_acq(count_ptr) ^ old_value) & kFinishSumTag) == 0);
}

// NVLink Barrier: 跨 rank 同步
void nvlink_barrier(workspace, sym_buffer, sm_idx, thread_idx, ...) {
    grid_sync();                          // 本地同步
    if (sm_idx == 0) {
        red_add_rel_sys(signal_ptr, ±1);  // 远程原子信号
        while (ld_acq_sys(signal_ptr) != target);  // 等待
    }
    grid_sync();                          // 本地同步
}

flowchart LR
    Q["Q: [BLOCK_Q, kNumHeads, kHeadDim]"] --> WGMMA["WGMMA<br>FP8×FP8→FP32"]
    KV["KV: [BLOCK_KV, kHeadDim]"] --> WGMMA
    WGMMA --> ACC["accum: [BLOCK_Q, BLOCK_KV]"]
    ACC --> REDUCE["Reduce over heads<br>+ weight + scale"]
    REDUCE --> LOGITS["Logits: [BLOCK_Q, BLOCK_KV]"]

// 矩阵乘法: [BLOCK_Q * kNumHeads, kHeadDim] @ [BLOCK_KV, kHeadDim]^T
WGMMA::wgmma(desc_kv, desc_q, accum, k);

// Head 维度归约 + 权重 + scaling
for (i = 0; i < BLOCK_Q; i++) {
    float sum[4] = {transform(0), transform(1), transform(2), transform(3)};
    for (j = 1; j < kNumHeads / 8; j++)
        for (k = 0; k < 4; k++) sum[k] += transform(j*4+k);
    v_0 = (sum[0] + sum[1]) * scale_kv_0;
    v_1 = (sum[2] + sum[3]) * scale_kv_1;
    // inter-thread reduction + store
}

// Block table 映射逻辑 block → 物理 block
const auto physical_block_idx = block_table[logical_block_idx];
tma::copy<kHeadDim, BLOCK_KV, kHeadDim>(&tensor_map_kv, ..., 0, physical_block_idx * BLOCK_KV);

// Metadata kernel: 计算每个 SM 的任务分配
template <uint32_t kAlignedBatchSize, uint32_t SPLIT_KV, uint32_t kNumSMs>
void smxx_paged_mqa_logits_metadata(batch_size, next_n, context_lens, schedule_metadata) {
    // 1. 计算每个 Q 的 KV segment 数
    // 2. Prefix sum 得到全局 segment 偏移
    // 3. 按 SM 均分 segments
    schedule_metadata[sm_idx * 2] = q_atom_idx;
    schedule_metadata[sm_idx * 2 + 1] = kv_split_idx;
}

// 1. 初始化 shared memory 为 -inf
for (i = threadIdx.x; i < BLOCK_KV; i += kNumWarps * 32)
    smem_buffer[i] = neg_inf;

// 2. 对无效区域 bulk copy -inf
if (right <= ks or ke <= left)
    SM90_BULK_COPY_S2G::copy(smem_buffer, logits + i * stride_logits + left, ...);

// 3. 处理对齐边界的逐元素写入
for (j = aligned_ks; j < ks; j++) logits[i * stride_logits + j] = neg_inf;

flowchart TB
    START["get_best_config()"] --> ENUM["枚举 Layout 候选<br>(block_m, block_n, cluster)"]

    ENUM --> FILTER["过滤不合法配置<br>- swizzle 对齐<br>- TMEM 容量<br>- multicast 整除"]

    FILTER --> MODEL["建模性能<br>- L1/L2 带宽<br>- wave 效率<br>- 计算周期"]

    MODEL --> COMPARE["比较排序<br>选择最优 Layout"]

    COMPARE --> STORAGE["推导 StorageConfig<br>(swizzle modes)"]

    STORAGE --> PIPELINE["推导 PipelineConfig<br>(num_stages from smem)"]

    PIPELINE --> LAUNCH["推导 LaunchConfig<br>(threads, cluster)"]

    LAUNCH --> OUTPUT["GemmConfig"]

// Block M 候选
if (Normal/Batched/KGrouped) block_m ∈ {64, 128, (256 if BF16)};
if (MGrouped) block_m = mk_alignment_for_contiguous_layout;

// Block N 候选
step = lcm(16, block_n_multiple_of);
start = (1D1D && FP32) ? 24 : step;  // 避免 bank conflict
end = (1D2D) ? 192 : (1D1D) ? 160 : 256;

// 关键过滤条件
- swizzle_mode % 64 == 0  (32B swizzle 性能差)
- num_stages >= 3 (小矩阵 >= 4)
- block_m > 128 && block_n > 128 → 拒绝（寄存器不足）

// M-Grouped: 强制 Swap AB
swap_ab = true;
block_n = 128;  // UMMA layout A/D 的 M 维度
block_m = mk_alignment_for_contiguous_layout;  // 最大 240

// Normal: 双向 Swap AB 探索
for (swap_ab in {0, 1}) {
    if (swap_ab) {
        block_m_candidates: step 16, end 256
        block_n = 128;
    } else {
        block_m ∈ {32, 64, 128}
        block_n: step 32, end (k≤256) ? 128 : 256
    }
}

// TMEM 容量检查
if (2 * umma_n + tmem_sf_cols > 512) reject;

// Swizzle 要求
if (FP4) swizzle_requirement = 128;
else swizzle_requirement = 64;

// 数据移动量
num_bytes_l2_ab = expected_k * (block_m/cluster_n + block_n/cluster_m) * elem_size;
num_bytes_l1_tc = expected_k * (max(64, block_m) + block_n) * elem_size
                + block_m * block_n * elem_size_cd;

// 周期估算
num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bw_per_cycle;
num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + num_bytes_l1_l2_cd) * num_blocks / l1_bw_per_cycle;
num_cycles = max(num_l1_cycles, num_l2_cycles) / wave_efficiency;

flowchart TB
    subgraph "SM90 寄存器分配"
        TMA_REG["TMA Warps<br>24-40 regs<br>(dealloc)"]
        MMA_REG["Math Warps<br>232-240 regs<br>(alloc)"]
    end

    subgraph "SM100 寄存器分配"
        TMA_REG2["TMA Warp<br>32 regs"]
        MMA_REG2["MMA Warps<br>128 regs"]
        EPI_REG["Epilogue Warps<br>128 regs"]
    end

// TMA warps 释放寄存器给 Math warps
cutlass::arch::warpgroup_reg_dealloc<40>();   // 仅保留 40 regs
cutlass::arch::warpgroup_reg_alloc<232>();    // 申请 232 regs

// TMEM 列数对齐到 32/64/128/256/512
template <uint32_t kNumCols>
constexpr uint32_t get_num_aligned_tmem_cols() {
    if (kNumCols <= 32) return 32;
    if (kNumCols <= 64) return 64;
    if (kNumCols <= 128) return 128;
    if (kNumCols <= 256) return 256;
    return 512;
}

flowchart LR
    subgraph "Stage 1: GMEM → SMEM"
        TMA_A["TMA Load A"] --> SMEM_A["smem_a[stage]"]
        TMA_B["TMA Load B"] --> SMEM_B["smem_b[stage]"]
        TMA_SF["TMA Load SF"] --> SMEM_SF["smem_sfa/sfb[stage]"]
    end

    subgraph "Stage 2: SMEM → MMA"
        SMEM_A -->|"SmemDescriptor"| MMA["WGMMA/UMMA"]
        SMEM_B -->|"SmemDescriptor"| MMA
        SMEM_SF -->|"UTCCP (SM100)"| TMEM_SF["TMEM SF"]
    end

    subgraph "Stage 3: Accumulate"
        MMA -->|"SM90: 寄存器"| ACC["accum"]
        MMA -->|"SM100: TMEM"| TMEM_ACC["TMEM accum"]
    end

    subgraph "Stage 4: Write-back"
        ACC -->|"SM90"| STSM["STSM → SMEM"]
        TMEM_ACC -->|"SM100"| TMEM_LOAD["TMEM Load → regs"]
        TMEM_LOAD --> STSM
        STSM --> TMA_STORE["TMA Store → GMEM"]
    end

// Pipeline stage 管理
auto get_pipeline = [=](iter_idx) -> tuple<stage, phase> {
    return {iter_idx % kNumStages, (iter_idx / kNumStages) & 1};
};

// full_barrier: TMA 写入完成 → MMA 可消费
// empty_barrier: MMA 消费完成 → TMA 可覆盖
full_barriers[stage]->wait(phase);     // MMA 等待 TMA
empty_barriers[stage]->wait(phase^1);  // TMA 等待 MMA

// SM90: 256-bit STSM (4 × float32)
SM90_U32x4_STSM_T<float4>::copy(r0, r1, r2, r3, smem_dst);

// SM100: 128-bit STSM (8 × float32 → BF16)
SM100_U8x8_STSM_T<uint2>::copy(pack(v0,v1), pack(v2,v3), smem_dst);

// 向量化 shared memory 写入
st_shared(smem_ptr, values[0], values[1], values[2], values[3]);

flowchart LR
    A["Block Size 选择"] --> B["Swizzle 模式"]
    B --> C["Stage 数量"]
    C --> D["TMA/MMA 重叠"]
    D --> E["寄存器压力"]
    E --> F["Occupancy"]
    F --> G["最终性能"]