> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.mimo_v2_flash.model

## Module Contents

### Classes

| Name                                                                                                             | Description                                                             |
| ---------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MiMoV2FlashAttention`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashAttention)             | MiMo-V2 attention with a Transformer Engine THD context-parallel path.  |
| [`MiMoV2FlashBlock`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashBlock)                     | Decoder block that alternates dense MLP and routed-MoE layers.          |
| [`MiMoV2FlashForCausalLM`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashForCausalLM)         | Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters. |
| [`MiMoV2FlashModel`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashModel)                     | Backbone model for Xiaomi MiMo-V2-Flash.                                |
| [`MiMoV2FlashRotaryEmbedding`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashRotaryEmbedding) | Rotary embedding module matching MiMo-V2-Flash partial-RoPE behavior.   |
| [`MiMoV2ForCausalLM`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2ForCausalLM)                   | MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.           |
| [`MiMoV2RMSNorm`](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2RMSNorm)                           | RMSNorm used by MiMo-V2-Flash decoder blocks.                           |

### Functions

| Name                                                                                                                           | Description                                                            |
| ------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------- |
| [`_apply_rotary_pos_emb`](#nemo_automodel-components-models-mimo_v2_flash-model-_apply_rotary_pos_emb)                         | -                                                                      |
| [`_convert_bool_4d_mask_to_additive`](#nemo_automodel-components-models-mimo_v2_flash-model-_convert_bool_4d_mask_to_additive) | -                                                                      |
| [`_derive_padding_mask`](#nemo_automodel-components-models-mimo_v2_flash-model-_derive_padding_mask)                           | -                                                                      |
| [`_eager_attention_forward`](#nemo_automodel-components-models-mimo_v2_flash-model-_eager_attention_forward)                   | -                                                                      |
| [`_ensure_additive_mask`](#nemo_automodel-components-models-mimo_v2_flash-model-_ensure_additive_mask)                         | -                                                                      |
| [`_fallback_additive_mask`](#nemo_automodel-components-models-mimo_v2_flash-model-_fallback_additive_mask)                     | -                                                                      |
| [`_initialize_mimo_rms_norm`](#nemo_automodel-components-models-mimo_v2_flash-model-_initialize_mimo_rms_norm)                 | Construct a MiMo RMSNorm from the configured backend.                  |
| [`_intersect_with_sliding_window`](#nemo_automodel-components-models-mimo_v2_flash-model-_intersect_with_sliding_window)       | Apply causal sliding-window bounds to an existing 4D document mask.    |
| [`_normalize_image_grid`](#nemo_automodel-components-models-mimo_v2_flash-model-_normalize_image_grid)                         | Accept the THW grids emitted by Qwen processors and PP's HW alias.     |
| [`_repeat_kv`](#nemo_automodel-components-models-mimo_v2_flash-model-_repeat_kv)                                               | -                                                                      |
| [`_replace_modal_embeddings`](#nemo_automodel-components-models-mimo_v2_flash-model-_replace_modal_embeddings)                 | Replace one kind of multimodal placeholder while preserving gradients. |
| [`_rotate_half`](#nemo_automodel-components-models-mimo_v2_flash-model-_rotate_half)                                           | -                                                                      |
| [`_validate_te_thd_sink_cudnn`](#nemo_automodel-components-models-mimo_v2_flash-model-_validate_te_thd_sink_cudnn)             | Reject the known pre-9.26 cuDNN ragged dSink kernel bug.               |

### Data

[`ModelClass`](#nemo_automodel-components-models-mimo_v2_flash-model-ModelClass)

### API

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    is_swa: bool,
    layer_idx: int
)
```

**Bases:** `Module`

MiMo-V2 attention with a Transformer Engine THD context-parallel path.

**`_te_v_head_dim`** `= self.v_head_dim`

---

**`attention_dropout`** `= float(config.attention_dropout or 0.0)`

---

**`head_dim`** `= int(config.swa_head_dim)`

---

**`k_proj`**

---

**`num_attention_heads`** `= int(config.swa_num_attention_heads)`

---

**`num_key_value_groups`**

---

**`num_key_value_heads`** `= int(config.swa_num_key_value_heads)`

---

**`o_proj`**

---

**`q_proj`**

---

