A warp MMA op describes the hardware ``mma.sync.aligned`` instruction to use,
it has parameters like data types and instruction shape. The operand layout is
fixed (A = row-major, B = col-major).

.. code-block:: python

   import cutlass
   import cutlass.cute as cute
   from cutlass.cute.nvgpu import warp

   op = warp.MmaF16BF16Op(
       cutlass.Float16,     # A/B element type
       cutlass.Float32,     # accumulator type
       (16, 8, 16),         # instruction shape (M, N, K)
   )

The key parameters are:

- **Instruction shape** ``(M, N, K)``: determines the size of one hardware MMA
  instruction. Valid shapes depend on the data type (see ops table below).
- **A/B element type** (``ab_dtype``) and **accumulator type** (``acc_dtype``):
  ``Float32`` is always a valid accumulator; ``Float16`` is only valid for F16
  inputs. Each op restricts ``ab_dtype`` to a specific family (F16/BF16, FP8,
  MXF4, etc.).
- **Operand layout**: fixed to A = row-major (K-major), B = col-major (K-major).
  Transpose is not supported. All 32 threads in a warp cooperate on each
  instruction.


CuTe DSL provides implementation of many warp-level MMA ops:

.. list-table:: warp-level MMA ops
   :header-rows: 1
   :widths: 34 22 34 10

   * - PTX name
     - Python class
     - Constructor parameters
     - SM Arch
   * - ``mma.sync.aligned.m16n8k{K}.row.col.{acc}.f16.f16`` / ``.bf16.bf16``
     - ``warp.MmaF16BF16Op``
     - ``ab_dtype, acc_dtype, shape_mnk``
     - ``sm_80+``
   * - ``mma.sync.aligned.m16n8k{K}.row.col.{acc}.{e4m3|e5m2}.{e4m3|e5m2}``
     - ``warp.MmaFP8Op``
     - ``ab_dtype, acc_dtype, shape_mnk``
     - ``sm_89+``
   * - ``mma.sync.aligned.kind::mxf4.block_scale.m16n8k64``
     - ``warp.MmaMXF4Op``
     - ``ab_dtype, acc_dtype, sf_type``
     - ``sm_120a+``
   * - ``mma.sync.aligned.kind::mxf4nvf4.block_scale.m16n8k64``
     - ``warp.MmaMXF4NVF4Op``
     - ``ab_dtype, acc_dtype, sf_type``
     - ``sm_120a+``

.. {$nv-internal-release begin}

Internal builds additionally provide:

.. list-table:: Internal warp-level MMA ops
   :header-rows: 1
   :widths: 34 22 34 10

   * - PTX name
     - Python class
     - Constructor parameters
     - SM Arch
   * - ``mma.sp.sync.aligned.m16n8k{K}.row.col.{acc}.f16.f16`` / ``.bf16.bf16``
     - ``warp.MmaF16BF16SparseOp``
     - ``ab_dtype, acc_dtype, shape_mnk, sparse_metadata_format``
     - ``sm_80+``

.. {$nv-internal-release end}

Creating a Tiled MMA
~~~~~~~~~~~~~~~~~~~~~

A ``TiledMma`` tiles the MMA atom across the thread block so that multiple
warps cooperate on a larger tile. You can pass the op directly or create an
explicit atom first:

.. code-block:: python

  # Option 1: directly from op (common shorthand)
  tiled_mma = cute.make_tiled_mma(op)

  # Option 2: explicit atom creation
  atom = cute.make_mma_atom(op)
  tiled_mma = cute.make_tiled_mma(atom)

With no extra arguments this wraps a single atom — one warp, one
``(16, 8, K)`` tile. The optional ``atom_layout_mnk`` and
``permutation_mnk`` parameters (described in the subsections below)
control multi-warp tiling and per-thread value layout respectively.

Spatial tiling with a repeat count

Warp MMA Atom (16x8x16)                make_tiled_mma(atom, (2, 2, 1))
+----------------+                     +----------------+----------------+
|                |                     |                |                | ^
|   16 x 8       |                     | Atom (0,0,0)   | Atom (0,1,0)   | |
|     x 16       |   --(2,2,1)-->      |   16 x 8       |   16 x 8       | | 2 x inst_M
|                |      repeat         |     x 16       |     x 16       | |  = 32
|                |                     | [Warp 0]       | [Warp 2]       | |
+----------------+                     +----------------+----------------+ |
                                      |                |                | |
                                      | Atom (1,0,0)   | Atom (1,1,0)   | |
                                      |   16 x 8       |   16 x 8       | |
                                      |     x 16       |     x 16       | |
                                      | [Warp 1]       | [Warp 3]       | v
                                      +----------------+----------------+
                                      <--- 2 x inst_N = 16 --->
                                      K unchanged = 16

``permutation_mnk`` is an optional third argument to ``make_tiled_mma``.
Each of its three entries is a **per-mode permutation** of the M, N, and
K coordinates inside the tiled MMA. In the common case shown in this
section, each entry is just a size, which is the identity permutation of
that size; in that case ``permutation_mnk`` simply sets the **total tile
footprint** of the tiled MMA along each dimension. When a mode's size is
larger than the atom layout's natural coverage
(``atom_layout x inst_shape``), each thread receives additional values to
fill the extended region — the thread count stays the same, but every
thread holds more data. The general form, where an entry is a
``Layout`` that reorders coordinates inside a mode, is covered in the
subsection below.

The standard convention for warp MMA (used in ``tensorop_gemm.py`` and
throughout the Ampere examples) doubles the N dimension:

