> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100

MiniMax M3 main-attention backward kernel for MSA on SM100.

KV-parallel: each CTA walks task rows (8 queries x 16 main heads = one 128-row
tile) bucketed by (batch, index\_head, key\_block), with K/V TMA-resident per
bucket and Q/dO TMA-loaded per tile (one load warp, 8-row gathers). One mma warp
issues all five tcgen05 GEMMs transposed (S^T, dP^T, dV, dK, dQ^T) over four
128-column TMEM allocations; dV/dK accumulate per bucket segment and are flushed
with fp32 vector atomics, dQ^T per tile with packed bf16 atomics into a
head-pair-interleaved pool that
`msa_backward_postprocess_sm100` casts to the bf16 gradient. The task tables come
from `msa_task_build_sm100`; `run_backward` is the host entry point.

## Module Contents

### Classes

| Name                                                                                                                            | Description |
| ------------------------------------------------------------------------------------------------------------------------------- | ----------- |
| [`_MSABackwardSm100Kernel`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_MSABackwardSm100Kernel) | -           |

### Functions

| Name                                                                                                                        | Description                                                                                       |
| --------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------- |
| [`_compile`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_compile)                           | Compile the backward once per process with dynamic token, workspace, and task counts.             |
| [`_contiguous_stride`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_contiguous_stride)       | -                                                                                                 |
| [`_head_major`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_head_major)                     | View a row-major `[T, H, ...]` tensor as the kernel's `[1, H, T, ...]` operand without a copy.    |
| [`_l2_policy_evict_last`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_l2_policy_evict_last) | 64-bit L2 cache policy: keep the whole line set as evict-last (fraction 1.0).                     |
| [`_operand`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_operand)                           | Return `tensor` detached, contiguous and on 16-byte-aligned storage, copying only when it is not. |
| [`_pack_bf16x2`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_pack_bf16x2)                   | Round two FP32 values to one bf16x2 word, low logical element in the low half.                    |
| [`_pack_f16x2`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_pack_f16x2)                     | Round two FP32 values to one f16x2 word, low logical element in the low half.                     |
| [`_red_add_16bitx2_hint`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_red_add_16bitx2_hint) | -                                                                                                 |
| [`_red_add_bf16x2_hint`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_red_add_bf16x2_hint)   | -                                                                                                 |
| [`_red_add_f16x2_hint`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_red_add_f16x2_hint)     | -                                                                                                 |
| [`_round_up`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-_round_up)                         | -                                                                                                 |
| [`run_backward`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-run_backward)                   | Run the SM100 main-attention backward on THD-contract tensors of one CUDA device.                 |

### Data

[`MAIN_HEADS_PER_INDEX`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-MAIN_HEADS_PER_INDEX)

[`TILE_M`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-TILE_M)

[`TILE_N`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_backward_sm100-TILE_N)

### API

```python
class nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel()
```

**`buffer_align_bytes`** `= 1024`

---

**`chunk_stage_starts`** `= (0, 2, 3)`

---

**`compute_chunk_cols`** `= 32`

---

**`compute_mma_chunk_stage`** `= len(self.chunk_stage_starts)`

---

**`compute_warp_id`** `= (4, 5, 6, 7)`

---

**`do_stage`** `= 2`

---

**`gather_mma_QdO_stage`** `= self.q_stage`

---

**`gather_row_stage`** `= self.row_stage`

---

**`index_heads_per_kv`** `= NUM_INDEX_HEADS // NUM_KV_HEADS`

---

**`load_mma_KV_stage`** `= 1`

---

**`load_warp_id`** `= 13`

---

**`mma_compute_S_stage`** `= 1`

---

**`mma_compute_dP_stage`** `= 1`

---

**`mma_reduce_dKV_stage`** `= 1`

---

**`mma_reduce_dQ_stage`** `= 1`

---

**`mma_warp_id`** `= 12`

---

**`num_compute_chunks`** `= TILE_N // self.compute_chunk_cols`

---

**`num_compute_warps`** `= 4`

---

**`num_reduce_warps`** `= 4`

---

**`num_regs_compute`** `= 184`

---

**`num_regs_empty`** `= 24`

---

