feat(kimi-vl): 3D MoonViT video support (image + video patch embedding)

[inureyes]

Summary

Add video input support for Kimi-VL 2.5 (model_type: "kimi_k25"). The model extends the MoonViT native-resolution image encoder to clips described by a 3-entry patch grid (t, h, w) (frames, patch rows, patch columns). The extension is not a temporal convolution: the patch projection stays the same per-frame 2D conv used for images, and the checkpoint contains no video-specific tensors. Video support consists of (1) a computed temporal position embedding, (2) tiling the 2D rotary tables across frames, (3) one attention segment spanning the whole clip, (4) temporal mean-pooling in the patch merger, and (5) frame sampling plus placeholder expansion in the processing path. All of the model-side pieces are additions to the existing MoonViT code; the frame decoding infrastructure already exists in src/multimodal/video.rs.

The former blocker (base MoonViT image support) has landed, so this issue is ready to implement.

Current state in mlxcel

The image path for both kimi_vl and kimi_k25 is in-tree and working:

• src/models/detection.rs (lines 211-212): "kimi_vl" => ModelType::KimiVL, "kimi_k25" => ModelType::KimiK25.
• src/loading/mod.rs (line 204): both route to load_kimi_vl_vlm.
• src/loading/vlm_kimi_vl.rs: loader; key remap (vision_tower.encoder.* strip, blocks.{i}.{wqkv,wo} to blocks.{i}.attn.*, patch-embed conv permutation, mm_projector.* vs multi_modal_projector.* prefix detection).
• src/vision/kimi_vl.rs: KimiVLModel (MoonViT tower + projector + DeepSeek-V3 text backbone). get_input_embeddings takes grid_shapes: &[(i32, i32)], i.e. 2D image grids only.
• src/vision/encoders/kimi_vl.rs + kimi_vl_pos_emb.rs + kimi_vl_rope.rs: encoder. forward_with_grid(pixel_values, shapes: &[(i32, i32)]) builds cu_seqlens as h*w per image, applies the learned 2D interpolated position embedding, 2D RoPE, block-diagonal attention, and a 2x2 spatial patch_merger. Everything is keyed on (h, w) pairs; there is no t anywhere.
• src/vision/processors/kimi_vl.rs: image-only preprocessing (preprocess_with_grid), producing (total_patches, C=3, pH=14, pW=14) channels-first patches plus per-image (grid_h, grid_w).
• src/multimodal/kimi_vl_prompt.rs: insert_kimi_vl_image_tokens expands each <|media_pad|> into (h/merge)*(w/merge) placeholder tokens; grids are (h, w) pairs.
• src/multimodal/vlm_runtime.rs (KimiVL arm, around line 962): image requests only.

What happens today when a video request hits a kimi_k25 model: the model loads and serves image and text requests normally, but video inputs are rejected before reaching the model. The CLI errors with --video input is currently only supported by Gemma4 VLMs (src/commands/generate_vlm.rs, around line 260), and the server returns HTTP 400 video inputs are only supported by Gemma 4 VLM models in this build (src/server/model_worker.rs, prepare_request_video_embeddings). No crash, no silent frame drop; a clean up-front rejection.

Video infrastructure that already exists and must be reused:

• src/multimodal/video.rs (+ video_tests.rs): ffmpeg-based probe and uniform frame extraction. load_video_source / load_videos return Vec<DynamicImage> per clip. smart_nframes implements the exact sampling rule this model needs: target fps default DEFAULT_FPS = 2.0, count clamped to [FPS_MIN_FRAMES = 4, min(FPS_MAX_FRAMES = 768, total_frames)], floored to a multiple of FRAME_FACTOR = 2. VideoSource::Fd provides the fd-backed anti-TOCTOU path used by the server media resolver. Gemma 4 is the existing in-tree consumer (compute_gemma4_video_embeddings in src/commands/generate_vlm.rs, prepare_request_video_embeddings in src/server/model_worker.rs); the KimiVL wiring should follow the same pattern.

Architecture

All shapes below use explicit named dimension order. patch_size = 14, embed_dim = 1152, num_heads = 16, head_dim = embed_dim / num_heads = 72, merge = spatial_merge_size = 2.

Media grids

Every media item carries a patch grid. Images use 2 entries (h, w); videos use 3 entries (t, h, w) where t is the number of sampled frames and (h, w) is the per-frame patch grid shared by all frames of the clip. A request may mix images and videos; grid entries appear in media order and the pixel patches are concatenated in the same order.

Patch embedding: same 2D weight, applied per frame

The checkpoint stores exactly one patch-projection tensor, vision_tower.patch_embed.proj.weight, with shape (out_ch = 1152, in_ch = 3, kH = 14, kW = 14) plus a bias (1152,). There is no kT axis and no 5D tensor; the vision config key temporal_patch_size (integer, value 2) exists but does not describe a convolution kernel depth. Its only observable effect is on frame sampling granularity: sampled frame counts are multiples of 2 (which is exactly FRAME_FACTOR = 2 in src/multimodal/video.rs).

Each frame is patchified independently, exactly like an image. For a clip the packed patch tensor is (t*h*w, C = 3, pH = 14, pW = 14) channels-first, ordered frame-major: frame index outermost, then row-major (row, col) within each frame. So patch index n = f*(h*w) + row*w + col.

MLX convolutions are channels-last, so the model transposes with axes [0, 2, 3, 1] to (t*h*w, pH = 14, pW = 14, C = 3) before the stride-14 conv (the existing image path in src/vision/kimi_vl.rs::get_input_embeddings already does this transpose). The conv output (t*h*w, 1, 1, 1152) is reshaped to (t*h*w, embed_dim = 1152).

Position embedding across frames

Two additive terms, both applied at the patch-embed stage:

1. Spatial (learned, stored): the (init_h = 64, init_w = 64, embed_dim = 1152) grid vision_tower.patch_embed.pos_emb.weight, bicubically resampled to (h, w, 1152) and flattened to (h*w, 1152), exactly as for images. For a video this per-frame table is tiled t times along axis 0, giving (t*h*w, 1152) in the same frame-major order as the patches.
2. Temporal (computed, NOT stored): a 1D sinusoid over the frame index. With half = embed_dim / 2 = 576:
   • pos has shape (t, 1), values 0..t-1
   • freq has shape (1, 576), freq[j] = exp(-ln(10000) * j / 576) for j in 0..576
   • ang = pos * freq has shape (t, 576) (broadcast of (t, 1) against (1, 576))
   • temporal = concat([sin(ang), cos(ang)], axis = 1) has shape (t, 1152) (if embed_dim were odd, zero-pad on axis 1 then truncate to embed_dim)
   • Row f of this table is repeated h*w times consecutively along axis 0 (frame f occupies rows f*h*w .. (f+1)*h*w), giving (t*h*w, 1152), and added to the tiled spatial term.

Important consequence: a t = 1 clip is not bit-identical to the same frame processed as an image. The frame-0 temporal row is sin(0) = 0 for columns 0..576 and cos(0) = 1 for columns 576..1152, a nonzero constant added to every token. Tests must assert this exact relationship, not naive equality (see acceptance criteria).

Rotary position embedding

The encoder's 2D RoPE (src/vision/encoders/kimi_vl_rope.rs) uses only spatial (row, col) positions; the frame index never enters the angles. For a video, the per-frame angle table of shape (h*w, head_dim/2 = 36) is tiled t times along axis 0 to (t*h*w, 36). Angle layout per token is unchanged from the image path: for j in 0..head_dim/4 = 18, column 2j is col * 10000^(-4j/head_dim) and column 2j+1 is row * 10000^(-4j/head_dim). There is no temporal rotary component.

Attention span

cu_seqlens segments cover the full token count per media item: t*h*w for a video, h*w for an image. All frames of one clip form a single bidirectional attention segment (every patch of every frame attends to every patch of every other frame of the same clip). Different media items never attend to each other, same as today.

Token packing and merging (temporal pooling)

After the final LayerNorm, the merger processes each media item's slice of the (total_tokens, 1152) sequence:

• For a video item, slice (t*h*w, 1152), reshape to (t, h*w, 1152) (axis 0 = frame, axis 1 = spatial position, axis 2 = channel), then mean over axis 0, producing (h*w, 1152). The whole clip collapses to one spatial token map before any spatial merging.
• Then the existing 2x2 spatial merge, identical to images: reshape (h*w, 1152) to (h/2, 2, w/2, 2, 1152), transpose axes (0, 2, 1, 3, 4) to (h/2, w/2, 2, 2, 1152), reshape to ((h/2)*(w/2), 4, 1152).

So a clip contributes (h/2)*(w/2) merged tokens regardless of t, the same count as a single image of that frame size. The projector is unchanged: (total_merged, 4, 1152) through pre_norm, flatten axes 1-2 to (total_merged, 4608), then Linear -> GELU -> Linear to (total_merged, text_hidden).

Frame sampling (processing)

Frames are sampled uniformly by index over [0, total_frames - 1]. The count is total_frames / native_fps * target_fps with target_fps defaulting to 2.0, clamped to [4, min(768, total_frames)], floored to a multiple of 2; an explicit frame-count override is rounded to a multiple of 2. This is exactly what smart_nframes in src/multimodal/video.rs already computes; no new sampling code is needed.

Each sampled frame is preprocessed exactly like a single image (resize under the patch budget, crop/pad to multiples of merge * patch_size = 28, normalize, patchify). All frames of a decoded clip have identical pixel dimensions, so computing the resize geometry once from frame 0 and applying it to every frame guarantees a single shared (h, w) for the clip.

Placeholder expansion

Videos use the same <|media_pad|> placeholder (media_placeholder_token_id, default 163606) as images; the chat template wraps the media span in <|media_start|> / <|media_end|>. One placeholder per media item expands to that item's merged-token count: (h/2)*(w/2) for a video, independent of t because of the temporal mean-pool. The expansion formula in src/multimodal/kimi_vl_prompt.rs is therefore already correct per item; it only needs to accept video grid entries.

Implementation plan

1. src/vision/encoders/kimi_vl.rs
   • Introduce a media-grid type, e.g. enum KimiMediaGrid { Image { h: i32, w: i32 }, Video { t: i32, h: i32, w: i32 } }. Do not normalize images to t = 1 videos: the two differ by the temporal embedding term (see above), so the distinction must be explicit.
   • forward_with_grid accepts &[KimiMediaGrid]. cu_seqlens becomes t*h*w per video item.
   • patch_merger gains the temporal pool: for video items, reshape the item slice (t*h*w, dim) to (t, h*w, dim) and mean over axis 0 before the existing 2x2 merge.
   • Parse temporal_patch_size (serde default 2) into KimiVLVisionConfig (currently the key is silently ignored) and use it as the frame-count rounding factor handed to sampling.
2. src/vision/encoders/kimi_vl_pos_emb.rs
   • Add the temporal sinusoid builder ((t, embed_dim) table as specified above, computed on the fly, never a weight).
   • Extend add_to for video items: tile pos_for(h, w) t times along axis 0, add the temporal table with each row repeated h*w times along axis 0.
3. src/vision/encoders/kimi_vl_rope.rs
   • cos_sin accepts the media-grid list; for video items emit the per-frame (row, col) position stream t times. Angle math is untouched.
4. src/vision/processors/kimi_vl.rs
   • Add preprocess_video_with_grid(clips: &[Vec<image::DynamicImage>]) -> (patches, Vec<(t, h, w)>): compute rescale geometry from frame 0, apply the identical rescale/crop/normalize/patchify to every frame, emit patches frame-major as (t*h*w, 3, 14, 14). Per-frame patch budget identical to the image path.
5. src/vision/kimi_vl.rs
   • get_input_embeddings takes the mixed media-grid list (images and videos in media order). The [0, 2, 3, 1] channels-last transpose and the projector stay as they are.
6. src/multimodal/kimi_vl_prompt.rs
   • Generalize insertion over media grids; per-item expansion count is (h/merge)*(w/merge) for both kinds. Keep the existing image-only tests green.