.. code-block:: python

   # From examples/cute/ampere/kernel/dense_gemm/tensorop_gemm.py
   permutation_mnk = (
       atom_layout_mnk[0] * mma_inst_shape[0],      # M: matches atom coverage
       atom_layout_mnk[1] * mma_inst_shape[1] * 2,   # N: 2x atom coverage
       atom_layout_mnk[2] * mma_inst_shape[2],        # K: matches atom coverage
   )

   tC = cute.make_layout(atom_layout_mnk)
   tiled_mma = cute.make_tiled_mma(
       op,
       tC,
       permutation_mnk=permutation_mnk,
   )

**Why double N?** The atom's N dimension is only 8 (inst_N = 8). Without
a permutation, each thread's B-operand values span a single 8-wide
N-range, which may not align well with SMEM load widths. The ``* 2``
on N gives each thread's B fragment two 8-wide N-ranges instead of one,
aligning the access pattern with wider contiguous SMEM regions for more
efficient loads.

For ``atom_layout_mnk = (2, 2, 1)`` and ``inst_shape = (16, 8, 16)``:

- Atom coverage = ``(2x16, 2x8, 1x16) = (32, 16, 16)``
- ``permutation_mnk = (32, 32, 16)`` — N extended from 16 to 32

.. code-block:: text

  Without permutation — natural atom coverage (M = 32, N = 16):

  C tile (M=32, N=16)
  +----------------+----------------+
  |                |                | ^
  |   [Warp 0]     |   [Warp 2]     | |
  |    16 x 8      |    16 x 8      | | 2 x inst_M
  |                |                | |  = 32
  +----------------+----------------+ |
  |                |                | |
  |   [Warp 1]     |   [Warp 3]     | |
  |    16 x 8      |    16 x 8      | |
  |                |                | v
  +----------------+----------------+
  <------------- N = 16 ---------->
  (each warp owns one (16, 8) atom;
   thread T0 of Warp 0 holds 4 C values in its 16x8 block)

  With permutation_mnk = (32, 32, 16) — N extended from 16 to 32:

  C tile (M=32, N=32)
  +----------------+----------------+----------------+----------------+
  |                |                |                |                | ^   N = 16 → 32:
  |   [Warp 0]     |   [Warp 2]     |   [Warp 0]     |   [Warp 2]     | |   atom pattern repeats
  |    16 x 8      |    16 x 8      |    16 x 8      |    16 x 8      | |   along N. Each thread
  |                |                |                |                | |   now holds 2x the
  +----------------+----------------+----------------+----------------+ |   values along N
  |                |                |                |                | |   (same threads, more
  |   [Warp 1]     |   [Warp 3]     |   [Warp 1]     |   [Warp 3]     | |   values per thread).
  |    16 x 8      |    16 x 8      |    16 x 8      |    16 x 8      | |
  |                |                |                |                | v
  +----------------+----------------+----------------+----------------+
  <---------------------------- N = 32 ---------------------------->
  |        atom coverage            |          value repeat           |


Reordering coordinates with a per-mode ``Layout``
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

So far each entry of ``permutation_mnk`` has been an integer, which is
shorthand for the identity layout ``Layout<Shape<S>, Stride<_1>>`` — the
atom pattern simply tiles to fill an ``S``-wide footprint. The general
form lets each entry be a ``Layout`` that **reorders coordinates inside
that mode** while keeping the same total size. That reordering is what
gives the parameter its name; the integer-only cases used earlier are
just the identity permutation.

The canonical illustration is the SM70 example from
`0t_mma_atom.md <../../cpp/cute/0t_mma_atom.md>`_. Take a 2x2 tiled MMA
of ``SM70_8x8x4_F32F16F16F32_NT`` atoms with a ``32x32x4`` footprint.
Without any M-mode permutation, thread ``T0``'s 8 A-values land at the
following ``(m, k)`` coordinates::

    T0V0 => (0, 0)     T0V4 => (16, 0)
    T0V1 => (1, 0)     T0V5 => (17, 0)
    T0V2 => (2, 0)     T0V6 => (18, 0)
    T0V3 => (3, 0)     T0V7 => (19, 0)

— two separate runs of 4 along M, with a gap from m=4 to m=15. We may
prefer those 8 values to sit in **one contiguous run** in the logical
M-coordinates (e.g. so register or SMEM layouts pack cleanly). Passing
the M-mode layout ``(4, 4, 2):(1, 8, 4)`` does exactly that: it is a
scatter permutation telling each old m-coord where to go in the new
image.

.. code-block:: text

    old m-coord:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
    new m-coord:  0  1  2  3  8  9 10 11 16 17 18 19 24 25 26 27  4  5  6  7 12 13 14 15 20 21 22 23 28 29 30 31

After the permutation, ``T0``'s 8 A-values occupy ``m = 0..7`` — one
contiguous run — and every other thread's M-values become equally
contiguous. Thread-data ownership and value counts are unchanged; only
the **mapping from values to m-coordinates** is permuted.

In CuTeDSL the permuted entry is built with ``cute.make_layout``;
identity entries stay as integers:

.. code-block:: python

   m_perm = cute.make_layout((4, 4, 2), stride=(1, 8, 4))
   tiled_mma = cute.make_tiled_mma(
       op,                                # SM70_8x8x4 NT atom
       atom_layout_mnk=(2, 2, 1),
       permutation_mnk=(m_perm, 32, 4),   # M: scatter, N/K: identity sizes
   )

The same mechanism applies to the N and K modes — any subset of the
three entries can be an integer (identity) or a ``Layout`` (real
permutation). For warp MMAs the most common case in practice is still
the integer-only form shown earlier in this section; the ``Layout`` form
is the tool you reach for when a register or SMEM layout wants each
thread's fragment to be contiguous in logical coordinates.


Partitioning Tensors
---------------------

Before computing, partition the CTA-tiled tensors according to the
tiled MMA layout. Warp MMA partitioning is **per-thread**: each of
the 32 threads in a warp (or 128 threads across 4 warps) receives
its own slice of the data, sized to match the register fragments
the MMA instruction expects.

