feat(09): gate register prefetch via REG_STAGES, fit two CTAs per SM

07-swizzling/cutedsl_swizzling.py

@@ -81,7 +81,6 @@ def swizzling_kernel(
     sB_layout: cute.ComposedLayout,
     sC_layout: cute.ComposedLayout,
     sO_layout: cute.ComposedLayout,
-    acc_dtype: cutlass.Constexpr,
     out_dtype: cutlass.Constexpr,
     is_gemm: cutlass.Constexpr[bool],
 ):
@@ -144,16 +143,16 @@ def swizzling_kernel(
     else:
         # Load C from smem in its native dtype (mC.element_type), then convert
         # into the accumulator fragment. The intermediate is required when
-        # acc_dtype differs from mC.element_type (e.g. fp16->fp32): the DSL's
-        # ``cute.copy`` requires source/destination bit widths to match, so
-        # the dtype change has to happen on a register-to-register ``.to()``
-        # step. When acc_dtype == mC.element_type this is a no-op cast.
+        # the accumulator dtype differs from mC.element_type (e.g. fp16->fp32):
+        # the DSL's ``cute.copy`` requires source/destination bit widths to
+        # match, so the dtype change has to happen on a register-to-register
+        # ``.to()`` step. When acc dtype == mC.element_type this is a no-op cast.
         thr_s2r_c = s2r_tiled_copy_c.get_slice(tid)
         tCrC_pre = cute.make_fragment_like(tCrC, mC.element_type)
         tCsC_s2r = thr_s2r_c.partition_S(sC)
         tCrC_s2r = thr_s2r_c.retile(tCrC_pre)
         cute.copy(s2r_tiled_copy_c, tCsC_s2r, tCrC_s2r)
-        tCrC.store(tCrC_pre.load().to(acc_dtype))
+        tCrC.store(tCrC_pre.load().to(tCrC.element_type))
 
     # ----- Phase 3: compute -----
     cute.gemm(tiled_mma, tCrC, tCrA, tCrB, tCrC)
@@ -328,7 +327,6 @@ def swizzling_gemm(
         sB_layout,
         sC_layout,
         sO_layout,
-        acc_dtype,
         out_dtype,
         is_gemm,
     ).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream)
@@ -428,7 +426,7 @@ def main() -> None:
     print(" Compiling fp16 in / fp32 acc / fp16 out ... ".center(PRINT_LENGTH, "-"))
     a_t = torch.empty(M, K, device="cuda", dtype=torch.float16)
     b_t = torch.empty(N, K, device="cuda", dtype=torch.float16)