**`num_regs_load`** `= 48`

---

**`num_regs_mma`** `= 48`

---

**`num_regs_reduce`** `= 128`

---

**`num_regs_scalar`** `= 48`

---

**`num_tmem_alloc_cols`** `= 512`

---

**`q_stage`** `= 2`

---

**`reduce_pass_cols`** `= 32`

---

**`reduce_sync_barrier`**

---

**`reduce_warp_id`** `= (8, 9, 10, 11)`

---

**`row_stage`** `= 2`

---

**`scalar_warp_id`** `= 14`

---

**`t2r_dQ_done_barrier`**

---

**`threads_per_cta`** `= 512`

---

**`threads_per_warp`** `= 32`

---

**`tmem_S_offset`** `= 0`

---

**`tmem_alloc_barrier`**

---

**`tmem_dK_offset`** `= 384`

---

**`tmem_dPdQ_offset`** `= 128`

---

**`tmem_dV_offset`** `= 256`

---

**`tmem_dealloc_barrier`**

---

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.__call__(
    mQ: cutlass.cute.Tensor,
    mK: cutlass.cute.Tensor,
    mV: cutlass.cute.Tensor,
    mdO: cutlass.cute.Tensor,
    mLSE: cutlass.cute.Tensor,
    mDelta: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    mTaskQRows: cutlass.cute.Tensor,
    mTaskQPos: cutlass.cute.Tensor,
    mdQ: cutlass.cute.Tensor,
    mdK: cutlass.cute.Tensor,
    mdV: cutlass.cute.Tensor,
    mDesc: cutlass.cute.Tensor,
    grid_launch: cutlass.Int32,
    softmax_scale: cutlass.Float32,
    stream: cuda.bindings.driver.CUstream
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel._quad_transpose4(
    blk,
    c0,
    c1
)
```

In-place 4x4 transpose of four F32x4 fragments across a quad's 4 lanes.

Lane q holds element (row q, block k) before and (row k, block q) after. `c0`/`c1`
are 4-wide lane-bit vectors, so each stage's slot choice lowers to FSEL, not a branch.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel._same_bucket(
    mTaskMeta: cutlass.cute.Tensor,
    row_a: cutlass.Int32,
    row_b: cutlass.Int32
) -> cutlass.Boolean
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel._task_fields(
    mTaskMeta: cutlass.cute.Tensor,
    row: cutlass.Int32
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.compute(
    tStS: cutlass.cute.Tensor,
    tdPtdP: cutlass.cute.Tensor,
    sPdS_store: cutlass.cute.Tensor,
    sLSE: cutlass.cute.Tensor,
    sDelta: cutlass.cute.Tensor,
    sQPos: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    row_lo: cutlass.Int32,
    row_hi: cutlass.Int32,
    scale_log2e: cutlass.Float32,
    log2_e: cutlass.Float32,
    pipelines
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.gather_scalars(
    mLSE: cutlass.cute.Tensor,
    mDelta: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    mTaskQRows: cutlass.cute.Tensor,
    mTaskQPos: cutlass.cute.Tensor,
    sLSE: cutlass.cute.Tensor,
    sDelta: cutlass.cute.Tensor,
    sQRows: cutlass.cute.Tensor,
    sQPos: cutlass.cute.Tensor,
    row_lo: cutlass.Int32,
    row_hi: cutlass.Int32,
    log2_e: cutlass.Float32,
    gather_row_pipeline
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.get_tmem_tensors(
    mma_S: cutlass.cute.TiledMma,
    mma_dP: cutlass.cute.TiledMma,
    mma_dQ: cutlass.cute.TiledMma,
    mma_dV: cutlass.cute.TiledMma,
    mma_dK: cutlass.cute.TiledMma,
    tmem_ptr_base: cutlass.cute.Pointer
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.kernel(
    mma_S: cutlass.cute.TiledMma,
    mma_dP: cutlass.cute.TiledMma,
    mma_dV: cutlass.cute.TiledMma,
    mma_dK: cutlass.cute.TiledMma,
    mma_dQ: cutlass.cute.TiledMma,
    tma_atom_K: cutlass.cute.CopyAtom,
    tma_tensor_K: cutlass.cute.Tensor,
    tma_atom_V: cutlass.cute.CopyAtom,
    tma_tensor_V: cutlass.cute.Tensor,
    tma_atom_Q: cutlass.cute.CopyAtom,
    tma_tensor_Q: cutlass.cute.Tensor,
    tma_atom_dO: cutlass.cute.CopyAtom,
    tma_tensor_dO: cutlass.cute.Tensor,
    mLSE: cutlass.cute.Tensor,
    mDelta: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    mTaskQRows: cutlass.cute.Tensor,
    mTaskQPos: cutlass.cute.Tensor,
    mdQ: cutlass.cute.Tensor,
    mdK: cutlass.cute.Tensor,
    mdV: cutlass.cute.Tensor,
    mDesc: cutlass.cute.Tensor,
    scale_log2e: cutlass.Float32,
    log2_e: cutlass.Float32,
    sK_layout: cutlass.cute.ComposedLayout,
    sV_layout: cutlass.cute.ComposedLayout,
    sQ_layout: cutlass.cute.ComposedLayout,
    sdO_layout: cutlass.cute.ComposedLayout,
    sPdS_layout: cutlass.cute.ComposedLayout,
    sPdS_store_layout: cutlass.cute.ComposedLayout,
    sdOb_layout: cutlass.cute.ComposedLayout,
    sQb_layout: cutlass.cute.ComposedLayout,
    sKt_layout: cutlass.cute.ComposedLayout,
    sPdSn_layout: cutlass.cute.ComposedLayout,
    tP_layout: cutlass.cute.ComposedLayout,
    sLSE_layout: cutlass.cute.Layout,
    sDelta_layout: cutlass.cute.Layout,
    sQRows_layout: cutlass.cute.Layout,
    sQPos_layout: cutlass.cute.Layout
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.load_kv_qdo(
    mma_S: cutlass.cute.TiledMma,
    mma_dP: cutlass.cute.TiledMma,
    tma_atom_K: cutlass.cute.CopyAtom,
    tma_tensor_K: cutlass.cute.Tensor,
    tma_atom_V: cutlass.cute.CopyAtom,
    tma_tensor_V: cutlass.cute.Tensor,
    tma_atom_Q: cutlass.cute.CopyAtom,
    tma_tensor_Q: cutlass.cute.Tensor,
    tma_atom_dO: cutlass.cute.CopyAtom,
    tma_tensor_dO: cutlass.cute.Tensor,
    sK: cutlass.cute.Tensor,
    sV: cutlass.cute.Tensor,
    sQ: cutlass.cute.Tensor,
    sdO: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    mTaskQRows: cutlass.cute.Tensor,
    row_lo: cutlass.Int32,
    row_hi: cutlass.Int32,
    load_mma_KV_pipeline,
    gather_mma_QdO_pipeline
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.mma(
    mma_S: cutlass.cute.TiledMma,
    mma_dP: cutlass.cute.TiledMma,
    mma_dV: cutlass.cute.TiledMma,
    mma_dK: cutlass.cute.TiledMma,
    mma_dQ: cutlass.cute.TiledMma,
    sK: cutlass.cute.Tensor,
    sV: cutlass.cute.Tensor,
    sQ: cutlass.cute.Tensor,
    sdO: cutlass.cute.Tensor,
    sdOb: cutlass.cute.Tensor,
    sQb: cutlass.cute.Tensor,
    sKt: cutlass.cute.Tensor,
    sPdSn: cutlass.cute.Tensor,
    tStS: cutlass.cute.Tensor,
    tdPtdP: cutlass.cute.Tensor,
    tdQtdQ: cutlass.cute.Tensor,
    tdVtdV: cutlass.cute.Tensor,
    tdKtdK: cutlass.cute.Tensor,
    tmem_ptr_base: cutlass.cute.Pointer,
    tP_layout: cutlass.cute.ComposedLayout,
    mTaskMeta: cutlass.cute.Tensor,
    row_lo: cutlass.Int32,
    row_hi: cutlass.Int32,
    pipelines
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.quantize(
    input: cutlass.cute.Tensor,
    frg_cnt: cutlass.Int32,
    scale: cutlass.Float32 | None = None
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._MSABackwardSm100Kernel.reduce(
    tdQtdQ: cutlass.cute.Tensor,
    tdVtdV: cutlass.cute.Tensor,
    tdKtdK: cutlass.cute.Tensor,
    mdQ: cutlass.cute.Tensor,
    mdK: cutlass.cute.Tensor,
    mdV: cutlass.cute.Tensor,
    sQRows: cutlass.cute.Tensor,
    mTaskMeta: cutlass.cute.Tensor,
    row_lo: cutlass.Int32,
    row_hi: cutlass.Int32,
    pipelines
)
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._compile() -> typing.Any
```

Compile the backward once per process with dynamic token, workspace, and task counts.

The fake tensors describe the head-major views `run_backward` builds; `stride_order[i]` is the
rank of mode `i`, `0` innermost. Returns an executable taking the kernel's positional arguments
minus the trailing stream.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._contiguous_stride(
    shape
) -> tuple[int, ...]
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._head_major(
    rows: torch.Tensor
) -> torch.Tensor
```

View a row-major `[T, H, ...]` tensor as the kernel's `[1, H, T, ...]` operand without a copy.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._l2_policy_evict_last(
    loc = None,
    ip = None
) -> cutlass.Uint64
```

64-bit L2 cache policy: keep the whole line set as evict-last (fraction 1.0).

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._operand(
    tensor: torch.Tensor
) -> torch.Tensor
```

Return `tensor` detached, contiguous and on 16-byte-aligned storage, copying only when it is not.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._pack_bf16x2(
    lo: cutlass.Float32,
    hi: cutlass.Float32,
    loc = None,
    ip = None
) -> cutlass.Uint32
```

Round two FP32 values to one bf16x2 word, low logical element in the low half.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._pack_f16x2(
    lo: cutlass.Float32,
    hi: cutlass.Float32,
    loc = None,
    ip = None
) -> cutlass.Uint32
```

Round two FP32 values to one f16x2 word, low logical element in the low half.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._red_add_16bitx2_hint(
    kind: str,
    destination: cutlass.cute.Pointer,
    word,
    policy,
    loc = None,
    ip = None
) -> None
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._red_add_bf16x2_hint(
    destination: cutlass.cute.Pointer,
    word: cutlass.Uint32,
    policy: cutlass.Uint64,
    loc = None,
    ip = None
) -> None
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._red_add_f16x2_hint(
    destination: cutlass.cute.Pointer,
    word: cutlass.Uint32,
    policy: cutlass.Uint64,
    loc = None,
    ip = None
) -> None
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100._round_up(
    n: int,
    m: int = 256
) -> int
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100.run_backward(
    q: torch.Tensor,
    k_aligned: torch.Tensor,
    v_aligned: torch.Tensor,
    grad_out: torch.Tensor,
    lse: torch.Tensor,
    out: torch.Tensor,
    schedule: nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.MSABackwardSchedule
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Run the SM100 main-attention backward on THD-contract tensors of one CUDA device.

**Parameters:**

**`q`** `torch.Tensor`

BF16 `[T, 64, 128]` compact queries.

---

**`k_aligned`** `torch.Tensor`

BF16 `[W, 4, 128]` keys in the 128-aligned workspace, `W` a positive multiple of 128.

---

**`v_aligned`** `torch.Tensor`

BF16 `[W, 4, 128]` values in the same workspace.

---

**`grad_out`** `torch.Tensor`

BF16 `[T, 64, 128]` output gradient.

---

**`lse`** `torch.Tensor`

FP32 `[T, 64]` forward log-sum-exp.

---

**`out`** `torch.Tensor`

BF16 `[T, 64, 128]` forward output.

---

**`schedule`** `MSABackwardSchedule`

Forward-derived task schedule of this call.

---

**Returns:** `torch.Tensor`

BF16 `(dq [T, 64, 128], dk_aligned [W, 4, 128], dv_aligned [W, 4, 128])`, zero outside the

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100.MAIN_HEADS_PER_INDEX = NUM_Q_HEADS // NUM_INDEX_HEADS
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100.TILE_M = 128
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_backward_sm100.TILE_N = 128
```