Example: ``GEMM (M, N, K) = (512, 512, 256)``,
``cta_tiler = (128, 128, 32)``, ``atom_layout_mnk = (2, 2, 1)``,
F16 atom = m16n8k16, ``permutation_mnk = (32, 32, 16)``,
``num_stages = 4``, 4 warps = 128 threads.

Global matrices:

.. code-block:: text

     mA: (M, K) = (512, 256)       mB: (N, K) = (512, 256)       mC: (M, N) = (512, 512)

          K=256                          K=256                       N=512
        |<--------->|                |<--------->|              |<---------------->|
        +-----------+                +-----------+              +----+----+----+---+
        |           | ^              |           | ^            |    |    |    |   | ^
        |    mA     | | M=512        |    mB     | | N=512      |    |    |    |   | | M=512
        |           | v              |           | v            |    |    |    |   | v
        +-----------+                +-----------+              +----+----+----+---+

Tiling with ``cta_tiler = (BM, BN, BK) = (128, 128, 32)`` gives
M/BM = 4 tiles, N/BN = 4 tiles, K/BK = 8 tiles:

.. code-block:: text

    mA tiled into (M/BM x K/BK)    mB tiled into (N/BN x K/BK)    mC tiled into (M/BM x N/BN)
    = (4 x 8) blocks               = (4 x 8) blocks               = (4 x 4) blocks

      BK=32  x8                       BK=32  x8                       BN=128  x4
    |<-->|                          |<-->|                          |<------>|
    +----+----+-- --+               +----+----+-- --+               +--------+--------+-- --+
    |    |    |..|  | ^  BM=128     |    |    |..|  | ^  BN=128     | (0,0)  | (0,1)  |..   | ^  BM=128
    +----+----+-- --+ v             +----+----+-- --+ v             +--------+--------+     + v
    |    |    |..|  | ^  BM=128     |    |    |..|  | ^  BN=128     | (1,0)  | (1,1)  |..   | ^  BM=128
    +----+----+-- --+ v             +----+----+-- --+ v             +--------+--------+     + v
    |    |    |..|  | ^             |    |    |..|  | ^             |  ...   |  ...   |..   | ^
    +----+----+-- --+ v             +----+----+-- --+ v             +--------+--------+-- --+ v
    |    |    |..|  | ^             |    |    |..|  | ^             | (3,0)  | (3,1)  |..   | ^
    +----+----+-- --+ v             +----+----+-- --+ v             +--------+--------+-- --+ v

Each CTA picks one (M-tile, N-tile) coordinate.
For example, CTA at ``tiler_coord = (0, 1, :)``.

After ``local_tile`` — one CTA's tile (``k = K/BK = 256/32 = 8``):

.. code-block:: text

     gA: (BM, BK, k) = (128, 32, 8)   gB: (BN, BK, k) = (128, 32, 8)   gC: (BM, BN) = (128, 128)

          BK=32                             BK=32                       BN=128
        |<----->|                        |<----->|                  |<-------->|
        +-------+--                      +-------+--                +----------+
        |       |..                      |       |..                |          | ^
   BM=  |  gA   | k=8                BN= |  gB   | k=8          BM= |    gC    | | 128
    128 |       |                    128 |       |              128 |          | v
        +-------+                        +-------+                  +----------+

SMEM tensors ``sA`` and ``sB`` have a pipeline staging dimension:

.. code-block:: text

     sA: (BM, BK, PIPE) = (128, 32, 4)         sB: (BN, BK, PIPE) = (128, 32, 4)

``get_slice(tidx)`` — each thread receives its own per-thread partition.
The tiled MMA footprint is ``permutation_mnk = (32, 32, 16)``, so BM,
BN, and BK are each subdivided into MMA-sized blocks:

.. code-block:: text

        sA: partition into (MMA, MMA_M, MMA_K, PIPE)

        Each SMEM stage (BM=128, BK=32):

        perm_K perm_K                               perm_M=32
          =16    =16                                |<---->|
        |<--->|<--->|                               +------+------+------+------+
        +-----+-----+  ^                            |      |      |      |      | ^
        |  0  |  1  |  |  perm_M=32                 |  0   |  1   |  2   |  3   | | perm_N
        +-----+-----+  v                            |      |      |      |      | v  =32
        |  0  |  1  |  ^                            +------+------+------+------+
        |     |     |  |  perm_M=32                   MMA_N = BN/perm_N = 4
        +-----+-----+  v
        |  0  |  1  |  ^                           sB: partition into (MMA, MMA_N, MMA_K, PIPE)
        |     |     |  |
        +-----+-----+  v                           gC: partition into (MMA, MMA_M, MMA_N)
        |  0  |  1  |  ^
        |     |     |  |
        +-----+-----+  v
          MMA_K = BK/perm_K = 2
          MMA_M = BM/perm_M = 4

After partition (per thread, e.g. thread ``tidx``):

- ``tCsA: (MMA, MMA_M, MMA_K, PIPE) = (MMA, 4, 2, 4)`` — MMA_M = BM/perm_M = 128/32 = 4, MMA_K = BK/perm_K = 32/16 = 2
- ``tCsB: (MMA, MMA_N, MMA_K, PIPE) = (MMA, 4, 2, 4)`` — MMA_N = BN/perm_N = 128/32 = 4, MMA_K = BK/perm_K = 32/16 = 2
- ``tCgC: (MMA, MMA_M, MMA_N) = (MMA, 4, 4)`` — MMA_M = 128/32 = 4, MMA_N = 128/32 = 4

The first mode ``MMA`` contains the atom's **thread × value** layout — it
encodes which registers within a single thread hold which matrix
elements. The remaining modes are repeat counts that tile the atom
across the full CTA tile.

