diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi.py b/scripts/checkpoint_conversion/numerical_tests_kimi.py
new file mode 100644
index 0000000000..ba4471144c
--- /dev/null
+++ b/scripts/checkpoint_conversion/numerical_tests_kimi.py
@@ -0,0 +1,317 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Full text+image e2e logit parity: torchtitan Kimi-VL vs the released HF model.
+
+Runs the released HF Kimi-VL (``trust_remote_code``) and torchtitan in ONE
+process on the same text+image prompt and compares last-token logits. torchtitan
+does its OWN Kimi image processing (``process_image`` + ``vision_to_patches``,
+raster order), so the full pipeline is exercised (preprocessing + vision +
+projector + scatter + DeepSeek-V3 text tower), not just the forward.
+
+The released remote code targets transformers ~4.50.x and does NOT import on 5.x,
+so run this in an env with ``transformers==4.50.3`` + ``tiktoken`` + ``blobfile``.
+
+Precision: the HF reference runs at ``--hf_dtype`` (default float32, the model's
+"true" output); torchtitan runs at ``--dtype`` (text) with ``--vision_dtype``
+overriding only its vision encoder. ``--dtype float32`` is the correctness gate;
+``--dtype bfloat16 --vision_dtype float16`` is the realistic config (~1e-2 KL).
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python -m \\
+        scripts.checkpoint_conversion.numerical_tests_kimi \\
+        --hf_model_path ~/hf_assets/moonshotai/Kimi-VL-A3B-Instruct \\
+        --tt_checkpoint_path outputs/kimi/kimi_vl_a3b_dcp --dtype float32
+
+Add ``--force-hf-routing`` to make titan use HF's exact per-token expert
+selections (diagnostic: removes the MoE routing-flip divergence).
+"""
+
+import argparse
+import os
+
+import torch
+import torch.distributed.checkpoint as dcp
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+from torchtitan.components.checkpoint import ModelWrapper
+from torchtitan.hf_datasets.multimodal.utils.image import (
+    process_image,
+    resize_to_patch_budget,
+    vision_to_patches,
+)
+from torchtitan.models.common.attention import ScaledDotProductAttention
+from torchtitan.models.kimi_k2_7 import model_registry
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+_MEDIA_TOKEN_ID = 163605
+_PATCH_SIZE = 14
+_MERGE_SIZE = 2
+_PROMPT = "<|media_pad|>\nWhat is shown in this image? Describe it briefly."
+
+
+class _VisionSDPA(nn.Module):
+    """Bidirectional SDPA for one image (no padding -> full attention is exact)."""
+
+    def forward(self, q, k, v, **kwargs):
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return out.transpose(1, 2)
+
+
+@torch.no_grad()
+def run_hf(hf_model_path, image_size, dtype, device):
+    """Released HF Kimi-VL (trust_remote_code) on a text+image prompt.
+
+    Returns the raw image, tokenized input_ids, last-token logits, the projector
+    output (scattered vision features), and per-MoE-layer top-k expert indices.
+    """
+    print(f"Loading released HF Kimi-VL on {device} ...")
+    proc = AutoProcessor.from_pretrained(hf_model_path, trust_remote_code=True)
+    model = (
+        AutoModelForCausalLM.from_pretrained(
+            hf_model_path, trust_remote_code=True, torch_dtype=dtype
+        )
+        .to(device)
+        .eval()
+    )
+
+    raw_image = (
+        torch.linspace(0, 255, image_size * image_size * 3)
+        .reshape(image_size, image_size, 3)
+        .to(torch.uint8)
+    )
+    # pyrefly: ignore [not-callable]
+    batch = proc(
+        text=[_PROMPT], images=[Image.fromarray(raw_image.numpy())], return_tensors="pt"
+    )
+
+    vis = {}
+    model.multi_modal_projector.register_forward_hook(
+        lambda m, i, o: vis.__setitem__("f", o.detach().float().cpu())
+    )
+    # MoE layers have a ``gate`` (MoEGate) whose forward returns
+    # (topk_idx, topk_weight, aux_loss); ``experts`` is a ModuleList (never
+    # called directly). Hook the gate to record the per-token expert selection.
+    expert_indices: dict[int, torch.Tensor] = {}
+    for i, layer in enumerate(model.language_model.model.layers):
+        if hasattr(layer.mlp, "gate"):
+            layer.mlp.gate.register_forward_hook(
+                lambda m, inp, out, i=i: expert_indices.__setitem__(
+                    i, out[0].detach().cpu()
+                )
+            )
+
+    inputs = {
+        k: (v.to(device) if isinstance(v, torch.Tensor) else v)
+        for k, v in batch.items()
+    }
+    inputs["pixel_values"] = inputs["pixel_values"].to(dtype)
+    out = model(**inputs)
+    ref = {
+        "input_ids": batch["input_ids"].cpu(),
+        "raw_image": raw_image,
+        "last_logits": out.logits[:, -1, :].float().cpu(),
+        "vision_features": vis["f"],
+        "expert_indices": expert_indices,
+    }
+    del model
+    torch.cuda.empty_cache()
+    return ref
+
+
+def _force_hf_routing(model, expert_indices, device):
+    """Monkeypatch each MoE router to use HF's recorded top-k experts.
+
+    Diagnostic: titan routes every token to exactly the experts HF chose (still
+    computing its own gating weights at those experts via the real scores), so the
+    discrete routing-flip divergence is removed and any residual is the non-routing
+    math (attention / expert FFN / fp). ``expert_indices`` maps the HF layer index
+    to a ``(num_tokens, top_k)`` LongTensor.
+    """
+    forced = 0
+    for key, layer in model.layers.items():
+        if not getattr(layer, "moe_enabled", False) or int(key) not in expert_indices:
+            continue
+        ids = expert_indices[int(key)]
+        ids = ids.view(1, ids.shape[0], ids.shape[-1]).to(device)  # (1, L, K)
+        router = layer.moe.router
+        orig = router.forward
+
+        def forced_forward(x_BLD, expert_bias_E=None, _r=router, _o=orig, _ids=ids):
+            # Reuse the real score computation; override only the selection.
+            _, _, scores_BLE = _o(x_BLD, expert_bias_E)
+            topk = scores_BLE.gather(dim=-1, index=_ids)
+            if _r.route_norm:
+                topk = topk / (topk.sum(dim=-1, keepdim=True) + 1e-20)
+            topk = topk * _r.route_scale
+            return topk, _ids, scores_BLE
+
+        router.forward = forced_forward
+        forced += 1
+    return forced
+
+
+@torch.no_grad()
+def run_tt(model_flavor, checkpoint_path, ref, dtype, vision_dtype, force_hf_routing):
+    """torchtitan Kimi-VL: its own image processing + forward on the same image."""
+    device = torch.device("cuda")
+    print(f"Loading torchtitan Kimi-VL ({model_flavor}) on {device} ...")
+    model_config = model_registry(model_flavor).model
+    with torch.device("meta"):
+        model = model_config.build()
+    model.to_empty(device="cpu")
+    # Cast before init_states so the complex64 ComplexRoPE cache survives.
+    model.to(dtype)
+    model.init_states(buffer_device=torch.device("cpu"))
+    state_dict = ModelWrapper(model)._get_state_dict()
+    dcp.load(state_dict, checkpoint_id=checkpoint_path)
+    model.to(device)
+
+    model.vision_encoder.to(vision_dtype)  # mixed precision: ViT in vision_dtype
+    for layer in model.layers.values():
+        layer.attention.inner_attention = ScaledDotProductAttention.Config().build()
+    for layer in model.vision_encoder.layers.values():
+        layer.attn.flex_attention = _VisionSDPA()
+    model.eval()
+
+    if force_hf_routing:
+        n = _force_hf_routing(model, ref["expert_indices"], device)
+        print(f"forced HF routing on {n} MoE layers")
+
+    # titan's OWN Kimi image processing on the raw image (the real data path).
+    img = process_image(
+        Image.fromarray(ref["raw_image"].numpy()),
+        patch_size=_PATCH_SIZE,
+        merge_size=_MERGE_SIZE,
+        resize_fn=resize_to_patch_budget,
+        max_patches=4096,
+        max_patches_per_side=512,
+        image_mean=(0.5, 0.5, 0.5),
+        image_std=(0.5, 0.5, 0.5),
+    )
+    # MoonViT3d consumes raster-order patches (matching the released processor).
+    patches, grid = vision_to_patches(
+        img,  # pyrefly: ignore [bad-argument-type]
+        patch_size=_PATCH_SIZE,
+        temporal_patch_size=1,
+        merge_size=_MERGE_SIZE,
+        patch_order="raster",
+    )
+    pixel_values = patches.unsqueeze(0).to(device=device, dtype=vision_dtype)
+    grid_thw = grid.unsqueeze(0).to(device)
+    tokens = ref["input_ids"].to(device)
+
+    # Sanity: titan's grid must yield the same vision-token count as HF produced
+    # (the placeholder run in input_ids), else the scatter compares different seqs.
+    n_titan = ((grid_thw[0, 1] // _MERGE_SIZE) * (grid_thw[0, 2] // _MERGE_SIZE)).item()
+    n_placeholders = (tokens == _MEDIA_TOKEN_ID).sum().item()
+    print(
+        f"tokens={tuple(tokens.shape)}  pixel_values={tuple(pixel_values.shape)}  "
+        f"grid_thw={grid_thw.tolist()}  titan_vis_tokens={n_titan}  "
+        f"hf_placeholders={n_placeholders}"
+    )
+    assert n_titan == n_placeholders, (
+        f"titan produced {n_titan} vision tokens but HF input_ids has "
+        f"{n_placeholders} placeholders -- preprocessing grids differ"
+    )
+
+    # Localize: titan's vision features (projector output) vs HF's, pre-scatter.
+    tt_feats = model.vision_encoder(pixel_values, grid_thw=grid_thw)
+    ref_feats = ref["vision_features"].float().reshape(-1, tt_feats.shape[-1])
+    tt_feats = tt_feats.float().cpu().reshape(-1, tt_feats.shape[-1])
+    vcos = F.cosine_similarity(ref_feats.flatten(), tt_feats.flatten(), dim=0).item()
+    vmax = (ref_feats - tt_feats).abs().max().item()
+    print(
+        f"vision features (pre-scatter): shape={tuple(tt_feats.shape)}  "
+        f"cos={vcos:.6f}  max_diff={vmax:.3e}"
+    )
+
+    logits = model(
+        tokens,
+        pixel_values=pixel_values,
+        grid_thw=grid_thw,
+        special_tokens={"image_id": _MEDIA_TOKEN_ID, "video_id": _MEDIA_TOKEN_ID},
+    )
+    return logits[:, -1, :].float().cpu().squeeze()
+
+
+def compare(ref_logits, tt_logits) -> bool:
+    ref, tt = ref_logits.squeeze(), tt_logits.squeeze()
+    pq = F.log_softmax(ref, dim=-1)
+    qq = F.log_softmax(tt, dim=-1)
+    kl = F.kl_div(qq, pq, log_target=True, reduction="sum").item()
+    cos = F.cosine_similarity(ref, tt, dim=-1).item()
+    max_diff = (ref - tt).abs().max().item()
+    top1 = (ref.argmax() == tt.argmax()).item()
+    ov5 = len(set(ref.topk(5).indices.tolist()) & set(tt.topk(5).indices.tolist())) / 5
+    print(f"\n{'=' * 60}\nFull MM logit parity (torchtitan vs HF Kimi-VL)\n{'=' * 60}")
+    print(
+        f"  KL={kl:.4e}  cos={cos:.6f}  max_diff={max_diff:.4e}  "
+        f"top1={'Y' if top1 else 'N'}  top5={ov5:.0%}"
+    )
+    passed = abs(kl) < 1e-3  # pyrefly: ignore [bad-argument-type]
+    print(
+        "RESULT: PASS (KL < 1e-3 -- fp noise)."
+        if passed
+        else "RESULT: FAIL (KL >= 1e-3)."
+    )
+    return passed
+
+
+@torch.no_grad()
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--hf_model_path",
+        default=os.path.expanduser("~/hf_assets/moonshotai/Kimi-VL-A3B-Instruct"),
+    )
+    p.add_argument("--tt_checkpoint_path", default="outputs/kimi/kimi_vl_a3b_dcp")
+    p.add_argument("--model_flavor", default="Kimi-VL-A3B")
+    p.add_argument("--image_size", type=int, default=336)
+    p.add_argument(
+        "--hf_dtype", default="float32", choices=["float32", "bfloat16", "float16"]
+    )
+    p.add_argument(
+        "--dtype", default="float32", choices=["float32", "bfloat16", "float16"]
+    )
+    # Mixed precision: bf16 is poor for the ViT's high-dynamic-range activations,
+    # so fp16 is the usual choice. Defaults to --dtype. (titan side only.)
+    p.add_argument(
+        "--vision_dtype", default=None, choices=["float32", "bfloat16", "float16"]
+    )
+    p.add_argument(
+        "--force-hf-routing",
+        action="store_true",
+        help="Force titan MoE routers to use HF's recorded expert selections "
+        "(diagnostic: removes routing-flip divergence to isolate the rest).",
+    )
+    args = p.parse_args()
+    device, dtype = torch.device("cuda"), getattr(torch, args.dtype)
+    vision_dtype = getattr(torch, args.vision_dtype) if args.vision_dtype else dtype
+    hf_dtype = getattr(torch, args.hf_dtype)
+    print(
+        f"hf_dtype={args.hf_dtype}  titan text={args.dtype}  "
+        f"titan vision={args.vision_dtype or args.dtype}"
+    )
+
+    ref = run_hf(args.hf_model_path, args.image_size, hf_dtype, device)
+    tt_logits = run_tt(
+        args.model_flavor,
+        args.tt_checkpoint_path,
+        ref,
+        dtype,
+        vision_dtype,
+        args.force_hf_routing,
+    )
+    if not compare(ref["last_logits"], tt_logits):
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/assets/tokenizer/tokenizer.json b/tests/assets/tokenizer/tokenizer.json
index 11fec5ce74..9dcddd757b 100644
--- a/tests/assets/tokenizer/tokenizer.json
+++ b/tests/assets/tokenizer/tokenizer.json
@@ -2047,7 +2047,11 @@
       "</tool_call>": 2012,
       "<tool_response>": 2013,
       "</tool_response>": 2014,
-      "</think>": 2015
+      "</think>": 2015,
+      "<|media_pad|>": 2016,
+      "<|media_begin|>": 2017,
+      "<|media_end|>": 2018,
+      "[PAD]": 2019
     },
     "merges": []
   }
diff --git a/tests/assets/tokenizer/tokenizer_config.json b/tests/assets/tokenizer/tokenizer_config.json
index 6fe58e88c8..5c269fdcb1 100644
--- a/tests/assets/tokenizer/tokenizer_config.json
+++ b/tests/assets/tokenizer/tokenizer_config.json
@@ -143,6 +143,38 @@
       "rstrip": false,
       "single_word": false,
       "special": true
+    },
+    "2016": {
+      "content": "<|media_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2017": {
+      "content": "<|media_begin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2018": {
+      "content": "<|media_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2019": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
     }
   },
   "bos_token": "<|begin_of_text|>",
diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py
index 5e9caed8a5..d950c01b83 100755
--- a/tests/integration_tests/models.py
+++ b/tests/integration_tests/models.py
@@ -13,7 +13,9 @@
 def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefinitions:
     """Use ``backend`` for every variant, or return an unsupported test unchanged."""
     if backend == "spmd_types" and any(
-        "--module qwen3_5" in arg for variant in t.override_args for arg in variant
+        "--module qwen3_5" in arg or "--module kimi_k2_7" in arg
+        for variant in t.override_args
+        for arg in variant
     ):
         return t
 
@@ -237,6 +239,22 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             "gpt_oss_pp+fsdp+cp+ep+sacop",
             ngpu=8,
         ),
+        # Integration Test Cases for Kimi K2.7
+        OverrideDefinitions(
+            [
+                [
+                    "--module kimi_k2_7 --config kimi_k2_5_debugmodel",
+                    "--training.local_batch_size 2",
+                    "--parallelism.data_parallel_shard_degree 2",
+                    "--parallelism.pipeline_parallel_degree 2",
+                    "--parallelism.tensor_parallel_degree 2",
+                    "--parallelism.expert_parallel_degree 2",
+                ],
+            ],
+            "Kimi K2.7 multimodal FSDP+TP+EP+PP",
+            "kimi_k2_5_mm_fsdp+tp+ep+pp",
+            ngpu=8,
+        ),
     ]
 
     return [_enable_spmd_backend(t, "spmd_types") for t in model_tests]
diff --git a/tests/unit_tests/test_mm_dataset_preprocessing.py b/tests/unit_tests/test_mm_dataset_preprocessing.py
new file mode 100644
index 0000000000..dff2b24f9b
--- /dev/null
+++ b/tests/unit_tests/test_mm_dataset_preprocessing.py
@@ -0,0 +1,164 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""CPU unit tests for the multimodal dataset image preprocessing.
+
+``resize_to_patch_budget`` is the NaViT patch-budget protocol: cap raw patches
+at the total and per-side limits, then pad to a ``patch_size * merge_size``
+multiple. These pin that pure-geometry behavior.
+"""
+
+import math
+import unittest
+
+import torch
+from PIL import Image
+
+from torchtitan.hf_datasets.multimodal.utils.image import (
+    process_image,
+    resize_to_patch_budget,
+    vision_to_patches,
+)
+
+
+def _patch_budget_geometry(h, w, *, patch_size, merge_size, max_patches):
+    """Reference geometry for the patch-budget resize: the final padded (H, W)
+    in pixels."""
+    nh, nw = h, w
+    if (nw // patch_size) * (nh // patch_size) > max_patches:
+        scale = math.sqrt(max_patches / ((nw // patch_size) * (nh // patch_size)))
+        nh, nw = int(nh * scale), int(nw * scale)
+    factor = merge_size * patch_size
+    pad_h = (factor - nh % factor) % factor
+    pad_w = (factor - nw % factor) % factor
+    return nh + pad_h, nw + pad_w
+
+
+class TestResizeToPatchBudget(unittest.TestCase):
+    PS, MERGE, LIMIT, SIDE = 14, 2, 4096, 512
+
+    def _final(self, h, w):
+        rh, rw, ph, pw = resize_to_patch_budget(
+            h,
+            w,
+            patch_size=self.PS,
+            merge_size=self.MERGE,
+            max_patches=self.LIMIT,
+            max_patches_per_side=self.SIDE,
+        )
+        return rh + ph, rw + pw
+
+    def test_matches_patch_budget_geometry(self):
+        for h, w in [(600, 800), (336, 336), (224, 448), (1000, 500), (101, 173)]:
+            got = self._final(h, w)
+            want = _patch_budget_geometry(
+                h,
+                w,
+                patch_size=self.PS,
+                merge_size=self.MERGE,
+                max_patches=self.LIMIT,
+            )
+            self.assertEqual(got, want, f"{h}x{w}: {got} != {want}")
+
+    def test_output_is_factor_multiple(self):
+        factor = self.PS * self.MERGE
+        for h, w in [(600, 800), (101, 173), (1400, 1400)]:
+            fh, fw = self._final(h, w)
+            self.assertEqual(fh % factor, 0)
+            self.assertEqual(fw % factor, 0)
+
+    def test_caps_patches_at_limit(self):
+        # 1400x1400 -> 100*100 = 10000 raw patches, well over the 4096 cap.
+        self.assertEqual((1400 // self.PS) * (1400 // self.PS), 10000)  # sanity
+        fh, fw = self._final(1400, 1400)
+        # Scaled down to a square grid at the cap; this size needs no padding,
+        # so the count must not exceed the limit at all.
+        self.assertEqual(fh % (self.PS * self.MERGE), 0)
+        self.assertLessEqual((fh // self.PS) * (fw // self.PS), self.LIMIT)
+
+    def test_small_image_not_upscaled(self):
+        # below the cap -> only padded, never scaled up.
+        rh, rw, ph, pw = resize_to_patch_budget(
+            30,
+            30,
+            patch_size=self.PS,
+            merge_size=self.MERGE,
+            max_patches=self.LIMIT,
+            max_patches_per_side=self.SIDE,
+        )
+        self.assertEqual((rh, rw), (30, 30))
+
+    def test_per_side_cap_scales_down(self):
+        # The per-side limit scales an extreme aspect ratio instead of dropping it.
+        final_h, final_w = self._final(self.PS * 600, self.PS * 2)
+        self.assertEqual(final_h // self.PS, self.SIDE)
+        self.assertEqual(final_w % (self.PS * self.MERGE), 0)
+
+    def test_per_side_cap_is_inclusive(self):
+        height, width = self.PS * self.SIDE, self.PS * self.MERGE
+        rh, rw, ph, pw = resize_to_patch_budget(
+            height,
+            width,
+            patch_size=self.PS,
+            merge_size=self.MERGE,
+            max_patches=self.LIMIT,
+            max_patches_per_side=self.SIDE,
+        )
+        self.assertEqual((rh, rw, ph, pw), (height, width, 0, 0))
+
+
+class TestProcessImagePatchBudget(unittest.TestCase):
+    def test_navit_pads_to_factor_multiple(self):
+        ps, merge = 14, 2
+        factor = ps * merge
+        # 100x173 is not a factor multiple -> navit must pad to one.
+        img = Image.fromarray((torch.rand(100, 173, 3) * 255).to(torch.uint8).numpy())
+        out = process_image(
+            img,
+            patch_size=ps,
+            merge_size=merge,
+            resize_fn=resize_to_patch_budget,
+            max_patches=4096,
+            image_mean=(0.5, 0.5, 0.5),
+            image_std=(0.5, 0.5, 0.5),
+        )
+        self.assertIsNotNone(out)
+        # (1, H, W, C)
+        _, H, W, C = out.shape
+        self.assertEqual(C, 3)
+        self.assertEqual(H % factor, 0)
+        self.assertEqual(W % factor, 0)
+        want_h, want_w = _patch_budget_geometry(
+            100, 173, patch_size=ps, merge_size=merge, max_patches=4096
+        )
+        self.assertEqual((H, W), (want_h, want_w))
+
+
+class TestVisionToPatchesOrder(unittest.TestCase):
+    """Patch sequence layout: 'block' vs 'raster'."""
+
+    def test_block_to_raster_permutation(self):
+        # 2x4 patch grid (h=2, w=4), merge_size=2. Distinct per-patch values so
+        # the two orderings are a pure permutation of each other.
+        img = torch.arange(1 * 28 * 56 * 3, dtype=torch.float32).reshape(1, 28, 56, 3)
+        block, grid = vision_to_patches(img, 14, 1, 2, patch_order="block")
+        raster, _ = vision_to_patches(img, 14, 1, 2, patch_order="raster")
+
+        self.assertEqual(grid.tolist(), [1, 2, 4])
+        self.assertEqual(block.shape, raster.shape)
+        # block slot b corresponds to raster slot block_to_raster_idx[b].
+        block_to_raster_idx = [0, 1, 4, 5, 2, 3, 6, 7]
+        for b, r in enumerate(block_to_raster_idx):
+            self.assertTrue(torch.equal(block[b], raster[r]))
+
+    def test_invalid_patch_order_raises(self):
+        img = torch.zeros(1, 28, 28, 3)
+        with self.assertRaises(ValueError):
+            vision_to_patches(img, 14, 1, 2, patch_order="bogus")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/torchtitan/distributed/fsdp.py b/torchtitan/distributed/fsdp.py
index 74be2265c6..9b2d47e093 100644
--- a/torchtitan/distributed/fsdp.py
+++ b/torchtitan/distributed/fsdp.py
@@ -78,6 +78,32 @@ def get_fsdp_reshard_after_forward_policy(
             )
 
 
+def apply_fsdp_to_vision_encoder(
+    vision_encoder: nn.Module,
+    dp_mesh: DeviceMesh,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
+    reshard_after_forward_policy: str = "default",
+    pp_enabled: bool = False,
+) -> None:
+    """FSDP a VLM vision encoder as a single unit.
+
+    One all-gather for all vision params is more efficient than per-layer sharding
+    (the vision encoder is small relative to the decoder). Call before
+    ``apply_fsdp_to_decoder`` so the encoder is already sharded.
+    """
+    mp_policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=reduce_dtype)
+    reshard_after_forward = get_fsdp_reshard_after_forward_policy(
+        reshard_after_forward_policy, pp_enabled=pp_enabled
+    )
+    fully_shard(
+        vision_encoder,
+        mesh=dp_mesh,
+        mp_policy=mp_policy,
+        reshard_after_forward=reshard_after_forward,
+    )
+
+
 def apply_fsdp_to_decoder(
     model: "Decoder",
     dp_mesh: DeviceMesh,
diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index b62cafb358..485cba6e69 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import copy
+import dataclasses
 import math
 import os
 from collections.abc import Callable
@@ -32,9 +33,9 @@
 from torchtitan.protocols.module import ModuleDict, ModuleList
 from torchtitan.tools.logging import logger
 
-# pipeline_llm is the public entrypoint for model-specific PP setup. Helpers in
-# this module are implementation details and should stay private.
-__all__ = ["pipeline_llm"]
+# pipeline_llm and pipeline_vlm are the public entrypoints for model-specific PP
+# setup. Helpers in this module are implementation details and stay private.
+__all__ = ["pipeline_llm", "pipeline_vlm"]
 
 
 def _build_get_mesh_callback(
@@ -140,6 +141,56 @@ def pipeline_llm(
     return pp_schedule, model_parts, has_first_stage, has_last_stage
 
 
+def pipeline_vlm(
+    model: nn.Module,
+    *,
+    parallel_dims: ParallelDims,
+    parallelism: ParallelismConfig,
+    model_config: BaseModel.Config,
+    **kwargs,
+) -> tuple[_PipelineSchedule, list[nn.Module], bool, bool]:
+    """PP entrypoint for vision-language models: co-locate the vision encoder
+    with the first stage, then delegate to ``pipeline_llm``.
+
+    The auto-generated LLM stage split only knows about decoder modules
+    (``tok_embeddings``, ``layers.*``, ``norm``, ``lm_head``). For a VLM we inject
+    ``vision_encoder`` into the first stage's FQN list so it runs alongside
+    ``tok_embeddings`` (vision features are scattered into the embedding sequence
+    before the decoder layers). On stages other than the first, ``tok_embeddings``
+    and ``vision_encoder`` are pruned to ``None``; each model's ``forward`` must
+    guard on ``self.tok_embeddings is not None`` so the multimodal logic is
+    skipped there.
+
+    NOTE: This adds load to stage 0 that the auto split does not model
+    (``input_weight`` only accounts for ``tok_embeddings``); for a heavy vision
+    encoder, bump ``parallelism.pipeline_parallel_first_stage_less_layers`` to
+    rebalance.
+    """
+    if parallelism.module_fqns_per_model_part is None:
+        (
+            num_virtual_stages,
+            num_layers,
+            input_weight,
+            output_weight,
+        ) = _get_pipeline_metadata(parallel_dims, parallelism, model_config)
+        fqn_per_part = _generate_llm_fqn_per_model_part(
+            num_virtual_stages, num_layers, input_weight, output_weight
+        )
+        if model.vision_encoder is not None:
+            fqn_per_part[0].insert(0, "vision_encoder")
+        parallelism = dataclasses.replace(
+            parallelism, module_fqns_per_model_part=fqn_per_part
+        )
+
+    return pipeline_llm(
+        model,
+        parallel_dims=parallel_dims,
+        parallelism=parallelism,
+        model_config=model_config,
+        **kwargs,
+    )
+
+
 def _get_pipeline_metadata(
     parallel_dims: ParallelDims,
     parallelism: ParallelismConfig,
diff --git a/torchtitan/hf_datasets/multimodal/mm_collator.py b/torchtitan/hf_datasets/multimodal/mm_collator.py
index 3350d89de8..c50d634d56 100644
--- a/torchtitan/hf_datasets/multimodal/mm_collator.py
+++ b/torchtitan/hf_datasets/multimodal/mm_collator.py
@@ -35,6 +35,7 @@ class MultiModalCollator:
     spatial_merge_size: int
     tokenizer: MultiModalTokenizer
     build_mrope_positions: bool
+    patch_order: str = "block"
 
     def collate_images(
         self, all_images: list[torch.Tensor]
@@ -52,7 +53,11 @@ def collate_images(
         """
         results = [
             vision_to_patches(
-                img, self.patch_size, self.temporal_patch_size, self.spatial_merge_size
+                img,
+                self.patch_size,
+                self.temporal_patch_size,
+                self.spatial_merge_size,
+                patch_order=self.patch_order,
             )
             for img in all_images
         ]
@@ -164,6 +169,13 @@ def _build_mrope_positions(
         Returns:
             (batch, seq_len, 3) MRoPE position IDs.
         """
+        # MRoPE position IDs are laid out in block order; a raster patch order
+        # would desync them from the patch sequence.
+        if self.patch_order != "block":
+            raise ValueError(
+                f"MRoPE requires patch_order='block', got {self.patch_order!r}."
+            )
+
         # Expand each video [T, H, W] into T rows of [1, H, W] so each frame is
         # treated like an image; temporal position comes from frame ordering.
         if grid_thw_videos is not None:
diff --git a/torchtitan/hf_datasets/multimodal/mm_datasets.py b/torchtitan/hf_datasets/multimodal/mm_datasets.py
index 71adda51a2..5d6e9d1e98 100644
--- a/torchtitan/hf_datasets/multimodal/mm_datasets.py
+++ b/torchtitan/hf_datasets/multimodal/mm_datasets.py
@@ -64,9 +64,10 @@
 import inspect
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from typing import Annotated, Any, Literal
 
 import torch
+import tyro
 from datasets import Dataset, load_dataset
 from datasets.distributed import split_dataset_by_node
 from torch.distributed.checkpoint.stateful import Stateful
@@ -79,7 +80,7 @@
 from torchtitan.hf_datasets import DatasetConfig
 from torchtitan.tools.logging import logger
 from .mm_collator import MultiModalCollator
-from .utils.image import calculate_vision_tokens, process_image
+from .utils.image import calculate_vision_tokens, process_image, resize_to_pixel_budget
 from .utils.packing import MMSamplePacker
 from .utils.text import insert_vision_placeholders
 
@@ -95,6 +96,9 @@ def _process_mm_sample(
     max_pixels: int,
     image_mean: tuple[float, ...],
     image_std: tuple[float, ...],
+    resize_fn: Callable[..., tuple[int, int, int, int]],
+    max_patches: int,
+    max_patches_per_side: int,
     **kwargs,
 ) -> dict[str, Any] | None:
     """Common processing logic for multimodal samples.
@@ -140,6 +144,9 @@ def _process_mm_sample(
                 max_pixels=max_pixels,
                 image_mean=image_mean,
                 image_std=image_std,
+                resize_fn=resize_fn,
+                max_patches=max_patches,
+                max_patches_per_side=max_patches_per_side,
             )
             if processed_img is not None:
                 num_tokens, _, _ = calculate_vision_tokens(
@@ -222,6 +229,7 @@ def _process_obelics_sample(
         max_pixels=max_pixels,
         image_mean=image_mean,
         image_std=image_std,
+        **kwargs,
     )
 
 
@@ -255,6 +263,7 @@ def _process_cc12_wd_sample(
         max_pixels=max_pixels,
         image_mean=image_mean,
         image_std=image_std,
+        **kwargs,
     )
 
 
@@ -313,6 +322,9 @@ def __init__(
         image_mean: tuple[float, ...],
         image_std: tuple[float, ...],
         packing_buffer_size: int,
+        resize_fn: Callable[..., tuple[int, int, int, int]],
+        max_patches: int,
+        max_patches_per_side: int,
         dp_rank: int = 0,
         dp_world_size: int = 1,
         infinite: bool = False,
@@ -346,6 +358,9 @@ def __init__(
         self.max_pixels = max_pixels
         self.image_mean = image_mean
         self.image_std = image_std
+        self.resize_fn = resize_fn
+        self.max_patches = max_patches
+        self.max_patches_per_side = max_patches_per_side
         self.video_dir = video_dir
         self.video_fps = video_fps
         self.video_min_frames = video_min_frames
@@ -376,6 +391,9 @@ def __iter__(self):
                     max_pixels=self.max_pixels,
                     image_mean=self.image_mean,
                     image_std=self.image_std,
+                    resize_fn=self.resize_fn,
+                    max_patches=self.max_patches,
+                    max_patches_per_side=self.max_patches_per_side,
                     video_dir=self.video_dir,
                     video_fps=self.video_fps,
                     video_min_frames=self.video_min_frames,
@@ -507,11 +525,31 @@ class Config(ParallelAwareDataloader.Config):
         spatial_merge_size: int
         """Spatially merge visual tokens after encoder. e.g. 2 means 2x2=4 patches merged."""
 
+        patch_order: Literal["block", "raster"] = "block"
+        """Patch sequence layout the collator emits: ``"block"`` (each
+        ``spatial_merge_size**2`` group contiguous, or ``"raster"`` (row-major).
+        Must be ``"block"`` when ``build_mrope_positions`` is set."""
+
+        resize_fn: Annotated[
+            Callable[..., tuple[int, int, int, int]], tyro.conf.Suppress
+        ] = resize_to_pixel_budget
+        """Image-resize strategy (a callable, like ``sample_processor``):
+        ``resize_to_pixel_budget`` or ``resize_to_patch_budget`` (cap patches at
+        ``max_patches``, pad to a ``patch_size * spatial_merge_size`` multiple).
+        Both share the signature ``(h, w, *, patch_size, merge_size,
+        **budget) -> (resize_h, resize_w, pad_h, pad_w)``."""
+
         min_pixels: int
-        """Minimum number of pixels for image resizing."""
+        """Minimum number of pixels for image resizing (pixel-budget strategy)."""
 
         max_pixels: int
-        """Maximum number of pixels for image resizing."""
+        """Maximum number of pixels for image resizing (pixel-budget strategy)."""
+
+        max_patches: int = 4096
+        """Max raw patches per image."""
+
+        max_patches_per_side: int = 512
+        """Per-side patch cap for the vision position-embedding grid (``navit``)."""
 
         image_mean: tuple[float, ...]
         """Per-channel mean for image normalization."""
@@ -561,6 +599,9 @@ def __init__(
             image_mean=config.image_mean,
             image_std=config.image_std,
             packing_buffer_size=config.packing_buffer_size,
+            resize_fn=config.resize_fn,
+            max_patches=config.max_patches,
+            max_patches_per_side=config.max_patches_per_side,
             dp_rank=dp_rank,
             dp_world_size=dp_world_size,
             infinite=config.infinite,
@@ -580,6 +621,7 @@ def __init__(
             spatial_merge_size=config.spatial_merge_size,
             tokenizer=tokenizer,
             build_mrope_positions=config.build_mrope_positions,
+            patch_order=config.patch_order,
         )
 
         dataloader_kwargs = {
diff --git a/torchtitan/hf_datasets/multimodal/utils/image.py b/torchtitan/hf_datasets/multimodal/utils/image.py
index 726d4d5703..b6325257f7 100644
--- a/torchtitan/hf_datasets/multimodal/utils/image.py
+++ b/torchtitan/hf_datasets/multimodal/utils/image.py
@@ -11,6 +11,7 @@
 """
 
 import math
+from collections.abc import Callable
 
 import einops as E
 import requests
@@ -46,6 +47,155 @@ def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor:
         return TVF.pil_to_tensor(image)
 
 
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int,
+    min_pixels: int,
+    max_pixels: int,
+) -> tuple[int, int]:
+    """Compute target (height, width) that satisfy per-frame pixel budget.
+
+    Both output dimensions are rounded to multiples of ``factor``.  The spatial
+    pixel count ``h * w`` is kept within [min_pixels, max_pixels].
+
+    Args:
+        height: Original height.
+        width: Original width.
+        factor: Spatial rounding factor (``patch_size * merge_size``).
+        min_pixels: Minimum spatial pixels per frame.
+        max_pixels: Maximum spatial pixels per frame.
+
+    Returns:
+        (resized_height, resized_width)
+    """
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"Absolute aspect ratio must be smaller than 200, "
+            f"got {max(height, width) / min(height, width):.1f}"
+        )
+
+    # Round spatial dims to nearest multiple of factor
+    h_bar = max(round(height / factor) * factor, factor)
+    w_bar = max(round(width / factor) * factor, factor)
+
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(math.floor(height / beta / factor) * factor, factor)
+        w_bar = max(math.floor(width / beta / factor) * factor, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+
+    return h_bar, w_bar
+
+
+def resize_to_pixel_budget(
+    height: int,
+    width: int,
+    *,
+    patch_size: int,
+    merge_size: int,
+    min_pixels: int,
+    max_pixels: int,
+    **_: object,
+) -> tuple[int, int, int, int]:
+    """Resize so the spatial pixel count lands in ``[min_pixels, max_pixels]``,
+    with both dims rounded to a ``patch_size * merge_size`` multiple (Qwen-VL
+    convention). Content is rescaled to the grid, so no padding is needed.
+
+    A resize strategy (``resize_fn``) for ``process_image`` -- extra budget
+    kwargs (e.g. ``max_patches``) are accepted and ignored via ``**_``.
+
+    Args:
+        height: Original height in pixels.
+        width: Original width in pixels.
+        patch_size: Spatial patch size.
+        merge_size: Spatial merge factor.
+        min_pixels: Lower spatial-pixel budget.
+        max_pixels: Upper spatial-pixel budget.
+
+    Returns:
+        ``(resize_h, resize_w, 0, 0)`` -- trailing zeros are padding (always 0
+        here), kept for a uniform interface with ``resize_to_patch_budget``.
+    """
+    factor = patch_size * merge_size
+    # Ensure both dims >= factor so the rounding has a valid starting point.
+    if height < factor or width < factor:
+        scale = max(factor / width, factor / height)
+        width, height = int(width * scale), int(height * scale)
+
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"Absolute aspect ratio must be smaller than 200, "
+            f"got {max(height, width) / min(height, width):.1f}"
+        )
+
+    h_bar = max(round(height / factor) * factor, factor)
+    w_bar = max(round(width / factor) * factor, factor)
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(math.floor(height / beta / factor) * factor, factor)
+        w_bar = max(math.floor(width / beta / factor) * factor, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar, 0, 0
+
+
+def resize_to_patch_budget(
+    height: int,
+    width: int,
+    *,
+    patch_size: int,
+    merge_size: int,
+    max_patches: int,
+    max_patches_per_side: int,
+    **_: object,
+) -> tuple[int, int, int, int]:
+    """Cap the total and per-side raw-patch counts (scale down,
+    aspect-preserving), then pad right/bottom to a ``patch_size * merge_size``
+    multiple. Small images within both limits are not upscaled.
+
+    A resize strategy (``resize_fn``) for ``process_image`` -- extra budget
+    kwargs (e.g. ``min_pixels`` / ``max_pixels``) are accepted and ignored via
+    ``**_``.
+
+    Args:
+        height: Original height in pixels.
+        width: Original width in pixels.
+        patch_size: Spatial patch size.
+        merge_size: Spatial merge factor.
+        max_patches: Max raw patches per image.
+        max_patches_per_side: Per-side patch cap (vision position-embedding limit).
+
+    Returns:
+        ``(resize_h, resize_w, pad_h, pad_w)`` -- resize to the first two, then
+        pad right/bottom by the last two.
+    """
+    num_patches_h = max(1.0, height // patch_size)
+    num_patches_w = max(1.0, width // patch_size)
+    num_patches = num_patches_h * num_patches_w
+
+    scale = min(
+        1.0,
+        math.sqrt(max_patches / num_patches),
+        max_patches_per_side * patch_size / height,
+        max_patches_per_side * patch_size / width,
+    )
+    h = max(1, int(height * scale))
+    w = max(1, int(width * scale))
+    h = min(h, max_patches_per_side * patch_size)
+    w = min(w, max_patches_per_side * patch_size)
+
+    factor = patch_size * merge_size
+    pad_h = (factor - h % factor) % factor
+    pad_w = (factor - w % factor) % factor
+    return h, w, pad_h, pad_w
+
+
 def process_image(
     image: str | bytes | Image.Image,
     patch_size: int = 16,
@@ -54,21 +204,31 @@ def process_image(
     min_pixels: int = 65536,
     image_mean: tuple[float, ...] = (0.5, 0.5, 0.5),
     image_std: tuple[float, ...] = (0.5, 0.5, 0.5),
+    resize_fn: Callable[..., tuple[int, int, int, int]] = resize_to_pixel_budget,
+    max_patches: int = 4096,
+    max_patches_per_side: int = 512,
 ) -> torch.Tensor | None:
     """Load and preprocess a single image for VLM training.
 
-    Uses torchvision APIs for decoding and resizing (faster uint8 SIMD paths).
-    Resizes to a pixel budget while keeping both dimensions multiples of
-    ``patch_size * merge_size``, then normalizes with the given mean/std.
+    Uses torchvision APIs for decoding and resizing (faster uint8 SIMD paths),
+    then normalizes with the given mean/std. ``resize_fn`` is the resize strategy
+    -- ``resize_to_pixel_budget`` (Qwen-VL, default) or
+    ``resize_to_patch_budget`` (Kimi-VL) -- with the uniform signature
+    ``(height, width, *, patch_size, merge_size, **budget) -> (resize_h,
+    resize_w, pad_h, pad_w)``. All budget kwargs are passed through; each
+    strategy uses the subset it needs.
 
     Args:
         image: Raw bytes, file path, HTTP(S) URL, or PIL Image.
         patch_size: Spatial patch size used by the vision encoder.
         merge_size: Spatial merge factor (patches merged per dimension).
-        max_pixels: Upper pixel budget for resizing.
-        min_pixels: Lower pixel budget for resizing.
+        max_pixels: Upper pixel budget (``resize_to_pixel_budget``).
+        min_pixels: Lower pixel budget (``resize_to_pixel_budget``).
         image_mean: Per-channel mean for normalization.
         image_std: Per-channel std for normalization.
+        resize_fn: Resize-strategy callable (see above).
+        max_patches: Max raw patches per image (``resize_to_patch_budget``).
+        max_patches_per_side: Per-side patch cap (``resize_to_patch_budget``).
 
     Returns:
         Tensor of shape (1, H, W, C) with a dummy temporal dim, or None on failure.
@@ -77,30 +237,29 @@ def process_image(
         # Decode to (C, H, W) uint8 tensor
         img_tensor = _decode_image(image)
         _, original_height, original_width = img_tensor.shape
-        factor = patch_size * merge_size
-
-        # Ensure both dimensions are at least ``factor`` so that
-        # smart_resize always has a valid starting point
-        if original_height < factor or original_width < factor:
-            scale = max(factor / original_width, factor / original_height)
-            original_width = int(original_width * scale)
-            original_height = int(original_height * scale)
 
-        resized_height, resized_width = smart_resize(
+        resize_h, resize_w, pad_h, pad_w = resize_fn(
             original_height,
             original_width,
-            factor=factor,
+            patch_size=patch_size,
+            merge_size=merge_size,
             min_pixels=min_pixels,
             max_pixels=max_pixels,
+            max_patches=max_patches,
+            max_patches_per_side=max_patches_per_side,
         )
-
-        # Bicubic resize on uint8 (leverages AVX2/NEON SIMD fast paths)
-        img_tensor = TVF.resize(
-            img_tensor,
-            [resized_height, resized_width],
-            interpolation=TVF.InterpolationMode.BICUBIC,
-            antialias=True,
-        )
+        if (resize_h, resize_w) != (original_height, original_width):
+            # Bicubic resize on uint8 (leverages AVX2/NEON SIMD fast paths)
+            img_tensor = TVF.resize(
+                img_tensor,
+                [resize_h, resize_w],
+                interpolation=TVF.InterpolationMode.BICUBIC,
+                antialias=True,
+            )
+        if pad_h or pad_w:
+            # Pad right/bottom with zeros (black); after normalization these are
+            # constant padding patches the encoder sees as part of the grid.
+            img_tensor = TVF.pad(img_tensor, [0, 0, pad_w, pad_h])
 
         # uint8 → float32 [0, 1] → normalize
         img_tensor = TVF.to_dtype(img_tensor, torch.float32, scale=True)
@@ -116,50 +275,6 @@ def process_image(
         return None
 
 
-def smart_resize(
-    height: int,
-    width: int,
-    factor: int,
-    min_pixels: int,
-    max_pixels: int,
-) -> tuple[int, int]:
-    """Compute target (height, width) that satisfy per-frame pixel budget.
-
-    Both output dimensions are rounded to multiples of ``factor``.  The spatial
-    pixel count ``h * w`` is kept within [min_pixels, max_pixels].
-
-    Args:
-        height: Original height.
-        width: Original width.
-        factor: Spatial rounding factor (``patch_size * merge_size``).
-        min_pixels: Minimum spatial pixels per frame.
-        max_pixels: Maximum spatial pixels per frame.
-
-    Returns:
-        (resized_height, resized_width)
-    """
-    if max(height, width) / min(height, width) > 200:
-        raise ValueError(
-            f"Absolute aspect ratio must be smaller than 200, "
-            f"got {max(height, width) / min(height, width):.1f}"
-        )
-
-    # Round spatial dims to nearest multiple of factor
-    h_bar = max(round(height / factor) * factor, factor)
-    w_bar = max(round(width / factor) * factor, factor)
-
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = max(math.floor(height / beta / factor) * factor, factor)
-        w_bar = max(math.floor(width / beta / factor) * factor, factor)
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = math.ceil(height * beta / factor) * factor
-        w_bar = math.ceil(width * beta / factor) * factor
-
-    return h_bar, w_bar
-
-
 def calculate_vision_tokens(
     num_frames: int,
     height: int,
@@ -193,18 +308,38 @@ def vision_to_patches(
     patch_size: int,
     temporal_patch_size: int,
     merge_size: int,
+    patch_order: str = "block",
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Convert an image/video tensor to flattened patches in block order.
+    """Convert an image/video tensor to flattened patches.
+
+    ``patch_order`` selects the sequence layout:
+
+    - ``"block"`` (default): each ``merge_size x merge_size`` spatial group is
+      contiguous (sequence ``(t, bh, bw, m, n)``), matching Qwen-style mergers
+      that fuse consecutive ``merge_size**2`` patches.
+    - ``"raster"``: plain row-major ``(t, h, w)`` order, matching encoders whose
+      RoPE / position embeddings index ``row = p // w, col = p % w`` (MoonViT3d).
 
-    Patches are ordered so that each ``merge_size × merge_size`` spatial group
-    is contiguous, matching the layout expected by the vision encoder's
-    spatial merge layer.
+    Example -- a 2x4 patch grid (h=2, w=4), ``merge_size=2``, patches labeled by
+    their raster index ``row * w + col``::
+
+        raster grid:   0 1 2 3
+                       4 5 6 7
+
+        raster order:  0 1 2 3 4 5 6 7
+        block order:   0 1 4 5  2 3 6 7
+                       \\_____/  \\_____/
+                       2x2 block  2x2 block
+
+    In block order the 4 patches that merge into one token are adjacent; in
+    raster order they are not.
 
     Args:
         img: (T, H, W, C) image or video tensor.
         patch_size: Spatial patch size (pixels per patch side).
         temporal_patch_size: Temporal patch size (frames per temporal patch).
         merge_size: Spatial merge size (patches merged per dimension).
+        patch_order: ``"block"`` or ``"raster"`` sequence layout.
 
     Returns:
         patches: (num_patches, patch_dim) flattened patch vectors in
@@ -227,21 +362,32 @@ def vision_to_patches(
     H_patches = H // ps
     W_patches = W // ps
 
-    # Reshape (T, H, W, C) → (num_patches, patch_dim) in block order:
-    #   T = t × pt        temporal patches × frames per temporal patch
-    #   H = bh × m × ph   block rows × merge patches × pixels per patch
-    #   W = bw × n × pw   block cols × merge patches × pixels per patch
-    # Sequence order:  (t, bh, bw, m, n) — merge group is contiguous
-    # Patch vector:    (c, pt, ph, pw)   — channel-first
-    patches = E.rearrange(
-        img,
-        "(t pt) (bh m ph) (bw n pw) c -> (t bh bw m n) (c pt ph pw)",
-        pt=ts,
-        ph=ps,
-        pw=ps,
-        m=merge_size,
-        n=merge_size,
-    )
+    # Reshape (T, H, W, C) -> (num_patches, patch_dim). The patch vector is
+    # always channel-first (c, pt, ph, pw); only the sequence ordering differs.
+    if patch_order == "block":
+        # (t, bh, bw, m, n) -- the merge group is contiguous.
+        patches = E.rearrange(
+            img,
+            "(t pt) (bh m ph) (bw n pw) c -> (t bh bw m n) (c pt ph pw)",
+            pt=ts,
+            ph=ps,
+            pw=ps,
+            m=merge_size,
+            n=merge_size,
+        )
+    elif patch_order == "raster":
+        # (t, h, w) -- plain row-major (merge_size unused in the layout).
+        patches = E.rearrange(
+            img,
+            "(t pt) (h ph) (w pw) c -> (t h w) (c pt ph pw)",
+            pt=ts,
+            ph=ps,
+            pw=ps,
+        )
+    else:
+        raise ValueError(
+            f"patch_order must be 'block' or 'raster', got {patch_order!r}."
+        )
 
     grid_thw = torch.tensor([T_patches, H_patches, W_patches])
     return patches, grid_thw
diff --git a/torchtitan/models/__init__.py b/torchtitan/models/__init__.py
index da1b5c6224..393e23f245 100644
--- a/torchtitan/models/__init__.py
+++ b/torchtitan/models/__init__.py
@@ -5,5 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 _supported_models = frozenset(
-    ["deepseek_v3", "flux", "gpt_oss", "llama3", "qwen3", "qwen3_5"]
+    [
+        "deepseek_v3",
+        "flux",
+        "gpt_oss",
+        "kimi_k2_7",
+        "llama3",
+        "qwen3",
+        "qwen3_5",
+    ]
 )
diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py
new file mode 100644
index 0000000000..0d5c4e1819
--- /dev/null
+++ b/torchtitan/models/common/multimodal.py
@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Model-agnostic vision<->text fusion for VLMs.
+
+The decoder embeds the full token sequence; the placeholder tokens
+get a throwaway text embedding that ``scatter_vision_embeds``
+overwrites with the vision encoder's per-item features at the positions
+``get_vision_positions`` locates.
+"""
+
+import torch
+
+
+def get_vision_positions(
+    tokens: torch.Tensor,
+    num_vision_tokens_per_item: torch.Tensor,
+    placeholder_id: int,
+) -> list[tuple[int, int, int, int]]:
+    """Locate each visual item's placeholder run in the token sequence.
+
+    Args:
+        tokens: (bsz, seq_len) token IDs.
+        num_vision_tokens_per_item: (num_items,) valid token count per visual item, in
+            the order the items appear in ``tokens``.
+        placeholder_id: token id whose contiguous runs mark vision spans.
+
+    Returns:
+        ``(item_idx, sample_idx, vision_start, n_tokens)`` per item, where
+        ``vision_start`` is the position of the run's first placeholder token
+        within its sample.
+
+    Raises:
+        ValueError: if the number of placeholder runs does not equal the number
+            of visual items, or a run's length does not match the item's token
+            count. Either mismatch means the text and vision streams are
+            misaligned; scattering anyway would silently corrupt the embeddings,
+            so fail loudly with the offending counts.
+    """
+    vision_mask = tokens == placeholder_id  # (bsz, seq_len)
+    # Shift within each row (row boundaries padded False) so a placeholder
+    # ending one sample and starting the next are NOT merged into one run across
+    # the flattened batch boundary.
+    prev_mask = torch.zeros_like(vision_mask)
+    prev_mask[:, 1:] = vision_mask[:, :-1]
+    next_mask = torch.zeros_like(vision_mask)
+    next_mask[:, :-1] = vision_mask[:, 1:]
+    flat_mask = vision_mask.view(-1)
+    region_starts = torch.where(flat_mask & ~prev_mask.view(-1))[0]
+    region_ends = torch.where(flat_mask & ~next_mask.view(-1))[0]
+    seq_len = tokens.shape[1]
+
+    num_items = int(num_vision_tokens_per_item.shape[0])
+    num_runs = int(region_starts.shape[0])
+    if num_runs != num_items:
+        raise ValueError(
+            f"Multimodal misalignment: found {num_runs} contiguous run(s) of "
+            f"placeholder id {placeholder_id} in the token sequence but received "
+            f"{num_items} visual item(s). Each visual item must correspond to "
+            f"exactly one placeholder run."
+        )
+
+    run_lengths = (region_ends - region_starts + 1).tolist()
+    positions: list[tuple[int, int, int, int]] = []
+    for i in range(num_items):
+        start = int(region_starts[i].item())
+        n_tokens = int(num_vision_tokens_per_item[i].item())
+        if run_lengths[i] != n_tokens:
+            raise ValueError(
+                f"Multimodal misalignment: placeholder run {i} spans "
+                f"{run_lengths[i]} token(s) but visual item {i} produced "
+                f"{n_tokens} embedding(s). The placeholder count in the prompt "
+                f"must match the vision token count for that item."
+            )
+        positions.append((i, start // seq_len, start % seq_len, n_tokens))
+    return positions
+
+
+def scatter_vision_embeds(
+    inputs_embeds: torch.Tensor,
+    *,
+    vision_embeds: torch.Tensor,
+    vision_positions: list[tuple[int, int, int, int]],
+) -> torch.Tensor:
+    """Copy padded vision features into the text sequence at placeholder runs.
+
+    Args:
+        inputs_embeds: (batch, seq_len, dim) text embeddings, modified in place.
+        vision_embeds: (num_items, max_tokens, dim) padded vision features.
+        vision_positions: from ``get_vision_positions``.
+    """
+    for item_idx, sample_idx, vision_start, n_tokens in vision_positions:
+        inputs_embeds[
+            sample_idx, vision_start : vision_start + n_tokens, :
+        ] = vision_embeds[item_idx, :n_tokens, :].to(inputs_embeds.dtype)
+    return inputs_embeds
diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py
new file mode 100644
index 0000000000..64ac8ec89a
--- /dev/null
+++ b/torchtitan/models/common/vision_encoder.py
@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Shared model-agnostic ViT building blocks for VLM vision encoders: a
+block-diagonal FlexAttention mask helper and the pre-norm transformer block
+(attention + MLP) over a padded ``(N, P, D)`` batch.
+
+RoPE differs per model, so each encoder passes it through the block to the
+attention as two per-forward args: ``rope_cache`` (a tensor, so config-based
+sharding can DTensor-wrap it before it meets the head-sharded q/k) and
+``rope_apply`` (a pass-through callable ``(q, k, rope_cache) -> (q, k)``).
+
+Shape suffixes:
+- N = num visual items
+- P = max patches per item (padded)
+- D = vision dim
+- H = num heads
+- Dh = head dim
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+
+from torchtitan.models.common import Linear
+from torchtitan.models.common.attention import FlexAttention
+from torchtitan.models.common.nn_modules import GELU, LayerNorm
+from torchtitan.protocols.module import Module
+
+compiled_create_block_mask = torch.compile(create_block_mask)
+
+# Applies rotary position embedding: (query, key, rope_cache) -> (query, key).
+RopeApply = Callable[
+    [torch.Tensor, torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]
+]
+
+
+def get_vision_block_mask_mod(num_patches: torch.Tensor) -> Callable:
+    """Block-diagonal mask: each visual item attends only to its own patches.
+
+    Args:
+        num_patches: (N,) real (non-padding) patch count per visual item (N is
+            the number of visual items, i.e. images/videos in the batch).
+    """
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        valid_q = q_idx < num_patches[b]
+        valid_kv = kv_idx < num_patches[b]
+        return valid_q & valid_kv
+
+    return mask_mod
+
+
+class VisionMLP(Module):
+    """Feed-forward network with GELU activation (fc1 -> act -> fc2)."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        fc1: Linear.Config
+        fc2: Linear.Config
+        act_fn: GELU.Config = field(
+            default_factory=lambda: GELU.Config(approximate="tanh")
+        )
+
+    def __init__(self, config: Config):
+        super().__init__()
+        self.linear_fc1 = config.fc1.build()
+        self.linear_fc2 = config.fc2.build()
+        self.act_fn = config.act_fn.build()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+
+
+class VisionAttention(Module):
+    """Multi-head self-attention with FlexAttention over a padded batch.
+
+    Separate q/k/v projections (clean per-head ColwiseParallel under TP). RoPE is
+    applied via the injected ``rope_apply`` callable so this class is reused
+    across models with different rotary formulations.
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        dim: int
+        num_heads: int
+        wq: Linear.Config
+        wk: Linear.Config
+        wv: Linear.Config
+        proj: Linear.Config
+        inner_attention: Module.Config = field(default_factory=FlexAttention.Config)
+
+    def __init__(self, config: Config):
+        super().__init__()
+        if config.dim % config.num_heads != 0:
+            raise ValueError(
+                f"VisionAttention dim ({config.dim}) must be divisible by "
+                f"num_heads ({config.num_heads})."
+            )
+        self.head_dim = config.dim // config.num_heads
+
+        self.wq = config.wq.build()
+        self.wk = config.wk.build()
+        self.wv = config.wv.build()
+        self.proj = config.proj.build()
+        self.flex_attention = config.inner_attention.build()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        rope_cache: torch.Tensor,
+        rope_apply: RopeApply,
+        attention_mask: BlockMask,
+    ) -> torch.Tensor:
+        N, P, _ = x.shape
+
+        # -1 infers the head count locally (= num_heads / TP under tensor
+        # parallelism, where wq/wk/wv are colwise-sharded).
+        q_NPHDh = self.wq(x).view(N, P, -1, self.head_dim)
+        k_NPHDh = self.wk(x).view(N, P, -1, self.head_dim)
+        v_NPHDh = self.wv(x).view(N, P, -1, self.head_dim)
+
+        q_NPHDh, k_NPHDh = rope_apply(q_NPHDh, k_NPHDh, rope_cache)
+
+        out_NPHDh = self.flex_attention(
+            q_NPHDh, k_NPHDh, v_NPHDh, attention_masks=attention_mask
+        )
+        out_NPD = out_NPHDh.reshape(N, P, -1)
+        return self.proj(out_NPD)
+
+
+class VisionTransformerBlock(Module):
+    """Pre-norm transformer block: norm -> attn -> residual -> norm -> mlp."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        norm1: LayerNorm.Config
+        norm2: LayerNorm.Config
+        attn: VisionAttention.Config
+        mlp: VisionMLP.Config
+
+    def __init__(self, config: Config):
+        super().__init__()
+        self.norm1 = config.norm1.build()
+        self.norm2 = config.norm2.build()
+        self.attn = config.attn.build()
+        self.mlp = config.mlp.build()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        rope_cache: torch.Tensor,
+        rope_apply: RopeApply,
+        attention_mask: BlockMask,
+    ) -> torch.Tensor:
+        x = x + self.attn(
+            self.norm1(x),
+            rope_cache=rope_cache,
+            rope_apply=rope_apply,
+            attention_mask=attention_mask,
+        )
+        x = x + self.mlp(self.norm2(x))
+        return x
diff --git a/torchtitan/models/deepseek_v3/__init__.py b/torchtitan/models/deepseek_v3/__init__.py
index 1b9f531678..e1d5759413 100644
--- a/torchtitan/models/deepseek_v3/__init__.py
+++ b/torchtitan/models/deepseek_v3/__init__.py
@@ -75,7 +75,7 @@ def _depth_experts_init(layer_id: int) -> dict[str, Callable]:
     }
 
 
-def _make_dsv3_attn_config(
+def make_mla_attention_config(
     *,
     layer_id: int,
     dim: int,
@@ -87,13 +87,16 @@ def _make_dsv3_attn_config(
     v_head_dim: int,
     mscale: float = 1.0,
     attn_backend: str,
+    linear_init: dict[str, Callable],
+    norm_init: dict[str, Callable],
+    depth_init: Callable[[int], dict[str, Callable]],
     rope: RoPE.Config,
 ) -> Attention.Config:
-    """Build a fully-specified DeepSeek V3 MLA Attention.Config.
+    """Build a fully-specified DeepSeek V3 MLA ``Attention.Config``.
 
-    All Linear and RMSNorm sub-configs have their dimensional fields set.
-    When q_lora_rank == 0, sets wq (not wq_a/wq_b).
-    When q_lora_rank > 0, sets wq_a/wq_b (not wq).
+    All Linear and RMSNorm sub-configs have their dimensional fields set. When
+    ``q_lora_rank == 0``, sets ``wq`` (not ``wq_a``/``wq_b``); when
+    ``q_lora_rank > 0``, sets ``wq_a``/``wq_b`` (not ``wq``).
     """
     inner_attention = get_attention_config(attn_backend)
     qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
@@ -102,26 +105,26 @@ def _make_dsv3_attn_config(
         wq = Linear.Config(
             in_features=dim,
             out_features=n_heads * qk_head_dim,
-            param_init=_LINEAR_INIT,
+            param_init=linear_init,
         )
         wq_a = None
         wq_b = None
         # q_norm is unused when q_lora_rank == 0 (never built), but the field is
         # required on Attention.Config so we supply a placeholder.
-        q_norm = RMSNorm.Config(normalized_shape=1, param_init=_NORM_INIT)
+        q_norm = RMSNorm.Config(normalized_shape=1, param_init=norm_init)
     else:
         wq = None
         wq_a = Linear.Config(
             in_features=dim,
             out_features=q_lora_rank,
-            param_init=_LINEAR_INIT,
+            param_init=linear_init,
         )
         wq_b = Linear.Config(
             in_features=q_lora_rank,
             out_features=n_heads * qk_head_dim,
-            param_init=_LINEAR_INIT,
+            param_init=linear_init,
         )
-        q_norm = RMSNorm.Config(normalized_shape=q_lora_rank, param_init=_NORM_INIT)
+        q_norm = RMSNorm.Config(normalized_shape=q_lora_rank, param_init=norm_init)
 
     return Attention.Config(
         dim=dim,
@@ -139,25 +142,25 @@ def _make_dsv3_attn_config(
         wkv_a=Linear.Config(
             in_features=dim,
             out_features=kv_lora_rank + qk_rope_head_dim,
-            param_init=_LINEAR_INIT,
+            param_init=linear_init,
         ),
-        kv_norm=RMSNorm.Config(normalized_shape=kv_lora_rank, param_init=_NORM_INIT),
+        kv_norm=RMSNorm.Config(normalized_shape=kv_lora_rank, param_init=norm_init),
         wkv_b=Linear.Config(
             in_features=kv_lora_rank,
             out_features=n_heads * (qk_nope_head_dim + v_head_dim),
-            param_init=_LINEAR_INIT,
+            param_init=linear_init,
         ),
         wo=Linear.Config(
             in_features=n_heads * v_head_dim,
             out_features=dim,
-            param_init=_depth_init(layer_id),
+            param_init=depth_init(layer_id),
         ),
         inner_attention=inner_attention,
         rope=dataclasses.replace(rope),
     )
 
 
-def _build_dsv3_layers(
+def build_mla_moe_layers(
     *,
     n_layers: int,
     n_dense_layers: int,
@@ -182,9 +185,13 @@ def _build_dsv3_layers(
     attn_backend: str,
     moe_comm_backend: str,
     non_blocking_capacity_factor: float | None,
+    linear_init: dict[str, Callable],
+    norm_init: dict[str, Callable],
+    depth_init: Callable[[int], dict[str, Callable]],
+    depth_experts_init: Callable[[int], dict[str, Callable]],
     rope: RoPE.Config,
 ) -> list[TransformerBlock.Config]:
-    """Build the list of per-layer TransformerBlock configs.
+    """Build the per-layer ``DeepSeekV3TransformerBlock`` configs (MLA + MoE).
 
     Layers with layer_id < n_dense_layers get a dense FeedForward and no MoE.
     Layers with layer_id >= n_dense_layers get a MoE and no FeedForward.
@@ -194,7 +201,7 @@ def _build_dsv3_layers(
     """
     layers = []
     for layer_id in range(n_layers):
-        attn_cfg = _make_dsv3_attn_config(
+        attn_cfg = make_mla_attention_config(
             layer_id=layer_id,
             dim=dim,
             n_heads=n_heads,
@@ -205,6 +212,9 @@ def _build_dsv3_layers(
             v_head_dim=v_head_dim,
             mscale=mscale,
             attn_backend=attn_backend,
+            linear_init=linear_init,
+            norm_init=norm_init,
+            depth_init=depth_init,
             rope=rope,
         )
 
@@ -212,8 +222,8 @@ def _build_dsv3_layers(
             ffn_cfg = make_ffn_config(
                 dim=dim,
                 hidden_dim=dense_hidden_dim,
-                w1_param_init=_LINEAR_INIT,
-                w2w3_param_init=_depth_init(layer_id),
+                w1_param_init=linear_init,
+                w2w3_param_init=depth_init(layer_id),
             )
             moe_cfg = None
         else:
@@ -223,7 +233,7 @@ def _build_dsv3_layers(
                 router=make_router_config(
                     dim=dim,
                     num_experts=num_experts,
-                    gate_param_init=_depth_init(layer_id),
+                    gate_param_init=depth_init(layer_id),
                     top_k=router_top_k,
                     score_func=router_score_func,
                     num_expert_groups=router_num_expert_groups,
@@ -236,15 +246,15 @@ def _build_dsv3_layers(
                     hidden_dim=moe_hidden_dim,
                     num_experts=num_experts,
                     top_k=router_top_k,
-                    param_init=_depth_experts_init(layer_id),
+                    param_init=depth_experts_init(layer_id),
                     comm_backend=moe_comm_backend,
                     non_blocking_capacity_factor=non_blocking_capacity_factor,
                 ),
                 shared_experts=make_ffn_config(
                     dim=dim,
                     hidden_dim=moe_hidden_dim * num_shared_experts,
-                    w1_param_init=_LINEAR_INIT,
-                    w2w3_param_init=_depth_init(layer_id),
+                    w1_param_init=linear_init,
+                    w2w3_param_init=depth_init(layer_id),
                 ),
             )
 
@@ -252,9 +262,9 @@ def _build_dsv3_layers(
             DeepSeekV3TransformerBlock.Config(
                 attention=attn_cfg,
                 attention_norm=RMSNorm.Config(
-                    normalized_shape=dim, param_init=_NORM_INIT
+                    normalized_shape=dim, param_init=norm_init
                 ),
-                ffn_norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+                ffn_norm=RMSNorm.Config(normalized_shape=dim, param_init=norm_init),
                 feed_forward=ffn_cfg,
                 moe=moe_cfg,
             )
@@ -262,6 +272,17 @@ def _build_dsv3_layers(
     return layers
 
 
+def _build_dsv3_layers(**kwargs) -> list[TransformerBlock.Config]:
+    """Thin wrapper: ``build_mla_moe_layers`` with DeepSeek V3's own inits."""
+    return build_mla_moe_layers(
+        **kwargs,
+        linear_init=_LINEAR_INIT,
+        norm_init=_NORM_INIT,
+        depth_init=_depth_init,
+        depth_experts_init=_depth_experts_init,
+    )
+
+
 def _debugmodel(
     attn_backend: str,
     moe_comm_backend: str,
diff --git a/torchtitan/models/kimi_k2_7/README.md b/torchtitan/models/kimi_k2_7/README.md
new file mode 100644
index 0000000000..2493dd58e5
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/README.md
@@ -0,0 +1,69 @@
+# Kimi K2.7
+
+Kimi K2.5, K2.6, and K2.7-Code share an architecture that pairs a
+**DeepSeek-V3-style** decoder (Multi-head Latent Attention + Mixture-of-Experts)
+with a **MoonViT3d** vision encoder.
+
+## Prerequisites
+
+Install the additional dependencies:
+
+```bash
+pip install av torchvision
+```
+
+## Architecture
+
+- **Decoder** — DeepSeek-V3 (MLA + MoE).
+- **Vision encoder** — MoonViT3d: linear patch embedding, learnable 2D spatial
+  position embeddings (plus a sinusoidal temporal term for video), 2D RoPE,
+  pre-norm transformer blocks, temporal mean-pool + 2x2 spatial merge, then a
+  2-layer MLP projector to the decoder hidden size.
+- **Multimodal forward** — projected vision embeddings are scattered into the
+  text embedding sequence at runs of the shared media placeholder token.
+
+## Model variants
+
+| Variant | LLM dim | Layers | Heads | Experts (top-k) | ViT dim | ViT layers | ViT heads |
+|---------|---------|--------|-------|-----------------|---------|------------|-----------|
+| debugmodel | 256 | 6 | 16 | 8 (top-3) | 256 | 4 | 4 |
+| moonlight-16B-A3B | 2048 | 27 | 16 | 64 (top-6) | — | — | — |
+| Kimi-VL-A3B | 2048 | 27 | 16 | 64 (top-6) | 1152 | 27 | 16 |
+| Kimi-K2.5 | 7168 | 61 | 64 | 384 (top-8) | 1152 | 27 | 16 |
+
+## Supported Parallelisms
+
+| Feature | Notes |
+|---------|-------|
+| FSDP / HSDP | Decoder sharded per-layer. Without PP, the vision encoder is a separate FSDP unit; with PP, it belongs to the first-stage root FSDP unit |
+| Tensor Parallelism (TP) | The token embedding and vision activations remain replicated for the vision scatter; decoder SP resumes at layer 0. Vision attention heads and linear layers are TP-sharded, without vision SP |
+| Expert Parallelism (EP) | DeepSeek-V3 routed + shared experts |
+| Pipeline Parallel (PP) | Vision encoder folded into the first stage; 1F1B and Interleaved1F1B schedules |
+
+## Numerical Checks
+
+The HuggingFace comparison covers the Kimi-VL compatibility flavor, not the 1T
+K2.x checkpoints. For full text+image float32 execution, it measures:
+
+- vision-feature cosine similarity: `0.999977`
+- normal end-to-end last-token logits: KL `4.3e-2`, top-1 match, top-5 4/5
+- with expert routing pinned to HuggingFace's selections: KL `5.3e-4`, top-1
+  match, top-5 5/5
+
+The routing-pinned result is a diagnostic that isolates non-routing math; it is
+not a normal end-to-end parity result.
+
+- **Parallelism correctness**: bit-identical logits (max diff `0.0`) for
+  no-parallel / FSDP / FSDP+EP; within bf16 tolerance for FSDP+EP+TP (with SP).
+
+Test scripts:
+- `scripts/checkpoint_conversion/numerical_tests_kimi.py`
+
+## TODO
+
+- Add a video dataset training pipeline.
+- Add INT4 (compressed-tensors) checkpoint loading. The released K2.5, K2.6,
+  and K2.7-Code 1T checkpoints are INT4 group-quantized; the inherited
+  DeepSeek-V3 adapter only handles FP8 block-scale, so the 1T config trains from
+  scratch but cannot load them yet.
+- Add Context Parallel (CP) support.
diff --git a/torchtitan/models/kimi_k2_7/__init__.py b/torchtitan/models/kimi_k2_7/__init__.py
new file mode 100644
index 0000000000..2373a5a0ef
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/__init__.py
@@ -0,0 +1,486 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections.abc import Callable
+from functools import partial
+
+import torch.nn as nn
+
+from torchtitan.components.optimizer import register_moe_load_balancing_hook
+
+from torchtitan.distributed.pipeline_parallel import pipeline_vlm
+from torchtitan.models.common import (
+    ComplexRoPE,
+    Embedding,
+    Linear,
+    RMSNorm,
+    TransformerBlock,
+)
+from torchtitan.models.common.nn_modules import LayerNorm
+from torchtitan.models.common.param_init import depth_scaled_std
+from torchtitan.models.common.vision_encoder import (
+    VisionAttention,
+    VisionMLP,
+    VisionTransformerBlock,
+)
+from torchtitan.models.deepseek_v3 import build_mla_moe_layers
+from torchtitan.models.utils import validate_converter_order
+from torchtitan.protocols.model import ModelConfigConverter
+from torchtitan.protocols.model_spec import ModelSpec
+
+from .model import KimiK25Model
+from .parallelize import parallelize_kimi_k2_5
+from .state_dict_adapter import KimiK25StateDictAdapter
+
+from .vision_encoder import (
+    KimiK25VisionEncoder,
+    VisionProjector,
+    VisionRotaryEmbedding2D,
+)
+
+__all__ = [
+    "parallelize_kimi_k2_5",
+    "KimiK25Model",
+    "KimiK25StateDictAdapter",
+    "KimiK25VisionEncoder",
+    "VisionProjector",
+    "VisionRotaryEmbedding2D",
+    "model_registry",
+    "kimi_k2_5_configs",
+    "KIMI_K2_5_SPECIAL_TOKENS",
+]
+
+
+KIMI_K2_5_SPECIAL_TOKENS = {
+    "image_token": "<|media_pad|>",
+    "video_token": "<|media_pad|>",
+    "vision_start_token": "<|media_begin|>",
+    "vision_end_token": "<|media_end|>",
+    "pad_token": "[PAD]",
+}
+
+
+_LINEAR_INIT = {
+    "weight": partial(nn.init.trunc_normal_, std=0.02),
+    "bias": nn.init.zeros_,
+}
+_NORM_INIT = {"weight": nn.init.ones_}
+_EMBEDDING_INIT = {"weight": partial(nn.init.normal_, std=1.0)}
+_POS_EMB_INIT = {"pos_embed": partial(nn.init.normal_, std=1.0)}
+
+
+def _output_linear_init(dim: int) -> dict[str, Callable]:
+    s = dim**-0.5
+    return {
+        "weight": partial(nn.init.trunc_normal_, std=s, a=-3 * s, b=3 * s),
+        "bias": nn.init.zeros_,
+    }
+
+
+def _depth_init(layer_id: int) -> dict[str, Callable]:
+    return {
+        "weight": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)),
+        "bias": nn.init.zeros_,
+    }
+
+
+def _depth_experts_init(layer_id: int) -> dict[str, Callable]:
+    return {
+        "w1_EFD": partial(nn.init.trunc_normal_, std=0.02),
+        "w2_EDF": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)),
+        "w3_EFD": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)),
+    }
+
+
+def _vl_linear(in_features: int, out_features: int) -> Linear.Config:
+    return Linear.Config(
+        in_features=in_features,
+        out_features=out_features,
+        bias=True,
+        param_init=_LINEAR_INIT,
+    )
+
+
+def _vl_layernorm(dim: int, eps: float = 1e-5) -> LayerNorm.Config:
+    return LayerNorm.Config(normalized_shape=dim, eps=eps)
+
+
+def _vision_encoder_config(
+    *,
+    dim: int,
+    ffn_dim: int,
+    num_layers: int,
+    num_heads: int,
+    patch_size: int = 14,
+    in_channels: int = 3,
+    init_pos_emb_height: int = 64,
+    init_pos_emb_width: int = 64,
+    rope_theta: float = 10000.0,
+    merge_kernel_size: list[int] | None = None,
+    text_hidden_size: int = 7168,
+) -> KimiK25VisionEncoder.Config:
+    """Build a fully-specified KimiK25VisionEncoder.Config (MoonViT3d)."""
+    if merge_kernel_size is None:
+        merge_kernel_size = [2, 2]
+    patch_dim = in_channels * patch_size * patch_size
+    head_dim = dim // num_heads
+    merged_dim = dim * merge_kernel_size[0] * merge_kernel_size[1]
+
+    block = VisionTransformerBlock.Config(
+        norm1=_vl_layernorm(dim),
+        norm2=_vl_layernorm(dim),
+        attn=VisionAttention.Config(
+            dim=dim,
+            num_heads=num_heads,
+            wq=_vl_linear(dim, dim),
+            wk=_vl_linear(dim, dim),
+            wv=_vl_linear(dim, dim),
+            proj=_vl_linear(dim, dim),
+        ),
+        mlp=VisionMLP.Config(
+            fc1=_vl_linear(dim, ffn_dim),
+            fc2=_vl_linear(ffn_dim, dim),
+        ),
+    )
+
+    return KimiK25VisionEncoder.Config(
+        dim=dim,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        patch_size=patch_size,
+        in_channels=in_channels,
+        merge_kernel_size=merge_kernel_size,
+        text_hidden_size=text_hidden_size,
+        init_pos_emb_height=init_pos_emb_height,
+        init_pos_emb_width=init_pos_emb_width,
+        param_init=_POS_EMB_INIT,
+        patch_embed_proj=_vl_linear(patch_dim, dim),
+        rotary_pos_emb=VisionRotaryEmbedding2D.Config(
+            head_dim=head_dim, theta=rope_theta
+        ),
+        block=block,
+        final_norm=_vl_layernorm(dim),
+        projector=VisionProjector.Config(
+            vt_hidden_size=dim,
+            merged_dim=merged_dim,
+            pre_norm=_vl_layernorm(dim),
+            linear_1=_vl_linear(merged_dim, merged_dim),
+            linear_2=_vl_linear(merged_dim, text_hidden_size),
+        ),
+    )
+
+
+def _build_kimi_layers(**kwargs) -> list[TransformerBlock.Config]:
+    """Build MLA/MoE layers with the Kimi-family parameter initializers."""
+    return build_mla_moe_layers(
+        **kwargs,
+        linear_init=_LINEAR_INIT,
+        norm_init=_NORM_INIT,
+        depth_init=_depth_init,
+        depth_experts_init=_depth_experts_init,
+    )
+
+
+def _debugmodel(
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None = None,
+) -> KimiK25Model.Config:
+    dim = 256
+    n_layers = 6
+    vocab_size = 2048
+    n_heads = 16
+    moe_hidden_dim = 256
+    num_shared_experts = 2
+    dense_hidden_dim = 1024
+    rope_dim = 64
+    num_experts = 8
+    n_dense_layers = 1
+
+    layers = _build_kimi_layers(
+        n_layers=n_layers,
+        n_dense_layers=n_dense_layers,
+        dim=dim,
+        n_heads=n_heads,
+        q_lora_rank=0,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=rope_dim,
+        v_head_dim=128,
+        mscale=0.70,
+        dense_hidden_dim=dense_hidden_dim,
+        moe_hidden_dim=moe_hidden_dim,
+        num_experts=num_experts,
+        num_shared_experts=num_shared_experts,
+        router_top_k=3,
+        router_score_func="softmax",
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope=ComplexRoPE.Config(
+            dim=rope_dim,
+            max_seq_len=4096 * 4,
+            theta=10000.0,
+            scaling="yarn",
+            rope_factor=40.0,
+            beta_fast=32.0,
+            beta_slow=1.0,
+            original_seq_len=4096,
+        ),
+    )
+    config = KimiK25Model.Config(
+        vocab_size=vocab_size,
+        dim=dim,
+        tok_embeddings=Embedding.Config(
+            num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT
+        ),
+        norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+        lm_head=Linear.Config(
+            in_features=dim,
+            out_features=vocab_size,
+            param_init=_output_linear_init(dim),
+        ),
+        layers=layers,
+        vision_encoder=_vision_encoder_config(
+            dim=256,
+            ffn_dim=512,
+            num_layers=4,
+            num_heads=4,
+            patch_size=14,
+            init_pos_emb_height=16,
+            init_pos_emb_width=16,
+            text_hidden_size=dim,
+        ),
+    )
+    return config
+
+
+def _moonlight_16b_a3b_config(
+    *,
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None,
+    rope_theta: float,
+    max_seq_len: int,
+    vision_encoder: "KimiK25VisionEncoder.Config | None" = None,
+) -> KimiK25Model.Config:
+    """Shared Moonshot 16B-A3B DeepSeekV3 text tower (MLA + sigmoid-routed MoE).
+
+    Used by both Moonlight (text-only) and Kimi-VL (with a vision encoder): they
+    share the architecture (no q-LoRA, no RoPE scaling, 64 experts top-6); only
+    the RoPE ``theta`` / ``max_seq_len`` and the vision tower differ.
+    """
+    dim = 2048
+    vocab_size = 163840
+    layers = _build_kimi_layers(
+        n_layers=27,
+        n_dense_layers=1,
+        dim=dim,
+        n_heads=16,
+        q_lora_rank=0,  # q_lora_rank null -> separate wq
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        mscale=1.0,
+        dense_hidden_dim=11264,
+        moe_hidden_dim=1408,
+        num_experts=64,
+        num_shared_experts=2,
+        router_top_k=6,
+        router_score_func="sigmoid",
+        router_num_expert_groups=None,
+        router_num_limited_groups=None,
+        router_route_scale=2.446,
+        router_route_norm=True,
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope=ComplexRoPE.Config(dim=64, max_seq_len=max_seq_len, theta=rope_theta),
+    )
+    return KimiK25Model.Config(
+        vocab_size=vocab_size,
+        dim=dim,
+        tok_embeddings=Embedding.Config(
+            num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT
+        ),
+        norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+        lm_head=Linear.Config(
+            in_features=dim,
+            out_features=vocab_size,
+            param_init=_output_linear_init(dim),
+        ),
+        layers=layers,
+        vision_encoder=vision_encoder,
+    )
+
+
+def _moonlight_16b_a3b(
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None = None,
+) -> KimiK25Model.Config:
+    """Build the text-only Moonlight 16B-A3B sibling without a vision tower."""
+    return _moonlight_16b_a3b_config(
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope_theta=50000.0,
+        max_seq_len=8192,
+        vision_encoder=None,
+    )
+
+
+def _kimi_vl_a3b(
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None = None,
+) -> KimiK25Model.Config:
+    """Kimi-VL 16B-A3B: Moonlight text tower plus a 2D MoonViT vision tower.
+
+    Kimi-VL's original 2D MoonViT is the ``t=1`` case of MoonViT3d, so it reuses
+    ``KimiK25VisionEncoder``. With image inputs, the temporal embedding is not
+    applied.
+    """
+    config = _moonlight_16b_a3b_config(
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope_theta=800000.0,
+        max_seq_len=131072,
+        vision_encoder=_vision_encoder_config(
+            dim=1152,
+            ffn_dim=4304,
+            num_layers=27,
+            num_heads=16,
+            patch_size=14,
+            init_pos_emb_height=64,
+            init_pos_emb_width=64,
+            text_hidden_size=2048,
+        ),
+    )
+    return config
+
+
+def _kimi_k2_5(
+    attn_backend: str,
+    moe_comm_backend: str,
+    non_blocking_capacity_factor: float | None = None,
+) -> KimiK25Model.Config:
+    """Architecture shared by Kimi K2.5, K2.6, and K2.7-Code: a ~1T-total /
+    ~32B-active DeepSeekV3-style text tower (384 routed experts, top-8) plus a
+    MoonViT3d vision tower.
+
+    All three checkpoints use the same parameterized architecture and tensor
+    schema. Their release-specific tokenizer and chat-template assets are
+    configured separately from this model definition.
+    """
+    dim = 7168
+    n_layers = 61
+    vocab_size = 163840
+    n_heads = 64
+    q_lora_rank = 1536
+    moe_hidden_dim = 2048
+    num_shared_experts = 1
+    dense_hidden_dim = 18432
+    rope_dim = 64  # qk_rope_head_dim
+    num_experts = 384
+    n_dense_layers = 1
+
+    layers = _build_kimi_layers(
+        n_layers=n_layers,
+        n_dense_layers=n_dense_layers,
+        dim=dim,
+        n_heads=n_heads,
+        q_lora_rank=q_lora_rank,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=rope_dim,
+        v_head_dim=128,
+        mscale=1.0,
+        dense_hidden_dim=dense_hidden_dim,
+        moe_hidden_dim=moe_hidden_dim,
+        num_experts=num_experts,
+        num_shared_experts=num_shared_experts,
+        router_top_k=8,
+        router_score_func="sigmoid",
+        router_num_expert_groups=None,
+        router_num_limited_groups=None,
+        router_route_scale=2.827,
+        router_route_norm=True,
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+        rope=ComplexRoPE.Config(
+            dim=rope_dim,
+            max_seq_len=262144,
+            theta=50000.0,
+            scaling="yarn",
+            rope_factor=64.0,
+            beta_fast=32.0,
+            beta_slow=1.0,
+            original_seq_len=4096,
+        ),
+    )
+    return KimiK25Model.Config(
+        vocab_size=vocab_size,
+        dim=dim,
+        tok_embeddings=Embedding.Config(
+            num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT
+        ),
+        norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT),
+        lm_head=Linear.Config(
+            in_features=dim,
+            out_features=vocab_size,
+            param_init=_output_linear_init(dim),
+        ),
+        layers=layers,
+        vision_encoder=_vision_encoder_config(
+            dim=1152,
+            ffn_dim=4304,
+            num_layers=27,
+            num_heads=16,
+            patch_size=14,
+            init_pos_emb_height=64,
+            init_pos_emb_width=64,
+            text_hidden_size=dim,
+        ),
+    )
+
+
+kimi_k2_5_configs = {
+    "debugmodel": _debugmodel,
+    "moonlight-16B-A3B": _moonlight_16b_a3b,
+    "Kimi-VL-A3B": _kimi_vl_a3b,
+    "Kimi-K2.5": _kimi_k2_5,
+}
+
+
+def model_registry(
+    flavor: str,
+    attn_backend: str = "flex",
+    moe_comm_backend: str = "standard",
+    non_blocking_capacity_factor: float | None = None,
+    converters: list[ModelConfigConverter.Config] | None = None,
+) -> ModelSpec:
+    config = kimi_k2_5_configs[flavor](
+        attn_backend=attn_backend,
+        moe_comm_backend=moe_comm_backend,
+        non_blocking_capacity_factor=non_blocking_capacity_factor,
+    )
+    if converters is not None:
+        validate_converter_order(converters)
+        for c in converters:
+            c.build().convert(config)
+    return ModelSpec(
+        name="kimi_k2_5",
+        flavor=flavor,
+        model=config,
+        parallelize_fn=parallelize_kimi_k2_5,
+        pipelining_fn=pipeline_vlm,
+        post_optimizer_build_fn=register_moe_load_balancing_hook,
+        state_dict_adapter=KimiK25StateDictAdapter,
+    )
diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py
new file mode 100644
index 0000000000..5dae9b381b
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/config_registry.py
@@ -0,0 +1,183 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
+from torchtitan.components.lr_scheduler import LRSchedulersContainer
+from torchtitan.components.metrics import MetricsProcessor
+from torchtitan.components.optimizer import default_adamw
+from torchtitan.components.tokenizer import MultiModalTokenizer
+from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig
+from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC
+from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader
+from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget
+from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader
+from torchtitan.models.common.config_utils import decoder_vocab_size
+from torchtitan.trainer import Trainer
+
+from . import KIMI_K2_5_SPECIAL_TOKENS, model_registry
+
+
+def _mm_dataloader(dataset: str, **kwargs) -> MMDataLoader.Config:
+    return MMDataLoader.Config(
+        dataset=dataset,
+        max_images_per_batch=128,
+        patch_size=14,
+        temporal_patch_size=1,
+        spatial_merge_size=2,
+        patch_order="raster",
+        resize_fn=resize_to_patch_budget,
+        max_patches=16384,
+        max_patches_per_side=512,
+        min_pixels=65536,
+        max_pixels=16777216,
+        image_mean=(0.5, 0.5, 0.5),
+        image_std=(0.5, 0.5, 0.5),
+        **kwargs,
+    )
+
+
+def kimi_k2_5_debugmodel() -> Trainer.Config:
+    model_spec = model_registry("debugmodel")
+    return Trainer.Config(
+        loss=ChunkedLossWrapper.Config(
+            loss_fn=CrossEntropyLoss.Config(
+                global_vocab_size=decoder_vocab_size(model_spec),
+            ),
+        ),
+        hf_assets_path="./tests/assets/tokenizer",
+        tokenizer=MultiModalTokenizer.Config(**KIMI_K2_5_SPECIAL_TOKENS),
+        metrics=MetricsProcessor.Config(log_freq=1),
+        model_spec=model_spec,
+        dataloader=_mm_dataloader("cc12m-test"),
+        optimizer=default_adamw(lr=8e-4),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=2,
+            decay_ratio=0.8,
+            decay_type="linear",
+            min_lr_factor=0.0,
+        ),
+        training=TrainingConfig(
+            local_batch_size=1,
+            seq_len=512,
+            steps=10,
+        ),
+        parallelism=ParallelismConfig(
+            expert_parallel_degree=1,
+        ),
+        checkpoint=CheckpointManager.Config(
+            interval=10,
+            last_save_model_only=False,
+        ),
+        activation_checkpoint=SelectiveAC.Config(),
+    )
+
+
+def moonlight_16b_a3b() -> Trainer.Config:
+    """Moonlight 16B-A3B: the text-only DeepSeekV3 sibling (no vision tower)."""
+    model_spec = model_registry("moonlight-16B-A3B", attn_backend="flex")
+    return Trainer.Config(
+        loss=ChunkedLossWrapper.Config(
+            loss_fn=CrossEntropyLoss.Config(
+                global_vocab_size=decoder_vocab_size(model_spec),
+            ),
+        ),
+        hf_assets_path="./assets/hf/Moonlight-16B-A3B",
+        model_spec=model_spec,
+        dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"),
+        optimizer=default_adamw(lr=3e-4),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=2000,
+            decay_ratio=0.8,
+            decay_type="cosine",
+            min_lr_factor=0.1,
+        ),
+        training=TrainingConfig(
+            local_batch_size=4,
+            seq_len=4096,
+            steps=10000,
+        ),
+        parallelism=ParallelismConfig(
+            expert_parallel_degree=8,
+        ),
+        checkpoint=CheckpointManager.Config(interval=500),
+        activation_checkpoint=FullAC.Config(),
+    )
+
+
+def kimi_vl_a3b() -> Trainer.Config:
+    """Kimi-VL A3B: Moonlight text tower + 2D MoonViT vision (image-text)."""
+    model_spec = model_registry("Kimi-VL-A3B", attn_backend="flex")
+    return Trainer.Config(
+        loss=ChunkedLossWrapper.Config(
+            loss_fn=CrossEntropyLoss.Config(
+                global_vocab_size=decoder_vocab_size(model_spec),
+            ),
+        ),
+        hf_assets_path="./assets/hf/Kimi-VL-A3B",
+        # Kimi-VL-A3B names the vision-start token <|media_start|>, whereas the
+        # K2.5 family uses <|media_begin|>; override just that one entry.
+        tokenizer=MultiModalTokenizer.Config(
+            **{**KIMI_K2_5_SPECIAL_TOKENS, "vision_start_token": "<|media_start|>"}
+        ),
+        model_spec=model_spec,
+        # Kimi-VL is a compatibility flavor; resizing intentionally follows
+        # Kimi-K2.5 per-side scaling instead of legacy Kimi-VL's side rejection.
+        dataloader=_mm_dataloader("cc12m"),
+        optimizer=default_adamw(lr=3e-4),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=2000,
+            decay_ratio=0.8,
+            decay_type="cosine",
+            min_lr_factor=0.1,
+        ),
+        training=TrainingConfig(
+            local_batch_size=1,
+            seq_len=4096,
+            steps=10000,
+        ),
+        parallelism=ParallelismConfig(
+            expert_parallel_degree=8,
+        ),
+        checkpoint=CheckpointManager.Config(interval=500),
+        activation_checkpoint=FullAC.Config(),
+    )
+
+
+def kimi_k2_5() -> Trainer.Config:
+    """Full Kimi K2.5 (~1T-total / ~32B-active)."""
+    compile_config = CompileConfig(enable=True, components=["loss"])
+    # The report uses BF16 compute; its FP8 path only compresses saved activations.
+    model_spec = model_registry("Kimi-K2.5", attn_backend="flex")
+    return Trainer.Config(
+        loss=ChunkedLossWrapper.Config(
+            loss_fn=CrossEntropyLoss.Config(
+                global_vocab_size=decoder_vocab_size(model_spec),
+            ),
+        ),
+        hf_assets_path="./assets/hf/Kimi-K2.5",
+        model_spec=model_spec,
+        dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"),
+        optimizer=default_adamw(lr=2.2e-4),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=2000,
+            decay_ratio=0.8,
+            decay_type="cosine",
+            min_lr_factor=0.1,
+        ),
+        training=TrainingConfig(
+            local_batch_size=4,
+            seq_len=4096,
+            steps=10000,
+        ),
+        parallelism=ParallelismConfig(
+            pipeline_parallel_schedule="Interleaved1F1B",
+            expert_parallel_degree=8,
+        ),
+        checkpoint=CheckpointManager.Config(interval=500),
+        activation_checkpoint=FullAC.Config(),
+        compile=compile_config,
+    )
diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py
new file mode 100644
index 0000000000..bfa6ffed65
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/model.py
@@ -0,0 +1,187 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Reference (SGLang):
+https://github.com/sgl-project/sglang/blob/e0c0c0a45cb1bda90392bfa2bba4184f5b0638a0/python/sglang/srt/models/kimi_k25.py
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+from torchtitan.models.common.attention import AttentionMasksType
+from torchtitan.models.common.decoder import Decoder
+from torchtitan.models.common.multimodal import (
+    get_vision_positions,
+    scatter_vision_embeds,
+)
+from torchtitan.models.deepseek_v3.model import DeepSeekV3Model
+
+from .sharding import set_kimi_k2_5_sharding_config
+from .vision_encoder import KimiK25VisionEncoder
+
+
+class KimiK25Model(DeepSeekV3Model):
+    """Kimi K2.5: DeepSeekV3 language model with a MoonViT3d vision encoder.
+
+    Forward pass flow::
+
+        forward(tokens, pixel_values[/videos], grid_thw, ...)
+          |
+          +-- tok_embeddings(tokens)               -> text embeddings
+          +-- vision_encoder(pixels)               -> padded vision features
+          +-- scatter at vision placeholder runs   -> multimodal embeddings
+          +-- decoder layers (MLA + MoE)           -> hidden states
+          +-- norm -> lm_head                      -> logits
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(DeepSeekV3Model.Config):
+        vision_encoder: KimiK25VisionEncoder.Config | None = None
+
+        def update_from_config(
+            self,
+            *,
+            config,
+            **kwargs,
+        ) -> None:
+            Decoder.Config.update_from_config(self, config=config, **kwargs)
+            parallelism = config.parallelism
+
+            # Decoder.Config validates the text attention heads. Vision attention
+            # is also head-sharded, so validate its head count independently.
+            tp = parallelism.tensor_parallel_degree
+            if (
+                tp > 1
+                and self.vision_encoder is not None
+                and self.vision_encoder.num_heads % tp != 0
+            ):
+                raise ValueError(
+                    f"tensor_parallel_degree ({tp}) must divide "
+                    f"vision num_heads ({self.vision_encoder.num_heads})."
+                )
+
+            set_kimi_k2_5_sharding_config(
+                self,
+                enable_sp=parallelism.enable_sequence_parallel,
+                enable_ep=parallelism.expert_parallel_degree > 1,
+            )
+
+    def __init__(self, config: Config):
+        super().__init__(config)
+        self.vision_encoder = (
+            config.vision_encoder.build() if config.vision_encoder is not None else None
+        )
+
+    def _prepare_multimodal_embeds(
+        self,
+        tokens: torch.Tensor,
+        *,
+        pixel_values: torch.Tensor | None,
+        grid_thw: torch.Tensor | None,
+        pixel_values_videos: torch.Tensor | None = None,
+        grid_thw_videos: torch.Tensor | None = None,
+        special_tokens: dict[str, int],
+    ) -> torch.Tensor:
+        """Embed tokens, run the vision encoder, scatter features into text.
+
+        With kimi's single unified placeholder, a one-modality batch's runs map
+        to visual items in order. Mixing images and videos in one batch is not
+        yet supported (see the TODO below).
+        """
+        inputs_embeds = (
+            self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens
+        )
+
+        modalities = []
+        if pixel_values is not None and grid_thw is not None:
+            modalities.append((pixel_values, grid_thw))
+        if pixel_values_videos is not None and grid_thw_videos is not None:
+            modalities.append((pixel_values_videos, grid_thw_videos))
+
+        if not modalities:
+            return inputs_embeds
+        # TODO: support mixed image+video batches. Upstream fix: when
+        # image_id == video_id, emit one document-ordered vision stream so the
+        # runs stay modality-agnostic and this branch goes away.
+        assert len(modalities) == 1, "mixed image+video batches not yet supported"
+        pixels, grid = modalities[0]
+        # A non-empty modalities list means a multimodal run, so the vision
+        # encoder is present (text-only configs never populate pixels).
+        assert self.vision_encoder is not None
+
+        placeholder_id = special_tokens["image_id"]
+        assert placeholder_id == special_tokens["video_id"]
+
+        # Patches arrive float32; match the encoder's compute dtype for the matmul.
+        pixels = pixels.to(self.vision_encoder.patch_embed.weight.dtype)
+        vision_embeds = self.vision_encoder(pixels, grid_thw=grid)
+        # MoonViT collapses time (temporal pooling) and merges 2x2 spatially, so
+        # the token count is (h/kh)*(w/kw), independent of t.
+        kh, kw = self.vision_encoder.merge_kernel_size
+        num_tokens_per_item = (grid[:, 1] // kh) * (grid[:, 2] // kw)
+        vision_positions = get_vision_positions(
+            tokens, num_tokens_per_item, placeholder_id
+        )
+        if vision_positions:
+            inputs_embeds = scatter_vision_embeds(
+                inputs_embeds,
+                vision_embeds=vision_embeds,
+                vision_positions=vision_positions,
+            )
+        return inputs_embeds
+
+    def forward(  # pyrefly: ignore [bad-override]
+        self,
+        tokens: torch.Tensor,
+        *,
+        pixel_values: torch.Tensor | None = None,
+        grid_thw: torch.Tensor | None = None,
+        pixel_values_videos: torch.Tensor | None = None,
+        grid_thw_videos: torch.Tensor | None = None,
+        special_tokens: dict[str, int] | None = None,
+        attention_masks: AttentionMasksType | None = None,
+        positions: torch.Tensor | None = None,
+    ):
+        """Forward pass for Kimi K2.5.
+
+        Images and videos share one unified ``<|media_pad|>`` placeholder.
+
+        Args:
+            tokens: (batch, seq_len) token IDs.
+            pixel_values: (num_images, max_num_patch, patch_dim) padded image
+                patches, or None for text-only / video-only batches.
+            grid_thw: (num_images, 3) patch counts ``[t, h, w]`` per image.
+            pixel_values_videos: padded video patches, or None (mixing with
+                ``pixel_values`` in one batch is not yet supported).
+            grid_thw_videos: (num_videos, 3) patch counts per video.
+            special_tokens: tokenizer-resolved ``image_id``/``video_id``;
+                required for image/video batches, None for text-only.
+            attention_masks: Decoder attention masks.
+            positions: Per-token position IDs for packed sequences.
+
+        Returns:
+            (batch, seq_len, vocab_size) logits.
+        """
+        if self.tok_embeddings is not None:
+            x = self._prepare_multimodal_embeds(
+                tokens,
+                pixel_values=pixel_values,
+                grid_thw=grid_thw,
+                pixel_values_videos=pixel_values_videos,
+                grid_thw_videos=grid_thw_videos,
+                special_tokens=special_tokens,  # pyrefly: ignore [bad-argument-type]
+            )
+        else:
+            x = tokens
+
+        for layer in self.layers.values():
+            x = layer(x, attention_masks, positions)
+
+        x = self.norm(x) if self.norm is not None else x
+        if self._skip_lm_head:
+            return x
+        return self.lm_head(x) if self.lm_head is not None else x
diff --git a/torchtitan/models/kimi_k2_7/parallelize.py b/torchtitan/models/kimi_k2_7/parallelize.py
new file mode 100644
index 0000000000..d02d0b359c
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/parallelize.py
@@ -0,0 +1,131 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Parallelization for Kimi K2.5 (MoonViT3d vision encoder + DeepSeekV3 decoder).
+
+TP/SP/EP is applied by ``model.parallelize(parallel_dims)`` from the
+``ShardingConfig`` that ``set_kimi_k2_5_sharding_config`` sets on every
+sub-config (see ``sharding.py``). FSDP is then applied in two parts: the vision
+encoder as a single unit (its compute is small, so one all-gather beats many
+per-layer ones), then the decoder with MoE-aware per-layer wrapping. Context
+Parallel and ``full_dtensor`` are not supported.
+"""
+
+import torch.nn as nn
+
+from torchtitan.config import (
+    CompileConfig,
+    ParallelismConfig,
+    TORCH_DTYPE_MAP,
+    TrainingConfig,
+)
+from torchtitan.distributed import ParallelDims
+from torchtitan.distributed.activation_checkpoint import ActivationCheckpointingConfig
+from torchtitan.distributed.compile import apply_compile
+from torchtitan.distributed.fsdp import (
+    apply_fsdp_to_decoder,
+    apply_fsdp_to_vision_encoder,
+)
+from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
+
+
+def parallelize_kimi_k2_5(
+    model: nn.Module,
+    *,
+    parallel_dims: ParallelDims,
+    training: TrainingConfig,
+    parallelism: ParallelismConfig,
+    compile_config: CompileConfig,
+    ac_config: ActivationCheckpointingConfig,
+    dump_folder: str,
+):
+    """Apply TP/EP, activation checkpointing, ``torch.compile``, and FSDP.
+
+    Order: config-based TP/EP (``model.parallelize``) -> async-TP ->
+    activation checkpointing -> compile -> FSDP (vision encoder as a single
+    unit, then the MoE-aware decoder).
+
+    NOTE: the passed-in model should preferably be on meta device; otherwise it
+    must fit in GPU or CPU memory.
+    """
+    if parallelism.spmd_backend == "full_dtensor":
+        raise NotImplementedError("full_dtensor is not supported yet.")
+
+    if parallel_dims.cp_enabled:
+        raise NotImplementedError(
+            "Context Parallel is not yet supported for Kimi K2.5: vision scatter "
+            "needs the full sequence before CP would shard it."
+        )
+
+    model_compile_enabled = (
+        compile_config.enable and "model" in compile_config.components
+    )
+
+    if parallel_dims.tp_enabled or parallel_dims.ep_enabled:
+        if parallelism.enable_async_tensor_parallel and not model_compile_enabled:
+            raise RuntimeError("Async TP requires torch.compile")
+        model.parallelize(parallel_dims)  # pyrefly: ignore [not-callable]
+
+    if parallel_dims.tp_enabled:
+        maybe_enable_async_tp(parallelism, compile_config, parallel_dims.get_mesh("tp"))
+
+    if ac_config is not None:
+        ac_policy = ac_config.build(dump_folder=dump_folder)
+        ac_policy.apply(model)
+        if model.vision_encoder is not None:
+            ac_policy.apply(model.vision_encoder)
+
+    if model_compile_enabled:
+        apply_compile(model, compile_config)
+        if model.vision_encoder is not None:
+            # pyrefly: ignore [bad-argument-type]
+            apply_compile(model.vision_encoder, compile_config)
+
+    dp_mesh_names = (
+        ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"]
+    )
+    dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
+
+    # FSDP the vision encoder as a single unit, before the decoder's FSDP.
+    #
+    # Skipped under PP: as its own fully_shard root, the encoder's forward may
+    # not run on a given microbatch (a text-only batch, or a stage that holds the
+    # encoder but gets no pixels), so its params are never all-gathered -- which
+    # trips FSDP's per-microbatch post-backward callback. Under PP we leave it
+    # unwrapped so its params fall into the root ``fully_shard(model)`` in
+    # ``apply_fsdp_to_decoder``.
+    if model.vision_encoder is not None and not parallel_dims.pp_enabled:
+        apply_fsdp_to_vision_encoder(
+            model.vision_encoder,  # pyrefly: ignore [bad-argument-type]
+            dp_mesh,
+            param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param],
+            reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce],
+            reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward,
+            pp_enabled=parallel_dims.pp_enabled,
+        )
+
+    edp_mesh = None
+    if parallel_dims.ep_enabled:
+        edp_mesh_names = (
+            ["dp_replicate", "efsdp"]
+            if parallel_dims.dp_replicate_enabled
+            else ["efsdp"]
+        )
+        edp_mesh = parallel_dims.get_optional_mesh(edp_mesh_names)
+
+    apply_fsdp_to_decoder(
+        model,  # pyrefly: ignore [bad-argument-type]
+        dp_mesh,
+        param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param],
+        reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce],
+        pp_enabled=parallel_dims.pp_enabled,
+        cpu_offload=training.enable_cpu_offload,
+        reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward,
+        ep_degree=parallel_dims.ep,
+        edp_mesh=edp_mesh,
+    )
+
+    return model
diff --git a/torchtitan/models/kimi_k2_7/sharding.py b/torchtitan/models/kimi_k2_7/sharding.py
new file mode 100644
index 0000000000..058e061006
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/sharding.py
@@ -0,0 +1,143 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Config-based sharding for Kimi K2.5 (MoonViT3d + DeepSeekV3).
+
+Sets ``ShardingConfig`` on every sub-config so ``model.parallelize()`` applies
+TP/EP/SP uniformly via the Module protocol.
+
+- Decoder (MLA + MoE): reuses ``set_deepseek_v3_sharding_config``. Multimodal
+  configs keep the token embedding ``Replicate`` for the vision scatter and
+  resume SP at layer 0 (see ``_shard_decoder_after_embedding_scatter``).
+- Vision encoder: activations flow ``Replicate`` (no SP -- the patch sequence is
+  short, so sequence-sharding would add gather/scatter around the block-diagonal
+  attention for little memory gain). Only the linear layers are Colwise/Rowwise
+  sharded for memory; norms and position embeddings stay ``Replicate``.
+"""
+
+from typing import TYPE_CHECKING
+
+import spmd_types as spmd
+
+from torchtitan.models.common.decoder_sharding import (
+    colwise_config,
+    dense_activation_placement,
+    dense_param_placement,
+    rowwise_config,
+    set_gqa_inner_attention_local_map,
+)
+from torchtitan.models.deepseek_v3.sharding import set_deepseek_v3_sharding_config
+from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig
+
+if TYPE_CHECKING:
+    from torchtitan.models.kimi_k2_7.model import KimiK25Model
+
+_REPLICATE_PARAM = dense_param_placement(tp=spmd.R)
+_REPLICATE_ACT = dense_activation_placement(tp=spmd.R)
+
+_REPLICATE_NORM = ShardingConfig(
+    state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM},
+    in_src_shardings={"input": _REPLICATE_ACT},
+    in_dst_shardings={"input": _REPLICATE_ACT},
+    out_dst_shardings=_REPLICATE_ACT,
+)
+
+
+def set_kimi_k2_5_sharding_config(
+    config: "KimiK25Model.Config",
+    *,
+    enable_sp: bool,
+    enable_ep: bool,
+) -> None:
+    set_deepseek_v3_sharding_config(
+        config,
+        enable_sp=enable_sp,
+        enable_ep=enable_ep,
+    )
+    if config.vision_encoder is not None:
+        if enable_sp:
+            _shard_decoder_after_embedding_scatter(config)
+        _set_vision_encoder_sharding(config.vision_encoder)
+
+
+def _shard_decoder_after_embedding_scatter(config: "KimiK25Model.Config") -> None:
+    """Keep ``tok_embeddings`` ``Replicate`` and resume SP at layer 0's output.
+
+    The vision scatter writes features at arbitrary sequence positions, so it
+    needs the full (``Replicate``) embedding -- a ``Shard(1)`` one cannot be
+    indexed by sequence position locally. Layer 0 then takes a ``Replicate``
+    input and its rowwise ``wo`` reduce-scatters back to ``Shard(1)``, so the
+    residual is sequence-parallel from layer 0's output and layers ``1..N-1``
+    are unchanged full SP.
+    """
+    config.tok_embeddings.sharding_config = ShardingConfig(
+        state_shardings={"weight": dense_param_placement(tp=spmd.S(0))},
+        in_src_shardings={"input": _REPLICATE_ACT},
+        in_dst_shardings={"input": _REPLICATE_ACT},
+        out_src_shardings=dense_activation_placement(tp=spmd.P),
+        out_dst_shardings=_REPLICATE_ACT,
+        local_map=LocalMapConfig(in_grad_placements=None),
+    )
+
+    layer0 = config.layers[0]
+    layer0.attention_norm.sharding_config = ShardingConfig(
+        state_shardings={"weight": _REPLICATE_PARAM},
+        in_src_shardings={"input": _REPLICATE_ACT},
+        out_src_shardings=_REPLICATE_ACT,
+    )
+    layer0.attention.sharding_config = ShardingConfig(
+        in_src_shardings={"x": _REPLICATE_ACT},
+        in_dst_shardings={"x": _REPLICATE_ACT},
+    )
+
+
+def _set_vision_encoder_sharding(ve_cfg) -> None:
+    """Replicate-activation TP plan for the MoonViT3d vision encoder.
+
+    Linear layers are Colwise/Rowwise sharded for memory; norms and the
+    learnable position table are Replicate. ``patch_embed`` wraps the plain
+    ``pixel_values`` input as ``DTensor(Replicate)`` so the rest of the encoder
+    runs in DTensor space.
+    """
+    # The encoder's own ``pos_embed`` table is Replicate (F.interpolate runs on it).
+    ve_cfg.sharding_config = ShardingConfig(
+        state_shardings={"pos_embed": _REPLICATE_PARAM},
+    )
+
+    # patch_embed (Linear): receives plain pixel_values -> wrap as Replicate.
+    ve_cfg.patch_embed_proj.sharding_config = ShardingConfig(
+        state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM},
+        in_src_shardings={"input": _REPLICATE_ACT},
+        in_dst_shardings={"input": _REPLICATE_ACT},
+        out_dst_shardings=_REPLICATE_ACT,
+    )
+
+    # Transformer block sub-modules (shared VisionTransformerBlock: norm1/norm2).
+    block = ve_cfg.block
+    block.norm1.sharding_config = _REPLICATE_NORM
+    block.norm2.sharding_config = _REPLICATE_NORM
+
+    # The stacked 2D rope_cache enters the attention as a plain (Replicate)
+    # tensor input so it is DTensor-wrapped before meeting head-sharded q/k.
+    block.attn.sharding_config = ShardingConfig(
+        in_src_shardings={"rope_cache": _REPLICATE_ACT},
+        in_dst_shardings={"rope_cache": _REPLICATE_ACT},
+    )
+    block.attn.wq.sharding_config = colwise_config()
+    block.attn.wk.sharding_config = colwise_config()
+    block.attn.wv.sharding_config = colwise_config()
+    block.attn.proj.sharding_config = rowwise_config(output_sp=False)
+    set_gqa_inner_attention_local_map(block.attn.inner_attention)
+
+    block.mlp.fc1.sharding_config = colwise_config()
+    block.mlp.fc2.sharding_config = rowwise_config(output_sp=False)
+
+    # Final norm + projector.
+    ve_cfg.final_norm.sharding_config = _REPLICATE_NORM
+    proj = ve_cfg.projector
+    proj.pre_norm.sharding_config = _REPLICATE_NORM
+    proj.linear_1.sharding_config = colwise_config()
+    proj.linear_2.sharding_config = rowwise_config(output_sp=False)
diff --git a/torchtitan/models/kimi_k2_7/state_dict_adapter.py b/torchtitan/models/kimi_k2_7/state_dict_adapter.py
new file mode 100644
index 0000000000..fc9d4b4f55
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/state_dict_adapter.py
@@ -0,0 +1,221 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""State dict adapter for Kimi K2.5 (MoonViT3d + DeepSeekV3).
+
+The language model is DeepSeekV3, so this subclasses
+``DeepSeekV3StateDictAdapter`` and delegates every LM key to it; only the vision
+tower / projector are handled here.
+
+Vision name/shape mappings (``HF -> torchtitan``; reversed on save):
+
+- attention qkv:  ``wqkv`` (fused)                       ->  ``attn.wq``/``wk``/``wv`` (split)
+- attention proj: ``wo``                                 ->  ``attn.proj``
+- projector mlp:  ``mm_projector.proj.0``/``2``            ->  ``projector.linear_1``/``2``
+- patch embed:    ``Conv2d`` weight                      ->  ``Linear`` weight (reshape)
+"""
+
+import re
+from typing import Any
+
+import torch
+
+from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter
+
+from .model import KimiK25Model
+
+
+class KimiK25StateDictAdapter(DeepSeekV3StateDictAdapter):
+    def __init__(
+        self,
+        model_config: KimiK25Model.Config,
+        hf_assets_path: str | None,
+    ):
+        super().__init__(model_config, hf_assets_path)
+
+        self.vision_encoder = model_config.vision_encoder
+        if self.vision_encoder is None:
+            return
+
+        self.patch_size = self.vision_encoder.patch_size
+        self.in_channels = self.vision_encoder.in_channels
+
+        # Vision tower: HF name -> torchtitan name (fused qkv handled separately).
+        self.vision_from_hf_map = {
+            # Patch embedding (Conv2d weight reshaped to Linear on load).
+            "vision_tower.patch_embed.proj.weight": "vision_encoder.patch_embed.weight",
+            "vision_tower.patch_embed.proj.bias": "vision_encoder.patch_embed.bias",
+            # Learnable spatial position embedding.
+            "vision_tower.patch_embed.pos_emb.weight": "vision_encoder.pos_embed",
+            # Block norms: HF norm0/norm1 (pre-attn / pre-mlp) -> tt norm1/norm2.
+            "vision_tower.encoder.blocks.{}.norm0.weight": "vision_encoder.layers.{}.norm1.weight",
+            "vision_tower.encoder.blocks.{}.norm0.bias": "vision_encoder.layers.{}.norm1.bias",
+            "vision_tower.encoder.blocks.{}.norm1.weight": "vision_encoder.layers.{}.norm2.weight",
+            "vision_tower.encoder.blocks.{}.norm1.bias": "vision_encoder.layers.{}.norm2.bias",
+            # Attention output projection: HF wo -> tt attn.proj.
+            "vision_tower.encoder.blocks.{}.wo.weight": "vision_encoder.layers.{}.attn.proj.weight",
+            "vision_tower.encoder.blocks.{}.wo.bias": "vision_encoder.layers.{}.attn.proj.bias",
+            # Block MLP: HF fc0/fc1 -> tt linear_fc1/linear_fc2.
+            "vision_tower.encoder.blocks.{}.mlp.fc0.weight": "vision_encoder.layers.{}.mlp.linear_fc1.weight",
+            "vision_tower.encoder.blocks.{}.mlp.fc0.bias": "vision_encoder.layers.{}.mlp.linear_fc1.bias",
+            "vision_tower.encoder.blocks.{}.mlp.fc1.weight": "vision_encoder.layers.{}.mlp.linear_fc2.weight",
+            "vision_tower.encoder.blocks.{}.mlp.fc1.bias": "vision_encoder.layers.{}.mlp.linear_fc2.bias",
+            # Final encoder norm.
+            "vision_tower.encoder.final_layernorm.weight": "vision_encoder.final_norm.weight",
+            "vision_tower.encoder.final_layernorm.bias": "vision_encoder.final_norm.bias",
+            # Multimodal projector (1T Kimi-K2.5 spelling; Kimi-VL's
+            # multi_modal_projector.linear_1/2 is normalized to this in from_hf).
+            "mm_projector.pre_norm.weight": "vision_encoder.projector.pre_norm.weight",
+            "mm_projector.pre_norm.bias": "vision_encoder.projector.pre_norm.bias",
+            "mm_projector.proj.0.weight": "vision_encoder.projector.linear_1.weight",
+            "mm_projector.proj.0.bias": "vision_encoder.projector.linear_1.bias",
+            "mm_projector.proj.2.weight": "vision_encoder.projector.linear_2.weight",
+            "mm_projector.proj.2.bias": "vision_encoder.projector.linear_2.bias",
+        }
+
+    def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]:
+        if self.vision_encoder is None:
+            return super().from_hf(
+                {
+                    key: value
+                    for key, value in hf_state_dict.items()
+                    if not key.endswith("rotary_emb.inv_freq")
+                }
+            )
+
+        lm_hf: dict[str, Any] = {}
+        vision: dict[str, Any] = {}
+        unmapped: list[str] = []
+        for key, value in hf_state_dict.items():
+            # RoPE inv_freq is recomputed at runtime -- no target to map.
+            if key.endswith("rotary_emb.inv_freq"):
+                continue
+
+            if not (
+                key.startswith("vision_tower.")
+                or key.startswith("mm_projector.")
+                or key.startswith("multi_modal_projector.")
+            ):
+                # LM key. The inherited DeepSeek-V3 adapter (super) is keyed on
+                # bare ``model.*`` names, but the released checkpoints nest the
+                # text tower under ``language_model.`` -- strip that prefix so super
+                # recognizes the keys. (The ``language_model.layers.`` variant,
+                # which omits the ``model.`` segment, first gets ``model.`` inserted.)
+                if key.startswith("language_model.layers."):
+                    key = key.replace(
+                        "language_model.layers.", "language_model.model.layers.", 1
+                    )
+                if key.startswith("language_model."):
+                    key = key.replace("language_model.", "", 1)
+                lm_hf[key] = value
+                continue
+
+            # Projector: the two released checkpoints disagree -- Kimi-VL spells it
+            # ``multi_modal_projector.linear_1/2``, the 1T Kimi-K2.5 spells it
+            # ``mm_projector.proj.0/2``. Normalize to the latter (the map/save form).
+            key = key.replace("multi_modal_projector.", "mm_projector.")
+            key = key.replace("mm_projector.linear_1", "mm_projector.proj.0")
+            key = key.replace("mm_projector.linear_2", "mm_projector.proj.2")
+
+            if re.search(r"\.wqkv\.(weight|bias)$", key):
+                # Split fused HF vision qkv -> separate wq/wk/wv.
+                # pyrefly: ignore [missing-attribute]
+                layer_num = re.search(r"\d+", key).group(0)
+                kind = "weight" if key.endswith("weight") else "bias"
+                if value.shape[0] % 3 != 0:
+                    raise ValueError(
+                        f"fused vision QKV '{key}' has first dim "
+                        f"{value.shape[0]}, not divisible by 3 (q|k|v)."
+                    )
+                q, k, v = torch.chunk(value, 3, dim=0)
+                base = f"vision_encoder.layers.{layer_num}.attn"
+                vision[f"{base}.wq.{kind}"] = q
+                vision[f"{base}.wk.{kind}"] = k
+                vision[f"{base}.wv.{kind}"] = v
+            elif key in self.vision_from_hf_map:
+                new_key = self.vision_from_hf_map[key]
+                if new_key == "vision_encoder.patch_embed.weight":
+                    # HF Conv2d (out, C, kH, kW) -> Linear (out, C*kH*kW).
+                    value = value.reshape(value.shape[0], -1)
+                vision[new_key] = value
+            else:
+                abstract_key = re.sub(r"(\d+)", "{}", key, count=1)
+                if abstract_key in self.vision_from_hf_map:
+                    # pyrefly: ignore [missing-attribute]
+                    layer_num = re.search(r"\d+", key).group(0)
+                    vision[
+                        self.vision_from_hf_map[abstract_key].format(layer_num)
+                    ] = value
+                else:
+                    unmapped.append(key)
+
+        if unmapped:
+            raise ValueError(
+                f"KimiK25StateDictAdapter: {len(unmapped)} vision key(s) have no "
+                f"mapping: {unmapped}. Add them to vision_from_hf_map, or filter "
+                f"them above if they are disposable buffers."
+            )
+
+        # DeepSeekV3 handles the LM keys (incl. RoPE validation + experts).
+        state_dict = super().from_hf(lm_hf)
+        state_dict.update(vision)
+        return state_dict
+
+    def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
+        if self.vision_encoder is None:
+            return super().to_hf(state_dict)
+
+        to_hf_map = {v: k for k, v in self.vision_from_hf_map.items()}
+        use_kimi_vl_projector_names = (
+            self.fqn_to_index_mapping is not None
+            and "multi_modal_projector.pre_norm.weight" in self.fqn_to_index_mapping
+        )
+        lm_titan: dict[str, Any] = {}
+        hf_state_dict: dict[str, Any] = {}
+        # Buffer separate vision q/k/v per layer to re-fuse into one HF tensor.
+        vision_qkv: dict[tuple[str, str], dict[str, torch.Tensor]] = {}
+
+        for key, value in state_dict.items():
+            if not key.startswith("vision_encoder."):
+                lm_titan[key] = value
+            elif re.search(r"vision_encoder\.layers\.\d+\.attn\.w[qkv]\.", key):
+                # pyrefly: ignore [missing-attribute]
+                layer_num = re.search(r"\d+", key).group(0)
+                proj = re.search(r"attn\.(w[qkv])\.(weight|bias)", key)
+                # pyrefly: ignore [missing-attribute]
+                which, kind = proj.group(1), proj.group(2)
+                vision_qkv.setdefault((layer_num, kind), {})[which] = value
+            elif "patch_embed.weight" in key:
+                # Linear (out, C*kH*kW) -> HF Conv2d (out, C, kH, kW).
+                hf_state_dict[to_hf_map[key]] = value.reshape(
+                    -1, self.in_channels, self.patch_size, self.patch_size
+                )
+            elif "vision_encoder.layers" in key:
+                abstract_key = re.sub(r"(\d+)", "{}", key, count=1)
+                # pyrefly: ignore [missing-attribute]
+                layer_num = re.search(r"\d+", key).group(0)
+                hf_state_dict[to_hf_map[abstract_key].format(layer_num)] = value
+            else:
+                hf_key = to_hf_map[key]
+                # Normalize the K2.5 and Kimi-VL projector naming conventions.
+                if use_kimi_vl_projector_names and hf_key.startswith("mm_projector."):
+                    hf_key = hf_key.replace("mm_projector.", "multi_modal_projector.")
+                    hf_key = hf_key.replace("proj.0", "linear_1")
+                    hf_key = hf_key.replace("proj.2", "linear_2")
+                hf_state_dict[hf_key] = value
+
+        # Fuse vision q/k/v -> single HF qkv tensor per (layer, weight|bias).
+        for (layer_num, kind), parts in vision_qkv.items():
+            fused = torch.cat([parts["wq"], parts["wk"], parts["wv"]], dim=0)
+            hf_state_dict[
+                f"vision_tower.encoder.blocks.{layer_num}.wqkv.{kind}"
+            ] = fused
+
+        # Re-add the ``language_model.`` nesting that super (DeepSeek-V3) drops.
+        hf_state_dict.update(
+            {f"language_model.{k}": v for k, v in super().to_hf(lm_titan).items()}
+        )
+        return hf_state_dict
diff --git a/torchtitan/models/kimi_k2_7/vision_encoder.py b/torchtitan/models/kimi_k2_7/vision_encoder.py
new file mode 100644
index 0000000000..5accd91f85
--- /dev/null
+++ b/torchtitan/models/kimi_k2_7/vision_encoder.py
@@ -0,0 +1,467 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""MoonViT3d vision encoder + multimodal projector for Kimi K2.5.
+
+Reference (SGLang):
+https://github.com/sgl-project/sglang/blob/e0c0c0a45cb1bda90392bfa2bba4184f5b0638a0/python/sglang/srt/models/kimi_k25.py
+
+Shape suffixes:
+- N = num visual items
+- P = max num of patches per visual item (padded)
+- D = vision dim
+- M = merged tokens
+- K = merged feature dim (kh*kw*D)
+"""
+
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torchtitan.models.common import Linear
+from torchtitan.models.common.nn_modules import GELU, LayerNorm
+from torchtitan.models.common.rope import ComplexRoPE
+from torchtitan.models.common.vision_encoder import (
+    compiled_create_block_mask,
+    get_vision_block_mask_mod,
+    VisionTransformerBlock,
+)
+from torchtitan.protocols.module import Module, ModuleDict
+
+
+def _get_temporal_pos_embed(
+    num_frames: int,
+    embed_dim: int,
+    *,
+    base: float = 10000.0,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    """Fixed 1D sinusoidal embeddings for the temporal axis (video frames).
+
+    Returns ``(num_frames, embed_dim)`` float32; the standard 1D sincos formula
+    over frame indices.
+
+    Args:
+        num_frames: Number of video frames (temporal positions).
+        embed_dim: Embedding width per frame.
+        base: Sinusoid base (longest wavelength); the conventional PE constant.
+        device: Device for the returned tensor.
+    """
+    grid = torch.arange(num_frames, dtype=torch.float32, device=device)
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / (
+        embed_dim / 2.0
+    )
+    omega = 1.0 / base**omega
+    out = torch.outer(grid, omega)
+    return torch.cat([out.sin(), out.cos()], dim=1)
+
+
+def _compute_learned_pos_embeds(
+    pos_embed: torch.Tensor,
+    grids: list[list[int]],
+    max_num_patch: int,
+    interpolation_mode: str,
+) -> torch.Tensor:
+    """Interpolated learnable 2D spatial pos-emb + fixed sinusoidal temporal.
+
+    The learnable ``(height, width, dim)`` spatial table is interpolated to each
+    item's ``(h, w)`` patch grid (raster order). For video (``t > 1``) the
+    spatial embedding is repeated per frame and summed with a fixed 1D sinusoidal
+    temporal embedding (``_get_temporal_pos_embed``; only video hits this path).
+
+    Args:
+        pos_embed: (height, width, dim) learnable spatial position table.
+        grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw``
+            read to CPU once by the caller, so the per-item loop adds no syncs).
+        max_num_patch: Padded sequence length.
+        interpolation_mode: ``F.interpolate`` mode (e.g. ``"bicubic"``).
+
+    Returns:
+        (N, max_num_patch, dim) padded position embeddings to add to the patches.
+    """
+    height, width, dim = pos_embed.shape
+    pos = pos_embed.new_zeros(len(grids), max_num_patch, dim)
+
+    # (dim, height, width) for F.interpolate; .float() for bicubic.
+    grid_table = pos_embed.permute(2, 0, 1).unsqueeze(0).float()
+
+    hw_to_indices: dict[tuple[int, int], list[int]] = {}
+    for i, (_, h, w) in enumerate(grids):
+        hw_to_indices.setdefault((h, w), []).append(i)
+
+    for (h, w), indices in hw_to_indices.items():
+        if (h, w) == (height, width):
+            pos_hw = pos_embed.flatten(end_dim=1)
+        else:
+            pos_hw = (
+                F.interpolate(grid_table, size=(h, w), mode=interpolation_mode)
+                .squeeze(0)
+                .permute(1, 2, 0)
+                .reshape(h * w, dim)
+                .to(pos_embed.dtype)
+            )
+        for i in indices:
+            t = grids[i][0]
+            seq_len = t * h * w
+            if t == 1:
+                pos[i, :seq_len] = pos_hw
+            else:
+                # (t, 1, dim) to broadcast the per-frame term over h*w patches.
+                time_weight = _get_temporal_pos_embed(
+                    t, dim, device=pos.device
+                ).unsqueeze(1)
+                frames = pos_hw.unsqueeze(0).repeat(t, 1, 1)
+                frames = frames + time_weight.to(frames.dtype)
+                pos[i, :seq_len] = frames.reshape(seq_len, dim)
+
+    return pos
+
+
+def _compute_2d_rope_cache(
+    freq_table: torch.Tensor,
+    grids: list[list[int]],
+    max_num_patch: int,
+    head_dim: int,
+) -> torch.Tensor:
+    """Compute the padded 2D-RoPE complex ``freqs_cis`` cache in raster order.
+
+    For head-dim pair index ``k`` (``k`` in ``[0, head_dim/4)``), even output
+    pairs are rotated by the *column* (x) position and odd pairs by the *row*
+    (y) position. The per-axis angle for a position ``p`` is ``p * inv_freq[k]``;
+    this looks it up by gathering row ``p`` of ``freq_table`` (built once by
+    ``VisionRotaryEmbedding2D`` and cached by the encoder) rather than
+    recomputing ``p * inv_freq`` each call. Frames repeat the spatial pattern.
+
+    Returns a complex cache consumed by ``ComplexRoPE.apply_rotary_emb``; only
+    the cache is 2D/per-grid, which is why it is built here rather than by the
+    1D ``ComplexRoPE`` cache machinery.
+
+    Args:
+        freq_table: ``(max_hw, head_dim/4)`` position-to-frequency table, where
+            ``freq_table[p, k] = p * inv_freq[k]``.
+        grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw``
+            read to CPU once by the caller, so the per-item loop adds no syncs).
+        max_num_patch: Padded sequence length.
+        head_dim: Attention head dim (must be divisible by 4).
+
+    Returns:
+        ``(N, max_num_patch, 1, head_dim/2)`` complex64 (head axis = 1 to
+        broadcast over the heads).
+    """
+    device = freq_table.device
+
+    angles = torch.zeros(
+        len(grids), max_num_patch, head_dim // 2, device=device, dtype=freq_table.dtype
+    )
+
+    # Group by (h, w) so the per-resolution angle grid is built once.
+    hw_to_indices: dict[tuple[int, int], list[int]] = {}
+    for i, (_, h, w) in enumerate(grids):
+        hw_to_indices.setdefault((h, w), []).append(i)
+
+    for (h, w), indices in hw_to_indices.items():
+        # Raster order: position p -> (row = p // w, col = p % w). Gather each
+        # axis's angles from the precomputed table (freq_table[pos] = pos*inv_freq).
+        flat = torch.arange(h * w, device=device)
+        x_ang = freq_table[flat % w]  # (h*w, head_dim/4) column
+        y_ang = freq_table[flat // w]  # (h*w, head_dim/4) row
+        # Interleave x/y so pair 2k uses x-position, pair 2k+1 uses y-position.
+        ang = torch.stack([x_ang, y_ang], dim=-1).reshape(h * w, head_dim // 2)
+        for i in indices:
+            t = grids[i][0]
+            seq_len = t * h * w
+            angles[i, :seq_len] = ang.repeat(t, 1)
+
+    # Complex unit-modulus cache; unsqueeze the head axis for broadcast.
+    return torch.polar(torch.ones_like(angles), angles).unsqueeze(2)
+
+
+def _tpool_patch_merger(
+    hidden_NPD: torch.Tensor,
+    grids: list[list[int]],
+    merge_kernel_size: tuple[int, int],
+) -> torch.Tensor:
+    """Temporal mean pooling + spatial merge over the padded batch.
+
+    For each item ``(t, h, w)``: reshape its valid patches to
+    ``(t, h, w, D)``, mean over the temporal axis, then group spatial
+    ``kh x kw`` neighbors and concatenate them along the feature axis, yielding
+    ``(h/kh * w/kw)`` merged tokens of dim ``kh*kw*D``.
+
+    Args:
+        hidden_NPD: (N, P, D) padded patch features.
+        grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw``
+            read to CPU once by the caller, so the per-item loop adds no syncs).
+        merge_kernel_size: ``(kh, kw)`` spatial merge factors.
+
+    Returns:
+        merged: ``(N, max_merged, kh*kw*D)`` padded merged tokens, where
+        ``max_merged = max_i (h_i/kh) * (w_i/kw)``. The valid token count per
+        item is ``(h/kh) * (w/kw)`` (recomputed by the caller from ``grids``
+        for the scatter).
+    """
+    num_vision, _, d_model = hidden_NPD.shape
+    kh, kw = merge_kernel_size
+    merged_dim = kh * kw * d_model
+
+    max_merged = max((h // kh) * (w // kw) for _, h, w in grids)
+    merged = hidden_NPD.new_zeros(num_vision, max_merged, merged_dim)
+
+    for i, (t, h, w) in enumerate(grids):
+        seq = hidden_NPD[i, : t * h * w]
+        new_h, new_w = h // kh, w // kw
+        # (t, new_h, kh, new_w, kw, D) -> mean over t -> group spatial kernel.
+        seq = seq.view(t, new_h, kh, new_w, kw, d_model)
+        seq = seq.permute(0, 1, 3, 2, 4, 5).mean(dim=0)
+        seq = seq.reshape(new_h * new_w, merged_dim)
+        merged[i, : new_h * new_w] = seq
+
+    return merged
+
+
+class VisionRotaryEmbedding2D(Module):
+    """2D rotary position embedding for the vision tower.
+
+    Holds the per-axis frequencies ``inv_freq`` (``head_dim/4`` of them, shared
+    by the row and column axes). ``forward(seqlen)`` returns the
+    position-to-frequency table ``freq_table[p, k] = p * inv_freq[k]`` for
+    positions up to ``seqlen``; ``_compute_2d_rope_cache`` gathers per-patch
+    row/col angles from it, and ``ComplexRoPE.apply_rotary_emb`` applies them.
+    ``head_dim`` must be divisible by 4.
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        head_dim: int
+        theta: float = 10000.0
+
+    def __init__(self, config: Config):
+        super().__init__()
+        if config.head_dim % 4 != 0:
+            raise ValueError(
+                f"2D RoPE requires head_dim divisible by 4, got {config.head_dim}."
+            )
+        self.head_dim = config.head_dim
+        self.theta = config.theta
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor:
+        # inv_freq[k] = theta**(-4k/head_dim) for k in [0, head_dim/4); the
+        # step of 4 leaves room for the row/col split of the 2D rotation.
+        return 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, self.head_dim, 4, dtype=torch.float32, device=device)
+                / self.head_dim
+            )
+        )
+
+    def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None:
+        """Re-compute inv_freq on the target device after to_empty()."""
+        device = buffer_device or self.inv_freq.device
+        self.inv_freq = self._compute_inv_freq(device=device)
+
+    def forward(self, seqlen: int) -> torch.Tensor:
+        """Frequency table ``(seqlen, head_dim/4)`` for positions ``[0, seqlen)``."""
+        seq = torch.arange(
+            seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
+        )
+        return torch.outer(seq, self.inv_freq)
+
+
+class VisionProjector(Module):
+    """Project merged vision features to the language-model hidden size.
+
+    Applies a per-patch LayerNorm (on ``vt_hidden_size``), then a 2-layer MLP
+    over the concatenated spatial-merge features. This is a unimodal vision
+    projection head; the cross-modal fusion (scattering these features into the
+    text sequence) happens later in the model, not here.
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        vt_hidden_size: int
+        merged_dim: int
+        pre_norm: LayerNorm.Config
+        linear_1: Linear.Config
+        linear_2: Linear.Config
+        act_fn: GELU.Config = field(default_factory=GELU.Config)
+
+    def __init__(self, config: Config):
+        super().__init__()
+        self.vt_hidden_size = config.vt_hidden_size
+        self.merged_dim = config.merged_dim
+        self.pre_norm = config.pre_norm.build()
+        self.linear_1 = config.linear_1.build()
+        self.linear_2 = config.linear_2.build()
+        self.act_fn = config.act_fn.build()
+
+    def forward(self, merged_NMK: torch.Tensor) -> torch.Tensor:
+        """Args: ``(N, M, kh*kw*vt_hidden_size)`` padded merged tokens.
+
+        The pre-norm runs per-patch on ``vt_hidden_size``; the merged kernel
+        features are then flattened for the projection MLP.
+        """
+        n, m, _ = merged_NMK.shape
+        x = merged_NMK.view(n, m, -1, self.vt_hidden_size)
+        x = self.pre_norm(x).view(n, m, self.merged_dim)
+        x = self.linear_1(x)
+        x = self.act_fn(x)
+        x = self.linear_2(x)
+        return x
+
+
+class KimiK25VisionEncoder(Module):
+    """MoonViT3d vision tower + multimodal projector for Kimi K2.5."""
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Module.Config):
+        dim: int = 1152
+        num_layers: int = 27
+        num_heads: int = 16
+
+        patch_size: int = 14
+        in_channels: int = 3
+        merge_kernel_size: list[int] = field(default_factory=lambda: [2, 2])
+        text_hidden_size: int = 7168
+
+        # Learnable 2D spatial position table, shape (height, width, dim).
+        init_pos_emb_height: int = 64
+        init_pos_emb_width: int = 64
+        interpolation_mode: str = "bicubic"
+
+        # Sub-modules.
+        patch_embed_proj: Linear.Config
+        rotary_pos_emb: VisionRotaryEmbedding2D.Config
+        block: VisionTransformerBlock.Config
+        final_norm: LayerNorm.Config
+        projector: VisionProjector.Config
+
+    def __init__(self, config: Config):
+        super().__init__()
+        self.config = config
+        self.merge_kernel_size = tuple(config.merge_kernel_size)
+        self.interpolation_mode = config.interpolation_mode
+
+        self.patch_embed = config.patch_embed_proj.build()
+        # Learnable 2D spatial position table.
+        self.pos_embed = nn.Parameter(
+            torch.empty(
+                config.init_pos_emb_height, config.init_pos_emb_width, config.dim
+            )
+        )
+        self.rotary_pos_emb = config.rotary_pos_emb.build()
+        self._cached_freq_table: torch.Tensor | None = None
+        self.layers = ModuleDict(
+            {str(idx): config.block.build() for idx in range(config.num_layers)}
+        )
+        self.final_norm = config.final_norm.build()
+        self.projector = config.projector.build()
+
+    def compute_position_embeddings(
+        self, grids: list[list[int]], max_num_patch: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute both position embeddings for the padded batch.
+
+        Delegates to two standalone helpers:
+        - ``_compute_learned_pos_embeds``: interpolated learnable spatial table
+          (plus a runtime sinusoidal temporal term for videos).
+        - ``_compute_2d_rope_cache``: 2D RoPE complex cache, gathering from the
+          RoPE frequency table (cached here, regrown only when a larger side
+          appears).
+
+        Args:
+            grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw``
+                read to CPU once by ``forward``).
+            max_num_patch: Padded sequence length.
+
+        Returns:
+            learned_pos: ``(N, max_num_patch, dim)`` additive position embeddings.
+            rope_cache: ``(N, max_num_patch, 1, head_dim/2)`` complex RoPE cache
+                for ``ComplexRoPE.apply_rotary_emb``.
+        """
+        max_hw = max(max(h, w) for _, h, w in grids)
+        if self._cached_freq_table is None or self._cached_freq_table.shape[0] < max_hw:
+            self._cached_freq_table = self.rotary_pos_emb(max_hw)
+
+        learned_pos = _compute_learned_pos_embeds(
+            self.pos_embed, grids, max_num_patch, self.interpolation_mode
+        )
+        rope_cache = _compute_2d_rope_cache(
+            self._cached_freq_table,
+            grids,
+            max_num_patch,
+            self.rotary_pos_emb.head_dim,
+        )
+        return learned_pos, rope_cache
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        *,
+        grid_thw: torch.Tensor,
+    ) -> torch.Tensor:
+        """Encode a padded batch of visual items.
+
+        Args:
+            pixel_values: ``(N, P, patch_dim)`` padded flattened patches.
+            grid_thw: ``(N, 3)`` patch counts ``[t, h, w]`` per item.
+
+        Returns:
+            ``(N, max_merged, text_hidden_size)`` padded projected features.
+
+        Each item's ``(h, w)`` must be divisible by ``merge_kernel_size`` and its
+        ``t*h*w`` must fit in ``P`` (both dataloader-guaranteed; asserted below).
+        Patches must arrive in raster order ``(t, h, w)`` (the dataloader's
+        ``patch_order="raster"``); the 2D RoPE / position embeddings index
+        ``row = p // w, col = p % w``.
+        """
+        num_vision, max_num_patch, _ = pixel_values.shape
+        # One host sync for the whole forward: read the (N, 3) grid to CPU ints.
+        grids = grid_thw.tolist()  # [[t, h, w], ...]
+
+        # Grid contract (dataloader-guaranteed); assert so a bad batch fails
+        # clearly instead of with a cryptic view/index error in the helpers.
+        kh, kw = self.merge_kernel_size
+        for t, h, w in grids:
+            assert (
+                h % kh == 0 and w % kw == 0
+            ), f"grid {h}x{w} indivisible by {(kh, kw)}"
+            assert t * h * w <= max_num_patch, f"t*h*w={t * h * w} > P={max_num_patch}"
+
+        num_patch = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).to(
+            torch.long
+        )  # (N,)
+
+        learned_pos, rope_cache = self.compute_position_embeddings(grids, max_num_patch)
+        x = self.patch_embed(pixel_values) + learned_pos
+
+        mask_mod = get_vision_block_mask_mod(num_patch)
+        attention_mask = compiled_create_block_mask(
+            mask_mod,
+            num_vision,
+            None,
+            max_num_patch,
+            max_num_patch,
+            device=x.device,
+        )
+
+        for block in self.layers.values():
+            x = block(
+                x,
+                rope_cache=rope_cache,
+                rope_apply=ComplexRoPE.apply_rotary_emb,
+                attention_mask=attention_mask,
+            )
+
+        x = self.final_norm(x)
+
+        # Temporal pool + spatial merge, then project to the LLM hidden size.
+        # pyrefly: ignore [bad-argument-type]
+        merged = _tpool_patch_merger(x, grids, self.merge_kernel_size)
+        return self.projector(merged)
diff --git a/torchtitan/models/qwen3_5/__init__.py b/torchtitan/models/qwen3_5/__init__.py
index 4488d9c568..363cbdcc20 100644
--- a/torchtitan/models/qwen3_5/__init__.py
+++ b/torchtitan/models/qwen3_5/__init__.py
@@ -10,6 +10,7 @@
 import torch.nn as nn
 
 from torchtitan.components.optimizer import register_moe_load_balancing_hook
+from torchtitan.distributed.pipeline_parallel import pipeline_vlm
 
 from torchtitan.models.common import (  # noqa: F401
     Conv1d,
@@ -26,6 +27,11 @@
 )
 from torchtitan.models.common.nn_modules import LayerNorm
 from torchtitan.models.common.param_init import depth_scaled_std  # noqa: F401
+from torchtitan.models.common.vision_encoder import (
+    VisionAttention,
+    VisionMLP,
+    VisionTransformerBlock,
+)
 from torchtitan.models.utils import validate_converter_order
 from torchtitan.protocols.model import ModelConfigConverter
 
@@ -41,17 +47,12 @@
     Qwen35TransformerBlock,
     RMSNormGated,
 )