-    c_t = torch.empty(M, N, device="cuda", dtype=torch.float16)
+    c_t = torch.empty(M, N, device="cuda", dtype=torch.float32)
     o_t = torch.empty(M, N, device="cuda", dtype=torch.float16)
     fp16f32_clear, fp16f32_accum = _compile_pair(
         a_t,
@@ -470,22 +468,22 @@ def main() -> None:
     print(f" M={M}, N={N}, K={K} ".center(PRINT_LENGTH, "-"))
 
     # ----- Spec 1: fp16 = fp16 * fp16 + fp32 (validated) -----
-    print(" fp16 = fp32_acc(fp16 * fp16) + fp16 ".center(PRINT_LENGTH, "="))
+    print(" fp16 = fp32_acc(fp16 * fp16) + fp32 ".center(PRINT_LENGTH, "="))
     torch.cuda.manual_seed_all(9527)
     a = torch.randn(M, K, device="cuda", dtype=torch.float16)
     b = torch.randn(N, K, device="cuda", dtype=torch.float16)
     c = torch.randn(M, N, device="cuda", dtype=torch.float32)
 
     # Case 1: MM (fp16 = fp16 * fp16)
     out = torch.empty(M, N, device="cuda", dtype=torch.float16)
-    fp16f32_clear(a, b, c.clone().half(), out)
+    fp16f32_clear(a, b, c.clone(), out)
     torch.cuda.synchronize()
     # For fp16 input, torch.matmul uses fp32 as the accumulator precision
-    compare_matrix(out, torch.matmul(a.float(), b.T.float()).half(), counters)
+    compare_matrix(out, torch.matmul(a, b.T), counters)
 
     # Case 2: MMA (fp16 = fp16 * fp16 + fp16)
     out = torch.empty(M, N, device="cuda", dtype=torch.float16)
-    fp16f32_accum(a, b, c.clone().half(), out)
+    fp16f32_accum(a, b, c.clone(), out)
     torch.cuda.synchronize()
     compare_matrix(out, torch.addmm(c, a.float(), b.T.float()).half(), counters)
 

08-dynamic-mma/cutedsl_dynamic_mma.py

@@ -70,6 +70,9 @@ def dynamic_mma_kernel(
     g2s_tiled_copy_a: cute.TiledCopy,
     g2s_tiled_copy_b: cute.TiledCopy,
     g2s_tiled_copy_c: cute.TiledCopy,
+    s2r_tiled_copy_a: cute.TiledCopy,
+    s2r_tiled_copy_b: cute.TiledCopy,
+    s2r_tiled_copy_c: cute.TiledCopy,
     r2s_tiled_copy_o: cute.TiledCopy,
     s2g_tiled_copy_o: cute.TiledCopy,
     sA_layout: cute.ComposedLayout,
@@ -133,12 +136,12 @@ def dynamic_mma_kernel(
     tBcB = thr_g2s_b.partition_S(cB)  # (CPY, CPY_N, CPY_K)
     tCcC = thr_g2s_c.partition_S(cC)  # (CPY, CPY_M, CPY_N)
 
-    # ----- M / N predicate tensors (independent of ik) -----
-    # Layout mirrors tensorop_gemm.py: the first mode is the CPY sub-mode
-    # (``tAgA.shape[0][1]`` — separate booleans for elements *within* a
-    # copy atom), the second mode is the tile's CPY_M, and the third is
-    # CPY_K with stride 0 so the same M/N predicate is replayed across
-    # every K coordinate in the tile.
+    # ----- M / N predicate (used for ik >= 1, i.e. all but the residue tile) -----
+    # Three-mode layout (rest_v, CPY_M, CPY_K) with the K mode broadcast at
+    # stride 0 so the same M/N predicate is replayed across every K iter.
+    # K-bound is not needed here because the domain_offset shift guarantees
+    # every K-tile from ik=1 onward is in-bounds; only the residue tile
+    # (handled by ``tApA_first`` below) needs the per-element K check.
     tApA = cute.make_rmem_tensor(
         cute.make_layout(
             (
@@ -178,7 +181,7 @@ def dynamic_mma_kernel(
     # No dynamic loop here so we can inline a per-(m, n) bound check.
     # Out-of-bounds slots are left at zero from the smem fill.
     if cutlass.const_expr(not is_gemm):
-        thr_g2s_c.partition_D(sC).fill(0)
+        tCsC.fill(0)
         for m in cutlass.range_constexpr(cute.size(tCgC, mode=[1])):
             for n in cutlass.range_constexpr(cute.size(tCgC, mode=[2])):
                 if cute.elem_less(tCcC[0, m, n][0], m_max) and cute.elem_less(
@@ -197,37 +200,19 @@ def dynamic_mma_kernel(
     tCrB = tiled_mma.make_fragment_B(thr_mma.partition_B(sB))
     tCrC = tiled_mma.make_fragment_C(thr_mma.partition_C(gC))
 
-    # S2R partitions reuse the same tiled_mma so the fragment layouts match
-    # the MMA atom directly (no explicit ldmatrix TiledCopy needed for this
-    # demo — the compiler will lower the universal copy onto LDS / LDSM).
-    universal = cute.nvgpu.CopyUniversalOp()
-    s2r_atom_a = cute.make_copy_atom(universal, mA.element_type)
-    s2r_atom_b = cute.make_copy_atom(universal, mB.element_type)
-    s2r_atom_c = cute.make_copy_atom(universal, mC.element_type)
-    s2r_tiled_copy_a = cute.make_tiled_copy_A(s2r_atom_a, tiled_mma)
-    s2r_tiled_copy_b = cute.make_tiled_copy_B(s2r_atom_b, tiled_mma)
-    s2r_tiled_copy_c = cute.make_tiled_copy_C(s2r_atom_c, tiled_mma)
-
     thr_s2r_a = s2r_tiled_copy_a.get_slice(tid)
     thr_s2r_b = s2r_tiled_copy_b.get_slice(tid)
     thr_s2r_c = s2r_tiled_copy_c.get_slice(tid)
 
     # ----- Mainloop: per K-tile predicated G2S, then S2R + MMA -----
     num_k_tiles = cute.size(tAgA, mode=[3])
 
-    # First K-tile (ik=0) needs per-element K-bound checks because the
-    # domain_offset shift can make some elements have negative gmem K
-    # coords. Build a per-element pred (rest_v, CPY_M, CPY_K with all
-    # strides distinct) that combines the M-bound and the K-bound check
-    # ``tile_k + k_residue >= 0`` for each per-thread CPY element.
-    #
-    # This per-element pred intentionally deviates from the C++ iteration-
-    # level gate ``if (get<1>(tAcA_g2s(0,0,k)) >= -k_residue)`` because
-    # that gate silently drops valid threads when K < BLK_K — a single ik
-    # can carry threads whose individual K-positions span the entire
-    # BLK_K range, some valid and some not. The C++ harness would fail
-    # tiny-K shapes for the same reason; we replace with a per-element
-    # pred (M-bound AND K-bound) so every valid thread fires its cp.async.
+    # Residue tile (ik = 0) needs a per-element predicate combining M-bound
+    # AND K-bound. The per-element K check (rather than the C++ iteration-
+    # level ``thread-0 K coord >= -k_residue`` gate) is required because
+    # with CPY_K = 1 a single iter carries lanes spread across the whole
+    # BLK_K range — some at valid in-bounds K-positions, others not — so the
+    # gate must fire per lane / per val element rather than per iter.
     tApA_first = cute.make_rmem_tensor(
         cute.make_layout(
             (
@@ -279,10 +264,19 @@ def dynamic_mma_kernel(
                     tBcB[(0, rest_v), n, k][1] + k_residue,
                 )
 
-    # Zero the smem tile so any predicated-off slot reads as zero.
+    # Zero the A/B smem ONCE before the K-loop (hoisted out of the mainloop).
+    # The G2S copy is predicated, so predicated-off slots are never written by
+    # cp.async; pre-zeroing lets them read as 0. A single clear suffices: the
+    # M-boundary rows masked off by the M-only predicate are never written by
+    # any K-tile, and the ik=0 K-residue columns are overwritten by every
+    # later (full) K-tile.
+    #
+    # No sync between the fill and the cp.async: each thread's fill targets the
+    # same g2s partition slots that its cp.async then writes, so program order
+    # within a thread suffices. The post-wait_group sync_threads below is what
+    # publishes both the zeros and the cp.async data to the s2r readers.
     tAsA.fill(0)
     tBsB.fill(0)
-    cute.arch.sync_threads()
     cute.copy(
         g2s_tiled_copy_a,
         tAgA[None, None, None, 0],
@@ -332,8 +326,6 @@ def dynamic_mma_kernel(
     # Remaining K-tiles (ik = 1..num_k_tiles-1) — full M/N predication only,
     # K is guaranteed in-bounds by the residue shift.
     for ik in cutlass.range(num_k_tiles - 1, unroll_full=False):
-        tAsA.fill(0)
-        tBsB.fill(0)
         cute.copy(
             g2s_tiled_copy_a,
             tAgA[None, None, None, ik + 1],
@@ -506,14 +498,32 @@ def dynamic_mma_gemm(
     g2s_atom_c = cute.make_copy_atom(g2s_op, mC.element_type, num_bits_per_copy=128)
     g2s_tiled_copy_c = cute.make_tiled_copy_tv(g2s_atom_c, tlC_thr, tlC_val)
 
+    universal = cute.nvgpu.CopyUniversalOp()
+
+    # ----- S2R copies (ldmatrix for 16-bit operands) -----
+    # A/B own 4 32-bit packets per thread (VAL_EXPAND_K=2), so the x4 ldmatrix
+    # variant fits exactly. C has no K val-expand, so each thread only owns
+    # 2 32-bit packets — use the x2 ldmatrix variant when C is 16-bit. For
+    # wider C (e.g. fp32) fall back to a universal copy.
+    ldm_op_ab = cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4)
+    ldm_op_c = cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2)
+    s2r_atom_a = cute.make_copy_atom(ldm_op_ab, mA.element_type)
+    s2r_atom_b = cute.make_copy_atom(ldm_op_ab, mB.element_type)
+    s2r_tiled_copy_a = cute.make_tiled_copy_A(s2r_atom_a, tm)
+    s2r_tiled_copy_b = cute.make_tiled_copy_B(s2r_atom_b, tm)
+    if cutlass.const_expr(mC.element_type.width == 16):
+        s2r_atom_c = cute.make_copy_atom(ldm_op_c, mC.element_type)
+    else:
+        s2r_atom_c = cute.make_copy_atom(universal, mC.element_type)
+    s2r_tiled_copy_c = cute.make_tiled_copy_C(s2r_atom_c, tm)
+
     # ----- R2S + S2G copies for the smem-staged epilogue -----
     # R2S: MMA-derived TiledCopy so the per-thread register fragment lands
     # at its natural smem position under the swizzled sO layout. Pick the
     # op the same way dynamic_mma.cu does: SM90+ with a 16-bit output uses
     # ``stmatrix`` (x2 because the C/O fragment has no K val-expand and so
     # owns 2 32-bit packets per thread); otherwise fall back to a universal
     # STS lowering.
-    universal = cute.nvgpu.CopyUniversalOp()
     sm_major, _ = torch.cuda.get_device_capability()
     if cutlass.const_expr(sm_major >= 9 and mO.element_type.width == 16):
         stm_op = cute.nvgpu.warp.StMatrix8x8x16bOp(False, 2)
@@ -550,6 +560,9 @@ def dynamic_mma_gemm(
         g2s_tiled_copy_a,
         g2s_tiled_copy_b,
         g2s_tiled_copy_c,
+        s2r_tiled_copy_a,
+        s2r_tiled_copy_b,
+        s2r_tiled_copy_c,
         r2s_tiled_copy_o,
         s2g_tiled_copy_o,
         sA_layout,
@@ -562,6 +575,15 @@ def dynamic_mma_gemm(
         grid=(grid_n, grid_m, 1),
         block=(NUM_THREADS, 1, 1),
         stream=stream,
+        # nvvm.minctasm hint: require >= 2 resident blocks per SM. Left
+        # unconstrained the CuTeDSL register allocator is ILP-greedy and
+        # inflates to ~150 registers/thread (not a spill -- just a loose
+        # allocation), which caps occupancy at one block per SM on
+        # H100/H200. The hint tightens the budget so the kernel fits at
+        # 128 registers/thread with zero spill, roughly doubling achieved
+        # occupancy. 2 is the largest spill-free value here -- 3+ would
+        # force register spills to local memory.
+        min_blocks_per_mp=2,
     )
 
 
@@ -684,7 +706,7 @@ def main() -> None:
     print(" Compiling fp16 in / fp32 acc / fp16 out ... ".center(PRINT_LENGTH, "-"))
     a_t = torch.empty(M0, K0, device="cuda", dtype=torch.float16)
     b_t = torch.empty(N0, K0, device="cuda", dtype=torch.float16)
-    c_t = torch.empty(M0, N0, device="cuda", dtype=torch.float16)
+    c_t = torch.empty(M0, N0, device="cuda", dtype=torch.float32)
     o_t = torch.empty(M0, N0, device="cuda", dtype=torch.float16)
     fp16f32_clear, fp16f32_accum = _compile_pair(
         a_t,
@@ -715,6 +737,7 @@ def main() -> None:
     # NOT validated against torch — fp16 accumulation drops too many bits
     # for large K to match torch's fp32-accumulated reference.
     print(" fp16 = fp16 * fp16 + fp16 (exercise only) ".center(PRINT_LENGTH, "="))
+    torch.cuda.manual_seed_all(9527)
     for m, n, k in exps:
         print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-"))
         a = torch.randn(m, k, device="cuda", dtype=torch.float16)
@@ -728,12 +751,13 @@ def main() -> None:
         torch.cuda.synchronize()
 
     # ----- Sweep: fp16 in, fp32 acc, fp16 out (Spec 2) -----
-    print(" fp16 = fp32_acc(fp16 * fp16) + fp16 ".center(PRINT_LENGTH, "="))
+    print(" fp16 = fp32_acc(fp16 * fp16) + fp32 ".center(PRINT_LENGTH, "="))
+    torch.cuda.manual_seed_all(9527)
     for m, n, k in exps:
         print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-"))
         a = torch.randn(m, k, device="cuda", dtype=torch.float16)
         b = torch.randn(n, k, device="cuda", dtype=torch.float16)
-        c = torch.randn(m, n, device="cuda", dtype=torch.float16)
+        c = torch.randn(m, n, device="cuda", dtype=torch.float32)
 
         out = torch.empty(m, n, device="cuda", dtype=torch.float16)
         fp16f32_clear(a, b, c.clone(), out)
@@ -743,10 +767,11 @@ def main() -> None:
         out = torch.empty(m, n, device="cuda", dtype=torch.float16)
         fp16f32_accum(a, b, c.clone(), out)
         torch.cuda.synchronize()
-        compare_matrix(out, torch.addmm(c.float(), a.float(), b.T.float()).half(), counters)
+        compare_matrix(out, torch.addmm(c, a.float(), b.T.float()).half(), counters)
 
     # ----- Sweep: bf16 in, fp32 acc, bf16 out (Spec 3) -----
     print(" bf16 = fp32_acc(bf16 * bf16) + fp32 ".center(PRINT_LENGTH, "="))
+    torch.cuda.manual_seed_all(9527)
     for m, n, k in exps:
         print(f" M={m}, N={n}, K={k} ".center(PRINT_LENGTH, "-"))
         a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)

08-dynamic-mma/dynamic_mma.cu

@@ -170,10 +170,20 @@ __global__ __launch_bounds__(Spec::kThreadNum) void dynamic_mma(void *__restrict
 
   int NTilesK = ceil_div(K, kBlockK);
 
+  // Zero the A/B smem ONCE before the K-loop (hoisted out of the mainloop).
+  // The predicated cp.async never writes masked-off slots, so they must read
+  // as 0. A single clear suffices: the M-boundary rows masked off by the
+  // M-only predicate are never written by any K-tile, and the ik=0 K-residue
+  // columns are overwritten by every later (full) K-tile.
+  //
+  // No sync between the clear and the cp.async: each thread's clear targets
+  // the same g2s partition slots that its cp.async then writes, so program
+  // order within a thread suffices. The __syncthreads() after cp_async_wait
+  // below publishes both the zeros and the cp.async data to the s2r readers.
+  clear(tAsA_g2s);
+  clear(tBsB_g2s);
+
   for (int ik = 0; ik < NTilesK; ++ik) {
-    // Clear the smem tiles to account for predicated off loads
-    clear(tAsA_g2s);
-    clear(tBsB_g2s);
     if (ik == 0) {
 #pragma unroll
       for (int k = 0; k < size<2>(tAsA_g2s); ++k) {
@@ -352,11 +362,53 @@ struct KernelSpec {
 
   using TiledMMA = decltype(make_tiled_mma(MMA_op{}, MMAThrLayout{}, MMATileLayout{}));
 
-  // 为什么直接用 AutoVectorizingCopy 会出现 RuntimeError: CUDA error: misaligned address？
-  // using Copy_G2S_op = AutoVectorizingCopy; -> cp.async.ca.shared.global.L2::128B
-  // 这样也 OK！
-  // using Copy_G2S_op = SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>;
-  // 但还是 cp.async.cg 最好
+  // Why AutoVectorizingCopy faults under copy_if (CUDA error 716, "misaligned
+  // address"):
+  //
+  //   `AutoVectorizingCopyWithAssumedAlignment<MaxBits>` inherits from
+  //   `UniversalCopy<uint_bit_t<MaxBits>>`, so `AutoVectorizingCopy` (=
+  //   `<128>`) is structurally a 128-bit atom applied to whatever element
+  //   type the tensor has.
+  //
+  //   In `cute/algorithm/copy.hpp`, the `copy()` overload for AutoVec does a
+  //   `recast<uint_bit_t<vec_bits>>` of src/dst BEFORE issuing the atom:
+  //
+  //       copy(AutoVec<N>, src, dst)
+  //         -> recast<uintN>(src/dst)        // fp16 -> uint128_t view
+  //         -> copy_if(true, src_v, dst_v)   // 1 atom call per iter
+  //
+  //   But `copy_if(Copy_Atom<AutoVec<N>, T>, pred, src, dst)` has NO matching
+  //   recast specialization. It falls through to the generic Copy_Atom
+  //   path that unrolls atom calls over the per-thread tile WITHOUT
+  //   recasting. With val (1, 8) of fp16 we get 8 atom calls at stride
+  //   1 fp16 (= 2 B), each invoking the 128-bit atom:
+  //
+  //       ld.global.nc.v2.u64 [base + 0]    aligned
+  //       ld.global.nc.v2.u64 [base + 2]    misaligned -> fault
+  //       ld.global.nc.v2.u64 [base + 4]    misaligned
+  //       ...
+  //       ld.global.nc.v2.u64 [base + 14]   misaligned
+  //
+  //   These loads also have NO @p in PTX: copy_if's `if (pred)` becomes one
+  //   coarse `setp + @p bra` gating the whole block. NVCC faithfully lowers
+  //   exactly what CUTLASS asked for; the broken stride is at the CUTLASS
+  //   layer, not at NVCC.
+  //
+  // Tracked upstream: NVIDIA/cutlass#2354 ("Missing copy_if implementation
+  // for AutoVectorizingCopyWithAssumedAlignment"). As of CUTLASS main the
+  // bug is still present -- no copy_if(AutoVec<N>, ...) specialization has
+  // been merged. The official workaround in
+  // cutlass/examples/cute/tutorial/tiled_copy_if.cu is to bake the width
+  // into the atom type explicitly:
+  //
+  //   using CopyOp = UniversalCopy<uint_byte_t<sizeof(T) * size(val_layout)>>;
+  //
+  // SM80_CP_ASYNC_CACHEGLOBAL<uint128_t> follows the same principle: its
+  // atom is 16 B intrinsically, so val (1, 8) collapses to one atom call per
+  // iter -> one aligned `cp.async.cg.shared.global [smem], [gmem], 16` per
+  // thread. It also drives the cp_async_fence / cp_async_wait pipeline
+  // (AutoVec is sync ld+st, the fences become no-ops with respect to data
+  // motion).
   using Copy_G2S_op = SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>;
 
   // Note: ldmatrix only support 16-bit data type (or below)