.. code-block:: python

   @cute.kernel
   def kernel(tiled_mma: cute.TiledMma, ...):
       tidx, _, _ = cute.arch.thread_idx()

       # CTA-tiled global tensors
       gA = cute.local_tile(mA, cta_tiler, tiler_coord, proj=(1, None, 1))
       gB = cute.local_tile(mB, cta_tiler, tiler_coord, proj=(None, 1, 1))
       gC = cute.local_tile(mC, cta_tiler, tiler_coord, proj=(1, 1, None))

       # Per-thread partition via the thread index
       thr_mma = tiled_mma.get_slice(tidx)

       # SMEM partitions (used by make_fragment_A/B and ldmatrix retiling)
       tCsA = thr_mma.partition_A(sA)   # (MMA, MMA_M, MMA_K, PIPE)
       tCsB = thr_mma.partition_B(sB)   # (MMA, MMA_N, MMA_K, PIPE)

       # C partitions for epilogue staging (SMEM) and destination (GMEM)
       tCsC = thr_mma.partition_C(sC)   # (MMA, MMA_M, MMA_N)
       tCgC = thr_mma.partition_C(gC)   # (MMA, MMA_M, MMA_N)

.. note:: The ``tCsA`` / ``tCsB`` SMEM partitions are not read directly
   by the GEMM — they establish the **shape** that
   ``make_fragment_A`` / ``make_fragment_B`` use to allocate register
   fragments. Actual SMEM→RMEM data movement goes through the S2R
   ``ldmatrix`` retiling path (see `Making Fragments`_).

Pre and Post-Conditions for Partitioning
-----------------------------------------

* The inputs of the partition should be at least rank-2 tensors.
* The output of the partition will have the layout that is compatible with the MMA atom's operand:

  - For A, the output will have the layout ``(MMA, MMA_M, MMA_K, ...)``.
  - For B, the output will have the layout ``(MMA, MMA_N, MMA_K, ...)``.
  - For C, the output will have the layout ``(MMA, MMA_M, MMA_N, ...)``.

* Note that the partition doesn't enforce any rules on the tensor's memory space or the tensor's data type. It only cares about the layout.


Making Fragments
-----------------

Fragments are the tensors that the warp MMA instruction operates on. For
warp MMA:

- **Fragment A**: per-thread register fragment holding one operand-A K-block.
- **Fragment B**: per-thread register fragment holding one operand-B K-block.
- **Fragment C (accumulator)**: per-thread register fragment that lives in
  RMEM and serves as both the input C and output D of ``cute.gemm()``.

Creating register fragments and ``ldmatrix`` copy views

The SMEM layouts define how A and B tiles are staged in shared memory before
the ``ldmatrix`` loads. For warp MMA, these layouts must satisfy two goals at
the same time:

- **Efficient GMEM -> SMEM copy:** ``cp.async`` should write contiguous
  16-byte regions for each thread.
- **Bank-conflict-free SMEM -> RMEM load:** the later ``ldmatrix`` loads
  should see a swizzled layout that matches the warp MMA operand access
  pattern.

The Ampere dense GEMM example
(``examples/cute/ampere/kernel/dense_gemm/tensorop_gemm.py``) builds these
layouts inline with a helper named ``_make_smem_layout_AB``.

**Host side** (``@cute.jit``):

.. code-block:: python

   # 16 bytes per thread for GMEM -> SMEM copies
   ab_copy_bits = 128

   sA_layout, sA_swizzle = self._make_smem_layout_AB(
       mA.element_type,       # dtype (e.g. Float16)
       self.a_major_mode,     # row-major or col-major
       ab_copy_bits,          # copy width in bits (128 = 16 bytes)
       (self.cta_tiler[0],    # BM
        self.cta_tiler[2],    # BK
        self.num_stages),     # PIPE
   )
   sB_layout, sB_swizzle = self._make_smem_layout_AB(
       mB.element_type,
       self.b_major_mode,
       ab_copy_bits,
       (self.cta_tiler[1],    # BN
        self.cta_tiler[2],    # BK
        self.num_stages),     # PIPE
   )

Here ``smem_tiler`` is ``(M_or_N, K, PIPE)``: ``(BM, BK, PIPE)`` for A and
``(BN, BK, PIPE)`` for B. The helper returns:

- ``sX_layout``: the logical SMEM layout with shape ``(BM_or_BN, BK, PIPE)``.
- ``sX_swizzle``: the swizzle applied when the tensor is materialized in SMEM.

The helper from ``tensorop_gemm.py`` implements the following four steps:

1. **Pick the major-mode size.** For a row-major operand, the contiguous
   dimension is K, so the helper uses ``smem_tiler[1]``. For a col-major
   operand, the contiguous dimension is M or N, so it uses ``smem_tiler[0]``.

2. **Cap the contiguous span at 128 bytes.** This keeps the layout atom within
   the swizzle span used by the example. The cap is 64 elements for F16/BF16
   and 128 elements for FP8.

3. **Build the swizzle.** With ``copy_bits = 128`` (16 bytes), the helper
   derives three arguments for ``make_swizzle``:

   - ``swizzle_bits = log2(major_mode_size * dtype.width / copy_bits)``,
     capped at 3. This is the number of address bits that get XOR'd.
   - ``base_bits = log2(copy_bits / 8)`` — log2 of the copy width in
     bytes (= 4 for 16-byte copies).
   - ``shift_bits = log2(copy_bits / dtype.width)`` — log2 of the copy
     width in elements (= 3 for F16 with 128-bit copies, i.e. 8 elements).

