nemo_automodel.components.models.minimax_m3_vl.kernels.msa_schedule
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
Functions
Data
API
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:
[4, rows + 1] CSR row offsets of the key-block to query map, one row per index head.
[4, edge_capacity] document-local query positions of that map.
[work_capacity, 6] forward work items in the column order above.
[1] number of valid work items.
[documents + 1] compact document offsets.
[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.
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.
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].
Select the CTA walk length; msa_task_build_sm100 mirrors this rule on the device.