**`rope_dim`** `= int(self.head_dim * config.partial_rotary_factor)`

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`sliding_window`** `= int(config.sliding_window) if is_swa else None`

---

**`v_head_dim`** `= int(config.swa_v_head_dim)`

---

**`v_proj`**

---

**`v_scale`** `= getattr(config, 'attention_value_scale', None)`

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._apply_te_rope(
    query_states: torch.Tensor,
    key_states: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    qkv_format: str
) -> tuple[torch.Tensor, torch.Tensor]
```

staticmethod

Apply partial RoPE to BSHD or flattened THD query/key tensors.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._forward_te(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run TE attention with native V dimensions, padding for legacy TE a2a CP.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden] for BSHD,
or \[tokens, hidden] for THD; a singleton batch is also accepted.

---

**`position_embeddings`** `tuple[torch.Tensor, torch.Tensor]`

Cosine and sine tensors aligned to the input
sequence, each ending in the rotary feature dimension.

---

**`attention_mask`** `torch.Tensor | None`

Optional mask of shape \[batch, sequence] for BSHD.

---

**`**kwargs`** `Any` — default: \{}

TE attention metadata, including one-dimensional cumulative
sequence-length tensors for THD inputs.

---

**Returns:** `torch.Tensor`

Attention output with the same shape as hidden\_states.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention._validate_a2a_cp_size(
    cp_size: int
) -> None
```

Validate the head partition required by TE's a2a CP transport.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashAttention.setup_cp_attention(
    cp_mesh,
    cp_stream: torch.cuda.Stream | None = None
) -> None
```

Validate and attach TE all-to-all context parallelism.

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashBlock(
    layer_idx: int,
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Decoder block that alternates dense MLP and routed-MoE layers.

**`attention_type`**

---

**`input_layernorm`**

---

**`mlp`** `= MoE(moe_config, backend)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashBlock.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashBlock.init_weights(
    buffer_device: torch.device
) -> None
```

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    revision: str | None = None,
    kwargs = {}
)
```

**Bases:** [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Causal LM wrapper for MiMo-V2-Flash with Automodel checkpoint adapters.

**`_keep_in_fp32_modules_strict`**

---

**`backend`** `= backend or BackendConfig()`

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM._get_multimodal_embeds(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    pixel_values: torch.Tensor | None = None,
    image_grid_thw: torch.Tensor | None = None,
    image_feature_indices: torch.Tensor | None = None,
    pixel_values_videos: torch.Tensor | None = None,
    video_grid_thw: torch.Tensor | None = None,
    video_feature_indices: torch.Tensor | None = None
) -> torch.Tensor
```

Encode global media and splice only this TE THD shard's feature rows.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM._local_modal_feature_indices(
    global_mask: torch.Tensor | None,
    local_indices: torch.Tensor | None
) -> torch.Tensor | None
```

staticmethod

Map local TE THD placeholders to rows in globally ordered media features.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM._pull_pipeline_media(
    input_ids: torch.Tensor | None,
    pixel_values: torch.Tensor | None,
    image_grid_thw: torch.Tensor | None,
    pixel_values_videos: torch.Tensor | None,
    video_grid_thw: torch.Tensor | None,
    global_image_mask: torch.Tensor | None = None,
    global_video_mask: torch.Tensor | None = None
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]
```

Retrieve PP media using global masks so every CP rank runs the same vision work.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]
```

Keep both rotary embeddings per stage and the vision tower on stage zero.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_grid_thw: torch.Tensor | None = None,
    pixel_values_videos: torch.Tensor | None = None,
    video_grid_thw: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor
```

Run text or VLM forward with TE-local packed media placement.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.from_config(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_input_embeddings()
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_output_embeddings()
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
```

Describe PP tensors using the already TE-sharded local token length.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]
```

Install MiMo's model-local adapter around framework TE THD sharding.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.set_input_embeddings(
    value
)
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashForCausalLM.set_output_embeddings(
    new_embeddings
)
```

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel(
    config: nemo_automodel.components.models.mimo_v2_flash.config.MiMoV2FlashConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)