4. **Build an 8-row layout atom and tile it.** The constant 8 comes from
   ``ldmatrix``: each warp-level load touches 8 rows of shared memory
   (32 threads, 4 matrices per load). Row-major uses an atom
   ``(8, major_mode_size):(major_mode_size, 1)`` — 8 rows of contiguous
   K-elements. Col-major uses
   ``(major_mode_size, 8):(1, major_mode_size)`` — contiguous MN-elements
   across 8 K-rows. ``tile_to_shape`` then broadcasts that atom across the
   full ``(M_or_N, K, PIPE)`` SMEM tensor.

For the running F16 example (``cta_tiler = (128, 128, 32)``,
``num_stages = 4``, ``copy_bits = 128``):

.. code-block:: text

   A operand (row-major, smem_tiler = (128, 32, 4)):
     major_mode_size = 32
     atom = (8, 32):(32, 1)
     swizzle = make_swizzle(2, 4, 3)
     tiled layout -> sA: (128, 32, 4)

   B operand (col-major, smem_tiler = (128, 32, 4)):
     major_mode_size = min(128, 64) = 64
     atom = (64, 8):(1, 64)
     swizzle = make_swizzle(3, 4, 3)
     tiled layout -> sB: (128, 32, 4)


**Kernel side** (``@cute.kernel``):

The layout and swizzle are passed to shared-memory allocation:

.. code-block:: python

   @cute.struct
   class SharedStorageAB:
       a: cute.struct.Align[
           cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)],
           16,
       ]
       b: cute.struct.Align[
           cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)],
           16,
       ]

   sA = SharedStorageAB(storage).a.get_tensor(sA_layout, swizzle=sA_swizzle)
   sB = SharedStorageAB(storage).b.get_tensor(sB_layout, swizzle=sB_swizzle)

After allocation:

- ``sA`` has shape ``(BM, BK, PIPE)``.
- ``sB`` has shape ``(BN, BK, PIPE)``.

These are the staged SMEM tensors written by ``cp.async`` and later consumed by
``partition_A`` / ``partition_B``, ``make_fragment_A`` / ``make_fragment_B``,
and the ``ldmatrix`` copy views described in `Making Fragments`_.


Executing the GEMM (Main Loop)
-------------------------------

The main loop iterates over K-tiles and, within each tile, over k-blocks
(``num_k_block = BK / perm_K``). Each k-block loads A and B from SMEM into
registers via ``ldmatrix``, then issues ``cute.gemm``.

.. code-block:: python

   tCrC.fill(0.0)

   for k_tile in range(k_tile_count):
       for k_block in cutlass.range(num_k_block, unroll_full=True):
           # Wait for next SMEM stage at the tile boundary
           if k_block == num_k_block - 1:
               cute.arch.cp_async_wait_group(num_smem_stages - 2)
               cute.arch.sync_threads()

           # ldmatrix: prefetch next k-block from SMEM → RMEM
           k_block_next = (k_block + 1) % num_k_block
           cute.copy(tiled_copy_s2r_A, tCsA_p[None, None, k_block_next],
                     tCrA_copy_view[None, None, k_block_next])
           cute.copy(tiled_copy_s2r_B, tCsB_p[None, None, k_block_next],
                     tCrB_copy_view[None, None, k_block_next])

           # cp.async: issue GMEM → SMEM for next K-tile
           # ... (see tensorop_gemm.py for pipeline pointer management)

           # MMA: tCrC += tCrA * tCrB
           cute.gemm(tiled_mma, tCrC, tCrA[None, None, k_block], tCrB[None, None, k_block], tCrC)

   cute.arch.cp_async_wait_group(0)
   cute.arch.sync_threads()

Key points:

- ``cute.gemm`` is **synchronous** — it emits ``mma.sync.aligned``
  instructions. There is no accumulate-mode flag; the accumulator
  (``tCrC``) is always read and written.
- All operands must be in **registers** before ``cute.gemm`` is called.
  The ``ldmatrix`` copies above prefetch the next k-block into
  ``tCrA`` / ``tCrB`` from SMEM each iteration.
- The ``cp.async`` / ``cp_async_wait_group`` calls manage the GMEM→SMEM
  pipeline; see ``tensorop_gemm.py`` for predication, K-residue handling,
  and pipeline pointer management.


Complete Workflow
------------------

Putting it all together, a typical Ampere warp MMA GEMM has this structure:

**Host function** (``@cute.jit``):

.. code-block:: python

   import cutlass
   import cutlass.cute as cute

   @cute.jit
   def host_function(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, stream):
       # 1. Create the MMA op and tiled MMA
       op = cute.nvgpu.warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16))
       atom_layout_mnk = (2, 2, 1)
       permutation_mnk = (
           atom_layout_mnk[0] * 16,
           atom_layout_mnk[1] * 8 * 2,
           atom_layout_mnk[2] * 16,
       )
       tC = cute.make_layout(atom_layout_mnk)
       tiled_mma = cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk)

       # 2. Create SMEM layouts
       ab_copy_bits = 128
       sA_layout, sA_swizzle = _make_smem_layout_AB(
           mA.element_type, a_major_mode, ab_copy_bits,
           (cta_tiler[0], cta_tiler[2], num_stages),
       )
       sB_layout, sB_swizzle = _make_smem_layout_AB(
           mB.element_type, b_major_mode, ab_copy_bits,
           (cta_tiler[1], cta_tiler[2], num_stages),
       )

       # 3. Launch the kernel
       kernel(mA, mB, mC, ..., tiled_mma, sA_layout, sA_swizzle,
              sB_layout, sB_swizzle).launch(
           grid=grid, block=[128, 1, 1], stream=stream,
       )

**Kernel function** (``@cute.kernel``):

.. code-block:: python

  @cute.kernel
  def kernel(mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor,
            ..., tiled_mma: cute.TiledMma):
      tidx, _, _ = cute.arch.thread_idx()
      bidx, bidy, bidz = cute.arch.block_idx()

      # -- CTA-tiled global tensors --
      gA = cute.local_tile(mA[None, None, bidz], cta_tiler, (bidx, bidy, None), proj=(1, None, 1))
      gB = cute.local_tile(mB[None, None, bidz], cta_tiler, (bidx, bidy, None), proj=(None, 1, 1))
      gC = cute.local_tile(mC[None, None, bidz], cta_tiler, (bidx, bidy, None), proj=(1, 1, None))

      # -- Allocate SMEM --
      @cute.struct
      class SharedStorageAB:
          a: cute.struct.Align[cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], 16]
          b: cute.struct.Align[cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], 16]

      smem = cutlass.utils.SmemAllocator()
      storage = smem.allocate(SharedStorageAB)
      sA = SharedStorageAB(storage).a.get_tensor(sA_layout, swizzle=sA_swizzle)  # (BM, BK, PIPE)
      sB = SharedStorageAB(storage).b.get_tensor(sB_layout, swizzle=sB_swizzle)  # (BN, BK, PIPE)
      sC = ...  # (BM, BN) SMEM for epilogue (non-MMA, see tensorop_gemm.py)

      # -- GMEM → SMEM copy partitions (cp.async) --
      # ... setup tAgA, tAsA, tBgB, tBsB (see tensorop_gemm.py)

      # -- MMA partitions and fragments --
      thr_mma = tiled_mma.get_slice(tidx)
      tCsA = thr_mma.partition_A(sA)   # (MMA, MMA_M, MMA_K, PIPE)
      tCsB = thr_mma.partition_B(sB)   # (MMA, MMA_N, MMA_K, PIPE)
      tCsC = thr_mma.partition_C(sC)   # (MMA, MMA_M, MMA_N)
      tCgC = thr_mma.partition_C(gC)   # (MMA, MMA_M, MMA_N)
      tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])  # (MMA, MMA_M, MMA_K)
      tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])  # (MMA, MMA_N, MMA_K)
      tCrC = tiled_mma.make_fragment_C(tCgC)                       # (MMA, MMA_M, MMA_N)
      tCrC.fill(0.0)

      # -- ldmatrix retiling (see "Making Fragments" § SMEM → RMEM load) --
      # ... build tiled_copy_s2r_A/B from LdMatrix8x8x16bOp + make_tiled_copy_A/B
      # ... then: tCsA_copy_view = partition_S(sA), tCrA_copy_view = retile(tCrA), etc.

      # -- Prologue: cp.async fills num_stages-1 SMEM buffers --
      # -- Prefetch first k-block into registers via ldmatrix --
      # ... (see tensorop_gemm.py for predication, residual_k, and pipeline setup)

      # -- Main loop --
      for k_tile in range(k_tile_count):
          for k_block in cutlass.range(num_k_block, unroll_full=True):
              if k_block == num_k_block - 1:
                  cute.arch.cp_async_wait_group(num_smem_stages - 2)
                  cute.arch.sync_threads()

              # ldmatrix: prefetch next k-block from SMEM → RMEM
              # tCsA_p / tCsB_p are per-pipeline-stage slices, e.g.:
              #   tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read]
              k_block_next = (k_block + 1) % num_k_block
              cute.copy(tiled_copy_s2r_A, tCsA_p[None, None, k_block_next],
                        tCrA_copy_view[None, None, k_block_next])
              cute.copy(tiled_copy_s2r_B, tCsB_p[None, None, k_block_next],
                        tCrB_copy_view[None, None, k_block_next])

              # cp.async: issue GMEM → SMEM for next K-tile
              # ... (see tensorop_gemm.py for pipeline pointer management)

              # MMA
              cute.gemm(tiled_mma, tCrC, tCrA[None, None, k_block],
                        tCrB[None, None, k_block], tCrC)

      # -- Epilogue: RMEM → SMEM → RMEM → GMEM --
      cute.arch.cp_async_wait_group(0)
      cute.arch.sync_threads()
      tCrD = cute.make_fragment_like(tCrC, c_dtype)
      tCrD[None] = epilogue_op(tCrC.load()).to(c_dtype)
      cute.autovec_copy(tCrD, tCsC)   # RMEM → SMEM
      cute.arch.sync_threads()
      # ... reload with epilogue thread layout, then vectorized store to GMEM


Beyond Simple Dense MMAs
------------------------

The warp MMA DSL supports more complex MMA operations beyond simple dense MMA:

- Block-scaled MMA

.. {$nv-internal-release begin}

Internal builds additionally provide:

- Sparse MMA

.. {$nv-internal-release end}

.. {$nv-internal-release begin}

Sparse MMA
~~~~~~~~~~

Sparse MMA exploits **2:4 structured sparsity** in operand A: out of every
4 consecutive K-elements, exactly 2 are non-zero. The hardware consumes a
compressed A operand together with a compact **metadata** tensor ``E`` that
encodes which 2 of 4 positions are non-zero.

Compared to dense MMA, the MMA API differences are:

**1. MMA op creation** — use ``MmaF16BF16SparseOp`` with an extra
``sparse_metadata_format`` parameter. The sparse instruction K is doubled
relative to dense (dense ``m16n8k8`` → sparse ``m16n8k16``, dense
``m16n8k16`` → sparse ``m16n8k32``) because operand A is 2:4 compressed:

.. code-block:: python

  from cutlass.cute.nvgpu.warp.mma import SparseMetadataFormat

  # Dense F16 (for comparison): inst_K = 16
  dense_op = cute.nvgpu.warp.MmaF16BF16Op(
      cutlass.Float16, cutlass.Float32, (16, 8, 16),
  )

  # Sparse F16: inst_K = 32 (2× dense, since A is 2:4 compressed)
  sparse_op = cute.nvgpu.warp.MmaF16BF16SparseOp(
      cutlass.Float16,                         # A/B element type
      cutlass.Float32,                         # accumulator type
      (16, 8, 32),                             # instruction shape (M, N, K)
      SparseMetadataFormat.TID,                # metadata format
  )
  tiled_mma = cute.make_tiled_mma(sparse_op, cute.make_layout((1, 1, 1)))

.. code-block:: text

  Supported instruction shapes for MmaF16BF16SparseOp:

  | A/B Type | Acc Type  | Inst Shape     |
  |----------|-----------|----------------|
  | F16      | F16, F32  | (16,8,16), (16,8,32) |
  | BF16     | F32       | (16,8,16), (16,8,32) |

**2. Compressed A tensor and metadata E** — operand A stores only the
two non-zero values per group of 4 K-elements (half the storage). The
metadata tensor ``E`` records which 2 of 4 positions are non-zero. The
exact bit encoding depends on ``SparseMetadataFormat`` and on how the
implementation packs metadata. In this repository, helper code that
generates 2:4 test inputs packs two 4-bit metadata entries into each
``uint8`` value:

.. code-block:: python

  # Example metadata values used by examples/CuTeDSL/helpers/sparse_utils.py
  # Each nibble selects which 2 of 4 positions are non-zero.
  metadata_values = [0x4, 0x8, 0x9, 0xC, 0xD, 0xE]

.. code-block:: text

  Dense A: (M, K)                    Sparse operands:
  +--+--+--+--+--+--+--+--+         +--+--+--+--+
  | a| 0| b| 0| c| 0| d| 0|   →     | a| b| c| d|   (compressed A values)
  +--+--+--+--+--+--+--+--+         +--+--+--+--+

                                    E stores the non-zero positions
                                    for each 2:4 group.

**3. Fragments** — the dense-style fragment APIs for A, B, and C still
apply to the sparse atom:

.. code-block:: python

  # A/B/C fragments — same public API shape as dense
  tCsA = thr_mma.partition_A(sA)
  tCsB = thr_mma.partition_B(sB)
  tCgC = thr_mma.partition_C(gC)

  tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
  tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
  tCrC = tiled_mma.make_fragment_C(tCgC)
  tCrC.fill(0.0)

Sparse metadata ``E`` is an auxiliary operand associated with A. The
public warp API and tests in this repository verify op construction and
the ``cute.gemm(..., [A, E], B, ...)`` calling convention, but they do
not provide an end-to-end warp sparse kernel showing the exact
``partition`` / ``copy`` / ``make_fragment`` sequence for ``E``. For
that reason, this document intentionally does not spell out an ``E``
fragment construction sequence that has no example backing it.

**4. Modified gemm call** — the metadata E is passed alongside operand A
as a list. This part of the API is verified by ``cutlass.cute.algorithm.gemm``:

.. code-block:: python

  # Schematic only: E_k is the metadata operand for the same k-slice as A_k.
  A_k = tCrA[None, None, k_block]
  E_k = metadata_k
  B_k = tCrB[None, None, k_block]

  cute.gemm(
      tiled_mma,
      tCrC,
      [A_k, E_k],   # [A, E]
      B_k,
      tCrC,
  )

.. code-block:: text

  Dense gemm call:
    cute.gemm(tiled_mma, tCrC, A_k, B_k, tCrC)

  Sparse gemm call:
    cute.gemm(tiled_mma, tCrC, [A_k, E_k], B_k, tCrC)
                              ^^^^  ^^^
                              A     metadata

The epilogue (RMEM → SMEM → GMEM) is identical to a dense kernel.

.. note:: An end-to-end warp sparse GEMM example is not yet available in the
  examples directory. The closest verified references in this repository are
  ``cutlass_ir/compiler/test/python/not_pytest/sm_80/test_mma_atom.py`` for
  op construction, ``cutlass_ir/compiler/test/python/api/sm_120a/test_nvgpu_warp_mma.py``
  for tiled sparse MMA construction, and
  ``examples/CuTeDSL/helpers/sparse_utils.py`` for
  2:4 metadata packing.

.. {$nv-internal-release end}


Block-scaled MMA
~~~~~~~~~~~~~~~~

Block-scaled MMA multiplies narrow-type matrices (FP4) while applying
**per-block scale factors** along the GEMM-K dimension. Each vector of
``sf_vec_size`` consecutive K-elements shares a single scale factor, so the
hardware computes ``D = (SFA · A) * (SFB · B) + C``. The scale factors live
in **registers** alongside the operands and must be loaded from SMEM before
each ``gemm`` call.

Supported ops: ``MmaMXF4Op`` (SM120a+), ``MmaMXF4NVF4Op`` (SM120a+).

Compared to a dense MMA kernel, a block-scaled kernel has four additional concerns:

**1. MMA op creation** — block-scaled ops fix the data type to FP4
(E2M1) and the accumulator to FP32. The scale-factor type and vector
size distinguish the two ops:

.. code-block:: python

  # MXF4: UE8M0 scales, sf_vec_size = 32
  op = cute.nvgpu.warp.MmaMXF4Op(
      cutlass.Float4E2M1FN,     # A/B element type (fixed: E2M1)
      cutlass.Float32,          # accumulator type (fixed: F32)
      cutlass.Float8E8M0FNU,    # scale-factor type
  )  # instruction shape = (16, 8, 64), sf_vec_size = 32

  # MXF4NVF4: UE4M3 scales, sf_vec_size = 16
  op = cute.nvgpu.warp.MmaMXF4NVF4Op(
      cutlass.Float4E2M1FN,     # A/B element type (fixed: E2M1)
      cutlass.Float32,          # accumulator type (fixed: F32)
      cutlass.Float8E4M3FN,     # scale-factor type
  )  # instruction shape = (16, 8, 64), sf_vec_size = 16

