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`	-

Functions

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

Data

MAIN_HEADS_PER_INDEX

TILE_M

TILE_N

API

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

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
)

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.

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

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

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
)

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
)

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
)

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
)

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
)

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
)

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
)

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
)

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.

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

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.

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).

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.

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.

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.

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

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

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

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

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:

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

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

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

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