Skip to content

feat(kimi): MoonViT-V2 vision encoder and patch merger for K3.1 - #795

Draft
kalebbroo wants to merge 1 commit into
mainfrom
feat/kimi-k31-vision
Draft

kalebbroo wants to merge 1 commit into
mainfrom
feat/kimi-k31-vision

Conversation

@kalebbroo

Copy link
Copy Markdown
Contributor

Summary

  • Adds MoonViTEncoder (IVlmImageEncoder, Family kimik3) under src/HartsyInference.LLM/Kimi/Vision/: 27 RMS-norm blocks, fused QKV split at load, interleaved 2D RoPE, bilinear position resampling, GELU (tanh) feed-forward, final RMS norm, and the patch merger (PatchMergerMLPV2).
  • KimiVisionWeightLoader reads vision_tower.* and mm_projector.* through IKimiK31WeightSource (the same index as the text decoder), converts to F32, and refuses a missing or misshapen tensor by name.
  • VisionPatchOps (shared, LLM/Multimodal) owns the merge-block patchify and token positions. QwenVlOps.Patchify moved there; Qwen25VlEncoder and Qwen3VlEncoder call it. Qwen's output is unchanged: this is a pure move, and QwenVlOps.BuildRope now takes its positions from the same helper.
  • Not wired: the engine route (TextService.LoadVisionInto, HfTextDirectoryLoader) and the placeholder-to-splice step are left for the route owner. VisionPlaceholderPrompt takes the Kimi literals <|media_begin|>, <|media_pad|>, <|media_end|> through its constructor; VisionTokenSplice takes the encoder output as is.

Source for the open questions (NVIDIA NeMo, main)

https://raw.githubusercontent.com/NVIDIA-NeMo/Automodel/main/nemo_automodel/components/models/kimi_k3/vision.py

  • RoPE is interleaved: xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)). Even pairs take the column (x_pos = flat % max_width), odd pairs the row (ret[h, w, 2*i+1] = cis(w * ...) in the docstring, freqs_cis = cat([x_cis, y_cis], -1)). Not rotate-half.
  • Merge order: reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, d_model) then permute(0, 1, 3, 2, 4, 5), so a merged token holds its 2x2 patches in row-major order. The encoder runs in merge-block order, which is the same permutation, so the merge is a contiguous group.
  • Blocks: norm_type rmsnorm, attn_bias/linear_bias false, activation gelu_pytorch_tanh. Projector: nn.GELU() (exact) between two bias-free linears, then RMSNorm(eps=projector_ln_eps).
  • Config (https://huggingface.co/moonshotai/Kimi-K3.1-Pro/resolve/main/config.json, vision_config): pos_emb_interpolation_mode bilinear, patch_embed_proj_bias false.
  • Checkpoint names and dtypes (model.safetensors.index.json; shard model-00096-of-000096 header): all vision tensors BF16; vision_tower.patch_embed.pos_emb.weight [64, 64, 1024]; vision_tower.encoder.blocks.N.wqkv.weight [4608, 1024]; mm_projector.proj.0.weight and proj.2.weight have no bias entries.
  • Preprocessing (kimi_k3_vision_processing.py, media_utils.py): images are resized, padded with zeros to a multiple of 28, then normalized with mean and std 0.5. Not implemented here; the route must reproduce it.

UNVERIFIED

  1. RMS-norm epsilon (BlockNormEpsilon). The checkpoint does not store it. NeMo calls cast_model_to_dtype(self.vision_tower, model_dtype), so a BF16 run makes nn.RMSNorm's default epsilon 2^-7 (0.0078125). This PR uses the F32 default, 1.19e-7, and the value is a config field. Confirm on real weights.
  2. Interleaved RoPE and merge order are from NeMo source only. No HF modeling_*.py exists for K3 (the repo file list has no modeling file). Nothing was compared against a Python reference run.
  3. Bilinear position resampling with align_corners=False. Taken from PyTorch's F.interpolate default, because NeMo does not pass it.
  4. Attention scale 1/sqrt(head_dim) with head dim 128 (NeMo's eager and flash paths).
  5. Preprocessing pad value. Whether the zero pad is applied before or after normalization (the route must match).
  6. Real weights. Nothing was run on the 1.6 GB checkpoint's vision shard. No GPU runs.
  7. Video (t > 1) is not supported; the encoder takes one image.
  8. Device placement. Weights and tensors are host F32, as in the Qwen-VL encoders. A CUDA or Vulkan route needs weight residency.
  9. Family tag kimik3 is unknown to MultimodalGenerator; the route must add a prompt format.

Tests (CPU lane)

  • dotnet test tests/HartsyInference.LLM.Tests -f net10.0 --filter "Category!=SyntheticSmoke&Category!=Integration&Category!=GpuIntegration&Category!=GpuBenchmark&Category!=Slow&Network!=Real": 196 passed, 0 failed (includes 18 new Kimi.Vision tests).
  • New tests: Patchify block order against a 4x4 pixel grid; bilinear resampling (corner average on a 2x2 to 1x1, identity on a native grid); projector golden on one merged token with g(1) and g(2) as constants; interleaved RoPE angles and the rotary table layout; loader refusals by name (missing tensor, wrong shape); BF16-to-F32 load through a real shard and index; encoder token count, finiteness, determinism, image dependence and size refusals.
  • Not run: the Qwen-VL LLM encoders have no tests in this repo, so the refactor is covered only by the pure-move argument and the merge-block order test. No Diffusion or Core tests were run; this PR does not touch them.

Changelog

One line under ## Unreleased, as requested in the task. No version bump (VersionSuffix unchanged).

Adds MoonViTEncoder (IVlmImageEncoder) and its loader for vision_tower.* and
mm_projector.* read through the same checkpoint index as the text decoder.

Token order follows NVIDIA NeMo kimi_k3/vision.py: patches in merge-block
order, interleaved 2D RoPE (even pairs take the column angle, odd pairs the
row angle), bilinear position resampling with align_corners off, GELU tanh in
blocks, exact GELU and RMS post-norm in the projector.

VisionPatchOps (shared) now owns the merge-block patchify and position order;
the Qwen-VL encoders call it instead of QwenVlOps.Patchify.

UNVERIFIED on real weights: the RMS-norm epsilon. The checkpoint does not
store it. NeMo casts the tower to the model dtype (BF16), which would make
PyTorch's default epsilon 2^-7; this uses the F32 default 1.19e-7.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant