From 181a8d95714d837bc4f53d43a867c85ae4428359 Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:10:09 +0000
Subject: [PATCH 1/9] config: the pipeline split field is
 pipeline_parallel_module_fqns_per_model_part

The field names the split that pipeline parallelism applies, so it takes the
prefix its neighbours carry, and it is exclusive with
pipeline_parallel_layers_per_stage: both describe the same split and a run
that sets each would have to pick one silently.
---
 tests/unit_tests/cpu/test_config_manager.py     | 6 +++---
 tests/unit_tests/cpu/test_no_new_cli_options.py | 2 +-
 torchtitan/config/configs.py                    | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/unit_tests/cpu/test_config_manager.py b/tests/unit_tests/cpu/test_config_manager.py
index 1174c1b5df5..8d90714d6b2 100644
--- a/tests/unit_tests/cpu/test_config_manager.py
+++ b/tests/unit_tests/cpu/test_config_manager.py
@@ -577,13 +577,13 @@ def test_cli_override_dump_folder(self):
         )
         assert config.dump_folder == "/tmp/test_tt/"
 
-    def test_parse_module_fqns_per_model_part(self):
-        """module_fqns_per_model_part defaults to None."""
+    def test_parse_pipeline_parallel_module_fqns_per_model_part(self):
+        """pipeline_parallel_module_fqns_per_model_part defaults to None."""
         config_manager = ConfigManager()
         config = config_manager.parse_args(
             ["--module", "llama3", "--config", "llama3_debugmodel"]
         )
-        assert config.parallelism.module_fqns_per_model_part is None
+        assert config.parallelism.pipeline_parallel_module_fqns_per_model_part is None
 
     def test_optional_component_configs_do_not_add_cli_subcommands(self):
         config_manager = ConfigManager()
diff --git a/tests/unit_tests/cpu/test_no_new_cli_options.py b/tests/unit_tests/cpu/test_no_new_cli_options.py
index 51dbed5e80a..fd5bf4a10f2 100644
--- a/tests/unit_tests/cpu/test_no_new_cli_options.py
+++ b/tests/unit_tests/cpu/test_no_new_cli_options.py
@@ -153,11 +153,11 @@
         "parallelism.enable_sequence_parallel",
         "parallelism.expert_parallel_degree",
         "parallelism.fsdp_reshard_after_forward",
-        "parallelism.module_fqns_per_model_part",
         "parallelism.pipeline_parallel_degree",
         "parallelism.pipeline_parallel_first_stage_less_layers",
         "parallelism.pipeline_parallel_last_stage_less_layers",
         "parallelism.pipeline_parallel_layers_per_stage",
+        "parallelism.pipeline_parallel_module_fqns_per_model_part",
         "parallelism.num_pp_microbatches",
         "parallelism.pipeline_parallel_schedule",
         "parallelism.pipeline_parallel_schedule_csv",
diff --git a/torchtitan/config/configs.py b/torchtitan/config/configs.py
index c0e0abc116c..e5696509463 100644
--- a/torchtitan/config/configs.py
+++ b/torchtitan/config/configs.py
@@ -185,7 +185,7 @@ class ParallelismConfig:
     of stages. Stages per rank are inferred from split points degree, and schedule.
     """
 
-    module_fqns_per_model_part: list[list[str]] | None = None
+    pipeline_parallel_module_fqns_per_model_part: list[list[str]] | None = None
     """
     Specify a list of lists containing the FQNs (Fully Qualified Names) of modules for each model chunk.
     Each inner list represents one model chunk and contains the module names that belong to that chunk.
@@ -209,7 +209,7 @@ class ParallelismConfig:
 
     pipeline_parallel_layers_per_stage: int | None = None
     """
-    The number of layers per (virtual) pipeline stage. If specified, the module_fqns_per_model_part will be
+    The number of layers per (virtual) pipeline stage. If specified, the pipeline_parallel_module_fqns_per_model_part will be
     calculated from the number of layers and pipeline_parallel_degree. If not specified, the
     layers per stage will be inferred from the model, schedule, and pipeline_parallel_degree.
     """

From 523d865e1222fcbbfd2ace9a29dae0cd03f6e06f Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:10:09 +0000
Subject: [PATCH 2/9] pipeline_parallel: pipeline_with_first_last_stage_modules
 pins both ends

The helper already co-located modules with the first stage; a model whose
output side has modules outside the decoder needs the same for the last one,
so it takes last_stage_module_fqns and a model declares both ends through
pipeline_first_stage_module_fqns and pipeline_last_stage_module_fqns.

It refuses pipeline_parallel_layers_per_stage instead of overriding it: the
split it derives is sized by the schedule's default stage count, so the two
cannot both hold. return_split hands the applied split back to a caller that
routes along it.
---
 .../unit_tests/cpu/test_pipeline_parallel.py  | 96 +++++++++++++++++--
 torchtitan/distributed/pipeline_parallel.py   | 88 +++++++++++++----
 .../graph_trainer/graph_pp/pipeline.py        |  2 +-
 .../transformers_modeling_backend/pipeline.py |  2 +-
 torchtitan/protocols/model.py                 | 11 ++-
 5 files changed, 170 insertions(+), 29 deletions(-)

diff --git a/tests/unit_tests/cpu/test_pipeline_parallel.py b/tests/unit_tests/cpu/test_pipeline_parallel.py
index 80d838669b9..71a4abed4c7 100644
--- a/tests/unit_tests/cpu/test_pipeline_parallel.py
+++ b/tests/unit_tests/cpu/test_pipeline_parallel.py
@@ -18,7 +18,7 @@
 )
 
 
-def test_pipeline_with_first_stage_modules_prepends_present_modules(monkeypatch):
+def test_pipeline_with_first_last_stage_modules_prepends_present_modules(monkeypatch):
     model = nn.Module()
     model.vision_encoder = nn.Linear(2, 2)
     model.vision_adapter = nn.Linear(2, 2)
@@ -32,7 +32,8 @@ def capture_pipeline_llm(model, **kwargs):
 
     monkeypatch.setattr(pipeline_parallel, "pipeline_llm", capture_pipeline_llm)
 
