Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsKimi K3nemo_automodel.components.models.kimi_k3.multimodal

nemo_automodel.components.models.kimi_k3.multimodal

View as Markdown

Native top-level Kimi K3 vision-language model.

Module Contents

Classes

NameDescription
KimiK31ForConditionalGenerationKimi K3.1 vision-language model: Kimi K3 with gated DSA on the full-attention layers.
KimiK3ForConditionalGenerationKimi K3 MoonViT3d tower, projector, and native KDA/MLA language model.

Functions

NameDescription
_positions_within_documentsReturn zero-based positions within each document, skipping masked tokens.
_projector_config-
_vision_tower_config-

Data

ModelClass

API

class nemo_automodel.components.models.kimi_k3.multimodal.KimiK31ForConditionalGeneration()

Bases: KimiK3ForConditionalGeneration

Kimi K3.1 vision-language model: Kimi K3 with gated DSA on the full-attention layers.

The architecture differs from Kimi K3 only through text_config.dsa_config, which KimiMLAAttention reads to attach a sparse-attention indexer to the listed layers.

ModelCapabilities
= KimiK3ForConditionalGeneration.ModelCapabilities
nemo_automodel.components.models.kimi_k3.multimodal.KimiK31ForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> 'KimiK31ForConditionalGeneration'
classmethod
class nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: KimiK3ForCausalLM

Kimi K3 MoonViT3d tower, projector, and native KDA/MLA language model.

ModelCapabilities
= KimiK3ForCausalLM.ModelCapabilities
_cp_batch_sharder
= staticmethod(shard_batch_for_kimi_vlm_cp)
cp_mesh
DeviceMesh | None = None
mm_projector
= projector_cls(projector_config)
state_dict_adapter
vision_tower
nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration._extract_image_features(
    pixel_values: torch.Tensor,
    grid_thws: torch.Tensor
) -> list[torch.Tensor]

Encode packed pixel_values using grid_thws image/video geometry.

nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration._merge_input_ids_with_image_features(
    image_features: list[torch.Tensor],
    inputs_embeds: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    labels: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]

Merge compact HF placeholders or pre-expanded training image tokens.

Parameters:

image_features
list[torch.Tensor]

Per-image tensors of shape [image_tokens, hidden], in batch order. Each image_tokens count follows spatial patch merging.

inputs_embeds
torch.Tensor

Tensor of shape [batch, sequence, hidden].

input_ids
torch.Tensor

Tensor of shape [batch, sequence], with either one media placeholder per image or one per projected image token.

attention_mask
torch.Tensor

Binary or indexed document mask of shape [batch, sequence].

labels
torch.Tensor | NoneDefaults to None

Optional unshifted targets of shape [batch, sequence].

Returns: torch.Tensor

Embeddings [batch, expanded_sequence, hidden], binary or indexed document mask,

nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    block_residual: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    grid_thws: torch.Tensor | None = None,
    image_grid_hws: torch.Tensor | None = None,
    image_grid_thw: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool | None = None,
    kimi_vlm_full_input: bool = False,
    kimi_vlm_token_order: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Run text-only or multimodal K3 inputs.

Parameters:

input_ids
torch.Tensor | NoneDefaults to None

Token IDs [batch, sequence], or PP activations [batch, sequence, hidden] after the first pipeline stage.

block_residual
torch.Tensor | NoneDefaults to None

PP residual starts [batch * sequence, blocks, hidden].

pixel_values
torch.Tensor | NoneDefaults to None

Flattened image patches [patches, channels, patch_h, patch_w].

grid_thws
torch.Tensor | NoneDefaults to None

Image geometry [images, 3], ordered time, height, width.

image_grid_thw
torch.Tensor | NoneDefaults to None

Canonical packed image geometry [images, 3].

image_grid_hws
torch.Tensor | NoneDefaults to None

Image geometry [images, 2] from the native VLM collator; a [images, 3] time-height-width grid is also accepted.

attention_mask
torch.Tensor | NoneDefaults to None

Binary or indexed document mask [batch, sequence].

position_ids
torch.Tensor | NoneDefaults to None

Optional positions [batch, sequence].

inputs_embeds
torch.Tensor | NoneDefaults to None

Optional embeddings [batch, sequence, hidden].

labels
torch.Tensor | NoneDefaults to None

Optional unshifted targets [batch, sequence], with ignore_index at unsupervised positions. Native recipes compute loss externally.

logits_to_keep
int | torch.TensorDefaults to 0

Trailing logit count or sequence indices [selected_tokens].

output_hidden_states
bool | NoneDefaults to None

Whether to return final hidden states.

kimi_vlm_full_input
boolDefaults to False

Whether CP left the primary token stream intact for image embedding before contiguous sequence sharding.

kimi_vlm_token_order
torch.Tensor | NoneDefaults to None

Optional source index of each CP sequence position [batch, padded_sequence], set when CP moved masked tokens behind valid ones.

**kwargs
AnyDefaults to {}

Text-backbone metadata, following KimiK3ForCausalLM.forward.

Returns: CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Causal LM output with logits [batch, expanded_sequence, vocab] and

nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration.from_config(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> 'KimiK3ForConditionalGeneration'
classmethod
nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> 'KimiK3ForConditionalGeneration'
classmethod
nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Describe global media token inputs and local CP stage activations.

Parameters:

is_first
bool

Whether this is the first pipeline stage.

microbatch_size
int

Number of sequences per pipeline microbatch.

seq_len
int

Global input token count before CP padding and slicing.

dtype
torch.dtype

Floating point dtype of pipeline activations.

Returns: tuple[torch.Tensor, ...]

Input and output meta tensors. Only first-stage token IDs retain

nemo_automodel.components.models.kimi_k3.multimodal.KimiK3ForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

Initialize text and vision after distributed materialization.

nemo_automodel.components.models.kimi_k3.multimodal._positions_within_documents(
    attention_mask: torch.Tensor
) -> torch.Tensor

Return zero-based positions within each document, skipping masked tokens.

Masked tokens inside a document, such as deliberately masked fake-image spans, do not advance positions or start a new document.

Parameters:

attention_mask
torch.Tensor

Binary or indexed document mask [batch, sequence].

Returns: torch.Tensor

Positions [batch, sequence], with masked positions set to one.

nemo_automodel.components.models.kimi_k3.multimodal._projector_config(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3Config
) -> types.SimpleNamespace
 
nemo_automodel.components.models.kimi_k3.multimodal._vision_tower_config(
    config: nemo_automodel.components.models.kimi_k3.config.KimiK3Config
)
 
nemo_automodel.components.models.kimi_k3.multimodal.ModelClass = KimiK3ForConditionalGeneration