bridge.models.kimi_vl.data.collate_fn#

Kimi VL collator implementations.

Module Contents#

Functions#

_kimi_assistant_mask_boundary_config

Build Kimi assistant loss boundaries and trim only empty thinking blocks.

_expand_image_tokens_and_aligned_mask

Expand image placeholder tokens and any aligned per-token loss mask.

_expand_image_tokens

Expand image placeholder tokens to the correct count based on grid_thws.

kimi_k25_vl_collate_fn

Collate function for Kimi K2.5 VL processors with pre-expanded image tokens.

Data#

KIMI_ASSISTANT_START

KIMI_ASSISTANT_END

KIMI_THINK_OPEN

KIMI_THINK_CLOSE

API#

bridge.models.kimi_vl.data.collate_fn.KIMI_ASSISTANT_START#

‘<|im_assistant|>assistant<|im_middle|>’

bridge.models.kimi_vl.data.collate_fn.KIMI_ASSISTANT_END#

‘<|im_end|>’

bridge.models.kimi_vl.data.collate_fn.KIMI_THINK_OPEN#

‘’

bridge.models.kimi_vl.data.collate_fn.KIMI_THINK_CLOSE#

‘’

bridge.models.kimi_vl.data.collate_fn._kimi_assistant_mask_boundary_config(
processor: Any,
) → megatron.bridge.data.vlm_processing.AssistantMaskBoundaryConfig#

Build Kimi assistant loss boundaries and trim only empty thinking blocks.

bridge.models.kimi_vl.data.collate_fn._expand_image_tokens_and_aligned_mask(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
loss_mask: torch.Tensor | None,
grid_thws: torch.Tensor,
media_token_id: int,
merge_kernel_size: tuple[int, int] = (2, 2),
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]#

Expand image placeholder tokens and any aligned per-token loss mask.

bridge.models.kimi_vl.data.collate_fn._expand_image_tokens(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
grid_thws: torch.Tensor,
media_token_id: int,
merge_kernel_size: tuple[int, int] = (2, 2),
) → tuple[torch.Tensor, torch.Tensor]#

Expand image placeholder tokens to the correct count based on grid_thws.

For PP, this ensures the sequence length is fixed BEFORE the model forward pass, eliminating dynamic sequence expansion inside the model.

Parameters:

  • input_ids – (seq_len,) tensor with one placeholder per image

  • attention_mask – (seq_len,) tensor

  • grid_thws – (num_images, 3) tensor with [t, h, w] for each image

  • media_token_id – Token ID of the image placeholder

  • merge_kernel_size – Vision tower’s patch merge kernel, default (2, 2)

Returns:

Input IDs with placeholder expanded to N tokens expanded_attention_mask: Attention mask expanded accordingly

Return type:

expanded_input_ids

bridge.models.kimi_vl.data.collate_fn.kimi_k25_vl_collate_fn(
examples: list[dict[str, Any]],
processor,
max_length: int | None = None,
*,
visual_keys: object = None,
min_pixels: int | None = None,
max_pixels: int | None = None,
sequence_length: int | None = None,
pad_to_max_length: bool = False,
pad_to_multiple_of: int = 128,
enable_in_batch_packing: bool = False,
in_batch_packing_pad_to_multiple_of: int = 1,
) → dict[str, torch.Tensor]#

Collate function for Kimi K2.5 VL processors with pre-expanded image tokens.

For pipeline parallelism, this function:

  1. Processes each sample to get input_ids with 1 placeholder per image

  2. Pre-expands each placeholder to N tokens (N = t*(h//2)*(w//2) from grid_thws)

  3. Pads all sequences to fixed max_length This ensures the model forward pass doesn’t change sequence length dynamically.