-    result = pipeline_parallel.pipeline_with_first_stage_modules(
+    parallelism = ParallelismConfig(pipeline_parallel_degree=2)
+    result = pipeline_parallel.pipeline_with_first_last_stage_modules(
         model,
         first_stage_module_fqns=(
             "vision_encoder",
@@ -41,23 +42,24 @@ def capture_pipeline_llm(model, **kwargs):
             "missing_module",
         ),
         parallel_dims=SimpleNamespace(pp=2),
-        parallelism=ParallelismConfig(pipeline_parallel_degree=2),
+        parallelism=parallelism,
         model_config=SimpleNamespace(layers=[None] * 4),
     )
 
     assert result is expected_result
-    assert captured["parallelism"].module_fqns_per_model_part == [
+    assert parallelism.pipeline_parallel_module_fqns_per_model_part is None
+    assert captured["parallelism"].pipeline_parallel_module_fqns_per_model_part == [
         ["vision_encoder", "vision_adapter", "tok_embeddings", "layers.0", "layers.1"],
         ["layers.2", "layers.3", "norm", "lm_head"],
     ]
 
 
-def test_pipeline_with_first_stage_modules_preserves_explicit_split(monkeypatch):
+def test_pipeline_with_first_last_stage_modules_preserves_explicit_split(monkeypatch):
     model = nn.Module()
     configured_fqns = [["input"], ["output"]]
     parallelism = ParallelismConfig(
         pipeline_parallel_degree=2,
-        module_fqns_per_model_part=configured_fqns,
+        pipeline_parallel_module_fqns_per_model_part=configured_fqns,
     )
     captured = {}
 
@@ -67,7 +69,7 @@ def capture_pipeline_llm(model, **kwargs):
 
     monkeypatch.setattr(pipeline_parallel, "pipeline_llm", capture_pipeline_llm)
 
-    pipeline_parallel.pipeline_with_first_stage_modules(
+    pipeline_parallel.pipeline_with_first_last_stage_modules(
         model,
         first_stage_module_fqns=("missing_module",),
         parallel_dims=SimpleNamespace(pp=2),
@@ -78,6 +80,36 @@ def capture_pipeline_llm(model, **kwargs):
     assert captured["parallelism"] is parallelism
 
 
+def test_pipeline_with_first_last_stage_modules_appends_present_last_stage_modules(
+    monkeypatch,
+):
+    model = nn.Module()
+    model.vision_encoder = nn.Linear(2, 2)
+    model.output_res_proj = nn.Linear(2, 2)
+    model.output_res_norm = None
+    captured = {}
+
+    def capture_pipeline_llm(model, **kwargs):
+        captured["parallelism"] = kwargs["parallelism"]
+        return object()
+
+    monkeypatch.setattr(pipeline_parallel, "pipeline_llm", capture_pipeline_llm)
+
+    pipeline_parallel.pipeline_with_first_last_stage_modules(
+        model,
+        first_stage_module_fqns=("vision_encoder",),
+        last_stage_module_fqns=("output_res_proj", "output_res_norm", "missing"),
+        parallel_dims=SimpleNamespace(pp=2),
+        parallelism=ParallelismConfig(pipeline_parallel_degree=2),
+        model_config=SimpleNamespace(layers=[None] * 4),
+    )
+
+    assert captured["parallelism"].pipeline_parallel_module_fqns_per_model_part == [
+        ["vision_encoder", "tok_embeddings", "layers.0", "layers.1"],
+        ["layers.2", "layers.3", "norm", "lm_head", "output_res_proj"],
+    ]
+
+
 def _assert_layer_assignment(module_names_per_stage: list[list[str]], num_layers: int):
     """Layers are assigned in order with no gaps or duplicates."""
     assigned = [
@@ -220,3 +252,53 @@ def test_pp_rank_to_stage_mapping_requires_even_division():
 def test_get_pipeline_metadata_requires_layers_attribute():
     with pytest.raises(ValueError, match="Model does not have layers attribute."):
         _get_pipeline_metadata(object(), ParallelismConfig(), object())
+
+
+def test_pipeline_with_first_last_stage_modules_hands_back_the_split(monkeypatch):
+    model = nn.Module()
+    model.vision_encoder = nn.Linear(2, 2)
+    model.output_res_norm = nn.Linear(2, 2)
+    fake_result = (object(), [], True, False)
+    captured = {}
+
+    def capture_pipeline_llm(model, **kwargs):
+        captured["parallelism"] = kwargs["parallelism"]
+        return fake_result
+
+    monkeypatch.setattr(pipeline_parallel, "pipeline_llm", capture_pipeline_llm)
+    common = dict(
+        first_stage_module_fqns=("vision_encoder",),
+        last_stage_module_fqns=("output_res_norm",),
+        parallel_dims=SimpleNamespace(pp=2),
+        model_config=SimpleNamespace(layers=[None] * 4),
+    )
+
+    plain = pipeline_parallel.pipeline_with_first_last_stage_modules(
+        model, parallelism=ParallelismConfig(pipeline_parallel_degree=2), **common
+    )
+    assert plain is fake_result
+
+    *head, split = pipeline_parallel.pipeline_with_first_last_stage_modules(
+        model,
+        parallelism=ParallelismConfig(pipeline_parallel_degree=2),
+        return_split=True,
+        **common,
+    )
+    assert tuple(head) == fake_result
+    assert split == [
+        ["vision_encoder", "tok_embeddings", "layers.0", "layers.1"],
+        ["layers.2", "layers.3", "norm", "lm_head", "output_res_norm"],
+    ]
+    assert split == captured["parallelism"].pipeline_parallel_module_fqns_per_model_part
+
+    explicit = [["tok_embeddings", "layers.0"], ["layers.1", "norm", "lm_head"]]
+    *_, split = pipeline_parallel.pipeline_with_first_last_stage_modules(
+        model,
+        parallelism=ParallelismConfig(
+            pipeline_parallel_degree=2,
+            pipeline_parallel_module_fqns_per_model_part=explicit,
+        ),
+        return_split=True,
+        **common,
+    )
+    assert split == explicit
diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index 79640dd0fd9..356ea155e46 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -9,6 +9,7 @@
 import math
 import os
 from collections.abc import Callable, Sequence
+from typing import Literal, overload
 
 import torch
 import torch.nn as nn
@@ -37,7 +38,7 @@
 logger = logging.getLogger(__name__)
 
 
-__all__ = ["pipeline_llm", "pipeline_with_first_stage_modules"]
+__all__ = ["pipeline_llm", "pipeline_with_first_last_stage_modules"]
 
 
 def _build_get_mesh_callback(
@@ -86,7 +87,7 @@ def pipeline_llm(
         output_weight,
     ) = _get_pipeline_metadata(parallel_dims, parallelism, model_config)
 
-    module_names_per_stage = parallelism.module_fqns_per_model_part
+    module_names_per_stage = parallelism.pipeline_parallel_module_fqns_per_model_part
     if module_names_per_stage is None:
         module_names_per_stage = _generate_llm_fqn_per_model_part(
             num_virtual_stages, num_layers, input_weight, output_weight
@@ -141,28 +142,73 @@ def pipeline_llm(
     return pp_schedule, model_parts, has_first_stage, has_last_stage
 
 
-def pipeline_with_first_stage_modules(
+@overload
+def pipeline_with_first_last_stage_modules(
     model: BaseModel,
     *,
     first_stage_module_fqns: Sequence[str],
     parallel_dims: ParallelDims,
     parallelism: ParallelismConfig,
     model_config: BaseModel.Config,
+    last_stage_module_fqns: Sequence[str] = ...,
+    return_split: Literal[False] = ...,
     **kwargs,
 ) -> tuple[_PipelineSchedule, list[BaseModel], bool, bool]:
-    """Co-locate additional model modules with the first pipeline stage.
+    ...
+
+
+@overload
+def pipeline_with_first_last_stage_modules(
+    model: BaseModel,
+    *,
+    first_stage_module_fqns: Sequence[str],
+    parallel_dims: ParallelDims,
+    parallelism: ParallelismConfig,
+    model_config: BaseModel.Config,
+    last_stage_module_fqns: Sequence[str] = ...,
+    return_split: Literal[True],
+    **kwargs,
+) -> tuple[_PipelineSchedule, list[BaseModel], bool, bool, list[list[str]]]:
+    ...
+
+
+def pipeline_with_first_last_stage_modules(
+    model: BaseModel,
+    *,
+    first_stage_module_fqns: Sequence[str],
+    parallel_dims: ParallelDims,
+    parallelism: ParallelismConfig,
+    model_config: BaseModel.Config,
+    last_stage_module_fqns: Sequence[str] = (),
+    return_split: bool = False,
+    **kwargs,
+) -> (
+    tuple[_PipelineSchedule, list[BaseModel], bool, bool]
+    | tuple[_PipelineSchedule, list[BaseModel], bool, bool, list[list[str]]]
+):
+    """Co-locate additional model modules with the first and last pipeline stages.
 
     The auto-generated LLM stage split only knows about decoder modules
     (``tok_embeddings``, ``layers.*``, ``norm``, ``lm_head``). This function
     prepends each present module from ``first_stage_module_fqns`` to the first
-    stage's FQN list before delegating to ``pipeline_llm``. On other stages, the
+    stage's FQN list and appends each from ``last_stage_module_fqns`` to the
+    last stage's before delegating to ``pipeline_llm``. On other stages, the
     modules are pruned to ``None``; the model's ``forward`` must tolerate that.
+    ``return_split=True`` also returns the split that was applied.
 
-    NOTE: This adds load to stage 0 that the auto split does not model
-    (``input_weight`` only accounts for ``tok_embeddings``). Use
-    ``parallelism.pipeline_parallel_first_stage_less_layers`` to rebalance.
+    NOTE: This adds load to the end stages that the auto split does not model
+    (``input_weight`` only accounts for ``tok_embeddings``, ``output_weight``
+    for ``norm`` and ``lm_head``). Use
+    ``parallelism.pipeline_parallel_first_stage_less_layers`` and
+    ``pipeline_parallel_last_stage_less_layers`` to rebalance.
     """
-    if parallelism.module_fqns_per_model_part is None:
+    if parallelism.pipeline_parallel_layers_per_stage is not None:
+        raise ValueError(
+            "pipeline_with_first_last_stage_modules derives the split, so "
+            "pipeline_parallel_layers_per_stage would be overridden; leave it unset."
+        )
+    fqn_per_part = parallelism.pipeline_parallel_module_fqns_per_model_part
+    if fqn_per_part is None:
         (
             num_virtual_stages,
             num_layers,
@@ -172,23 +218,31 @@ def pipeline_with_first_stage_modules(
         fqn_per_part = _generate_llm_fqn_per_model_part(
             num_virtual_stages, num_layers, input_weight, output_weight
         )
-        present_module_fqns = [
-            module_fqn
-            for module_fqn in first_stage_module_fqns
-            if getattr(model, module_fqn, None) is not None
-        ]
-        fqn_per_part[0][:0] = present_module_fqns
+
+        def present(module_fqns: Sequence[str]) -> list[str]:
+            return [
+                module_fqn
+                for module_fqn in module_fqns
+                if getattr(model, module_fqn, None) is not None
+            ]
+
+        fqn_per_part[0][:0] = present(first_stage_module_fqns)
+        fqn_per_part[-1].extend(present(last_stage_module_fqns))
+        # The caller's config is not touched.
         parallelism = dataclasses.replace(
-            parallelism, module_fqns_per_model_part=fqn_per_part
+            parallelism, pipeline_parallel_module_fqns_per_model_part=fqn_per_part
         )
 
-    return pipeline_llm(
+    result = pipeline_llm(
         model,
         parallel_dims=parallel_dims,
         parallelism=parallelism,
         model_config=model_config,
         **kwargs,
     )
+    if return_split:
+        return (*result, fqn_per_part)
+    return result
 
 
 def _get_pipeline_metadata(
diff --git a/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py b/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py
index 73d2be0885b..dbd24fc563f 100644
--- a/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py
+++ b/torchtitan/experiments/graph_trainer/graph_pp/pipeline.py
@@ -109,7 +109,7 @@ def graph_pipeline_llm(
         output_weight,
     ) = _get_pipeline_metadata(parallel_dims, parallelism, model_config)
 
-    module_names_per_stage = parallelism.module_fqns_per_model_part
+    module_names_per_stage = parallelism.pipeline_parallel_module_fqns_per_model_part
     if module_names_per_stage is None:
         module_names_per_stage = _generate_llm_fqn_per_model_part(
             num_virtual_stages,
diff --git a/torchtitan/experiments/transformers_modeling_backend/pipeline.py b/torchtitan/experiments/transformers_modeling_backend/pipeline.py
index 6290b9f1d30..3438dedb225 100644
--- a/torchtitan/experiments/transformers_modeling_backend/pipeline.py
+++ b/torchtitan/experiments/transformers_modeling_backend/pipeline.py
@@ -367,7 +367,7 @@ def pipeline_hf_transformers(
         stages_per_rank = 1 if is_single_stage_schedule else 2
         num_virtual_stages = parallel_dims.pp * stages_per_rank
 
-    module_names_per_stage = parallelism.module_fqns_per_model_part
+    module_names_per_stage = parallelism.pipeline_parallel_module_fqns_per_model_part
     if module_names_per_stage is None:
         module_names_per_stage = generate_llm_fqn_per_model_part(
             num_virtual_stages, num_layers, input_weight, output_weight
diff --git a/torchtitan/protocols/model.py b/torchtitan/protocols/model.py
index 4e1b3f72a76..233464a721c 100644
--- a/torchtitan/protocols/model.py
+++ b/torchtitan/protocols/model.py
@@ -103,6 +103,7 @@ def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]:
 
     state_dict_adapter_cls: ClassVar[type[BaseStateDictAdapter] | None] = None
     pipeline_first_stage_module_fqns: ClassVar[tuple[str, ...]] = ()
+    pipeline_last_stage_module_fqns: ClassVar[tuple[str, ...]] = ()
     supports_pipeline_parallel: ClassVar[bool] = True
 
     def pipeline(self, **kwargs: Any) -> tuple[Any, list[BaseModel], bool, bool]:
@@ -114,13 +115,17 @@ def pipeline(self, **kwargs: Any) -> tuple[Any, list[BaseModel], bool, bool]:
 
         from torchtitan.distributed.pipeline_parallel import (
             pipeline_llm,
-            pipeline_with_first_stage_modules,
+            pipeline_with_first_last_stage_modules,
         )
 
-        if self.pipeline_first_stage_module_fqns:
-            return pipeline_with_first_stage_modules(
+        if (
+            self.pipeline_first_stage_module_fqns
+            or self.pipeline_last_stage_module_fqns
+        ):
+            return pipeline_with_first_last_stage_modules(
                 self,
                 first_stage_module_fqns=self.pipeline_first_stage_module_fqns,
+                last_stage_module_fqns=self.pipeline_last_stage_module_fqns,
                 **kwargs,
             )
         return pipeline_llm(self, **kwargs)

From b01881e2ef48467cf25bb6d7e4df6e925ca94da5 Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:10:22 +0000
Subject: [PATCH 3/9] kimi_k3: the pipeline lives in kimi_k3/pipeline_parallel,
 and the model owns it

The model declares the modules that belong with the first and the last
stage and overrides pipeline() to route the attention-residual blocks along
the split it gets back. The package holds the stage subclass, whose forward
carries the blocks a receiving rank lacks and whose backward reads the
delta's gradient need off the receive metadata, and the routing tables,
which derive the block count from the layer count and assume the loop-style
stage-to-rank assignment the contiguity check enforces.

A block's first layer closes the previous block and joins the stack, so a
stage that opens inside a block carries a partial sum rather than a closed
block, and a stage without the head hands the stack on unchanged.
---
 .../cpu/test_kimi_k3_pp_block_grads.py        | 275 +++++++++++++
 .../unit_tests/cpu/test_kimi_k3_pp_layout.py  | 192 +++++++++
 tests/unit_tests/cpu/test_kimi_k3_pp_stage.py | 162 ++++++++
 torchtitan/models/kimi_k3/model.py            |  70 ++--
 .../kimi_k3/pipeline_parallel/__init__.py     | 121 ++++++
 .../kimi_k3/pipeline_parallel/layout.py       | 159 ++++++++
 .../models/kimi_k3/pipeline_parallel/stage.py | 364 ++++++++++++++++++
 7 files changed, 1314 insertions(+), 29 deletions(-)
 create mode 100644 tests/unit_tests/cpu/test_kimi_k3_pp_block_grads.py
 create mode 100644 tests/unit_tests/cpu/test_kimi_k3_pp_layout.py
 create mode 100644 tests/unit_tests/cpu/test_kimi_k3_pp_stage.py
 create mode 100644 torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
 create mode 100644 torchtitan/models/kimi_k3/pipeline_parallel/layout.py
 create mode 100644 torchtitan/models/kimi_k3/pipeline_parallel/stage.py

diff --git a/tests/unit_tests/cpu/test_kimi_k3_pp_block_grads.py b/tests/unit_tests/cpu/test_kimi_k3_pp_block_grads.py
new file mode 100644
index 00000000000..ec7ae784bd2
--- /dev/null
+++ b/tests/unit_tests/cpu/test_kimi_k3_pp_block_grads.py
@@ -0,0 +1,275 @@
+# 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.
+
+import unittest
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributed.pipelining.schedules import ScheduleInterleaved1F1B
+from torch.distributed.pipelining.stage import _PipelineStageBase
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
+
+from torchtitan.models.kimi_k3.pipeline_parallel.layout import infer_block_layout_tables
+from torchtitan.models.kimi_k3.pipeline_parallel.stage import (
+    AttnResPipelineStage,
+    PPRankLocalCache,
+)
+
+NUM_LAYERS, LAYERS_PER_BLOCK = 16, 4
+NUM_BLOCKS = NUM_LAYERS // LAYERS_PER_BLOCK
+# Channels: one per layer's reads, the head's reads, the loss readout, the input.
+HEAD, READOUT, INPUT = NUM_LAYERS, NUM_LAYERS + 1, NUM_LAYERS + 2
+DIM = NUM_LAYERS + 3
+TOKENS, MICROBATCHES, STEPS = 1, 4, 3
+# pp4 x vp4 with the head alone on the last stage, two layers per stage, and uneven stages
+SPLITS = {
+    "pp4 x vp4, head alone": [[0], [1, 2]] + [[s + 1] for s in range(2, 15)] + [[]],
+    "pp4 x vp2": [[2 * s, 2 * s + 1] for s in range(8)],
+    "pp4 x vp2, blocks open inside stages": [
+        [0, 1],
+        [2, 3, 4],
+        [5],
+        [6, 7, 8, 9],
+        [10],
+        [11, 12, 13],
+        [14],
+        [15],
+    ],
+}
+
+
+class _ExactStage(nn.Module):
+    """A Kimi K3-shaped pipeline stage whose block gradients are small integers."""
+
+    def __init__(
+        self, layers: list[int], *, first: bool, last: bool, dtype: torch.dtype
+    ):
+        super().__init__()
+        self.layers, self.first, self.last = layers, first, last
+        self.blocks = nn.ParameterDict(
+            {
+                str(layer // LAYERS_PER_BLOCK): nn.Parameter(
+                    torch.arange(DIM, dtype=dtype) % 3 + 1 + layer // LAYERS_PER_BLOCK
+                )
+                for layer in layers
+                if layer % LAYERS_PER_BLOCK == 0
+            }
+        )
+
+    def forward(self, hidden: torch.Tensor, stack: torch.Tensor | None = None):
+        if self.first:
+            hidden = F.pad(hidden, (INPUT, 0))
+            stack = hidden.new_zeros(hidden.shape[0], 0, DIM)
+        assert stack is not None
+        for layer in self.layers:
+            if layer % LAYERS_PER_BLOCK == 0:
+                block = self.blocks[str(layer // LAYERS_PER_BLOCK)] * hidden[:, INPUT:]
+                stack = torch.cat((stack, block.unsqueeze(1)), dim=1)
+            for _ in range(2):
+                read = _read(stack, layer).unsqueeze(1)
+                hidden = hidden + F.pad(read, (READOUT, DIM - READOUT - 1))
+        if self.last:
+            return hidden[:, READOUT] + _read(stack, HEAD)
+        return hidden, stack
+
+
+def _read(stack: torch.Tensor, channel: int) -> torch.Tensor:
+    weights = torch.arange(1, stack.shape[1] + 1, dtype=stack.dtype)
+    return (stack[:, :, channel] * weights).sum(1)
+
+
+def _loss(output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return (output * target).sum()
+
+
+def _batch(dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+    weights = torch.arange(1, MICROBATCHES + 1, dtype=dtype).repeat_interleave(TOKENS)
+    return weights.unsqueeze(1), weights
+
+
+def _expected_block_grad(block: int, dtype: torch.dtype) -> torch.Tensor:
+    total = sum(TOKENS * (mb + 1) ** 2 for mb in range(MICROBATCHES))
+    reads = torch.zeros(DIM, dtype=dtype)
+    reads[block * LAYERS_PER_BLOCK : NUM_LAYERS] = 2
+    reads[HEAD] = 1
+    return total * (block + 1) * reads
+
+
+def _block_grads(modules) -> dict[int, torch.Tensor]:
+    return {
+        int(name.split(".")[1]): param.grad.clone()
+        for module in modules
+        for name, param in module.named_parameters()
+    }
+
+
+def _train(modules, step_fn, dtype: torch.dtype):
+    params = [p for module in modules for p in module.parameters()]
+    optimizer = torch.optim.SGD(params, lr=1.0) if params else None
+    history = []
+    for _ in range(STEPS):
+        losses = step_fn(*_batch(dtype))
+        history.append((_block_grads(modules), losses))
+        if optimizer is not None:
+            optimizer.step()
+            optimizer.zero_grad()
+    return history
+
+
+def _run_single_device(split: list[list[int]], dtype: torch.dtype):
+    last = len(split) - 1
+    modules = [
+        _ExactStage(layers, first=s == 0, last=s == last, dtype=dtype)
+        for s, layers in enumerate(split)
+    ]
+
+    def step(inputs, targets):
+        losses = []
+        for x, y in zip(
+            inputs.chunk(MICROBATCHES), targets.chunk(MICROBATCHES), strict=True
+        ):
+            out = modules[0](x)
+            for module in modules[1:]:
+                out = module(*out)
+            loss = _loss(out, y)
+            loss.backward()
+            losses.append(loss.detach())
+        return losses
+
+    return _train(modules, step, dtype)
+
+
+class TestKimiK3PipelineExactBlockGradients(DTensorTestBase):
+    @property
+    def device_type(self) -> str:
+        return "cpu"
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    def _run_pipeline(
+        self,
+        split: list[list[int]],
+        dtype: torch.dtype,
+        cache: bool,
+        eval_losses: list | None = None,
+    ):
+        num_stages, last = len(split), len(split) - 1
+        mine = range(self.rank, num_stages, self.world_size)
+        modules = [
+            _ExactStage(split[s], first=s == 0, last=s == last, dtype=dtype)
+            for s in mine
+        ]
+        stages = [
+            AttnResPipelineStage(module, s, num_stages, torch.device("cpu"))
+            for module, s in zip(modules, mine, strict=True)
+        ]
+        schedule_stages: list[_PipelineStageBase] = list(stages)
+        schedule = ScheduleInterleaved1F1B(
+            schedule_stages,
+            n_microbatches=MICROBATCHES,
+            loss_fn=_loss,
+            scale_grads=False,
+        )
+        layout = infer_block_layout_tables(
+            stage_to_rank=dict(stages[0].stage_index_to_group_rank),
+            n_layers=NUM_LAYERS,
+            layers_per_block=LAYERS_PER_BLOCK,
+            layer_to_stage={
+                layer: s for s, layers in enumerate(split) for layer in layers
+            },
+            cache=cache,
+        )
+        store = PPRankLocalCache()
+        for stage in stages:
+            stage.set_routing(layout, store)
+
+        def step(inputs, targets):
+            losses: list[torch.Tensor] = []
+            args = (inputs,) if 0 in mine else ()
+            if last in mine:
+                schedule.step(*args, target=targets, losses=losses)
+            else:
+                schedule.step(*args)
+            if eval_losses is not None:
+                evaluated: list[torch.Tensor] = []
+                with torch.no_grad():
+                    if last in mine:
+                        schedule.eval(*args, target=targets, losses=evaluated)
+                    else:
+                        schedule.eval(*args)
+                if store.blocks(0) or any(
+                    store.blocks(mb) for mb in range(MICROBATCHES)
+                ):
+                    raise AssertionError("eval left blocks in the rank store")
+                eval_losses.append(
+                    (
+                        [loss.detach() for loss in evaluated],
+                        sum(len(s._order) + len(s._delta_in) for s in stages),
+                    )
+                )
+            return [loss.detach() for loss in losses]
+
+        sent = sum(len(layout.delta_to_send(s)) for s in range(num_stages))
+        return _train(modules, step, dtype), sent
+
+    @with_comms
+    def test_rank_cache_matches_whole_stack_and_single_device(self):
+        for dtype in (torch.bfloat16, torch.float32):
+            for name, split in SPLITS.items():
+                with self.subTest(dtype=dtype, split=name):
+                    reference = _run_single_device(split, dtype)
+                    cached, cached_sent = self._run_pipeline(split, dtype, cache=True)
+                    naive, naive_sent = self._run_pipeline(split, dtype, cache=False)
+                    self.assertLess(cached_sent, naive_sent)
+                    for step in range(STEPS):
+                        ref_grads, ref_losses = reference[step]
+                        for grads, losses in (cached[step], naive[step]):
+                            for block, grad in grads.items():
+                                expected = _expected_block_grad(block, dtype)
+                                torch.testing.assert_close(
+                                    grad, ref_grads[block], rtol=0, atol=0
+                                )
+                                torch.testing.assert_close(
+                                    grad, expected, rtol=0, atol=0
+                                )
+                            if losses:
+                                torch.testing.assert_close(
+                                    losses, ref_losses, rtol=0, atol=0
+                                )
+
+    @with_comms
+    def test_forward_only_eval_between_steps(self):
+        for name, split in SPLITS.items():
+            with self.subTest(split=name):
+                reference = _run_single_device(split, torch.float32)
+                evaluated: list = []
+                history, _ = self._run_pipeline(
+                    split, torch.float32, cache=True, eval_losses=evaluated
+                )
+                for step in range(STEPS):
+                    grads, losses = history[step]
+                    ref_grads, ref_losses = reference[step]
+                    for block, grad in grads.items():
+                        torch.testing.assert_close(
+                            grad, ref_grads[block], rtol=0, atol=0
+                        )
+                    eval_step_losses, leftover = evaluated[step]
+                    if losses:
+                        torch.testing.assert_close(losses, ref_losses, rtol=0, atol=0)
+                        torch.testing.assert_close(
+                            eval_step_losses, ref_losses, rtol=0, atol=0
+                        )
+                    self.assertEqual(leftover, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit_tests/cpu/test_kimi_k3_pp_layout.py b/tests/unit_tests/cpu/test_kimi_k3_pp_layout.py
new file mode 100644
index 00000000000..84ed428a54c
--- /dev/null
+++ b/tests/unit_tests/cpu/test_kimi_k3_pp_layout.py
@@ -0,0 +1,192 @@
+# 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.
+
+"""The block routing tables, on CPU: uneven splits, the store, the deltas."""
+
+import unittest
+
+from torchtitan.models.kimi_k3.pipeline_parallel import _require_loop_style
+from torchtitan.models.kimi_k3.pipeline_parallel.layout import (
+    BlockLayoutTables,
+    infer_block_layout_tables,
+    layer_to_stage_from_split,
+)
+
+
+def _uneven_map() -> dict[int, int]:
+    # 24 layers over 4 stages as 5 / 7 / 6 / 6.
+    ranges = [(0, 5), (5, 12), (12, 18), (18, 24)]
+    return {
+        layer: stage for stage, (lo, hi) in enumerate(ranges) for layer in range(lo, hi)
+    }
+
+
+# Two ranks, two stages each, interleaved: stage s runs on rank s % 2.
+_STAGE_TO_RANK = {0: 0, 1: 1, 2: 0, 3: 1}
+
+# The split the B200 pp4 x vp4 cell spells out, repeated here so the test needs
+# no recipe module: 16 stages, one layer per stage from layer 5 on.
+_PP4_VP4_SPLIT = [
+    ["vision_encoder", "tok_embeddings", "layers.0"],
+    ["layers.1", "layers.2"],
+    ["layers.3", "layers.4"],
+    *[[f"layers.{i}"] for i in range(5, 17)],
+    ["norm", "lm_head", "output_res_proj", "output_res_norm"],
+]
+
+
+def _tables(cache: bool = True) -> BlockLayoutTables:
+    return BlockLayoutTables(
+        stage_to_rank=_STAGE_TO_RANK,
+        n_layers=24,
+        layers_per_block=12,
+        layer_to_stage=_uneven_map(),
+        cache=cache,
+    )
+
+
+_EIGHT_STAGES = [
+    ["vision_encoder", "tok_embeddings", "layers.0", "layers.1"],
+    ["layers.2", "layers.3", "layers.4"],
+    ["layers.5", "layers.6", "layers.7"],
+    ["layers.8", "layers.9"],
+    ["layers.10", "layers.11"],
+    ["layers.12", "layers.13"],
+    ["layers.14", "layers.15"],
+    ["layers.16", "norm", "lm_head", "output_res_proj", "output_res_norm"],
+]
+
+
+class TestRouting(unittest.TestCase):
+    def test_tables_follow_the_map_not_an_equal_split(self):
+        tables = _tables()
+        self.assertEqual(tables.producer_stage_of_block(0), 0)
+        self.assertEqual(tables.producer_stage_of_block(1), 2)
+        self.assertEqual(tables.commits_at(1), [])
+        self.assertEqual(tables.delta_to_send(0), [0])
+        self.assertEqual(tables.delta_to_send(1), [])
+        self.assertEqual(tables.delta_to_send(2), [1])
+        self.assertEqual(tables.delta_to_send(3), [])
+        self.assertEqual(tables.cache_at_entry(2), frozenset({0}))
+        self.assertEqual(tables.cache_at_entry(3), frozenset({0}))
+        self.assertEqual(tables.cache_readers_of_block(0), [2, 3])
+        self.assertEqual(tables.deposits_expected(0, 0), 1)
+        self.assertEqual(tables.deposits_expected(0, 1), 1)
+        self.assertEqual(tables.deposits_expected(1, 2), 0)
+
+    def test_without_the_cache_every_hop_carries_everything(self):
+        tables = _tables(cache=False)
+        self.assertEqual(tables.delta_to_send(0), [0])
+        self.assertEqual(tables.delta_to_send(1), [0])
+        self.assertEqual(tables.delta_to_send(2), [0, 1])
+        for stage in range(4):
+            self.assertEqual(tables.cache_at_entry(stage), frozenset())
+        self.assertEqual(tables.deposits_expected(0, 0), 0)
+
+    def test_infer_accepts_the_map_and_rejects_a_broken_one(self):
+        common = dict(
+            stage_to_rank=_STAGE_TO_RANK,
+            n_layers=24,
+            layers_per_block=12,
+        )
+        tables = infer_block_layout_tables(layer_to_stage=_uneven_map(), **common)
+        self.assertEqual(tables.producer_stage_of_block(1), 2)
+        incomplete = _uneven_map()
+        del incomplete[7]
+        with self.assertRaisesRegex(ValueError, "exactly once"):
+            infer_block_layout_tables(layer_to_stage=incomplete, **common)
+        scrambled = _uneven_map()
+        scrambled[7], scrambled[20] = scrambled[20], scrambled[7]
+        with self.assertRaisesRegex(ValueError, "non-contiguous"):
+            infer_block_layout_tables(layer_to_stage=scrambled, **common)
+
+    def test_the_map_is_read_off_the_split(self):
+        split = [
+            ["tok_embeddings", "vision_encoder", "layers.0", "layers.1"],
+            ["layers.2"],
+            ["layers.3", "norm", "lm_head", "output_res_proj"],
+        ]
+        self.assertEqual(layer_to_stage_from_split(split), {0: 0, 1: 0, 2: 1, 3: 2})
+
+
+class TestSplit(unittest.TestCase):
+    def _tables(self, split, pp):
+        layer_to_stage = layer_to_stage_from_split(split)
+        self.assertEqual(sorted(layer_to_stage), list(range(17)))
+        return BlockLayoutTables(
+            stage_to_rank={s: s % pp for s in range(len(split))},
+            n_layers=17,
+            layers_per_block=4,
+            layer_to_stage=layer_to_stage,
+            cache=True,
+        )
+
+    def test_the_pp2_vp2_cell_has_every_transport_path(self):
+        from torchtitan.distributed.pipeline_parallel import (
+            _generate_llm_fqn_per_model_part,
+        )
+
+        split = _generate_llm_fqn_per_model_part(4, 17)
+        self.assertEqual([len(stage) for stage in split], [5, 5, 5, 5])
+        tables = self._tables(split, pp=2)
+        self.assertEqual(
+            [tables.producer_stage_of_block(b) for b in range(5)], [0, 1, 1, 2, 3]
+        )
+        self.assertEqual(
+            [tables.delta_to_send(s) for s in range(3)], [[0], [1, 2], [3]]
+        )
+        self.assertEqual(tables.cache_at_entry(2), frozenset({0}))
+        self.assertEqual(tables.cache_at_entry(3), frozenset({0, 1, 2}))
+
+    def test_an_eight_stage_split_collects_three_deposits_per_block(self):
+        split = _EIGHT_STAGES
+        self.assertEqual(len(split), 8)
+        self.assertEqual(split[0][:2], ["vision_encoder", "tok_embeddings"])
+        self.assertEqual(split[-1][-2:], ["output_res_proj", "output_res_norm"])
+        tables = self._tables(split, pp=2)
+        self.assertEqual(
+            [tables.producer_stage_of_block(b) for b in range(5)], [0, 1, 3, 5, 7]
+        )
+        self.assertEqual(
+            [tables.delta_to_send(s) for s in range(7)],
+            [[0], [1], [], [2], [], [3], []],
+        )
+        self.assertEqual(tables.cache_at_entry(7), frozenset({0, 1, 2, 3}))
+        self.assertEqual(tables.deposits_expected(0, 0), 3)
+        self.assertEqual(tables.deposits_expected(0, 1), 3)
+
+    def test_the_pp4_vp4_cell_runs_one_layer_per_stage(self):
+        split = _PP4_VP4_SPLIT
+        self.assertEqual(len(split), 16)
+        self.assertEqual(sum(n.startswith("layers.") for s in split for n in s), 17)
+        self.assertEqual(
+            split[-1], ["norm", "lm_head", "output_res_proj", "output_res_norm"]
+        )
+        tables = self._tables(split, pp=4)
+        self.assertEqual(
+            [tables.producer_stage_of_block(b) for b in range(5)], [0, 2, 6, 10, 14]
+        )
+        self.assertEqual(tables.delta_to_send(2), [0, 1])
+        self.assertEqual(tables.delta_to_send(5), [])
+        self.assertEqual(tables.delta_to_send(14), [4])
+        self.assertEqual(tables.cache_at_entry(15), frozenset({0, 1, 2, 3}))
+        self.assertEqual(tables.deposits_expected(0, 0), 3)
+
+
+class TestLoopStyleGuard(unittest.TestCase):
+    def test_a_loop_style_assignment_passes(self):
+        _require_loop_style(object(), {0: 0, 1: 1, 2: 0, 3: 1}, pp=2)
+
+    def test_any_other_assignment_is_refused(self):
+        # Stages 0..3 on ranks 0, 1, 1, 0, the v shape: stage 2 revisits rank 1
+        # out of index order, so the rank store would be asked for blocks it
+        # never saw. The map is what is checked, whatever schedule produced it.
+        with self.assertRaisesRegex(ValueError, "unsupported with attn_res_cache"):
+            _require_loop_style(object(), {0: 0, 1: 1, 2: 1, 3: 0}, pp=2)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/unit_tests/cpu/test_kimi_k3_pp_stage.py b/tests/unit_tests/cpu/test_kimi_k3_pp_stage.py
new file mode 100644
index 00000000000..d77d06e5ba4
--- /dev/null
+++ b/tests/unit_tests/cpu/test_kimi_k3_pp_stage.py
@@ -0,0 +1,162 @@
+# 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.
+
+"""The stage's carrier handling, on CPU: assembly, routing, the gradient split, and
+the swap of core's stages for AttnRes ones."""
+
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.distributed.pipelining import PipelineStage
+from torch.distributed.pipelining.schedules import Schedule1F1B, ScheduleInterleaved1F1B
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
+
+from torchtitan.models.kimi_k3.pipeline_parallel import _swap_in_attn_res_stages
+from torchtitan.models.kimi_k3.pipeline_parallel.stage import (
+    assemble_stack,
+    AttnResPipelineStage,
+    pack_outgoing_delta,
+    PPRankLocalCache,
+    split_stack_grad,
+)
+
+
+class TestCarrier(unittest.TestCase):
+    def test_assembly_orders_blocks_and_hands_back_a_leaf(self):
+        T, D = 4, 8
+        hidden = torch.randn(T, D)
+        delta = torch.randn(T, 1, D, requires_grad=True)  # block 2 on the wire
+        store = {0: torch.randn(T, D), 1: torch.randn(T, D)}
+        stack, order = assemble_stack(hidden, delta, [2], store)
+        self.assertEqual(order, [0, 1, 2])
+        self.assertTrue(stack.is_leaf and stack.requires_grad)
+        self.assertTrue(torch.equal(stack[:, 0], store[0]))
+        self.assertTrue(torch.equal(stack[:, 2], delta[:, 0]))
+        empty, order = assemble_stack(hidden, hidden.new_zeros(T, 0, D), [], {})
+        self.assertEqual((tuple(empty.shape), order), ((T, 0, D), []))
+        with self.assertRaisesRegex(ValueError, "routing expects"):
+            assemble_stack(hidden, delta, [2, 3], store)
+
+    def test_payload_is_the_routed_columns_of_the_model_stack(self):
+        T, D = 4, 8
+        stack_out = torch.randn(T, 3, D, requires_grad=True)
+        payload = pack_outgoing_delta(stack_out, [0, 1, 2], [1, 2])
+        self.assertEqual(tuple(payload.shape), (T, 2, D))
+        self.assertTrue(torch.equal(payload[:, 0], stack_out[:, 1]))
+        self.assertTrue(payload.requires_grad)
+        self.assertEqual(
+            tuple(pack_outgoing_delta(stack_out, [0, 1, 2], []).shape), (T, 0, D)
+        )
+
+    def test_gradient_split_sends_the_received_and_deposits_the_stored(self):
+        T, D = 4, 8
+        grad_stack = torch.randn(T, 3, D)
+        like = torch.zeros(T, D)
+        grad_delta, deposits = split_stack_grad(grad_stack, [0, 1, 2], [2], like)
+        self.assertEqual(tuple(grad_delta.shape), (T, 1, D))
+        self.assertTrue(grad_delta.is_contiguous())
+        self.assertTrue(torch.equal(grad_delta[:, 0], grad_stack[:, 2]))
+        self.assertEqual(set(deposits), {0, 1})
+        self.assertTrue(torch.equal(deposits[1], grad_stack[:, 1]))
+        grad_delta, deposits = split_stack_grad(None, [0], [0], like)
+        self.assertTrue(torch.equal(grad_delta, torch.zeros(T, 1, D)))
+        self.assertEqual(deposits, {})
+
+    def test_store_accumulates_deposits_and_releases_blocks_separately(self):
+        store = PPRankLocalCache()
+        store.put(0, 0, torch.zeros(4, 2))
+        store.deposit(0, 0, torch.ones(4, 2))
+        store.deposit(0, 0, torch.ones(4, 2))
+        store.release(0)
+        self.assertEqual(store.blocks(0), {})
+        self.assertTrue(store.has_deposits(0))
+        grad, count = store.collect(0, 0)
+        self.assertEqual(count, 2)
+        assert grad is not None
+        self.assertTrue(torch.equal(grad, torch.full((4, 2), 2.0)))
+        self.assertFalse(store.has_deposits(0))
+        self.assertEqual(store.collect(0, 0), (None, 0))
+
+
+class TestForwardOnly(unittest.TestCase):
+    def test_eval_backward_routes_nothing_and_forgets_the_chunk(self):
+        stage = AttnResPipelineStage.__new__(AttnResPipelineStage)
+        stage._has_backward = False
+        stage.fwd_cache = {0: ((torch.zeros(1),), [])}
+        stage.bwd_cache = {}
+        stage._order = {0: [0, 1]}
+        stage._delta_in = {0: [1]}
+        AttnResPipelineStage.backward_one_chunk(stage, 0)
+        self.assertEqual((stage.fwd_cache, stage._order, stage._delta_in), ({}, {}, {}))
+
+
+def _loss(output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return output.sum()
+
+
+def _get_mesh(*args, **kwargs):
+    return None
+
+
+class TestStageSwap(DTensorTestBase):
+    """Core's stages are rebuilt as AttnRes stages on the schedule, one rank on gloo."""
+
+    @property
+    def device_type(self) -> str:
+        return "cpu"
+
+    @property
+    def world_size(self) -> int:
+        return 1
+
+    @staticmethod
+    def _stage(index: int, num_stages: int) -> PipelineStage:
+        return PipelineStage(
+            nn.Linear(2, 2),
+            index,
+            num_stages,
+            torch.device("cpu"),
+            get_mesh=_get_mesh,
+        )
+
+    def _assert_rebuilt(self, old: PipelineStage, new: AttnResPipelineStage) -> None:
+        self.assertIsInstance(new, AttnResPipelineStage)
+        self.assertIsNot(new, old)
+        # The same module object, so the model parts core returned stay valid.
+        self.assertIs(new.submod, old.submod)
+        self.assertEqual(
+            (new.stage_index, new.num_stages, new.device),
+            (old.stage_index, old.num_stages, old.device),
+        )
+        self.assertIs(new.group, old.group)
+        self.assertIs(new._mesh_cache._get_mesh_cb, old._mesh_cache._get_mesh_cb)
+        self.assertEqual(new.stage_index_to_group_rank, old.stage_index_to_group_rank)
+
+    @with_comms
+    def test_single_stage_schedule(self):
+        old = self._stage(0, 1)
+        schedule = Schedule1F1B(old, n_microbatches=1, loss_fn=_loss)
+        (new,) = _swap_in_attn_res_stages(schedule)
+        self._assert_rebuilt(old, new)
+        self.assertIs(schedule._stage, new)
+
+    @with_comms
+    def test_multi_stage_schedule(self):
+        old = [self._stage(index, 2) for index in range(2)]
+        schedule = ScheduleInterleaved1F1B(old, n_microbatches=2, loss_fn=_loss)
+        new = _swap_in_attn_res_stages(schedule)
+        self.assertEqual(len(new), 2)
+        for old_stage, new_stage in zip(old, new, strict=True):
+            self._assert_rebuilt(old_stage, new_stage)
+        self.assertEqual(schedule._stages, new)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py
index eed581d6616..e82156a3fa3 100644
--- a/torchtitan/models/kimi_k3/model.py
+++ b/torchtitan/models/kimi_k3/model.py
@@ -157,7 +157,7 @@ def forward(
 
 
 def _apply_attention_residual(
-    prefix_sum_TD: torch.Tensor,
+    partial_block_TD: torch.Tensor | None,
     block_residual_TND: torch.Tensor,
     projection: Linear,
     norm: RMSNorm,
@@ -165,7 +165,11 @@ def _apply_attention_residual(
     """Apply Kimi's block-level attention residual in FP32."""
     assert norm.eps is not None
 
-    values_TND = torch.cat((block_residual_TND, prefix_sum_TD.unsqueeze(1)), dim=1)
+    values_TND = (
+        block_residual_TND
+        if partial_block_TD is None
+        else torch.cat((block_residual_TND, partial_block_TD.unsqueeze(1)), dim=1)
+    )
     values_float = values_TND.float()
     variance = values_float.pow(2).mean(dim=-1, keepdim=True)
     keys_TND = values_float * torch.rsqrt(variance + norm.eps)
@@ -204,6 +208,8 @@ def __init__(self, config: Config):
             raise ValueError("Exactly one of feed_forward or moe must be configured.")
         self.layer_id = config.layer_id
         self.attn_res_block_size = config.attn_res_block_size
+        # A block's first layer closes the previous block and joins the stack.
+        self.first_layer_in_block = self.layer_id % self.attn_res_block_size == 0
         self.attention = (
             config.attention.build() if config.attention is not None else None
         )
@@ -244,29 +250,25 @@ def forward(
         *,
         padding_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        prefix_sum_TD = x_TD
+        if self.first_layer_in_block:
+            block_residual_TND = torch.cat(
+                (block_residual_TND, x_TD.unsqueeze(1)), dim=1
+            )
+            partial_block_TD = None
+        else:
+            partial_block_TD = x_TD
 
-        if block_residual_TND.shape[1] > 0:
-            assert self.attention_res_proj is not None
+        if self.attention_res_proj is None:
+            h_TD = x_TD
+        else:
             assert self.attention_res_norm is not None
-            x_TD = _apply_attention_residual(
-                prefix_sum_TD,
+            h_TD = _apply_attention_residual(
+                partial_block_TD,
                 block_residual_TND,
                 self.attention_res_proj,
                 self.attention_res_norm,
             )
-
-        opens_block = self.layer_id % self.attn_res_block_size == 0
-        if opens_block:
-            block_residual_TND = torch.cat(
-                (
-                    block_residual_TND,
-                    prefix_sum_TD.unsqueeze(1),
-                ),
-                dim=1,
-            )
-
-        h_TD = self.attention_norm(x_TD)
+        h_TD = self.attention_norm(h_TD)
         layer_mask = (
             attention_masks[self.attn_mask_key] if attention_masks is not None else None
         )
@@ -275,7 +277,7 @@ def forward(
         else:
             assert self.delta_attention is not None
             h_TD = self.delta_attention(h_TD, layer_mask, positions)
-        prefix_sum_TD = h_TD if opens_block else prefix_sum_TD + h_TD
+        prefix_sum_TD = h_TD if self.first_layer_in_block else x_TD + h_TD
 
         h_TD = _apply_attention_residual(
             prefix_sum_TD,
@@ -302,7 +304,14 @@ def _register_optimizer_hooks(cls, optimizers, model_parts, parallel_dims) -> No
 
         register_moe_quantile_balancing_hook(optimizers, model_parts, parallel_dims)
 
-    supports_pipeline_parallel = False
+    pipeline_first_stage_module_fqns = ("vision_encoder",)
+    pipeline_last_stage_module_fqns = ("output_res_proj", "output_res_norm")
+
+    def pipeline(self, **kwargs):
+        """Partition the model, then route the attention-residual blocks along the split."""
+        from .pipeline_parallel import pipeline_kimi_k3
+
+        return pipeline_kimi_k3(self, **kwargs)
 
     @dataclass(kw_only=True, slots=True)
     class Config(Decoder.Config):
@@ -384,15 +393,13 @@ def parallelize(
     ) -> KimiK3Model:
         unsupported = [
             name
-            for name, enabled in (
-                ("pipeline parallel", parallel_dims.pp_enabled),
-                ("context parallel", parallel_dims.cp_enabled),
-            )
+            for name, enabled in (("context parallel", parallel_dims.cp_enabled),)
             if enabled
         ]
         if unsupported:
             raise NotImplementedError(
-                "Kimi K3 currently supports FSDP2 data parallelism only; "
+                "Kimi K3 currently supports FSDP2 data parallelism and pipeline "
+                "parallelism only; "
                 f"disable {', '.join(unsupported)}."
             )
         if compile_config is not None and "model" in compile_config.components:
@@ -526,9 +533,10 @@ def _prepare_multimodal_embeds(
             vision_positions=vision_positions,
         )
 
-    def forward(  # pyrefly: ignore [bad-override]
+    def forward(  # pyrefly: ignore[bad-param-name-override, bad-override]
         self,
         tokens: torch.Tensor,
+        block_residual_TND: torch.Tensor | None = None,
         *,
         pixel_values: torch.Tensor | None = None,
         grid_thw: torch.Tensor | None = None,
@@ -538,9 +546,10 @@ def forward(  # pyrefly: ignore [bad-override]
         positions: torch.Tensor | None = None,
         attention_masks: KimiK3AttentionMaskDict | None = None,
         padding_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if pixel_values_videos is not None or grid_thw_videos is not None:
             raise NotImplementedError("Kimi K3 v1 supports images but not videos.")
+
         if self.tok_embeddings is not None:
             with spmd_local_context("dp"):
                 h_TD = self._prepare_multimodal_embeds(
@@ -555,7 +564,8 @@ def forward(  # pyrefly: ignore [bad-override]
         if spmd.is_type_checking():
             spmd.assert_type(h_TD, {MeshAxisName.DP: spmd.S(0)})
 
-        block_residual_TND = h_TD.unsqueeze(1)[:, :0]
+        if block_residual_TND is None:
+            block_residual_TND = h_TD.unsqueeze(1)[:, :0]
         for layer in self.layers.values():
             h_TD, block_residual_TND = layer(
                 h_TD,
@@ -565,6 +575,8 @@ def forward(  # pyrefly: ignore [bad-override]
                 padding_mask=padding_mask,
             )
 
+        if self.output_res_proj is None:
+            return h_TD, block_residual_TND
         h_TD = _apply_attention_residual(
             h_TD,
             block_residual_TND,
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py b/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
new file mode 100644
index 00000000000..8f3f5409846
--- /dev/null
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
@@ -0,0 +1,121 @@
+# 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.
+
+"""Pipeline parallelism for Kimi K3: core's split with the tower and the aggregation
+pinned to its ends, AttnRes stages, and the block routing tables."""
+
+import logging
+
+from torch.distributed.pipelining.schedules import (
+    _PipelineSchedule,
+    PipelineScheduleMulti,
+    PipelineScheduleSingle,
+)
+from torch.distributed.pipelining.stage import _PipelineStageBase, PipelineStage
+
+from torchtitan.distributed.pipeline_parallel import (
+    pipeline_with_first_last_stage_modules,
+)
+from torchtitan.protocols.model import BaseModel
+
+from .layout import infer_block_layout_tables, layer_to_stage_from_split
+from .stage import AttnResPipelineStage, PPRankLocalCache
+
+__all__ = ["pipeline_kimi_k3"]
+
+logger = logging.getLogger(__name__)
+
+
+def _as_attn_res_stage(stage: _PipelineStageBase) -> AttnResPipelineStage:
+    assert isinstance(stage, PipelineStage)
+    rebuilt = AttnResPipelineStage(
+        stage.submod,
+        stage.stage_index,
+        stage.num_stages,
+        stage.device,
+        group=stage.group,
+        dw_builder=stage.dw_builder,
+        get_mesh=stage._mesh_cache._get_mesh_cb,
+    )
+    # The schedule wrote its stage-to-rank map onto the stage it was handed.
+    rebuilt.stage_index_to_group_rank = stage.stage_index_to_group_rank
+    return rebuilt
+
+
+def _swap_in_attn_res_stages(
+    schedule: _PipelineSchedule,
+) -> list[AttnResPipelineStage]:
+    if isinstance(schedule, PipelineScheduleSingle):
+        rebuilt = _as_attn_res_stage(schedule._stage)
+        schedule._stage = rebuilt
+        return [rebuilt]
+    if isinstance(schedule, PipelineScheduleMulti):
+        rebuilt_stages = [_as_attn_res_stage(s) for s in schedule._stages]
+        held: list[_PipelineStageBase] = list(rebuilt_stages)
+        schedule._stages = held
+        return rebuilt_stages
+    raise RuntimeError(f"Unexpected pipeline schedule class {type(schedule).__name__}.")
+
+
+def _require_loop_style(
+    schedule: _PipelineSchedule, stage_to_rank: dict[int, int], pp: int
+) -> None:
+    """The rank store keeps a block for the rank's later stages, so the cached transport
+    needs the loop-style assignment, stage s on rank s % pp; any other one is refused."""
+    for stage in sorted(stage_to_rank):
+        rank = stage_to_rank[stage]
+        if rank != stage % pp:
+            raise ValueError(
+                f"{type(schedule).__name__} is unsupported with attn_res_cache: stage "
+                f"{stage} sits on rank {rank}, not on rank {stage % pp} of the "
+                "loop-style assignment the rank store assumes. Use a looped schedule "
+                "such as Interleaved1F1B, or turn attn_res_cache off."
+            )
+
+
+def pipeline_kimi_k3(model: BaseModel, *, attn_res_cache: bool = True, **kwargs):
+    """pipelining_fn for Kimi K3; with attn_res_cache a hop carries only the blocks the
+    receiving rank lacks, without it the whole stack, and every rank must agree."""
+    (
+        pp_schedule,
+        model_parts,
+        has_first_stage,
+        has_last_stage,
+        split,
+    ) = pipeline_with_first_last_stage_modules(
+        model,
+        first_stage_module_fqns=model.pipeline_first_stage_module_fqns,
+        last_stage_module_fqns=model.pipeline_last_stage_module_fqns,
+        return_split=True,
+        **kwargs,
+    )
+
+    stages = _swap_in_attn_res_stages(pp_schedule)
+    stage_to_rank = dict(stages[0].stage_index_to_group_rank)
+    if attn_res_cache:
+        _require_loop_style(pp_schedule, stage_to_rank, kwargs["parallel_dims"].pp)
+    model_config = kwargs["model_config"]
+    layer_cfgs = model_config.layers
+    n_layers = len(layer_cfgs)
+    layers_per_block = layer_cfgs[0].attn_res_block_size
+    layer_to_stage = layer_to_stage_from_split(split)
+    layout = infer_block_layout_tables(
+        stage_to_rank=stage_to_rank,
+        n_layers=n_layers,
+        layers_per_block=layers_per_block,
+        layer_to_stage=layer_to_stage,
+        cache=attn_res_cache,
+    )
+    store = PPRankLocalCache()
+    for stage in stages:
+        stage.set_routing(layout, store)
+    logger.info(
+        "Kimi K3 pipeline: %d stage(s) on this rank %s, block transport %s",
+        len(stages),
+        [s.stage_index for s in stages],
+        "delta with rank store" if attn_res_cache else "whole stack every hop",
+    )
+    return pp_schedule, model_parts, has_first_stage, has_last_stage
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/layout.py b/torchtitan/models/kimi_k3/pipeline_parallel/layout.py
new file mode 100644
index 00000000000..3fdabd9bff0
--- /dev/null
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/layout.py
@@ -0,0 +1,159 @@
+# 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.
+
+"""Routing tables for the block attention residual across pipeline stages."""
+
+from collections.abc import Sequence
+
+
+class BlockLayoutTables:
+    """Per-stage routing of the block stack for one micro-batch.
+
+    Stages are walked in index order and a rank keeps what its earlier stages
+    committed, which is the loop-style stage-to-rank assignment; the v-shaped
+    schedules are not supported, and ``pipeline_kimi_k3`` refuses them.
+    """
+
+    def __init__(
+        self,
+        *,
+        stage_to_rank: dict[int, int],
+        n_layers: int,
+        layers_per_block: int,
+        layer_to_stage: dict[int, int],
+        cache: bool = True,
+    ) -> None:
+        if n_layers <= 0 or layers_per_block <= 0:
+            raise ValueError("n_layers and layers_per_block must be positive")
+        self.num_stages = len(stage_to_rank)
+        if sorted(stage_to_rank) != list(range(self.num_stages)):
+            raise ValueError(
+                f"stage_to_rank must cover stages 0..{self.num_stages - 1}; "
+                f"got {sorted(stage_to_rank)}"
+            )
+        self.stage_to_rank = dict(stage_to_rank)
+        # a ceiling: the last block may be partial
+        self.num_blocks = -(-n_layers // layers_per_block)
+        self.n_layers = n_layers
+        self.layers_per_block = layers_per_block
+        self.cache = cache
+        self._layer_to_stage = dict(layer_to_stage)
+        self._commits_at: dict[int, list[int]] = {}
+        self._producer_stage_of_block: dict[int, int] = {}
+        self._cache_at_entry: dict[int, frozenset[int]] = {}
+        self._delta_to_send: dict[int, list[int]] = {}
+        self._cache_readers: dict[int, list[int]] = {}
+        self._build()
+
+    def commits_at(self, stage_id: int) -> list[int]:
+        return list(self._commits_at.get(stage_id, ()))
+
+    def cache_at_entry(self, stage_id: int) -> frozenset[int]:
+        return self._cache_at_entry[stage_id]
+
+    def delta_to_send(self, stage_id: int) -> list[int]:
+        return list(self._delta_to_send.get(stage_id, ()))
+
+    def producer_stage_of_block(self, block_idx: int) -> int:
+        return self._producer_stage_of_block[block_idx]
+
+    def cache_readers_of_block(self, block_idx: int) -> list[int]:
+        return list(self._cache_readers.get(block_idx, ()))
+
+    def deposits_expected(self, block_idx: int, owner_stage: int) -> int:
+        """Deposits the stage that brought ``block_idx`` onto the rank collects: one per later reader there."""
+        rank = self.stage_to_rank[owner_stage]
+        return sum(
+            1
+            for reader in self._cache_readers.get(block_idx, ())
+            if self.stage_to_rank[reader] == rank and reader > owner_stage
+        )
+
+    def _build(self) -> None:
+        for stage_id in range(self.num_stages):
+            self._commits_at[stage_id] = []
+        for ell in range(self.n_layers):
+            if ell % self.layers_per_block != 0:
+                continue
+            block_idx = ell // self.layers_per_block
+            stage_id = self._layer_to_stage[ell]
+            self._commits_at[stage_id].append(block_idx)
+            self._producer_stage_of_block[block_idx] = stage_id
+
+        held: dict[int, set[int]] = {r: set() for r in set(self.stage_to_rank.values())}
+        accumulated: set[int] = set()
+        for stage_id in range(self.num_stages):
+            rank = self.stage_to_rank[stage_id]
+            self._cache_at_entry[stage_id] = frozenset(held[rank])
+            accumulated.update(self._commits_at[stage_id])
+            if self.cache:
+                held[rank].update(accumulated)
+            next_stage = stage_id + 1
+            if next_stage < self.num_stages:
+                receiver = held[self.stage_to_rank[next_stage]]
+                self._delta_to_send[stage_id] = sorted(accumulated - receiver)
+            else:
+                self._delta_to_send[stage_id] = []
+
+        readers: dict[int, list[int]] = {b: [] for b in range(self.num_blocks)}
+        for stage_id in range(self.num_stages):
+            for b in sorted(self._cache_at_entry[stage_id]):
+                readers[b].append(stage_id)
+        self._cache_readers = readers
+
+
+def infer_block_layout_tables(
+    *,
+    stage_to_rank: dict[int, int],
+    n_layers: int,
+    layers_per_block: int,
+    layer_to_stage: dict[int, int],
+    cache: bool = True,
+) -> BlockLayoutTables:
+    """Build the tables; every layer must sit on one stage, in contiguous runs."""
+    num_stages = len(stage_to_rank)
+    if sorted(layer_to_stage) != list(range(n_layers)):
+        raise ValueError(
+            f"layer_to_stage must cover layers 0..{n_layers - 1} exactly once; "
+            f"got {sorted(layer_to_stage)}"
+        )
+    previous = -1
+    for layer_id in range(n_layers):
+        stage_idx = layer_to_stage[layer_id]
+        if not 0 <= stage_idx < num_stages:
+            raise ValueError(
+                f"layer {layer_id} sits on stage {stage_idx}, outside the "
+                f"{num_stages} stages of this pipeline"
+            )
+        if stage_idx < previous:
+            raise ValueError(
+                f"layer {layer_id} sits on stage {stage_idx} after layer "
+                f"{layer_id - 1} on stage {previous}. A non-contiguous "
+                "pipeline split is not supported: the block routing would "
+                "carry deltas to the wrong stages."
+            )
+        previous = stage_idx
+    return BlockLayoutTables(
+        stage_to_rank=stage_to_rank,
+        n_layers=n_layers,
+        layers_per_block=layers_per_block,
+        layer_to_stage=layer_to_stage,
+        cache=cache,
+    )
+
+
+def layer_to_stage_from_split(
+    module_fqns_per_model_part: Sequence[Sequence[str]],
+) -> dict[int, int]:
+    """The layer-to-stage map, read off the split core applies."""
+    layer_to_stage: dict[int, int] = {}
+    for stage_idx, names in enumerate(module_fqns_per_model_part):
+        for name in names:
+            prefix, _, layer = name.partition(".")
+            if prefix != "layers" or not layer.isdigit():
+                continue
+            layer_to_stage[int(layer)] = stage_idx
+    return layer_to_stage
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/stage.py b/torchtitan/models/kimi_k3/pipeline_parallel/stage.py
new file mode 100644
index 00000000000..c16855a87d9
--- /dev/null
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/stage.py
@@ -0,0 +1,364 @@
+# 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.
+
+"""Pipeline stage carrying the block attention residual across hops.
+
+Suffixes: T tokens, N blocks, D model dim.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch.distributed.pipelining import PipelineStage
+from torch.distributed.pipelining._utils import flatten_args
+
+from .layout import BlockLayoutTables
+
+
+class PPRankLocalCache:
+    """Blocks a rank holds per micro-batch and the gradient deposits, shared by its stages."""
+
+    def __init__(self) -> None:
+        self._blocks: dict[int, dict[int, torch.Tensor]] = {}
+        self._deposits: dict[tuple[int, int], torch.Tensor] = {}
+        self._counts: dict[tuple[int, int], int] = {}
+
+    def put(self, mb: int, block_idx: int, block_TD: torch.Tensor) -> None:
+        self._blocks.setdefault(mb, {})[block_idx] = block_TD
+
+    def blocks(self, mb: int) -> dict[int, torch.Tensor]:
+        return self._blocks.get(mb, {})
+
+    def release(self, mb: int) -> None:
+        """Free the blocks of ``mb``; the deposits stay until collected."""
+        self._blocks.pop(mb, None)
+
+    def deposit(self, mb: int, block_idx: int, grad_TD: torch.Tensor) -> None:
+        key = (mb, block_idx)
+        prior = self._deposits.get(key)
+        self._deposits[key] = grad_TD.clone() if prior is None else prior + grad_TD
+        self._counts[key] = self._counts.get(key, 0) + 1
+
+    def collect(self, mb: int, block_idx: int) -> tuple[torch.Tensor | None, int]:
+        key = (mb, block_idx)
+        return self._deposits.pop(key, None), self._counts.pop(key, 0)
+
+    def has_deposits(self, mb: int) -> bool:
+        return any(key[0] == mb for key in self._deposits)
+
+
+def assemble_stack(
+    hidden_TD: torch.Tensor,
+    delta_TND: torch.Tensor,
+    delta_blocks: list[int],
+    store_blocks: dict[int, torch.Tensor],
+) -> tuple[torch.Tensor, list[int]]:
+    """The block stack as a fresh autograd leaf, with each column's block index."""
+    if delta_TND.shape[1] != len(delta_blocks):
+        raise ValueError(
+            f"received {delta_TND.shape[1]} block(s) but the routing expects "
+            f"{delta_blocks}"
+        )
+    order = sorted(set(delta_blocks) | set(store_blocks))
+    pieces = [
+        store_blocks[b] if b in store_blocks else delta_TND[:, delta_blocks.index(b)]
+        for b in order
+    ]
+    if pieces:
+        stack_TND = torch.stack(pieces, dim=1)
+    else:
+        stack_TND = hidden_TD.new_zeros(hidden_TD.shape[0], 0, hidden_TD.shape[-1])
+    return stack_TND.detach().requires_grad_(True), order
+
+
+def pack_outgoing_delta(
+    stack_out_TND: torch.Tensor,
+    order_out: list[int],
+    out_blocks: list[int],
+) -> torch.Tensor:
+    """The blocks the next hop carries, as views of the model's stack."""
+    if stack_out_TND.shape[1] != len(order_out):
+        raise ValueError(
+            f"the model returned {stack_out_TND.shape[1]} block(s); the routing "
+            f"expects {len(order_out)} ({order_out})"
+        )
+    pieces = [stack_out_TND[:, order_out.index(b)] for b in out_blocks]
+    if pieces:
+        return torch.stack(pieces, dim=1)
+    num_tokens, _, dim = stack_out_TND.shape
+    return stack_out_TND.new_zeros(num_tokens, 0, dim)
+
+
+def split_stack_grad(
+    grad_stack_TND: torch.Tensor | None,
+    order: list[int],
+    delta_blocks: list[int],
+    like_TD: torch.Tensor,
+) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
+    """Split a stack gradient into the received columns, in wire order, and the stored blocks' deposits."""
+    num_tokens, dim = like_TD.shape[0], like_TD.shape[-1]
+    grad_delta = like_TD.new_zeros(num_tokens, len(delta_blocks), dim)
+    deposits: dict[int, torch.Tensor] = {}
+    if grad_stack_TND is None:
+        return grad_delta, deposits
+    for col, b in enumerate(order):
+        if b in delta_blocks:
+            grad_delta[:, delta_blocks.index(b)] = grad_stack_TND[:, col]
+        else:
+            deposits[b] = grad_stack_TND[:, col]
+    return grad_delta, deposits
+
+
+class AttnResPipelineStage(PipelineStage):
+    """``PipelineStage`` whose hops carry the block residual's delta."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._layout: BlockLayoutTables | None = None
+        self._store: PPRankLocalCache | None = None
+        # per micro-batch: the stack's block order and the blocks the delta carried in
+        self._order: dict[int, list[int]] = {}
+        self._delta_in: dict[int, list[int]] = {}
+
+    def set_routing(self, layout: BlockLayoutTables, store: PPRankLocalCache) -> None:
+        self._layout = layout
+        self._store = store
+
+    def layout(self) -> BlockLayoutTables:
+        if self._layout is None:
+            raise RuntimeError(
+                f"stage {self.stage_index}: set_routing() must run first"
+            )
+        return self._layout
+
+    def store(self) -> PPRankLocalCache:
+        if self._store is None:
+            raise RuntimeError(
+                f"stage {self.stage_index}: set_routing() must run first"
+            )
+        return self._store
+
+    def _is_first_on_rank(self) -> bool:
+        layout = self.layout()
+        mine = [s for s, r in layout.stage_to_rank.items() if r == self.group_rank]
+        return self.stage_index == min(mine)
+
+    def _is_last_on_rank(self) -> bool:
+        layout = self.layout()
+        mine = [s for s, r in layout.stage_to_rank.items() if r == self.group_rank]
+        return self.stage_index == max(mine)
+
+    def _assemble(
+        self, mb: int, hidden_TD: torch.Tensor, delta_TND: torch.Tensor
+    ) -> torch.Tensor:
+        layout, store = self.layout(), self.store()
+        delta_blocks = layout.delta_to_send(self.stage_index - 1)
+        expected = layout.cache_at_entry(self.stage_index)
+        held = store.blocks(mb)
+        if set(held) != set(expected):
+            raise RuntimeError(
+                f"stage {self.stage_index} micro-batch {mb}: the store holds "
+                f"blocks {sorted(held)} but the routing expects {sorted(expected)}"
+            )
+        stack_TND, order = assemble_stack(hidden_TD, delta_TND, delta_blocks, held)
+        if layout.cache:
+            # Keep what arrived for the rank's later stages.
+            for i, b in enumerate(delta_blocks):
+                store.put(mb, b, delta_TND[:, i].detach())
+        self._order[mb] = order
+        self._delta_in[mb] = delta_blocks
+        return stack_TND
+
+    def _commit_and_route(
+        self, mb: int, stack_out_TND: torch.Tensor, order_in: list[int]
+    ) -> torch.Tensor:
+        layout, store = self.layout(), self.store()
+        my_commits = layout.commits_at(self.stage_index)
+        order_out = order_in + my_commits
+        if layout.cache:
+            for i, b in enumerate(my_commits):
+                store.put(mb, b, stack_out_TND[:, len(order_in) + i].detach())
+        return pack_outgoing_delta(
+            stack_out_TND, order_out, layout.delta_to_send(self.stage_index)
+        )
+
+    def forward_one_chunk(
+        self,
+        fwd_chunk_id: int,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any] | None = None,
+        save_forward_output: bool = True,
+    ):
+        store = self.store()
+        if self.is_first:
+            composite_args: tuple[Any, ...] = args
+            order_in: list[int] = []
+        else:
+            hidden_TD, delta_TND = self._retrieve_recv_activations(fwd_chunk_id)
+            stack_TND = self._assemble(fwd_chunk_id, hidden_TD, delta_TND)
+            composite_args = (hidden_TD, stack_TND)
+            order_in = self._order[fwd_chunk_id]
+        composite_kwargs = kwargs or {}
+
+        output = self.forward_maybe_with_nosync(*composite_args, **composite_kwargs)
+
+        if self.is_last:
+            output_tuple = (
+                (output,) if isinstance(output, torch.Tensor) else tuple(output)
+            )
+            if save_forward_output:
+                self.output_chunks.append(output)
+        else:
+            hidden_out_TD, stack_out_TND = output
+            payload_TND = self._commit_and_route(fwd_chunk_id, stack_out_TND, order_in)
+            output_tuple = (hidden_out_TD, payload_TND)
+
+        # flatten_args returns a list with detach=False; lists keep the checker on that overload.
+        flatten_input_tensors: list[torch.Tensor] = list(
+            flatten_args(composite_args)
+        ) + list(flatten_args(composite_kwargs))
+        self.fwd_cache[fwd_chunk_id] = (output_tuple, flatten_input_tensors)
+
+        if self._is_last_on_rank():
+            store.release(fwd_chunk_id)
+        return output
+
+    def _retrieve_recv_grads(self, bwd_chunk_id: int):
+        grads = super()._retrieve_recv_grads(bwd_chunk_id)
+        if self.is_last:
+            return grads
+        layout = self.layout()
+        grad_hidden, grad_delta = grads
+        mine = set(layout.commits_at(self.stage_index))
+        out_blocks = layout.delta_to_send(self.stage_index)
+        committed = [j for j, b in enumerate(out_blocks) if b in mine]
+        if not committed:
+            return (grad_hidden, grad_delta)
+        if grad_delta is None:
+            outputs_meta = self._stage_meta.outputs
+            if outputs_meta is not None and not outputs_meta[1].requires_grad:
+                # Nothing upstream of the payload's blocks is trainable (a frozen embedding under LoRA):
+                # no gradient channel, nowhere for the deposits to go.
+                for j in committed:
+                    self._collect_into(None, bwd_chunk_id, out_blocks[j])
+                return (grad_hidden, None)
+            raise RuntimeError(
+                f"stage {self.stage_index} micro-batch {bwd_chunk_id}: no gradient "
+                f"arrived for the payload carrying its own blocks "
+                f"{[out_blocks[j] for j in committed]}"
+            )
+        grad_delta = grad_delta.clone()
+        for j in committed:
+            self._collect_into(grad_delta[:, j], bwd_chunk_id, out_blocks[j])
+        return (grad_hidden, grad_delta)
+
+    def _collect_into(self, grad_col_TD: torch.Tensor | None, mb: int, b: int) -> None:
+        layout, store = self.layout(), self.store()
+        deposit, count = store.collect(mb, b)
+        expected = layout.deposits_expected(b, self.stage_index)
+        if count != expected:
+            raise RuntimeError(
+                f"stage {self.stage_index} micro-batch {mb} block {b}: "
+                f"{count} gradient deposit(s) but {expected} expected; a "
+                "later stage on this rank did not run its backward"
+            )
+        if deposit is not None and grad_col_TD is not None:
+            grad_col_TD.add_(deposit)
+
+    def backward_one_chunk(
+        self,
+        bwd_chunk_id: int,
+        loss=None,
+        full_backward: bool = True,
+        last_backward=False,
+    ):
+        super().backward_one_chunk(
+            bwd_chunk_id,
+            loss=loss,
+            full_backward=full_backward,
+            last_backward=last_backward,
+        )
+        if not self.has_backward:
+            # Forward-only pass (schedule.eval): no backward ran; drop the forward's bookkeeping.
+            self.fwd_cache.pop(bwd_chunk_id, None)
+            self._order.pop(bwd_chunk_id, None)
+            self._delta_in.pop(bwd_chunk_id, None)
+            return
+        if self.is_first:
+            return
+        store = self.store()
+        grad_hidden, grad_stack = self.bwd_cache[bwd_chunk_id]
+        order = self._order.pop(bwd_chunk_id)
+        delta_blocks = self._delta_in.pop(bwd_chunk_id)
+        like = grad_hidden if grad_hidden is not None else grad_stack
+        if like is None:
+            raise RuntimeError(
+                f"stage {self.stage_index}: backward produced no gradient for "
+                "either input"
+            )
+        grad_delta, deposits = split_stack_grad(grad_stack, order, delta_blocks, like)
+        for b, grad_TD in deposits.items():
+            store.deposit(bwd_chunk_id, b, grad_TD)
+        for j, b in enumerate(delta_blocks):
+            self._collect_into(grad_delta[:, j], bwd_chunk_id, b)
+        # Whether the previous stage expects a delta gradient comes from the receive metadata:
+        # the assembled stack is a detached leaf, so autograd cannot tell.
+        inputs_meta = self._stage_meta.inputs
+        if (
+            inputs_meta is None
+            or len(inputs_meta) != 2
+            or inputs_meta[1] is None
+            or len(inputs_meta[1].shape) != 3
+        ):
+            raise RuntimeError(
+                f"stage {self.stage_index}: the receive metadata should describe "
+                f"(hidden, delta) with a [T, N, D] delta; got {inputs_meta}"
+            )
+        delta_needs_grad = inputs_meta[1].requires_grad
+        self.bwd_cache[bwd_chunk_id] = (
+            grad_hidden.contiguous() if grad_hidden is not None else None,
+            grad_delta if delta_needs_grad else None,
+        )
+        if self._is_first_on_rank() and store.has_deposits(bwd_chunk_id):
+            raise RuntimeError(
+                f"stage {self.stage_index} micro-batch {bwd_chunk_id}: gradient "
+                "deposits left uncollected after the rank's last backward"
+            )
+
+    def _compute_outputs(
+        self, *args: torch.Tensor, module: torch.nn.Module, **kwargs: Any
+    ):
+        layout = self.layout()
+        if self.is_first:
+            output = module(*args, **kwargs)
+            order_in: list[int] = []
+        else:
+            hidden_TD, delta_TND = args
+            delta_blocks = layout.delta_to_send(self.stage_index - 1)
+            held = {
+                b: hidden_TD.new_zeros(hidden_TD.shape)
+                for b in layout.cache_at_entry(self.stage_index)
+            }
+            stack_TND, order_in = assemble_stack(
+                hidden_TD, delta_TND, delta_blocks, held
+            )
+            output = module(hidden_TD, stack_TND, **kwargs)
+        if self.is_last:
+            return output
+        hidden_out_TD, stack_out_TND = output
+        order_out = order_in + layout.commits_at(self.stage_index)
+        payload_TND = pack_outgoing_delta(
+            stack_out_TND, order_out, layout.delta_to_send(self.stage_index)
+        )
+        return hidden_out_TD, payload_TND
+
+    def _compute_input_grads(self, outputs, all_fwd_inputs, grad_outputs=None):
+        grads = super()._compute_input_grads(outputs, all_fwd_inputs, grad_outputs)
+        return tuple(
+            g.contiguous() if isinstance(g, torch.Tensor) else g for g in grads
+        )

From f14dbb29f9b01be3eddb92af29ea1960b1dad430 Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:10:22 +0000
Subject: [PATCH 4/9] kimi_k3: the debug model is 17 layers at debug widths
 with a 2048 vocabulary

The flavor carried the released vocabulary and 24 layers at model widths,
which is neither cheap nor a debug shape. 17 is the least depth in the
model's own (3 KDA + 1 MLA) pattern whose units split into the 16 stages a
pp4 x vp4 cell asks for.
---
 torchtitan/models/kimi_k3/__init__.py | 36 +++++++++++++--------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py
index c9202f9ec7c..9b614fa8d37 100644
--- a/torchtitan/models/kimi_k3/__init__.py
+++ b/torchtitan/models/kimi_k3/__init__.py
@@ -490,36 +490,36 @@ def _debugmodel(
     *,
     seq_len: int,
 ) -> KimiK3Model.Config:
-    dim = 1024
+    dim = 256
     return _kimi_k3_config(
         max_context_length=seq_len,
         dim=dim,
         moe_comm_backend=moe_comm_backend,
-        vocab_size=163840,
-        num_layers=24,
-        full_attention_layers={3, 7, 11, 15, 19, 23},
-        attn_res_block_size=12,
-        num_heads=16,
-        q_lora_rank=512,
-        kv_lora_rank=256,
+        vocab_size=2048,
+        num_layers=17,
+        full_attention_layers={3, 7, 11, 15, 16},
+        attn_res_block_size=4,
+        num_heads=4,
+        q_lora_rank=128,
+        kv_lora_rank=64,
         qk_nope_head_dim=64,
         qk_rope_head_dim=32,
         v_head_dim=64,
         kda_head_dim=128,
         conv_kernel_size=4,
-        dense_hidden_dim=4096,
-        latent_dim=512,
-        expert_hidden_dim=384,
-        num_experts=32,
-        top_k=4,
+        dense_hidden_dim=512,
+        latent_dim=128,
+        expert_hidden_dim=128,
+        num_experts=8,
+        top_k=2,
         num_shared_experts=2,
         vision_encoder=_vision_encoder_config(
             text_dim=dim,
-            dim=512,
-            qkv_dim=768,
-            hidden_dim=2048,
-            num_layers=8,
-            num_heads=6,
+            dim=256,
+            qkv_dim=512,
+            hidden_dim=512,
+            num_layers=2,
+            num_heads=4,
             init_pos_emb_height=32,
             init_pos_emb_width=32,
         ),

From fb019c17e5566b095acb923d36c1e69639fab7fe Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:10:22 +0000
Subject: [PATCH 5/9] kimi_k3: the pipeline cell in the B200 suite is pp4 x vp4

One recipe and one cell, as the llama3 pipeline cells are, with the split
spelled out because the debug model's units do not divide into 16 stages.
Type checking is off, which the pipeline recipes upstream also do.
---
 tests/integration_tests/b200.py               |  6 ++++++
 .../cpu/test_integration_test_definitions.py  |  1 +
 torchtitan_recipes/tests/b200.py              | 19 +++++++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/tests/integration_tests/b200.py b/tests/integration_tests/b200.py
index 4bfd566592b..93439d1699b 100644
--- a/tests/integration_tests/b200.py
+++ b/tests/integration_tests/b200.py
@@ -24,6 +24,12 @@ def build_b200_tests_list() -> list[OverrideDefinitions]:
             test_name="kimi_k3_mm_muon",
             ngpu=2,
         ),
+        OverrideDefinitions(
+            configs=[recipes.kimi_k3_debugmodel_pp4_vp4],
+            test_descr="Kimi K3 pipeline parallel pp4 x vp4",
+            test_name="kimi_k3_pp4_vp4",
+            ngpu=4,
+        ),
         OverrideDefinitions(
             configs=[recipes.llama3_debugmodel_mxfp8_fsdp2],
             test_descr="MXFP8 linear with an FSDP-managed weight cache",
diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py
index eb12468539a..a5bb5c558bb 100644
--- a/tests/unit_tests/cpu/test_integration_test_definitions.py
+++ b/tests/unit_tests/cpu/test_integration_test_definitions.py
@@ -131,6 +131,7 @@ def test_b200_tests_are_registered_in_separate_suite() -> None:
     assert {test.test_name for test in build_b200_tests_list()} == {
         "kimi_k3_mm",
         "kimi_k3_mm_muon",
+        "kimi_k3_pp4_vp4",
         "mxfp8_linear_fsdp",
         "nvfp4_linear_fsdp",
     }
diff --git a/torchtitan_recipes/tests/b200.py b/torchtitan_recipes/tests/b200.py
index 01eac209408..b9ffdf8120a 100644
--- a/torchtitan_recipes/tests/b200.py
+++ b/torchtitan_recipes/tests/b200.py
@@ -55,3 +55,22 @@ def llama3_debugmodel_nvfp4_fsdp2() -> Trainer.Config:
     config.parallelism.data_parallel_shard_degree = 2
     config.training.num_tokens_per_microbatch_per_dp_rank = 2048
     return config
+
+
+def kimi_k3_debugmodel_pp4_vp4() -> Trainer.Config:
+    from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel
+
+    config = kimi_k3_debugmodel()
+    _set_spmd_typechecking(config, typechecking=False)
+    config.parallelism.pipeline_parallel_degree = 4
+    config.parallelism.pipeline_parallel_schedule = "Interleaved1F1B"
+    config.parallelism.num_pp_microbatches = 4
+    # 16 stages: one layer per stage from layer 5 on, the head alone on the last.
+    config.parallelism.pipeline_parallel_module_fqns_per_model_part = [
+        ["vision_encoder", "tok_embeddings", "layers.0"],
+        ["layers.1", "layers.2"],
+        ["layers.3", "layers.4"],
+        *[[f"layers.{i}"] for i in range(5, 17)],
+        ["norm", "lm_head", "output_res_proj", "output_res_norm"],
+    ]
+    return config

From 267997a2bc5486ce6123ba5d7690411388ffb76a Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Wed, 23 Sep 2026 05:42:28 +0000
Subject: [PATCH 6/9] kimi_k3: the composability cell in the B200 suite is
 fsdp2 x tp2 x ep2 x pp2

One eight-GPU cell with the four parallelisms together, as the reviewer asked
for: FSDP 2, TP 2 with sequence parallel, EP 2 and PP 2 under 1F1B with four
micro-batches. Type checking is off, which the pipeline recipes have in common.
---
 tests/integration_tests/b200.py                 |  6 ++++++
 .../cpu/test_integration_test_definitions.py    |  1 +
 torchtitan_recipes/tests/b200.py                | 17 +++++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/tests/integration_tests/b200.py b/tests/integration_tests/b200.py
index 93439d1699b..847a63d8919 100644
--- a/tests/integration_tests/b200.py
+++ b/tests/integration_tests/b200.py
@@ -30,6 +30,12 @@ def build_b200_tests_list() -> list[OverrideDefinitions]:
             test_name="kimi_k3_pp4_vp4",
             ngpu=4,
         ),
+        OverrideDefinitions(
+            configs=[recipes.kimi_k3_debugmodel_fsdp2_tp2_ep2_pp2],
+            test_descr="Kimi K3 FSDP, TP, EP and pipeline parallel together",
+            test_name="kimi_k3_fsdp2_tp2_ep2_pp2",
+            ngpu=8,
+        ),
         OverrideDefinitions(
             configs=[recipes.llama3_debugmodel_mxfp8_fsdp2],
             test_descr="MXFP8 linear with an FSDP-managed weight cache",
diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py
index a5bb5c558bb..fb5308369b1 100644
--- a/tests/unit_tests/cpu/test_integration_test_definitions.py
+++ b/tests/unit_tests/cpu/test_integration_test_definitions.py
@@ -129,6 +129,7 @@ def test_h100_tests_are_registered_in_separate_suite() -> None:
 
 def test_b200_tests_are_registered_in_separate_suite() -> None:
     assert {test.test_name for test in build_b200_tests_list()} == {
+        "kimi_k3_fsdp2_tp2_ep2_pp2",
         "kimi_k3_mm",
         "kimi_k3_mm_muon",
         "kimi_k3_pp4_vp4",
diff --git a/torchtitan_recipes/tests/b200.py b/torchtitan_recipes/tests/b200.py
index b9ffdf8120a..97e0420b549 100644
--- a/torchtitan_recipes/tests/b200.py
+++ b/torchtitan_recipes/tests/b200.py
@@ -74,3 +74,20 @@ def kimi_k3_debugmodel_pp4_vp4() -> Trainer.Config:
         ["norm", "lm_head", "output_res_proj", "output_res_norm"],
     ]
     return config
+
+
+def kimi_k3_debugmodel_fsdp2_tp2_ep2_pp2() -> Trainer.Config:
+    from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel
+
+    config = kimi_k3_debugmodel()
+    # Type checking stays off under pipeline parallelism, as the other pipeline
+    # recipes have it.
+    _set_spmd_typechecking(config, typechecking=False)
+    config.parallelism.data_parallel_shard_degree = 2
+    config.parallelism.tensor_parallel_degree = 2
+    config.parallelism.enable_sequence_parallel = True
+    config.parallelism.expert_parallel_degree = 2
+    config.parallelism.pipeline_parallel_degree = 2
+    config.parallelism.pipeline_parallel_schedule = "1F1B"
+    config.parallelism.num_pp_microbatches = 4
+    return config

From 99894da64b3597acd89a4f9e31ffdadd50e7139a Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Wed, 23 Sep 2026 08:09:44 +0000
Subject: [PATCH 7/9] kimi_k3: the cache transport drawn, the way
 MOE_SHARDING.md draws MoE sharding

One note in the pipeline package with one vector figure: a stage on one rank
with the rank store, forward and backward, and the same stage with the
whole-stack transport, with h, B and delta and their gradients on the arrows
and the words kept for a legend. The rest of the note is the two transports
side by side and what the stage and the tables do, so a reader has the idea
before the code.
---
 assets/images/kimi_k3_pp_attn_res_cache.svg   | 151 ++++++++++++++++++
 .../pipeline_parallel/PP_ATTN_RES_CACHE.md    |  47 ++++++
 2 files changed, 198 insertions(+)
 create mode 100644 assets/images/kimi_k3_pp_attn_res_cache.svg
 create mode 100644 torchtitan/models/kimi_k3/pipeline_parallel/PP_ATTN_RES_CACHE.md

diff --git a/assets/images/kimi_k3_pp_attn_res_cache.svg b/assets/images/kimi_k3_pp_attn_res_cache.svg
new file mode 100644
index 00000000000..a308b6b4222
--- /dev/null
+++ b/assets/images/kimi_k3_pp_attn_res_cache.svg
@@ -0,0 +1,151 @@
+<svg xmlns="http://www.w3.org/2000/svg" font-family="IBM Plex Sans, Helvetica Neue, Arial, sans-serif" font-size="12.5" viewBox="0 0 1460 690" role="img" aria-label="Forward and backward data flow of one Kimi K3 pipeline stage with the block cache on, next to the same stage with the cache off">
+    <defs>
+      <marker id="mf" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="8" markerHeight="8" orient="auto-start-reverse">
+        <path d="M0,0 L10,5 L0,10 z" class="mk-fwd" fill="#1f2328"/>
+      </marker>
+      <marker id="mb" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="8" markerHeight="8" orient="auto-start-reverse">
+        <path d="M0,0 L10,5 L0,10 z" class="mk-bwd" fill="#d92d20"/>
+      </marker>
+      <pattern id="hatch" width="6" height="6" patternUnits="userSpaceOnUse" patternTransform="rotate(45)">
+        <rect width="6" height="6" fill="#eef9f0"/>
+        <line x1="0" y1="0" x2="0" y2="6" class="hatch-line" stroke="#2f9e44" stroke-width="1.2"/>
+      </pattern>
+    </defs>
+
+    <!-- ===================== panel 1: cache on ===================== -->
+    <text x="20" y="30" class="title" fill="#1f2328" font-size="15" font-weight="600">attn_res_cache on</text>
+    <text x="180" y="30" class="sub" fill="#5f6b76" font-size="12">stage s on rank r; a hop carries the blocks the next rank has not seen</text>
+
+    <!-- input / output labels -->
+    <text x="365" y="72" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h, Δ  from stage s−1</text>
+    <text x="365" y="604" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h′, Δ′  to stage s+1</text>
+    <text x="665" y="604" text-anchor="middle" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h′, ∇Δ′  from stage s+1</text>
+    <text x="665" y="72" text-anchor="middle" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h, ∇Δ  to stage s−1</text>
+
+    <!-- forward column -->
+    <rect x="260" y="90" width="210" height="40" rx="8" class="comm" fill="#c6efce" stroke="#2f9e44"/>
+    <text x="365" y="115" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P recv</text>
+
+    <rect x="260" y="190" width="210" height="50" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="365" y="211" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">assemble</text>
+    <text x="365" y="228" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">B = store ⊕ Δ  (leaf)</text>
+
+    <rect x="260" y="300" width="210" height="60" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="365" y="324" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">layers l..m</text>
+    <text x="365" y="343" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">first layer of a block: B ← B ⊕ h</text>
+
+    <rect x="260" y="420" width="210" height="50" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="365" y="441" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">pack Δ′</text>
+    <text x="365" y="458" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">columns the next rank lacks</text>
+
+    <rect x="260" y="530" width="210" height="40" rx="8" class="comm" fill="#c6efce" stroke="#2f9e44"/>
+    <text x="365" y="555" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P send</text>
+
+    <text x="260" y="660" class="cap f" font-size="13" font-weight="600" letter-spacing="0.08em" fill="#d9760a">FORWARD</text>
+
+    <!-- forward arrows -->
+    <line x1="365" y1="130" x2="365" y2="188" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="373" y="164" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h, Δ</text>
+    <line x1="365" y1="240" x2="365" y2="298" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="373" y="274" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h, B</text>
+    <line x1="365" y1="360" x2="365" y2="418" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="373" y="394" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h′, B′</text>
+    <line x1="365" y1="470" x2="365" y2="528" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="373" y="516" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h′, Δ′</text>
+
+    <!-- rank store -->
+    <rect x="20" y="190" width="170" height="170" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="105" y="214" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">rank store</text>
+    <text x="105" y="236" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">blocks B ≤ s−P</text>
+    <text x="105" y="254" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">deposits ∇B</text>
+    <text x="105" y="278" text-anchor="middle" class="sub" fill="#5f6b76" font-size="12">keyed (micro-batch, block)</text>
+    <text x="105" y="296" text-anchor="middle" class="sub" fill="#5f6b76" font-size="12">one per rank, shared by</text>
+    <text x="105" y="312" text-anchor="middle" class="sub" fill="#5f6b76" font-size="12">the rank's stages</text>
+    <text x="105" y="340" text-anchor="middle" class="note" fill="#1f6fb2" font-size="11.5">freed after the rank's last stage</text>
+
+    <line x1="190" y1="208" x2="258" y2="208" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="224" y="203" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">held B</text>
+    <line x1="260" y1="232" x2="192" y2="232" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="224" y="247" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">cache Δ</text>
+    <polyline points="260,330 224,330 224,300 192,300" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="222" y="346" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">commit new B</text>
+
+    <!-- backward column -->
+    <rect x="560" y="530" width="210" height="40" rx="8" class="commh" fill="url(#hatch)" stroke="#2f9e44"/>
+    <text x="665" y="555" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P recv</text>
+
+    <rect x="560" y="420" width="210" height="50" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="665" y="441" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">collect</text>
+    <text x="665" y="458" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">∇Δ′ += deposits of own B</text>
+
+    <rect x="560" y="190" width="210" height="50" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="665" y="211" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">split ∇B</text>
+    <text x="665" y="228" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">Δ part → wire, stored → deposit</text>
+
+    <rect x="560" y="90" width="210" height="40" rx="8" class="commh" fill="url(#hatch)" stroke="#2f9e44"/>
+    <text x="665" y="115" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P send</text>
+
+    <text x="560" y="660" class="cap b" font-size="13" font-weight="600" letter-spacing="0.08em" fill="#d92d20">BACKWARD</text>
+
+    <!-- backward arrows -->
+    <line x1="665" y1="530" x2="665" y2="472" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="673" y="506" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h′, ∇Δ′</text>
+    <line x1="560" y1="445" x2="472" y2="352" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="500" y="418" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">backward</text>
+    <line x1="470" y1="312" x2="558" y2="220" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="478" y="258" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h, ∇B</text>
+    <line x1="665" y1="190" x2="665" y2="132" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="673" y="164" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h, ∇Δ</text>
+
+    <!-- store <-> backward: deposits in (from split), deposits out (to collect) -->
+    <polyline points="560,200 520,200 520,160 105,160 105,188" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="480" y="154" text-anchor="end" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">deposit stored ∇B</text>
+    <polyline points="105,360 105,500 600,500 600,472" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="230" y="494" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">deposits, counted against the tables</text>
+
+    <!-- notes -->
+    <text x="790" y="323" class="note" fill="#1f6fb2" font-size="11.5">last stage: aggregation</text>
+    <text x="790" y="339" class="note" fill="#1f6fb2" font-size="11.5">on (h, B), no send</text>
+    <text x="790" y="212" class="note" fill="#1f6fb2" font-size="11.5">tables: the same on</text>
+    <text x="790" y="228" class="note" fill="#1f6fb2" font-size="11.5">every rank, no collective</text>
+
+    <!-- divider -->
+    <line x1="960" y1="20" x2="960" y2="670" stroke="#d9d4c7" stroke-dasharray="4 6"/>
+
+    <!-- ===================== panel 2: cache off ===================== -->
+    <text x="1000" y="30" class="title" fill="#1f2328" font-size="15" font-weight="600">attn_res_cache off</text>
+    <text x="1160" y="30" class="sub" fill="#5f6b76" font-size="12">the whole stack on every hop</text>
+
+    <text x="1145" y="72" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h, B  from stage s−1</text>
+    <text x="1145" y="604" text-anchor="middle" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h′, B′  to stage s+1</text>
+    <text x="1370" y="604" text-anchor="middle" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h′, ∇B′  from stage s+1</text>
+    <text x="1370" y="72" text-anchor="middle" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h, ∇B  to stage s−1</text>
+
+    <rect x="1040" y="90" width="210" height="40" rx="8" class="comm" fill="#c6efce" stroke="#2f9e44"/>
+    <text x="1145" y="115" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P recv</text>
+    <rect x="1040" y="300" width="210" height="60" rx="8" class="ent" fill="#ffffff" stroke="#1f2328"/>
+    <text x="1145" y="324" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">layers l..m</text>
+    <text x="1145" y="343" text-anchor="middle" class="sym" fill="#1f2328" font-family="IBM Plex Mono, Consolas, monospace">B is the leaf, nothing assembled</text>
+    <rect x="1040" y="530" width="210" height="40" rx="8" class="comm" fill="#c6efce" stroke="#2f9e44"/>
+    <text x="1145" y="555" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P send</text>
+    <text x="1040" y="660" class="cap f" font-size="13" font-weight="600" letter-spacing="0.08em" fill="#d9760a">FORWARD</text>
+
+    <line x1="1145" y1="130" x2="1145" y2="298" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="1153" y="220" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h, B</text>
+    <line x1="1145" y1="360" x2="1145" y2="528" class="fwd" stroke="#1f2328" marker-end="url(#mf)"/>
+    <text x="1153" y="450" class="lf" fill="#d9760a" font-family="IBM Plex Mono, Consolas, monospace">h′, B′</text>
+
+    <rect x="1290" y="530" width="160" height="40" rx="8" class="commh" fill="url(#hatch)" stroke="#2f9e44"/>
+    <text x="1370" y="555" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P recv</text>
+    <rect x="1290" y="90" width="160" height="40" rx="8" class="commh" fill="url(#hatch)" stroke="#2f9e44"/>
+    <text x="1370" y="115" text-anchor="middle" class="t" fill="#1f2328" font-size="13.5" font-weight="600">P2P send</text>
+    <text x="1290" y="660" class="cap b" font-size="13" font-weight="600" letter-spacing="0.08em" fill="#d92d20">BACKWARD</text>
+
+    <polyline points="1370,530 1370,345 1252,345" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="1378" y="445" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h′, ∇B′</text>
+    <polyline points="1250,315 1330,315 1330,132" class="bwd" stroke="#d92d20" marker-end="url(#mb)"/>
+    <text x="1338" y="225" class="lb" fill="#d92d20" font-family="IBM Plex Mono, Consolas, monospace">∇h, ∇B</text>
+
+    <text x="1040" y="392" class="note" fill="#1f6fb2" font-size="11.5">no store, no deposits;</text>
+    <text x="1040" y="408" class="note" fill="#1f6fb2" font-size="11.5">B grows with the stage index</text>
+  </svg>
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/PP_ATTN_RES_CACHE.md b/torchtitan/models/kimi_k3/pipeline_parallel/PP_ATTN_RES_CACHE.md
new file mode 100644
index 00000000000..24518bfb0ce
--- /dev/null
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/PP_ATTN_RES_CACHE.md
@@ -0,0 +1,47 @@
+# Pipeline parallelism for the block attention residual
+
+Kimi K3 carries a stack of block residuals alongside the hidden state, and a
+pipeline hop has to carry that stack too. The split and the schedule are
+core's, through `pipeline_with_first_last_stage_modules`
+([`__init__.py`](__init__.py)); the stage that moves the stack is
+[`stage.py`](stage.py), and the tables that decide what each hop carries are
+[`layout.py`](layout.py).
+
+## Overview
+
+The figure shows one stage on one rank with the cache on, forward and backward,
+and the same stage with the cache off. `h` is the hidden state, `B` the block
+stack, `Δ` the blocks a hop carries, `∇` a gradient; P2P boxes are green, solid
+forward and hatched backward; forward arrows are black and backward arrows red.
+
+![Kimi K3 pipeline with the attention residual cache](../../../../assets/images/kimi_k3_pp_attn_res_cache.svg)
+
+## The two transports
+
+| `attn_res_cache` | a hop carries | the rank keeps | against a single device |
+|---|---|---|---|
+| on (default) | hidden `[T, D]` and the blocks the receiving rank has not seen, `[T, Nd, D]` | every block its earlier stages committed or received, per micro-batch, in `PPRankLocalCache`, released after its last stage's forward | the same values; the cached blocks' gradients are summed in another order, so not bitwise |
+| off | hidden `[T, D]` and the whole stack `[T, N, D]` | nothing between hops | bitwise |
+
+Plain `1F1B` has one stage per rank, so the rank store never holds anything
+and the two transports are the same hop.
+
+## What the stage does
+
+- `assemble_stack`: the received delta and the held blocks become one leaf
+  `[T, N, D]` in block order, and the stage's model part runs on it.
+- `pack_outgoing_delta`: the columns the next rank lacks, `delta_to_send`,
+  as views of the model's stack.
+- Backward: `split_stack_grad` returns the received columns' gradient to the
+  previous stage in wire order and deposits the held blocks' gradients in the
+  store; the stage that committed a block collects them in
+  `_retrieve_recv_grads`, `deposits_expected` many, else it raises.
+
+## The routing tables
+
+`BlockLayoutTables` is a pure function of the split and the stage-to-rank map,
+the same on every rank: `commits_at(stage)`, `cache_at_entry(stage)`,
+`delta_to_send(stage)` and `deposits_expected(block, stage)`. It walks stages
+in index order and keeps a block on the rank that sees the stage next, the
+loop-style assignment, stage `s` on rank `s % pp`; with the cache on,
+`pipeline_kimi_k3` refuses any other map.

From 3b2746ce0c0c1402cf39f2f6c1878b96959bdb89 Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Wed, 23 Sep 2026 10:12:13 +0000
Subject: [PATCH 8/9] kimi_k3: the pipeline cells train with AdamW under the
 per-head DistMuon recipe

---
 torchtitan_recipes/tests/b200.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/torchtitan_recipes/tests/b200.py b/torchtitan_recipes/tests/b200.py
index 97e0420b549..691d4ff10b1 100644
--- a/torchtitan_recipes/tests/b200.py
+++ b/torchtitan_recipes/tests/b200.py
@@ -73,6 +73,10 @@ def kimi_k3_debugmodel_pp4_vp4() -> Trainer.Config:
         *[[f"layers.{i}"] for i in range(5, 17)],
         ["norm", "lm_head", "output_res_proj", "output_res_norm"],
     ]
+    # DistMuon refuses a stage where a param-group pattern matches nothing (the
+    # head-only stage here) and refuses tensor parallel; the pipeline cells train
+    # with AdamW.
+    config.optimizer = default_adamw(lr=8e-4)
     return config
 
 
@@ -90,4 +94,8 @@ def kimi_k3_debugmodel_fsdp2_tp2_ep2_pp2() -> Trainer.Config:
     config.parallelism.pipeline_parallel_degree = 2
     config.parallelism.pipeline_parallel_schedule = "1F1B"
     config.parallelism.num_pp_microbatches = 4
+    # DistMuon refuses a stage where a param-group pattern matches nothing (the
+    # head-only stage here) and refuses tensor parallel; the pipeline cells train
+    # with AdamW.
+    config.optimizer = default_adamw(lr=8e-4)
     return config

From 949bc6c1c2912dc5d34225c084054c4cf82ae003 Mon Sep 17 00:00:00 2001
From: QIU023 <yiqiaoqiu@hotmail.com>
Date: Tue, 22 Sep 2026 22:13:44 +0000
Subject: [PATCH 9/9] kimi_k3: a pipeline rank can park its committed blocks on
 host memory

The rank store holds every block a rank has committed for a micro-batch
until the stage that reads them runs. With attn_res_cache_offload the store
takes the device it belongs to and parks each block on pinned host memory
on the way in, bringing it back on the way out, which trades the hop's
device memory for two copies.

The device is given rather than inferred from the first block, so a store
that has been handed a host tensor does not start moving every later one.
The switch lives on the model config because the model owns its pipelining,
and it refuses to pair with the whole-stack transport, which keeps nothing
between hops for it to park.
---
 .../cpu/test_kimi_k3_pp_store_offload.py      | 38 +++++++++++++++++++
 torchtitan/models/kimi_k3/model.py            |  2 +
 .../kimi_k3/pipeline_parallel/__init__.py     |  9 ++++-
 .../models/kimi_k3/pipeline_parallel/stage.py | 15 ++++++--
 4 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit_tests/cpu/test_kimi_k3_pp_store_offload.py

diff --git a/tests/unit_tests/cpu/test_kimi_k3_pp_store_offload.py b/tests/unit_tests/cpu/test_kimi_k3_pp_store_offload.py
new file mode 100644
index 00000000000..06d54c60b9c
--- /dev/null
+++ b/tests/unit_tests/cpu/test_kimi_k3_pp_store_offload.py
@@ -0,0 +1,38 @@
+# 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.
+
+"""The rank store's offload path: what it parks and what it hands back."""
+
+import unittest
+
+import torch
+
+from torchtitan.models.kimi_k3.pipeline_parallel.stage import PPRankLocalCache
+
+
+class TestRankStoreOffload(unittest.TestCase):
+    def test_without_a_device_the_store_hands_back_what_it_was_given(self):
+        store = PPRankLocalCache()
+        block = torch.randn(4, 8)
+        store.put(0, 0, block)
+        self.assertIs(store.blocks(0)[0], block)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "the park is a device to host copy")
+    def test_with_a_device_the_block_parks_on_the_host_and_comes_back(self):
+        device = torch.device("cuda", torch.cuda.current_device())
+        store = PPRankLocalCache(device=device)
+        block = torch.randn(4, 8, device=device)
+        store.put(0, 0, block)
+        parked = store._blocks[0][0]
+        self.assertEqual(parked.device.type, "cpu")
+        self.assertTrue(parked.is_pinned())
+        back = store.blocks(0)[0]
+        self.assertEqual(back.device, device)
+        torch.testing.assert_close(back, block)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py
index e82156a3fa3..cd1c1c44cb6 100644
--- a/torchtitan/models/kimi_k3/model.py
+++ b/torchtitan/models/kimi_k3/model.py
@@ -319,6 +319,8 @@ class Config(Decoder.Config):
         output_res_norm: RMSNorm.Config
         output_res_proj: Linear.Config
         vision_encoder: KimiK3VisionEncoder.Config | None = None
+        attn_res_cache_offload: bool = False
+        """Park a pipeline rank's committed blocks on pinned host memory."""
 
         def update_from_config(self, *, config, **kwargs) -> None:
             Decoder.Config.update_from_config(self, config=config, **kwargs)
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py b/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
index 8f3f5409846..c1ee70afc07 100644
--- a/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/__init__.py
@@ -109,7 +109,14 @@ def pipeline_kimi_k3(model: BaseModel, *, attn_res_cache: bool = True, **kwargs)
         layer_to_stage=layer_to_stage,
         cache=attn_res_cache,
     )
-    store = PPRankLocalCache()
+    offload = kwargs["model_config"].attn_res_cache_offload
+    if offload and not attn_res_cache:
+        raise NotImplementedError(
+            "attn_res_cache_offload parks the blocks a rank holds between hops, "
+            "which only the cached transport keeps; it has no effect when every "
+            "hop carries the whole stack."
+        )
+    store = PPRankLocalCache(device=kwargs["device"] if offload else None)
     for stage in stages:
         stage.set_routing(layout, store)
     logger.info(
diff --git a/torchtitan/models/kimi_k3/pipeline_parallel/stage.py b/torchtitan/models/kimi_k3/pipeline_parallel/stage.py
index c16855a87d9..a7ecf82bbba 100644
--- a/torchtitan/models/kimi_k3/pipeline_parallel/stage.py
+++ b/torchtitan/models/kimi_k3/pipeline_parallel/stage.py
@@ -21,18 +21,27 @@
 
 
 class PPRankLocalCache:
-    """Blocks a rank holds per micro-batch and the gradient deposits, shared by its stages."""
+    """Blocks a rank holds per micro-batch and the gradient deposits, shared by its stages;
+    given a device, a stored block parks on pinned host memory between commit and read."""
 
-    def __init__(self) -> None:
+    def __init__(self, *, device: torch.device | None = None) -> None:
+        self._device = device
         self._blocks: dict[int, dict[int, torch.Tensor]] = {}
         self._deposits: dict[tuple[int, int], torch.Tensor] = {}
         self._counts: dict[tuple[int, int], int] = {}
 
     def put(self, mb: int, block_idx: int, block_TD: torch.Tensor) -> None:
+        if self._device is not None and block_TD.is_cuda:
+            host_TD = torch.empty_like(block_TD, device="cpu", pin_memory=True)
+            host_TD.copy_(block_TD, non_blocking=True)
+            block_TD = host_TD
         self._blocks.setdefault(mb, {})[block_idx] = block_TD
 
     def blocks(self, mb: int) -> dict[int, torch.Tensor]:
-        return self._blocks.get(mb, {})
+        held = self._blocks.get(mb, {})
+        if self._device is None:
+            return held
+        return {b: t.to(self._device, non_blocking=True) for b, t in held.items()}
 
     def release(self, mb: int) -> None:
         """Free the blocks of ``mb``; the deposits stay until collected."""