7. src/multimodal/vlm_runtime.rs
   • Extend the KimiVL arm to accept videos: decode via crate::multimodal::video (reuse load_video_source / load_videos and per-video fps overrides exactly as the Gemma 4 arm does), preprocess clips, build the mixed grid list, run insertion and embeddings. Add a KimiVL video variant to VlmPreparationSummary for logging.
8. src/commands/generate_vlm.rs
   • Lift the Gemma4-only --video gate (around line 260) to also route KimiVL-loaded models; default fps 2.0 via the existing flag.
9. src/server/model_worker.rs
   • Add a KimiVL branch to prepare_request_video_embeddings (keep the fd-backed VideoSource path and the allowlist behavior); remove kimi_k25 from the 400-rejection set.

Weight loading and sanitize

The video path uses exactly the image-path tensor set; the checkpoint contains no video-specific weights.

• vision_tower.patch_embed.proj.weight is stored as (out_ch = 1152, in_ch = 3, kH = 14, kW = 14). The loader (transform_vision_key in src/loading/vlm_kimi_vl.rs) already permutes it with axes [0, 2, 3, 1] to the MLX channels-last layout (out_ch = 1152, kH = 14, kW = 14, in_ch = 3), guarded by conv2d_weight_is_channel_last so an already-permuted checkpoint is passed through. This rank-4 handling is complete for video: since there is no kT axis, the video path reuses this weight verbatim, one conv application per frame. Nothing is currently mis-loaded or skipped, and no loader change is required for weights.
• The temporal position embedding is a computed sinusoid, not a checkpoint parameter; RoPE has no weights; the temporal mean-pool has no weights.
• The only load-time delta is config parsing: add temporal_patch_size: usize (default 2) to KimiVLVisionConfig so the value present in kimi_k25 checkpoints is read instead of ignored.

Validation and acceptance criteria

Unit tests (synthetic weights, src/vision/encoders/kimi_vl_tests.rs and module tests):

• Temporal sinusoid: shape (t, 1152); row 0 equals zeros in columns 0..576 and ones in columns 576..1152; spot-check a t > 1 row against the closed form.
• Position embedding relationship: for any (h, w), the video table for grid (1, h, w) equals the image table for (h, w) plus the frame-0 temporal row broadcast to all h*w rows.
• RoPE tiling: the (t*h*w, 36) video table equals the (h*w, 36) image table repeated t times along axis 0.
• Merger temporal pool: construct distinct per-frame inputs and assert the output equals the frame mean followed by the existing 2x2 merge; output shape ((h/2)*(w/2), 4, dim) for any t.
• cu_seqlens: one segment of t*h*w per clip, mixed with h*w image segments in media order.
• Placeholder expansion: a video grid (t, h, w) expands its <|media_pad|> to (h/2)*(w/2) tokens independent of t; mixed image+video prompts expand in media order.

Single-frame equivalence (the guaranteed relationship, not naive equality): a t = 1 clip through the full tower must produce exactly the image-path output once the frame-0 temporal constant is accounted for (either by adding it to the image path's patch embeddings in the test, or via a test-only hook that zeroes the temporal term). Token counts must match exactly.

Regression: image-only requests through a kimi_vl and a kimi_k25 checkpoint produce token-identical outputs before and after this change.

Real-checkpoint smoke: load a kimi_k25 checkpoint and run a short video (a few seconds) end-to-end through both mlxcel generate-vlm --video <file> and a server chat completion with a video_url part. Assertions: the number of expanded placeholder tokens equals the number of projected feature rows, generation completes without shape errors, and the output plausibly describes the clip's temporal content.

Merge gate: cargo test, cargo clippy, cargo fmt --check clean (CI does not build this workspace's tests; local runs are the gate).

Out of scope

• Temporal attention or temporal RoPE components: the architecture has none; do not invent them.
• Video support for any other model family, and any change to the Gemma 4 video paths beyond shared-code reuse.
• Audio, mixed video+audio requests, streaming video, and long-video chunking policies.
• Aligning the image-path normalization constants with per-checkpoint preprocessor configuration; the video path must simply match whatever the image path does today.
• Vision-feature caching for video and any batching/scheduler work.