-from .parallelize import parallelize_qwen3_5, pipeline_qwen3_5
+
+from .parallelize import parallelize_qwen3_5
 from .rope import MRoPE
 from .state_dict_adapter import Qwen35StateDictAdapter
-from .vision_encoder import (
-    PatchMerger,
-    Qwen35VisionEncoder,
-    VisionAttention,
-    VisionMLP,
-    VisionRotaryEmbedding,
-    VisionTransformerBlock,
-)
+
+from .vision_encoder import PatchMerger, Qwen35VisionEncoder, VisionRotaryEmbedding
 
 __all__ = [
     "parallelize_qwen3_5",
@@ -1102,7 +1103,7 @@ def model_registry(
         flavor=flavor,
         model=config,
         parallelize_fn=parallelize_qwen3_5,
-        pipelining_fn=pipeline_qwen3_5,
+        pipelining_fn=pipeline_vlm,
         post_optimizer_build_fn=register_moe_load_balancing_hook,
         state_dict_adapter=Qwen35StateDictAdapter,
     )
diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py
index 2d247c37a8..64f001841f 100644
--- a/torchtitan/models/qwen3_5/model.py
+++ b/torchtitan/models/qwen3_5/model.py
@@ -29,6 +29,10 @@
     VarlenMetadata,
 )
 from torchtitan.models.common.decoder import Decoder
