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

CuTe-free contract of the MSA kernels.

The one home of the attention topology the kernels are compiled for, the schedule the forward saves
for the backward (validated once, when it is built), the CTA-walk rule that `msa_task_build_sm100`
mirrors on the device, and the grid bound the main kernel is launched with. The CPU tests exhaust the
walk rule from here without importing CuTe.

## Module Contents

### Classes

| Name                                                                                                              | Description                                                                                      |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------ |
| [`MSABackwardSchedule`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-MSABackwardSchedule) | Forward-derived int32 metadata the backward task build reads; save with `ctx.save_for_backward`. |

### Functions

| Name                                                                                                          | Description                                                                         |
| ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- |
| [`chunk_map`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-chunk_map)                 | Return `(num_full_ctas, tail_rows, grid_ctas)` for a walk covering every row once.  |
| [`grid_launch_bound`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-grid_launch_bound) | Bound every count up to capacity: full CTAs plus at most one tail CTA per SM.       |
| [`rows_per_cta`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-rows_per_cta)           | Select the CTA walk length; `msa_task_build_sm100` mirrors this rule on the device. |

### Data

[`BLOCK_SIZE`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-BLOCK_SIZE)

[`DQ_ACCUM_DTYPE`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-DQ_ACCUM_DTYPE)

[`HEAD_DIM`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-HEAD_DIM)

[`INDEX_DIM`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-INDEX_DIM)

[`NUM_INDEX_HEADS`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-NUM_INDEX_HEADS)

[`NUM_KV_HEADS`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-NUM_KV_HEADS)

[`NUM_Q_HEADS`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-NUM_Q_HEADS)

[`QUERY_CHUNK`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-QUERY_CHUNK)

[`ROWS_PER_CTA_LARGE`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-ROWS_PER_CTA_LARGE)

[`ROWS_PER_CTA_SMALL`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-ROWS_PER_CTA_SMALL)

[`ROWS_PER_CTA_SWITCH`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-ROWS_PER_CTA_SWITCH)

[`SOFTMAX_SCALE`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-SOFTMAX_SCALE)

[`TOPK_BLOCKS`](#nemo_automodel-components-models-minimax_m3_vl-kernels-msa_schedule-TOPK_BLOCKS)

### API

```python
class nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.MSABackwardSchedule(
    row_ptr: torch.Tensor,
    q_indices: torch.Tensor,
    scheduler_metadata: torch.Tensor,
    work_count: torch.Tensor,
    cu_seqlens: torch.Tensor,
    document_workspace_starts: torch.Tensor
)
```

Dataclass

Forward-derived int32 metadata the backward task build reads; save with `ctx.save_for_backward`.

`scheduler_metadata` columns are
`(index_head, row_linear, q_begin, q_count, document_ordinal, document_local_kblock)`,
valid only up to `work_count`. Shapes and dtypes are checked once, here, so every kernel
wrapper downstream can take the schedule as given.

**Parameters:**

**`row_ptr`** `torch.Tensor`

`[4, rows + 1]` CSR row offsets of the key-block to query map, one row per index head.

---

**`q_indices`** `torch.Tensor`

`[4, edge_capacity]` document-local query positions of that map.

---

**`scheduler_metadata`** `torch.Tensor`

`[work_capacity, 6]` forward work items in the column order above.

---

**`work_count`** `torch.Tensor`

`[1]` number of valid work items.

---

**`cu_seqlens`** `torch.Tensor`

`[documents + 1]` compact document offsets.

---

**`document_workspace_starts`** `torch.Tensor`

`[documents]` 128-aligned workspace row of each document.

---

**Raises:**

* `TypeError`: If a field is not int32.
* `ValueError`: If a field does not have the shape stated above.

**`cu_seqlens`** `Tensor`

---

**`document_workspace_starts`** `Tensor`

---

**`q_indices`** `Tensor`

---

**`row_ptr`** `Tensor`

---

**`scheduler_metadata`** `Tensor`

---

**`work_count`** `Tensor`

---

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.MSABackwardSchedule.__post_init__() -> None
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.chunk_map(
    num_rows: int,
    rows_per_cta: int,
    num_sms: int
) -> tuple[int, int, int]
```

Return `(num_full_ctas, tail_rows, grid_ctas)` for a walk covering every row once.

The tables kernel of `msa_task_build_sm100` mirrors this rule on the device.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.grid_launch_bound(
    capacity: int,
    num_sms: int
) -> int
```

Bound every count up to capacity: full CTAs plus at most one tail CTA per SM.

`chunk_map(n, r, s)[2] <= n // r + s` for every count, so the bound takes the larger of the
small and the large walk regime over `[0, capacity]`.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.rows_per_cta(
    num_rows: int
) -> int
```

Select the CTA walk length; `msa_task_build_sm100` mirrors this rule on the device.

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.BLOCK_SIZE = 128
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.DQ_ACCUM_DTYPE = torch.bfloat16
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.HEAD_DIM = 128
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.INDEX_DIM = 128
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.NUM_INDEX_HEADS = 4
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.NUM_KV_HEADS = 4
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.NUM_Q_HEADS = 64
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.QUERY_CHUNK = 8
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.ROWS_PER_CTA_LARGE = 64
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.ROWS_PER_CTA_SMALL = 4
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.ROWS_PER_CTA_SWITCH = 2400
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.SOFTMAX_SCALE = HEAD_DIM ** -0.5
```

```python
nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule.TOPK_BLOCKS = 16
```