```

**Bases:** `Module`

Backbone model for Xiaomi MiMo-V2-Flash.

**`embed_tokens`**

---

**`has_sliding_layers`**

---

**`layers`**

---

**`moe_config`** `= moe_config or MoEConfig(**moe_defaults)`

---

**`norm`**

---

**`rotary_emb`**

---

**`swa_rotary_emb`**

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel._build_causal_mask_mapping(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None,
    position_ids: torch.Tensor,
    cache_position: torch.Tensor
) -> dict[str, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    cache_position: torch.Tensor | None = None,
    _packed_seq_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    seq_lens_padded: torch.Tensor | None = None,
    cu_seqlens: torch.Tensor | None = None,
    cu_seqlens_padded: torch.Tensor | None = None,
    max_seqlen: int | torch.Tensor | None = None,
    qkv_format: str | None = None,
    cp_size: int = 1,
    cp_rank: int = 0,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the text backbone in BSHD or genuine TE THD layout.

THD batches must already have passed through the framework TE sharder.
Their cumulative lengths delimit packed documents and are forwarded to
every attention layer, including non-first pipeline stages.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashModel.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding(
    rope_theta: float,
    head_dim: int,
    partial_rotary_factor: float = 1.0,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** `Module`

Rotary embedding module matching MiMo-V2-Flash partial-RoPE behavior.

**`_inv_freq_initialized`** `= not inv_freq.is_meta`

---

**`attention_scaling`** `= 1.0`

---

**`inv_freq`** `Tensor`

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding._build_inv_freq(
    device: torch.device | None = None
) -> torch.Tensor
```

Build the FP32 inverse frequencies on the requested device.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding._materialize_inv_freq(
    device: torch.device
) -> None
```

Restore a non-persistent RoPE buffer after meta-model materialization.

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2FlashRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2ForCausalLM()
```

**Bases:** [MiMoV2FlashForCausalLM](#nemo_automodel-components-models-mimo_v2_flash-model-MiMoV2FlashForCausalLM)

MiMo-V2.5/V2.6 wrapper using the `mimo_v2` checkpoint config.

**`ModelCapabilities`** `= MiMoV2FlashForCausalLM.ModelCapabilities`

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> 'MiMoV2ForCausalLM'
```

classmethod

Resolve a MiMo-V2 config while checkpoint loading remains framework-owned.

```python
class nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm(
    hidden_size: int,
    eps: float = 1e-06,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** `Module`

RMSNorm used by MiMo-V2-Flash decoder blocks.

**`weight`** `= nn.Parameter(torch.ones(hidden_size, dtype=dtype))`

---

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model.MiMoV2RMSNorm.reset_parameters() -> None
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._convert_bool_4d_mask_to_additive(
    attention_mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._derive_padding_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._eager_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0
) -> tuple[torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._ensure_additive_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None,
    sliding_window: int | None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._fallback_additive_mask(
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._initialize_mimo_rms_norm(
    backend: nemo_automodel.components.models.common.BackendConfig,
    hidden_size: int,
    eps: float,
    dtype: torch.dtype
) -> torch.nn.Module
```

Construct a MiMo RMSNorm from the configured backend.

```python
nemo_automodel.components.models.mimo_v2_flash.model._intersect_with_sliding_window(
    attention_mask: torch.Tensor,
    sliding_window: int,
    dtype: torch.dtype
) -> torch.Tensor
```

Apply causal sliding-window bounds to an existing 4D document mask.

```python
nemo_automodel.components.models.mimo_v2_flash.model._normalize_image_grid(
    grid: torch.Tensor | None
) -> torch.Tensor | None
```

Accept the THW grids emitted by Qwen processors and PP's HW alias.

```python
nemo_automodel.components.models.mimo_v2_flash.model._repeat_kv(
    hidden_states: torch.Tensor,
    n_rep: int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._replace_modal_embeddings(
    input_ids: torch.Tensor,
    inputs_embeds: torch.Tensor,
    token_id: int | None,
    modal_embeds: torch.Tensor | None
) -> torch.Tensor
```

Replace one kind of multimodal placeholder while preserving gradients.

```python
nemo_automodel.components.models.mimo_v2_flash.model._rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.mimo_v2_flash.model._validate_te_thd_sink_cudnn(
    head_dim: int
) -> None
```

Reject the known pre-9.26 cuDNN ragged dSink kernel bug.

```python
nemo_automodel.components.models.mimo_v2_flash.model.ModelClass = MiMoV2FlashForCausalLM
```