Module vision 

Kimi K3’s MoonViT-V2 vision encoder: patch embedding -> N transformer encoder layers (RMSNorm, self-attention with 2D RoPE, gated-MLP-free plain MLP2 feed-forward) -> a patch-merger projector into the text decoder’s embedding space. Genuinely new territory for ferrox, which has been text-only until now.

Transcribed directly from real reference source fetched live from the model repo (not guessed, not derived by analogy to other ViT/RoPE designs): moonshotai/Kimi-K3’s modeling_kimi_k3.py (MoonVision3dPatchEmbed, Rope2DPosEmbRepeated, apply_rope, MoonViTEncoderLayer, MoonViT3dEncoder, tpool_patch_merger, PatchMergerMLPV2) and its real config.json’s vision_config (norm_type='rmsnorm', activation_func='gelu_pytorch_tanh', attn_bias=False, linear_bias=False, patch_embed_proj_bias=False, patch_size=14, vt_hidden_size=1024, vt_num_attention_heads=12, qkv_hidden_size=1536, vt_intermediate_size=4096, vt_num_hidden_layers=27, merge_kernel_size=[2,2], mm_projector_type='patchmergerv2', projector_ln_eps=1e-5).

Deliberately scoped, and disclosed as such, to the common case rather than every real code path:

• Single image, single frame (t=1) – no multi-frame video input, so the real code’s temporal sincos position-embedding addition and the patch-merger’s temporal-mean pooling are both no-ops here, not implemented as general code paths.
• The caller’s patch grid is assumed to exactly match the positional-embedding grid it was constructed with, so the real code’s bicubic/bilinear interpolation branch (get_rope_shape, for when input resolution differs from the pretrained grid) is never exercised and not implemented.
• One image at a time – no variable-length multi-image batch packing (cu_seqlens); process a full patch sequence for one image per call.
• The real multimodal splicing of vision output into the text decoder’s input sequence (KimiK3ForConditionalGeneration.forward) is out of scope here – this module only implements the vision tower itself.

Two real, non-obvious facts confirmed by reading source rather than assuming standard ViT/RoPE conventions:

1. MoonVision3dPatchEmbed.forward reassigns x before calling x.size(0) in the same statement (x = self.proj(x).view(x.size(0), -1)), but Python evaluates the right-hand side of an assignment before rebinding the name, so x.size(0) refers to the original, pre-conv x – meaning the real input to this module is already a batch of N individually pre-extracted [C, patch_size, patch_size] patches (one “image” per patch), not one whole image run through a strided conv. A Conv2d(kernel_size=patch_size, stride=patch_size) applied to an input exactly patch_size wide/tall is mathematically just one linear projection of the flattened patch – so patch embedding here is a single shared Linear(in_dim*patch_size*patch_size, out_dim, bias=false) applied per patch (WeightMatrix::apply), not a general strided convolution.
2. The 2D RoPE (apply_rope) rotates consecutive real-value pairs ((x[2i], x[2i+1]), the original RoFormer/GPT-NeoX convention via view_as_complex), NOT the “rotate half” convention ((x[i], x[i+dim/2])) ferrox_core::attention::apply_rope implements for the text decoder’s RoPE – a distinct function is needed here, not a reuse of that one. Within each complex pair index j in [0, head_dim/2), _precompute_freqs_cis interleaves width- and height-based angles: even j uses the patch’s width coordinate, odd j uses its height coordinate (confirmed by reading the actual torch.cat(...).reshape(...) code, which contradicts that same function’s own docstring – code wins).

Not yet wired into Decoder or spliced into the text embedding sequence. Tested here against synthetic weights, cross-validated against an independent Python transcription of the same real algorithm.

Structs

VisionConfig

VisionEncoderLayerWeights

VisionEncoderWeights

VisionMergerWeights

Functions

embed_patches

Embeds a full image’s patches (already extracted as flattened in_dim*patch_size*patch_size pixel vectors, row-major (h, w) order) and adds the learned position embedding. Returns [n_patches, hidden_dim] flattened.

encoder_forward

Runs the full encoder (patch embed + every layer + final norm) for one image. Returns [n_patches, hidden_dim] flattened.

patch_merge

tpool_patch_merger for the single-frame (t=1) case: groups spatially-adjacent merge_kh x merge_kw patches (row-major (h, w)) into one merged patch each, without the temporal mean (a no-op at t=1). Returns [num_merged, merge_kh*merge_kw*hidden_dim] flattened – already contiguous per merged patch, ready for the projector.

project_merged_patches

PatchMergerMLPV2: per merged patch, Linear -> GELU (erf) -> Linear -> RMSNorm. merged is [num_merged, merge_kh*merge_kw*hidden_dim] flattened (patch_merge’s output); returns [num_merged, text_hidden] flattened.