.. code-block:: text

  | Op            | A/B Type | SF Type | Acc  | Inst Shape  | SF Vec Size |
  |---------------|----------|---------|------|-------------|-------------|
  | MmaMXF4Op     | E2M1     | UE8M0   | F32  | (16,8,64)   | 32          |
  | MmaMXF4NVF4Op | E2M1     | UE4M3   | F32  | (16,8,64)   | 16          |

**2. Extra global tensors and SMEM layouts for scale factors** — the host
function creates SFA/SFB tensors and allocates SMEM layouts for them
alongside A and B:

.. code-block:: python

  import cutlass.utils.blockscaled_layout as blockscaled_utils
  import cutlass.utils.blackwell_helpers as sm120_utils

  # Scale-factor global tensors (host side)
  sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, sf_vec_size)
  sfa_tensor = cute.make_tensor(sfa.iterator, sfa_layout)
  sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, sf_vec_size)
  sfb_tensor = cute.make_tensor(sfb.iterator, sfb_layout)

  # SMEM layouts for scale factors (SM120-specific helper)
  sfa_smem_layout = blockscaled_utils.sm120_make_smem_layout_sfa(
      tiled_mma, tile_shape_mnk, sf_vec_size, num_stages,
  )
  sfb_smem_layout = blockscaled_utils.sm120_make_smem_layout_sfb(
      tiled_mma, tile_shape_mnk, sf_vec_size, num_stages,
  )

**3. SF fragment creation and SMEM→RMEM retiling** — scale-factor
fragments use a ``CopyUniversalOp`` with thread-value layouts derived
from the tiled MMA, rather than the ``ldmatrix``-based path used for
A and B:

.. code-block:: python

  # A/B fragments (same as dense)
  tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
  tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])

  # SF fragments (SM120-specific partition helpers)
  tCrSFA = sm120_utils.partition_fragment_SFA(sSFA[None, None, 0], thr_mma, tidx)
  tCrSFB = sm120_utils.partition_fragment_SFB(sSFB[None, None, 0], thr_mma, tidx)

  # A/B: ldmatrix retiling (same as dense)
  atom_copy_A = cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(...), a_dtype)
  smem_tiled_copy_A = cute.make_tiled_copy_A(atom_copy_A, tiled_mma)

  # SF: CopyUniversal with SF-specific thread-value layout
  atom_copy_SF = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), sf_dtype)
  smem_tiled_copy_SFA = cute.make_tiled_copy(
      atom_copy_SF,
      sm120_utils.get_layoutSFA_TV(tiled_mma),
      (cute.size(tiled_mma.permutation_mnk[0]), cute.size(tiled_mma.permutation_mnk[2])),
  )
  smem_tiled_copy_SFB = cute.make_tiled_copy(
      atom_copy_SF,
      sm120_utils.get_layoutSFB_TV(tiled_mma),
      (cute.size(tiled_mma.permutation_mnk[1]), cute.size(tiled_mma.permutation_mnk[2])),
  )

**4. Modified main loop** — each k-block loads A, B, SFA, and SFB from
SMEM into registers. The ``cute.gemm`` call passes ``[A, SFA]`` and
``[B, SFB]`` as operand lists:

.. code-block:: python

  for k_block_idx in cutlass.range(num_k_blocks, unroll_full=True):
      # ldmatrix: load A and B from SMEM → RMEM (same as dense)
      cute.copy(smem_tiled_copy_A, tCsA_p[None, None, k_block_next],
                tCrA_copy_view[None, None, k_block_next])
      cute.copy(smem_tiled_copy_B, tCsB_p[None, None, k_block_next],
                tCrB_copy_view[None, None, k_block_next])

      # CopyUniversal: load SFA and SFB from SMEM → RMEM            # NEW
      cute.copy(smem_tiled_copy_SFA,
                cute.filter_zeros(tCsSFA_p)[None, None, k_block_next],
                cute.filter_zeros(tCrSFA_copy_view)[None, None, k_block_next])
      cute.copy(smem_tiled_copy_SFB,
                cute.filter_zeros(tCsSFB_p)[None, None, k_block_next],
                cute.filter_zeros(tCrSFB_copy_view)[None, None, k_block_next])

      # MMA with scale factors passed as [value, scale] pairs
      cute.gemm(
          tiled_mma,
          accumulators,
          [tCrA[None, None, k_block_idx], tCrSFA[None, None, k_block_idx]],  # [A, SFA]
          [tCrB[None, None, k_block_idx], tCrSFB[None, None, k_block_idx]],  # [B, SFB]
          accumulators,
      )

.. code-block:: text

  Dense gemm call:
    cute.gemm(tiled_mma, acc, tCrA[k], tCrB[k], acc)

  Block-scaled gemm call:
    cute.gemm(tiled_mma, acc, [tCrA[k], tCrSFA[k]], [tCrB[k], tCrSFB[k]], acc)
                              ^^^^^^^^  ^^^^^^^^^    ^^^^^^^^  ^^^^^^^^^
                              value     scale        value     scale
                              (RMEM)    (RMEM)       (RMEM)    (RMEM)

Note that ``cute.filter_zeros`` is applied to the SF copy views because
the scale-factor SMEM layouts may contain padding zeros from the TMA
tiling. This strips the padded entries so the copy operates only on
valid elements.

The epilogue (RMEM → SMEM → GMEM) is identical to a dense kernel.


See also:

- Dense GEMM example (Ampere): ``examples/cute/ampere/kernel/dense_gemm/tensorop_gemm.py``
- Block-scaled GEMM example (SM120a): ``examples/cute/blackwell_geforce/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent_pingpong.py``
- Block-scaled layout utilities: ``cutlass.utils.blockscaled_layout``
- SM120 helper utilities: ``cutlass.utils.blackwell_helpers``