> 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.kimi_k3.vision

## Module Contents

### Classes

| Name                                                                                                                            | Description                                                 |
| ------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------- |
| [`IdentityMap`](#nemo_automodel-components-models-kimi_k3-vision-IdentityMap)                                                   | -                                                           |
| [`Learnable2DInterpPosEmbDivided_fixed`](#nemo_automodel-components-models-kimi_k3-vision-Learnable2DInterpPosEmbDivided_fixed) | -                                                           |
| [`MLP`](#nemo_automodel-components-models-kimi_k3-vision-MLP)                                                                   | -                                                           |
| [`MLP2`](#nemo_automodel-components-models-kimi_k3-vision-MLP2)                                                                 | -                                                           |
| [`MoonViT3dEncoder`](#nemo_automodel-components-models-kimi_k3-vision-MoonViT3dEncoder)                                         | -                                                           |
| [`MoonViT3dPretrainedModel`](#nemo_automodel-components-models-kimi_k3-vision-MoonViT3dPretrainedModel)                         | -                                                           |
| [`MoonViTEncoderLayer`](#nemo_automodel-components-models-kimi_k3-vision-MoonViTEncoderLayer)                                   | -                                                           |
| [`MoonVision3dPatchEmbed`](#nemo_automodel-components-models-kimi_k3-vision-MoonVision3dPatchEmbed)                             | -                                                           |
| [`PatchMergerMLP`](#nemo_automodel-components-models-kimi_k3-vision-PatchMergerMLP)                                             | -                                                           |
| [`PatchMergerMLPV2`](#nemo_automodel-components-models-kimi_k3-vision-PatchMergerMLPV2)                                         | -                                                           |
| [`Rope2DPosEmbRepeated`](#nemo_automodel-components-models-kimi_k3-vision-Rope2DPosEmbRepeated)                                 | 2D rotary position embedding with multi-resolution support. |

### Functions

| Name                                                                                                                      | Description                                   |
| ------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------- |
| [`_apply_rope_input_validation`](#nemo_automodel-components-models-kimi_k3-vision-_apply_rope_input_validation)           | -                                             |
| [`apply_rope`](#nemo_automodel-components-models-kimi_k3-vision-apply_rope)                                               | Returns:                                      |
| [`eager_attention`](#nemo_automodel-components-models-kimi_k3-vision-eager_attention)                                     | -                                             |
| [`get_1d_sincos_pos_embed`](#nemo_automodel-components-models-kimi_k3-vision-get_1d_sincos_pos_embed)                     | t\_size: int of the temporal size             |
| [`get_1d_sincos_pos_embed_from_grid`](#nemo_automodel-components-models-kimi_k3-vision-get_1d_sincos_pos_embed_from_grid) | From:                                         |
| [`get_rope_shape`](#nemo_automodel-components-models-kimi_k3-vision-get_rope_shape)                                       | -                                             |
| [`get_rope_shape_decorate`](#nemo_automodel-components-models-kimi_k3-vision-get_rope_shape_decorate)                     | -                                             |
| [`multihead_attention`](#nemo_automodel-components-models-kimi_k3-vision-multihead_attention)                             | Multi-head attention using flash attention 2. |
| [`tpool_patch_merger`](#nemo_automodel-components-models-kimi_k3-vision-tpool_patch_merger)                               | -                                             |

### Data

[`VL_VISION_ATTENTION_FUNCTIONS`](#nemo_automodel-components-models-kimi_k3-vision-VL_VISION_ATTENTION_FUNCTIONS)

### API

```python
class nemo_automodel.components.models.kimi_k3.vision.IdentityMap()
```

**Bases:** `Module`

```python
nemo_automodel.components.models.kimi_k3.vision.IdentityMap.forward(
    x,
    args = (),
    kwargs = {}
)
```

```python
class nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed(
    height: int,
    width: int,
    num_frames: int,
    dim: int,
    interpolation_mode: str = 'bicubic'
)
```

**Bases:** `Module`

**`weight`** `= nn.Parameter(torch.empty(height, width, dim))`

---

```python
nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed.forward(
    x: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
```

```python
nemo_automodel.components.models.kimi_k3.vision.Learnable2DInterpPosEmbDivided_fixed.reset_parameters()
```

```python
class nemo_automodel.components.models.kimi_k3.vision.MLP(
    config
)
```

**Bases:** `Module`

**`pre_norm`** `= nn.LayerNorm(config.mm_hidden_size)`

---

**`proj`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.MLP.forward(
    x,
    args = (),
    kwargs = {}
)
```

```python
class nemo_automodel.components.models.kimi_k3.vision.MLP2(
    dims: list[int],
    activation,
    bias = True
)
```

**Bases:** `Module`

**Parameters:**

**`dims`** `list[int]`

\[in\_dim, hidden\_dim, out\_dim]

---

**`bias`** — default: True

whether to use bias in linear layer.

---

**`fc0`** `= nn.Linear(dims[0], dims[1], bias=bias)`

---

**`fc1`** `= nn.Linear(dims[1], dims[2], bias=bias)`

---

```python
nemo_automodel.components.models.kimi_k3.vision.MLP2.forward(
    x: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.kimi_k3.vision.MoonViT3dEncoder(
    hidden_dim: int,
    num_layers: int,
    block_cfg: dict,
    use_deterministic_attn: bool = False
)
```

**Bases:** `Module`

**`blocks`**

---

**`final_layernorm`** `= nn.LayerNorm(hidden_dim)`

---

**`rope_2d`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.MoonViT3dEncoder.forward(
    hidden_states: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.kimi_k3.vision.MoonViT3dPretrainedModel(
    config,
    inputs = (),
    kwargs = {}
)
```

**Bases:** `PreTrainedModel`

**`_no_split_modules`** `= ['MoonViTEncoderLayer']`

---

**`encoder`**

---

**`merge_kernel_size`** `= config.merge_kernel_size`

---

**`merge_type`** `= config.merge_type`

---

**`model_type`** `= 'moonvit3d'`

---

**`patch_embed`**

---

**`patch_size`** `= config.patch_size`

---

```python
nemo_automodel.components.models.kimi_k3.vision.MoonViT3dPretrainedModel.forward(
    pixel_values: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
```

**Parameters:**

**`pixel_values`** `torch.Tensor`

The input pixel values.

---

**`grid_thws`** `torch.Tensor`

Temporal, height and width.

---

**Returns:** `torch.Tensor`

torch.Tensor: The output tokens.

```python
class nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer(
    num_heads: int,
    hidden_dim: int,
    mlp_dim: int,
    qkv_hidden_size: int | None = None,
    norm_type: str = 'layernorm',
    mlp_type: str = 'mlp2',
    attn_implementation: str = 'flash_attention_2',
    activation = F.gelu,
    attn_bias: bool = False,
    linear_bias: bool = True,
    use_deterministic_attn: bool = False
)
```

**Bases:** `Module`

**`hidden_size_per_attention_head`** `= self.qkv_hidden_size // self.num_heads`

---

**`mlp`**

---

**`norm0`** `= nn.LayerNorm(hidden_dim)`

---

**`norm1`** `= nn.LayerNorm(hidden_dim)`

---

**`qkv_hidden_size`**

---

**`wo`**

---

**`wqkv`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer.attention_qkvpacked(
    x: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: torch.Tensor,
    rope_freqs_cis: torch.Tensor | None = None
)
```

**Parameters:**

**`x`** `torch.Tensor`

(batch\_size, seqlen, hidden\_dim)

---

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

---

```python
nemo_automodel.components.models.kimi_k3.vision.MoonViTEncoderLayer.forward(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    rope_freqs_cis: torch.Tensor | None = None
)
```

```python
class nemo_automodel.components.models.kimi_k3.vision.MoonVision3dPatchEmbed(
    out_dim: int,
    in_dim: int = 3,
    patch_size: int | tuple[int, int] = (14, 14),
    pos_emb_height: int = 14,
    pos_emb_width: int = 14,
    pos_emb_time: int = 4,
    pos_emb_type: str = 'divided_fixed',
    patch_embed_proj_bias: bool = True,
    pos_emb_interpolation_mode: str = 'bicubic'
)
```

**Bases:** `Module`

**`pos_emb`**

---

**`proj`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.MoonVision3dPatchEmbed.forward(
    x: torch.Tensor,
    grid_thws: torch.Tensor
) -> torch.Tensor
```

**Parameters:**

**`x`** `(L, Channels)`

input tensor

---

**`grid_hws`** `(N, 3)`

temporal, height and width

---

**Returns:** `torch.Tensor`

(L, Cout) tensor

```python
class nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLP(
    config
)
```

**Bases:** `Module`

**`hidden_size`**

---

**`pre_norm`** `= nn.LayerNorm(config.mm_hidden_size, eps=eps)`

---

**`proj`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLP.forward(
    x,
    args = (),
    kwargs = {}
)
```

```python
class nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLPV2(
    config
)
```

**Bases:** `Module`

**`hidden_size`**

---

**`post_norm`** `= nn.RMSNorm(config.hidden_size, eps=eps)`

---

**`proj`**

---

```python
nemo_automodel.components.models.kimi_k3.vision.PatchMergerMLPV2.forward(
    x,
    args = (),
    kwargs = {}
)
```

```python
class nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated(
    dim: int,
    max_height: int,
    max_width: int,
    theta_base = 10000
)
```

**Bases:** `Module`

2D rotary position embedding with multi-resolution support.

This class is intended to be used in the following way:

1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.
3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.
   The rope is shared across all attention layers and all heads.

Refs:

* RoFormer: [https://arxiv.org/abs/2104.09864](https://arxiv.org/abs/2104.09864)
* VisionLLaMA: [https://arxiv.org/abs/2403.00522](https://arxiv.org/abs/2403.00522)
* [https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py](https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py)

**Parameters:**

**`dim`** `int`

usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)

---

**`max_height`** `int`

the maximum height of the 2D grid

---

**`max_width`** `int`

the maximum width of the 2D grid

---

**`theta_base`** `float` — default: 10000

the base of the theta

---

**`device`** `str`

the device to store the precomputed cis

---

```python
nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated._precompute_freqs_cis(
    device: torch.device
) -> torch.Tensor
```

Calculate the cis(freqs) for each position in the 2D grid.

```python
nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated.extra_repr()
```

```python
nemo_automodel.components.models.kimi_k3.vision.Rope2DPosEmbRepeated.get_freqs_cis(
    grid_thws: torch.Tensor,
    device: torch.device
) -> torch.Tensor
```

**Parameters:**

**`grid_thws`** `torch.Tensor`

grid time, height and width

---

**Returns:** `torch.Tensor`

tensor of shape (sum(t \* height \* width), dim//2)

```python
nemo_automodel.components.models.kimi_k3.vision._apply_rope_input_validation(
    x,
    freqs_cis
)
```

```python
nemo_automodel.components.models.kimi_k3.vision.apply_rope(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Returns:
xq\_out, xk\_out: tensors of shape (..., num\_heads, head\_dim)

**Parameters:**

**`xq`** `torch.Tensor`

query, tensor of shape (..., num\_heads, head\_dim)

---

**`xk`** `torch.Tensor`

key, tensor of shape (..., num\_heads, head\_dim)

---

**`freqs_cis`** `torch.Tensor`

tensor of shape (..., head\_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.

---

```python
nemo_automodel.components.models.kimi_k3.vision.eager_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_cu_seqlens: typing.Optional[torch.Tensor] = None,
    k_cu_seqlens: typing.Optional[torch.Tensor] = None,
    kwargs = {}
) -> torch.Tensor
```

```python
nemo_automodel.components.models.kimi_k3.vision.get_1d_sincos_pos_embed(
    embed_dim,
    t_size,
    cls_token = False
)
```

t\_size: int of the temporal size
return:
pos\_embed: \[t\_size, embed\_dim] or \[1+t\_size, embed\_dim] (w/ or w/o cls\_token)

```python
nemo_automodel.components.models.kimi_k3.vision.get_1d_sincos_pos_embed_from_grid(
    embed_dim,
    pos
)
```

From:
[https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi\_modality/models/backbones/internvideo2/pos\_embed.py#L86](https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86)
embed\_dim: output dimension for each position
pos: a list of positions to be encoded: size (M,)
out: (M, D)

```python
nemo_automodel.components.models.kimi_k3.vision.get_rope_shape(
    org,
    interpolation_mode,
    shape
)
```

```python
nemo_automodel.components.models.kimi_k3.vision.get_rope_shape_decorate(
    func
)
```

```python
nemo_automodel.components.models.kimi_k3.vision.multihead_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    q_cu_seqlens: torch.Tensor | None = None,
    k_cu_seqlens: torch.Tensor | None = None,
    max_seqlen_q: int | None = None,
    max_seqlen_k: int | None = None,
    deterministic: bool = False
)
```

Multi-head attention using flash attention 2.

**Parameters:**

**`q, k, v`**

tensor of shape (batch\_size, seqlen, num\_heads, head\_dim),
or (tot\_seqlens, num\_heads, head\_dim) if packing.

---

**`q_cu_seqlens`** `torch.Tensor` — default: None

cumulative sequence lengths of q.
The first element should be 0 and the last element should be q.shape\[0].

---

**`k_cu_seqlens`** `torch.Tensor` — default: None

cumulative sequence lengths of k.
The first element should be 0 and the last element should be k.shape\[0].

---

**Returns:**

shape (batch\_size, seqlen, dim) or (tot\_seqlens, dim) if packing,
where dim = num\_heads \* head\_dim

```python
nemo_automodel.components.models.kimi_k3.vision.tpool_patch_merger(
    x: torch.Tensor,
    grid_thws: torch.Tensor,
    merge_kernel_size: tuple[int, int] = (2, 2)
) -> list[torch.Tensor]
```

```python
nemo_automodel.components.models.kimi_k3.vision.VL_VISION_ATTENTION_FUNCTIONS = {'flash_attention_2': multihead_attention, 'eager': eager_attention}
```