+from torchtitan.models.common.multimodal import (
+    get_vision_positions,
+    scatter_vision_embeds,
+)
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
 from torchtitan.protocols.module import Module
 
@@ -669,6 +673,9 @@ def update_from_config(
         def get_nparams_and_flops(
             self, model: nn.Module, seq_len: int
         ) -> tuple[int, int]:
+            # The shared helper excludes the vision encoder from the per-token
+            # FLOP term (ViT cost scales with patches, not seq_len), so this MFU
+            # is decoder-only. TODO: add a per-batch vision FLOP term for VLMs.
             attn_cfg = self.first_attention
             # pyrefly: ignore [missing-attribute]
             n_heads = attn_cfg.n_heads
@@ -754,42 +761,16 @@ def _get_vision_embeds(
             grid_thw: Grid dimensions (num_items, 3) for [t, h, w]
 
         Returns:
-            merged_embeds: (num_items, max_tokens, dim) padded vision embeddings
+            vision_embeds: (num_items, max_tokens, dim) padded vision embeddings
             num_tokens_per_item: (num_items,) actual token count per item
         """
         pixel_values = pixel_values.to(self.vision_encoder.patch_embed.weight.dtype)
-        merged_embeds = self.vision_encoder(pixel_values, grid_thw=grid_thw)
+        vision_embeds = self.vision_encoder(pixel_values, grid_thw=grid_thw)
 
         merge_unit = self.vision_encoder.spatial_merge_unit
         num_tokens_per_item = grid_thw.prod(-1) // merge_unit
 
-        return merged_embeds, num_tokens_per_item
-
-    def _scatter_vision_embeds(
-        self,
-        inputs_embeds: torch.Tensor,
-        *,
-        merged_embeds: torch.Tensor,
-        vision_positions: list[tuple[int, int, int, int]],
-    ) -> torch.Tensor:
-        """Scatter vision embeddings into text embeddings at placeholder positions.
-
-        Copies directly from the padded vision encoder output into the text
-        sequence.
-
-        Args:
-            inputs_embeds: Text embeddings (batch, seq_len, dim)
-            merged_embeds: Padded vision embeddings (num_items, max_tokens, dim)
-            vision_positions: List of (item_idx, sample_idx, vision_start, n_tokens)
-
-        Returns:
-            Updated embeddings
-        """
-        for item_idx, sample_idx, vision_start, n_tokens in vision_positions:
-            inputs_embeds[
-                sample_idx, vision_start : vision_start + n_tokens, :
-            ] = merged_embeds[item_idx, :n_tokens, :]
-        return inputs_embeds
+        return vision_embeds, num_tokens_per_item
 
     def _prepare_multimodal_embeds(
         self,
@@ -822,30 +803,26 @@ def _prepare_multimodal_embeds(
         )
 
         if pixel_values is not None and grid_thw is not None:
-            merged_embeds, num_tokens = self._get_vision_embeds(
+            vision_embeds, num_tokens = self._get_vision_embeds(
                 pixel_values, grid_thw=grid_thw
             )
-            image_positions = self._get_vision_positions(
-                tokens, num_tokens, image_token_id
-            )
+            image_positions = get_vision_positions(tokens, num_tokens, image_token_id)
             if image_positions:
-                inputs_embeds = self._scatter_vision_embeds(
+                inputs_embeds = scatter_vision_embeds(
                     inputs_embeds,
-                    merged_embeds=merged_embeds,
+                    vision_embeds=vision_embeds,
                     vision_positions=image_positions,
                 )
 
         if pixel_values_videos is not None and grid_thw_videos is not None:
-            merged_embeds, num_tokens = self._get_vision_embeds(
+            vision_embeds, num_tokens = self._get_vision_embeds(
                 pixel_values_videos, grid_thw=grid_thw_videos
             )
-            video_positions = self._get_vision_positions(
-                tokens, num_tokens, video_token_id
-            )
+            video_positions = get_vision_positions(tokens, num_tokens, video_token_id)
             if video_positions:
-                inputs_embeds = self._scatter_vision_embeds(
+                inputs_embeds = scatter_vision_embeds(
                     inputs_embeds,
-                    merged_embeds=merged_embeds,
+                    vision_embeds=vision_embeds,
                     vision_positions=video_positions,
                 )
 
diff --git a/torchtitan/models/qwen3_5/parallelize.py b/torchtitan/models/qwen3_5/parallelize.py
index 7c22ab83c6..666c22db48 100644
--- a/torchtitan/models/qwen3_5/parallelize.py
+++ b/torchtitan/models/qwen3_5/parallelize.py
@@ -11,10 +11,7 @@
 (activation checkpointing, compile, FSDP) to the Qwen3.5 model.
 """
 
-import torch
 import torch.nn as nn
-from torch.distributed._composable.fsdp import fully_shard
-from torch.distributed.fsdp import MixedPrecisionPolicy
 
 from torchtitan.config import (
     CompileConfig,
@@ -28,37 +25,11 @@
 from torchtitan.distributed.compile import apply_compile
 from torchtitan.distributed.fsdp import (
     apply_fsdp_to_decoder,
-    get_fsdp_reshard_after_forward_policy,
+    apply_fsdp_to_vision_encoder,
 )
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
 
 
-def _apply_fsdp_to_vision_encoder(
-    vision_encoder: nn.Module,
-    dp_mesh,
-    param_dtype: torch.dtype,
-    reduce_dtype: torch.dtype,
-    reshard_after_forward_policy: str = "default",
-    pp_enabled: bool = False,
-):
-    """FSDP the vision encoder as a single unit.
-
-    One AllGather for all vision params is more efficient than per-layer
-    sharding — the vision encoder is small relative to the decoder.
-    Must be called before apply_fsdp on the decoder.
-    """
-    mp_policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=reduce_dtype)
-    reshard_after_forward = get_fsdp_reshard_after_forward_policy(
-        reshard_after_forward_policy, pp_enabled=pp_enabled
-    )
-    fully_shard(
-        vision_encoder,
-        mesh=dp_mesh,
-        mp_policy=mp_policy,
-        reshard_after_forward=reshard_after_forward,
-    )
-
-
 def parallelize_qwen3_5(
     model: nn.Module,
     *,
@@ -118,7 +89,7 @@ def parallelize_qwen3_5(
     dp_mesh = parallel_dims.get_mesh(dp_mesh_names)
 
     if model.vision_encoder is not None:
-        _apply_fsdp_to_vision_encoder(
+        apply_fsdp_to_vision_encoder(
             model.vision_encoder,  # pyrefly: ignore [bad-argument-type]
             dp_mesh,
             param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param],
@@ -149,54 +120,3 @@ def parallelize_qwen3_5(
     )
 
     return model
-
-
-def pipeline_qwen3_5(
-    model: nn.Module,
-    *,
-    parallel_dims: ParallelDims,
-    parallelism: ParallelismConfig,
-    model_config,
-    **kwargs,
-):
-    """PP wrapper that assigns vision_encoder to the first pipeline stage.
-
-    Delegates to ``pipeline_llm`` after injecting ``vision_encoder`` into
-    the first stage's FQN list (the auto-generated LLM split doesn't know
-    about vision encoder modules).
-    """
-    import dataclasses
-
-    from torchtitan.distributed.pipeline_parallel import (
-        _generate_llm_fqn_per_model_part,
-        _get_pipeline_metadata,
-        pipeline_llm,
-    )
-
-    if parallelism.module_fqns_per_model_part is None:
-        (
-            num_virtual_stages,
-            num_layers,
-            input_weight,
-            output_weight,
-        ) = _get_pipeline_metadata(parallel_dims, parallelism, model_config)
-        fqn_per_part = _generate_llm_fqn_per_model_part(
-            num_virtual_stages, num_layers, input_weight, output_weight
-        )
-        # Vision encoder lives on the first stage alongside tok_embeddings. This
-        # adds load to stage 0 that the auto split doesn't model (input_weight
-        # only accounts for tok_embeddings); for a heavy vision encoder, bump
-        # parallelism.pipeline_parallel_first_stage_less_layers to rebalance.
-        if hasattr(model, "vision_encoder") and model.vision_encoder is not None:
-            fqn_per_part[0].insert(0, "vision_encoder")
-        parallelism = dataclasses.replace(
-            parallelism, module_fqns_per_model_part=fqn_per_part
-        )
-
-    return pipeline_llm(
-        model,
-        parallel_dims=parallel_dims,
-        parallelism=parallelism,
-        model_config=model_config,
-        **kwargs,
-    )
diff --git a/torchtitan/models/qwen3_5/vision_encoder.py b/torchtitan/models/qwen3_5/vision_encoder.py
index d4624f280e..2fd11120cd 100644
--- a/torchtitan/models/qwen3_5/vision_encoder.py
+++ b/torchtitan/models/qwen3_5/vision_encoder.py
@@ -4,7 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from collections.abc import Callable
 from dataclasses import dataclass, field
 
 import torch
@@ -12,37 +11,21 @@
 import torch.nn.functional as F
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.experimental import local_map
-from torch.nn.attention.flex_attention import BlockMask, create_block_mask
 
 from torchtitan.models.common import Linear
-from torchtitan.models.common.attention import FlexAttention
 from torchtitan.models.common.nn_modules import GELU, LayerNorm
 from torchtitan.models.common.rope import _maybe_wrap_positions, CosSinRoPE
+from torchtitan.models.common.vision_encoder import (
+    compiled_create_block_mask,
+    get_vision_block_mask_mod,
+    VisionTransformerBlock,
+)
 from torchtitan.protocols.module import Module, ModuleDict
 
-_compiled_create_block_mask = torch.compile(create_block_mask)
-
-
-def get_vision_block_mask_mod(num_patch: torch.Tensor) -> Callable:
-    """Create a mask modifier for block-diagonal attention.
-
-    Each image only attends to its own patches.
-
-    Args:
-        num_patch: (num_vision,) actual number of patches per visual item
-    """
-
-    def mask_mod(b, h, q_idx, kv_idx):
-        valid_q = q_idx < num_patch[b]
-        valid_kv = kv_idx < num_patch[b]
-        return valid_q & valid_kv
-
-    return mask_mod
-
 
 def _compute_learned_pos_embeds(
     learned_pos_embed: torch.Tensor,
-    grid_thw: torch.Tensor,
+    grids: list[list[int]],
     max_num_patch: int,
     num_grid_per_side: int,
     spatial_merge_size: int,
@@ -56,7 +39,7 @@ def _compute_learned_pos_embeds(
 
     Args:
         learned_pos_embed: (num_position_embeddings, dim) learnable position embeddings
-        grid_thw: (num_vision, 3) with patch counts [t, h, w] per visual item
+        grids: per-item ``[t, h, w]`` patch counts as host ints.
         max_num_patch: Maximum number of patches (for padding)
         num_grid_per_side: Side length of the square position embedding grid
         spatial_merge_size: Number of patches to merge per spatial dimension
@@ -65,16 +48,14 @@ def _compute_learned_pos_embeds(
     Returns:
         pos_embeds: (num_vision, max_num_patch, dim) interpolated position embeddings
     """
-    num_vision = grid_thw.shape[0]
     dtype = learned_pos_embed.dtype
     merge_size = spatial_merge_size
 
-    pos_embeds = learned_pos_embed.new_zeros(num_vision, max_num_patch, dim)
+    pos_embeds = learned_pos_embed.new_zeros(len(grids), max_num_patch, dim)
 
     # Group images by (h, w) to batch compute position embeddings
     hw_to_indices: dict[tuple[int, int], list[int]] = {}
-    for i in range(num_vision):
-        h, w = int(grid_thw[i, 1].item()), int(grid_thw[i, 2].item())
+    for i, (_, h, w) in enumerate(grids):
         key = (h, w)
         if key not in hw_to_indices:
             hw_to_indices[key] = []
@@ -120,7 +101,7 @@ def _compute_learned_pos_embeds(
         # For videos (t > 1), repeat spatial embeddings per frame;
         # temporal position encoding is handled by MRoPE in the LLM
         for i in indices:
-            t = int(grid_thw[i, 0].item())
+            t = grids[i][0]
             seq_len = t * h * w
             if t > 1:
                 pos_embeds[i, :seq_len] = pos_hw_block.repeat(t, 1)
@@ -132,7 +113,7 @@ def _compute_learned_pos_embeds(
 
 def _compute_2d_rope_cache(
     freq_table: torch.Tensor,
-    grid_thw: torch.Tensor,
+    grids: list[list[int]],
     max_num_patch: int,
     spatial_merge_size: int,
     head_dim: int,
@@ -145,7 +126,7 @@ def _compute_2d_rope_cache(
 
     Args:
         freq_table: (max_hw, head_dim//4) precomputed RoPE frequencies
-        grid_thw: (num_vision, 3) with patch counts [t, h, w] per visual item
+        grids: per-item ``[t, h, w]`` patch counts as host ints.
         max_num_patch: Maximum number of patches (for padding)
         spatial_merge_size: Number of patches to merge per spatial dimension
         head_dim: Attention head dimension
@@ -154,18 +135,16 @@ def _compute_2d_rope_cache(
         rope_cache: (num_vision, max_num_patch, 1, head_dim*2) float32 for
             VisionAttention
     """
-    num_vision = grid_thw.shape[0]
-    device = grid_thw.device
+    device = freq_table.device
     merge_size = spatial_merge_size
 
     rope_embeds = torch.zeros(
-        num_vision, max_num_patch, head_dim // 2, device=device, dtype=torch.float32
+        len(grids), max_num_patch, head_dim // 2, device=device, dtype=torch.float32
     )
 
     # Group images by (h, w) to batch compute RoPE embeddings
     hw_to_indices: dict[tuple[int, int], list[int]] = {}
-    for i in range(num_vision):
-        h, w = int(grid_thw[i, 1].item()), int(grid_thw[i, 2].item())
+    for i, (_, h, w) in enumerate(grids):
         key = (h, w)
         if key not in hw_to_indices:
             hw_to_indices[key] = []
@@ -212,7 +191,7 @@ def _compute_2d_rope_cache(
         # For videos (t > 1), repeat spatial embeddings per frame;
         # temporal position encoding is handled by MRoPE in the LLM
         for i in indices:
-            t = int(grid_thw[i, 0].item())
+            t = grids[i][0]
             seq_len = t * h * w
             if t > 1:
                 rope_embeds[i, :seq_len] = rope_2d.repeat(t, 1)
@@ -314,103 +293,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-class VisionAttention(Module):
-    """Multi-head attention with FlexAttention for efficient batched processing."""
-
-    @dataclass(kw_only=True, slots=True)
-    class Config(Module.Config):
-        dim: int
-        num_heads: int
-        wq: Linear.Config
-        wk: Linear.Config
-        wv: Linear.Config
-        proj: Linear.Config
-        inner_attention: Module.Config = field(default_factory=FlexAttention.Config)
-
-    def __init__(self, config: Config):
-        super().__init__()
-        self.dim = config.dim
-        self.num_heads = config.num_heads
-        self.head_dim = self.dim // self.num_heads
-
-        self.wq = config.wq.build()
-        self.wk = config.wk.build()
-        self.wv = config.wv.build()
-        self.proj = config.proj.build()
-        self.flex_attention = config.inner_attention.build()
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        *,
-        rope_cache: torch.Tensor,
-        attention_mask: BlockMask,
-    ) -> torch.Tensor:
-        bs, seqlen, _ = x.shape
-
-        xq = self.wq(x).view(bs, seqlen, -1, self.head_dim)
-        xk = self.wk(x).view(bs, seqlen, -1, self.head_dim)
-        xv = self.wv(x).view(bs, seqlen, -1, self.head_dim)
-
-        xq, xk = CosSinRoPE.apply_rotary_emb(xq, xk, rope_cache)
-
-        output = self.flex_attention(xq, xk, xv, attention_masks=attention_mask)
-        output = output.reshape(bs, seqlen, -1)
-        return self.proj(output)
-
-
-class VisionMLP(Module):
-    """Feed-forward network with GELU activation."""
-
-    @dataclass(kw_only=True, slots=True)
-    class Config(Module.Config):
-        fc1: Linear.Config
-        fc2: Linear.Config
-        act_fn: GELU.Config = field(
-            default_factory=lambda: GELU.Config(approximate="tanh")
-        )
-
-    def __init__(self, config: Config):
-        super().__init__()
-        self.linear_fc1 = config.fc1.build()
-        self.linear_fc2 = config.fc2.build()
-        self.act_fn = config.act_fn.build()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.linear_fc2(self.act_fn(self.linear_fc1(x)))
-
-
-class VisionTransformerBlock(Module):
-    """Single transformer block for vision encoder."""
-
-    @dataclass(kw_only=True, slots=True)
-    class Config(Module.Config):
-        norm1: LayerNorm.Config
-        norm2: LayerNorm.Config
-        attn: VisionAttention.Config
-        mlp: VisionMLP.Config
-
-    def __init__(self, config: Config):
-        super().__init__()
-        self.norm1 = config.norm1.build()
-        self.norm2 = config.norm2.build()
-        self.attn = config.attn.build()
-        self.mlp = config.mlp.build()
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        *,
-        rope_cache: torch.Tensor,
-        attention_mask: BlockMask,
-    ) -> torch.Tensor:
-        x = x + self.attn(
-            self.norm1(x), rope_cache=rope_cache, attention_mask=attention_mask
-        )
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-
 class Qwen35VisionEncoder(Module):
     """Qwen3.5 Vision Encoder with FlexAttention.
 
@@ -464,7 +346,7 @@ def __init__(self, config: Config):
         self.merger = config.merger.build()
 
     def compute_position_embeddings(
-        self, grid_thw: torch.Tensor, max_num_patch: int
+        self, grids: list[list[int]], max_num_patch: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute position embeddings for padded batch.
 
@@ -473,7 +355,7 @@ def compute_position_embeddings(
         - ``_compute_2d_rope_cache``: 2D RoPE cache
 
         Args:
-            grid_thw: (num_vision, 3) with patch counts [t, h, w] per visual item
+            grids: per-item ``[t, h, w]`` patch counts as host ints.
             max_num_patch: Maximum number of patches (for padding)
 
         Returns:
@@ -484,13 +366,13 @@ def compute_position_embeddings(
         head_dim = self.config.dim // self.config.num_heads
 
         # Get RoPE freq table, reusing cache when possible
-        max_hw = int(grid_thw[:, 1:].max().item())
+        max_hw = max(max(h, w) for _, h, w in grids)
         if self._cached_freq_table is None or self._cached_freq_table.shape[0] < max_hw:
             self._cached_freq_table = self.rotary_pos_emb(max_hw)
 
         learned_pos = _compute_learned_pos_embeds(
             self.pos_embed,
-            grid_thw,
+            grids,
             max_num_patch,
             self.num_grid_per_side,
             self.spatial_merge_size,
@@ -503,7 +385,7 @@ def compute_position_embeddings(
                 out_placements=(self._cached_freq_table.placements,),
             )(
                 self._cached_freq_table,
-                grid_thw,  # pyrefly: ignore [bad-argument-count]
+                grids,  # pyrefly: ignore [bad-argument-count]
                 max_num_patch,
                 self.spatial_merge_size,
                 head_dim,
@@ -511,7 +393,7 @@ def compute_position_embeddings(
         else:
             rope_cache = _compute_2d_rope_cache(
                 self._cached_freq_table,
-                grid_thw,
+                grids,
                 max_num_patch,
                 self.spatial_merge_size,
                 head_dim,
@@ -539,16 +421,18 @@ def forward(
         """
         num_vision, max_num_patch, _ = pixel_values.shape
 
+        # One host sync for the whole forward: read the (N, 3) grid to CPU ints
+        # so every per-item loop below builds shapes without a device sync. The
+        # GPU grid_thw is still used for num_patch (a pure tensor op, no sync).
+        grids = grid_thw.tolist()  # [[t, h, w], ...]
         num_patch = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).to(torch.long)
 
         x = self.patch_embed(pixel_values)  # (num_vision, max_num_patch, dim)
-        learned_pos, rope_cache = self.compute_position_embeddings(
-            grid_thw, max_num_patch
-        )
+        learned_pos, rope_cache = self.compute_position_embeddings(grids, max_num_patch)
         x = x + learned_pos
 
         mask_mod = get_vision_block_mask_mod(num_patch)
-        attention_mask = _compiled_create_block_mask(
+        attention_mask = compiled_create_block_mask(
             mask_mod,
             num_vision,
             None,
@@ -558,6 +442,11 @@ def forward(
         )
 
         for layer in self.layers.values():
-            x = layer(x, rope_cache=rope_cache, attention_mask=attention_mask)
+            x = layer(
+                x,
+                rope_cache=rope_cache,
+                rope_apply=CosSinRoPE.apply_rotary_emb,
+                attention_mask=attention_mask,
+            )
 
         return self.merger(x)
diff --git a/torchtitan/models/utils.py b/torchtitan/models/utils.py
index 7d4a2800ef..ba6da084bc 100644
--- a/torchtitan/models/utils.py
+++ b/torchtitan/models/utils.py
@@ -482,9 +482,13 @@ def get_moe_model_nparams_and_flops(
     nparams_shared_experts = 0
     nparams_experts = 0
     nparams_dense = 0
+    # TODO: add a per-batch vision encoder FLOP term for accurate VLM MFU.
+    nparams_vision = 0
 
     for name, p in model.named_parameters():
-        if "embedding" in name:
+        if "vision_encoder" in name:
+            nparams_vision += p.numel()
+        elif "embedding" in name:
             nparams_embedding += p.numel()
             nparams_dense += p.numel()
         elif "moe.shared_experts" in name:
@@ -497,7 +501,7 @@ def get_moe_model_nparams_and_flops(
             nparams_dense += p.numel()
 
     nparams_sparse = nparams_moe_router + nparams_shared_experts + nparams_experts
-    nparams = nparams_dense + nparams_sparse
+    nparams = nparams_dense + nparams_sparse + nparams_vision
 
     moe_config = next((l.moe for l in model_config.layers if l.moe is not None), None)
     if moe_config is not None:
@@ -511,7 +515,8 @@ def get_moe_model_nparams_and_flops(
 
     logger.info(
         f"Total parameter count: dense {nparams_dense:,}, "
-        f"sparse {nparams_sparse:,}, active {nparams_dense + nparams_sparse_active:,}"
+        f"sparse {nparams_sparse:,}, vision {nparams_vision:,}, "
+        f"active {nparams_dense + nparams_sparse_active:,}"
     )
 
     # With tied embeddings, PyTorch's parameter iterator already counts the