nemo_automodel.components.models.minimax_m3_vl.kernels

Model-private MSA kernels, bound lazily by msa_bindings.kernels; no eager CuTe imports.

The SM100 kernel modules are CuTe DSL source and call require_cute_dsl before their first DSL import, so a host without the msa extra sees UnavailableError instead of ModuleNotFoundError.

Submodules

Package Contents

Functions

Name	Description
`require_cute_dsl`	Refuse to bind the CuTe DSL in a kernel module on a host that cannot import it.
`require_sm100`	Reject any device the MSA kernels are not built for, before compiling or launching one.
`sm_capability`	Return the memoized CUDA compute capability of `device`.
`sm_count`	Return the memoized streaming-multiprocessor count of `device`; it sizes the backward CTA walk.

Data

MSA_KERNEL_IMPORT_ERROR

API

nemo_automodel.components.models.minimax_m3_vl.kernels.require_cute_dsl() -> None

Refuse to bind the CuTe DSL in a kernel module on a host that cannot import it.

The SM100 kernel modules are CuTe DSL source: their decorators, annotations and module constants need cutlass and cuda.bindings while the module loads, and the CI import walker imports every module of the package on hosts without the msa extra. Each kernel module calls this before its first DSL import so absence surfaces as UnavailableError, the signal both the walker and msa_bindings.kernels understand. It is the software twin of require_sm100: one gates the toolchain at import, the other the device at launch.

Raises:

UnavailableError: If cutlass or cuda.bindings.driver cannot be imported.

nemo_automodel.components.models.minimax_m3_vl.kernels.require_sm100(
    device: torch.device
) -> None

Reject any device the MSA kernels are not built for, before compiling or launching one.

Parameters:

device

torch.device

CUDA device the caller is about to run MSA kernels on.

Raises:

NotImplementedError: If device is not SM100.

nemo_automodel.components.models.minimax_m3_vl.kernels.sm_capability(
    device: torch.device
) -> tuple[int, int]

Return the memoized CUDA compute capability of device.

torch.cuda.get_device_capability costs about 1.6 us and require_sm100 runs on every MSA forward; a device’s capability cannot change, so look it up once per device.

Parameters:

device

torch.device

CUDA device taken from a tensor, which always carries an explicit index.

Returns: tuple[int, int]

The (major, minor) compute capability of that device.

nemo_automodel.components.models.minimax_m3_vl.kernels.sm_count(
    device: torch.device
) -> int

Return the memoized streaming-multiprocessor count of device; it sizes the backward CTA walk.

Parameters:

device

torch.device

CUDA device taken from a tensor, which always carries an explicit index.

Returns: int

The number of streaming multiprocessors of that device.

nemo_automodel.components.models.minimax_m3_vl.kernels.MSA_KERNEL_IMPORT_ERROR = "BackendConfig.sparse_attn='msa' backward requires nvidia-cutlass-dsl==4.6.2 and...