diff --git a/megatron/core/inference/disaggregation/__init__.py b/megatron/core/inference/disaggregation/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/megatron/core/inference/disaggregation/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/disaggregation/kv_reshard.py b/megatron/core/inference/disaggregation/kv_reshard.py
new file mode 100644
index 00000000000..7fa01488d1d
--- /dev/null
+++ b/megatron/core/inference/disaggregation/kv_reshard.py
@@ -0,0 +1,182 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""TP/PP/EP/ETP KV-shard layouts and the range-intersection reshard planner."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+from megatron.core.inference.disaggregation.utils import intersect
+
+
+@dataclass(frozen=True)
+class KVShardLayout:
+    """A worker's KV-cache ownership within the global model.
+
+    ``num_layers`` / ``num_heads`` are the *global* attention layer count
+    and KV-head count (for GQA, the number of KV heads). ``global_rank``
+    is the worker's torch rank (used as the transport peer id).
+    """
+
+    num_layers: int
+    num_heads: int
+    tp_size: int
+    tp_rank: int
+    pp_size: int
+    pp_rank: int
+    global_rank: int
+    # Expert dimensions. KV-replica dimensions only: they shard the MoE
+    # expert weights, never the attention KV cache, so they don't affect
+    # head_range/layer_range -- only representative (source) selection.
+    ep_size: int = 1
+    ep_rank: int = 0
+    etp_size: int = 1
+    etp_rank: int = 0
+    # Optional explicit PP layer window for this stage. When None, an even split
+    # of num_layers across pp_size is assumed -- correct for pure-attention
+    # models. Models that do NOT split attention layers evenly across PP stages
+    # (e.g. hybrid Mamba+attention) must pass an explicit (layer_start,
+    # num_local_layers); the even-split default would otherwise map the wrong
+    # global layer indices.
+    layer_start: Optional[int] = None
+    num_local_layers: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        # TP must divide heads (the head split is always even).
+        if self.num_heads % self.tp_size != 0:
+            raise ValueError(f"num_heads={self.num_heads} not divisible by tp_size={self.tp_size}")
+        # layer_start and num_local_layers are an all-or-nothing explicit window:
+        # setting only one would silently fall back to the even-split count and
+        # defeat the purpose (uneven stage with an even count).
+        if (self.layer_start is None) != (self.num_local_layers is None):
+            raise ValueError(
+                "layer_start and num_local_layers must be set together (or both omitted)"
+            )
+        # Only the even-split path requires PP to divide layers; an explicit
+        # window may be uneven across stages.
+        if self.layer_start is None and self.num_layers % self.pp_size != 0:
+            raise ValueError(
+                f"num_layers={self.num_layers} not divisible by pp_size={self.pp_size}; "
+                "pass an explicit (layer_start, num_local_layers) for uneven PP splits"
+            )
+
+    def kv_shard_key(self) -> Tuple[int, int]:
+        """The attention shard this rank holds: ``(tp_rank, pp_rank)``.
+        Ranks sharing a key hold identical KV (EP/ETP replicas of it)."""
+        return (self.tp_rank, self.pp_rank)
+
+    def layer_range(self) -> Tuple[int, int]:
+        """Global attention-layer range ``[lo, hi)`` owned by this rank."""
+        # num_local_layers is guaranteed set whenever layer_start is (see __post_init__).
+        if self.layer_start is not None:
+            return (self.layer_start, self.layer_start + self.num_local_layers)
+        per = self.num_layers // self.pp_size
+        return (self.pp_rank * per, (self.pp_rank + 1) * per)
+
+    def head_range(self) -> Tuple[int, int]:
+        """Global KV-head range ``[lo, hi)`` owned by this rank."""
+        per = self.num_heads // self.tp_size
+        return (self.tp_rank * per, (self.tp_rank + 1) * per)
+
+    def local_num_layers(self) -> int:
+        """Number of attention layers held locally by this rank."""
+        lo, hi = self.layer_range()
+        return hi - lo
+
+    def local_num_heads(self) -> int:
+        """Number of KV heads held locally by this rank."""
+        lo, hi = self.head_range()
+        return hi - lo
+
+
+@dataclass(frozen=True)
+class KVReshardTransfer:
+    """One sub-block exchange between a (src, dst) rank pair.
+
+    Global coords identify the intersection; the local-slice helpers
+    convert to each side's buffer offsets. There is at most one transfer
+    per (src, dst) pair (each owns a contiguous rectangle, so the
+    intersection is a single rectangle).
+    """
+
+    src_rank: int
+    dst_rank: int
+    # The transferred sub-block's GLOBAL bounds as half-open ranges:
+    # layers [global_layer_lo, global_layer_hi) x kv-heads [global_head_lo, global_head_hi).
+    global_layer_lo: int
+    global_layer_hi: int
+    global_head_lo: int
+    global_head_hi: int
+
+    def src_layer_slice(self, src: KVShardLayout) -> slice:
+        """Local layer slice on the source side for this transfer."""
+        off = src.layer_range()[0]
+        return slice(self.global_layer_lo - off, self.global_layer_hi - off)
+
+    def src_head_slice(self, src: KVShardLayout) -> slice:
+        """Local KV-head slice on the source side for this transfer."""
+        off = src.head_range()[0]
+        return slice(self.global_head_lo - off, self.global_head_hi - off)
+
+    def dst_layer_slice(self, dst: KVShardLayout) -> slice:
+        """Local layer slice on the destination side for this transfer."""
+        off = dst.layer_range()[0]
+        return slice(self.global_layer_lo - off, self.global_layer_hi - off)
+
+    def dst_head_slice(self, dst: KVShardLayout) -> slice:
+        """Local KV-head slice on the destination side for this transfer."""
+        off = dst.head_range()[0]
+        return slice(self.global_head_lo - off, self.global_head_hi - off)
+
+
+def plan_kv_reshard(
+    srcs: List[KVShardLayout], dsts: List[KVShardLayout]
+) -> List[KVReshardTransfer]:
+    """Full reshard plan: every sub-block that must move src -> dst.
+
+    Both sides compute the same plan from the same layouts and filter to
+    their own rank (``transfers_for_src`` / ``transfers_for_dst``).
+
+    KV is replicated across the EP and ETP dimensions, so each attention
+    shard ``(tp_rank, pp_rank)`` may be held by several source ranks. We
+    source each shard from exactly one of them -- the smallest
+    ``global_rank`` -- which avoids duplicate sends and is independent of
+    how EP/ETP map onto ranks.
+    """
+    if srcs and dsts:
+        if srcs[0].num_layers != dsts[0].num_layers or srcs[0].num_heads != dsts[0].num_heads:
+            raise ValueError("src and dst describe different global models")
+
+    # One representative source rank per attention shard (dedupe EP/ETP
+    # replicas that hold identical KV).
+    rep_rank: dict = {}
+    for s in srcs:
+        key = s.kv_shard_key()
+        if key not in rep_rank or s.global_rank < rep_rank[key]:
+            rep_rank[key] = s.global_rank
+    source_ranks = set(rep_rank.values())
+
+    transfers: List[KVReshardTransfer] = []
+    for d in dsts:
+        dl, dh = d.layer_range(), d.head_range()
+        for s in srcs:
+            if s.global_rank not in source_ranks:
+                continue
+            li = intersect(s.layer_range(), dl)
+            if li is None:
+                continue
+            hi = intersect(s.head_range(), dh)
+            if hi is None:
+                continue
+            transfers.append(
+                KVReshardTransfer(
+                    src_rank=s.global_rank,
+                    dst_rank=d.global_rank,
+                    global_layer_lo=li[0],
+                    global_layer_hi=li[1],
+                    global_head_lo=hi[0],
+                    global_head_hi=hi[1],
+                )
+            )
+    return transfers
diff --git a/megatron/core/inference/disaggregation/mamba_reshard.py b/megatron/core/inference/disaggregation/mamba_reshard.py
new file mode 100644
index 00000000000..8a23735154a
--- /dev/null
+++ b/megatron/core/inference/disaggregation/mamba_reshard.py
@@ -0,0 +1,222 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Heterogeneous TP/PP reshard of Mamba conv/ssm state between prefill and
+decode shard layouts (the Mamba analog of the attention KV reshard)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Tuple
+
+from megatron.core.inference.disaggregation.utils import intersect
+
+# Channel bands of a Mamba layer's state, in the order the conv state
+# concatenates them on its channel axis (x, B, C); ssm is the head axis.
+# (name, lives_in_conv). conv bands share one tensor; ssm is its own tensor.
+_CONV_BANDS = ("x", "B", "C")
+
+
+@dataclass(frozen=True)
+class MambaStateDims:
+    """The model's (global, unsharded) Mamba structural dims.
+
+    These belong to the MambaMixer / model config -- carried as one unit (rather
+    than loose constants spread across the layout) so there's a single source
+    and they can't drift apart. The producer should read them straight from the
+    model config (e.g. ``ngroups = config.mamba_num_groups``) rather than
+    reverse-deriving from tensor shapes. TP shards ``nheads``/``ngroups``; the
+    rest are unsharded.
+    """
+
+    nheads: int
+    headdim: int
+    d_state: int
+    ngroups: int
+    d_conv: int
+
+
+@dataclass(frozen=True)
+class MambaShardLayout:
+    """One rank's Mamba-state ownership: which global layers + TP rank, plus the
+    model's structural dims (:class:`MambaStateDims`). Per-rank locals follow by
+    dividing by ``tp_size``."""
+
+    global_rank: int
+    tp_size: int
+    tp_rank: int
+    layer_start: int  # global Mamba-layer index of this rank's first layer
+    num_layers: int  # Mamba layers held locally (this PP stage)
+    dims: MambaStateDims
+
+    def __post_init__(self) -> None:
+        # Wire reconstruction (MambaShardLayout(**dict)) hands ``dims`` as a
+        # plain dict; coerce it back to MambaStateDims.
+        if isinstance(self.dims, dict):
+            object.__setattr__(self, "dims", MambaStateDims(**self.dims))
+        # TP shards heads and groups; both must divide evenly or the local
+        # conv/ssm band sizes truncate to the wrong (or zero) width silently.
+        if self.dims.nheads % self.tp_size != 0:
+            raise ValueError(f"nheads={self.dims.nheads} not divisible by tp_size={self.tp_size}")
+        if self.dims.ngroups % self.tp_size != 0:
+            raise ValueError(f"ngroups={self.dims.ngroups} not divisible by tp_size={self.tp_size}")
+
+    # Convenience proxies onto the dims so callers read ``layout.headdim`` etc.
+    @property
+    def nheads(self) -> int:
+        """Global (unsharded) number of Mamba heads."""
+        return self.dims.nheads
+
+    @property
+    def headdim(self) -> int:
+        """Dimension of each Mamba head."""
+        return self.dims.headdim
+
+    @property
+    def d_state(self) -> int:
+        """SSM state size per head."""
+        return self.dims.d_state
+
+    @property
+    def ngroups(self) -> int:
+        """Global (unsharded) number of B/C groups."""
+        return self.dims.ngroups
+
+    @property
+    def d_conv(self) -> int:
+        """Convolution kernel width."""
+        return self.dims.d_conv
+
+    def mamba_shard_key(self) -> Tuple[int, int]:
+        """The Mamba shard this rank holds: ``(tp_rank, layer_start)``. Ranks
+        sharing a key hold identical state (e.g. EP/DP replicas of it)."""
+        return (self.tp_rank, self.layer_start)
+
+    @property
+    def d_inner(self) -> int:
+        """Global inner dimension (nheads * headdim)."""
+        return self.dims.nheads * self.dims.headdim
+
+    @property
+    def nheads_local(self) -> int:
+        """Number of Mamba heads held by this TP rank."""
+        return self.dims.nheads // self.tp_size
+
+    @property
+    def d_inner_local(self) -> int:
+        """Local inner dimension for this TP rank."""
+        return self.d_inner // self.tp_size
+
+    @property
+    def ngroups_local(self) -> int:
+        """Number of B/C groups held by this TP rank."""
+        return self.dims.ngroups // self.tp_size
+
+    @property
+    def conv_dim_local(self) -> int:
+        """Total local conv channel width (x + B + C bands)."""
+        return self.d_inner_local + 2 * self.ngroups_local * self.dims.d_state
+
+    def layer_range(self) -> Tuple[int, int]:
+        """Global Mamba-layer range ``[lo, hi)`` owned by this rank."""
+        return (self.layer_start, self.layer_start + self.num_layers)
+
+    def _band(self, name: str) -> Tuple[int, int, int]:
+        """``(global_total, local_size, conv_local_offset)`` for a band.
+
+        ``conv_local_offset`` is the band's start on the local conv channel
+        axis; for the ``ssm`` (head) band it is the start on the local head
+        axis (always 0, heads are the whole tensor)."""
+        if name == "x":
+            g = self.d_inner
+            return g, self.d_inner_local, 0
+        if name == "B":
+            g = self.dims.ngroups * self.dims.d_state
+            return g, self.ngroups_local * self.dims.d_state, self.d_inner_local
+        if name == "C":
+            g = self.dims.ngroups * self.dims.d_state
+            return (
+                g,
+                self.ngroups_local * self.dims.d_state,
+                self.d_inner_local + self.ngroups_local * self.dims.d_state,
+            )
+        if name == "ssm":
+            return self.dims.nheads, self.nheads_local, 0
+        raise KeyError(name)
+
+
+@dataclass(frozen=True)
+class MambaReshardTransfer:
+    """One sub-block move for the reshard.
+
+    ``band`` is ``"x"``/``"B"``/``"C"`` (conv channel axis) or ``"ssm"`` (head
+    axis). ``src_layer``/``dst_layer`` are local layer indices on each side;
+    ``*_lo``/``*_hi`` are the local channel/head slice bounds.
+    """
+
+    src_rank: int
+    dst_rank: int
+    band: str
+    global_layer: int
+    src_layer: int
+    dst_layer: int
+    src_lo: int
+    src_hi: int
+    dst_lo: int
+    dst_hi: int
+
+    @property
+    def is_conv(self) -> bool:
+        """True if this transfer targets the conv state; False for ssm."""
+        return self.band in _CONV_BANDS
+
+
+def plan_mamba_reshard(
+    src_layouts: List[MambaShardLayout], dst_layouts: List[MambaShardLayout]
+) -> List[MambaReshardTransfer]:
+    """Plan the conv/ssm sub-block moves from the prefill (src) layouts to the
+    decode (dst) layouts. One transfer per (src rank, dst rank, global layer,
+    band) where both the layer ranges and the channel ranges overlap."""
+    # Dedupe replica sources: ranks sharing (tp_rank, layer_start) hold identical
+    # Mamba state (e.g. EP/DP replicas), so source each shard from exactly one of
+    # them -- the smallest global_rank -- to avoid duplicate sends.
+    rep_rank: dict = {}
+    for s in src_layouts:
+        key = s.mamba_shard_key()
+        if key not in rep_rank or s.global_rank < rep_rank[key]:
+            rep_rank[key] = s.global_rank
+    source_ranks = set(rep_rank.values())
+
+    out: List[MambaReshardTransfer] = []
+    for s in src_layouts:
+        if s.global_rank not in source_ranks:
+            continue
+        s_lr = s.layer_range()
+        for d in dst_layouts:
+            layer_ov = intersect(s_lr, d.layer_range())
+            if layer_ov is None:
+                continue
+            for band in (*_CONV_BANDS, "ssm"):
+                _, s_size, s_off = s._band(band)
+                _, d_size, d_off = d._band(band)
+                s_glo = (s.tp_rank * s_size, s.tp_rank * s_size + s_size)
+                d_glo = (d.tp_rank * d_size, d.tp_rank * d_size + d_size)
+                chan_ov = intersect(s_glo, d_glo)
+                if chan_ov is None:
+                    continue
+                lo, hi = chan_ov
+                for g in range(layer_ov[0], layer_ov[1]):
+                    out.append(
+                        MambaReshardTransfer(
+                            src_rank=s.global_rank,
+                            dst_rank=d.global_rank,
+                            band=band,
+                            global_layer=g,
+                            src_layer=g - s.layer_start,
+                            dst_layer=g - d.layer_start,
+                            src_lo=s_off + (lo - s_glo[0]),
+                            src_hi=s_off + (hi - s_glo[0]),
+                            dst_lo=d_off + (lo - d_glo[0]),
+                            dst_hi=d_off + (hi - d_glo[0]),
+                        )
+                    )
+    return out
diff --git a/megatron/core/inference/disaggregation/utils.py b/megatron/core/inference/disaggregation/utils.py
new file mode 100644
index 00000000000..9b5e153b443
--- /dev/null
+++ b/megatron/core/inference/disaggregation/utils.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Shared helpers for the disaggregation modules."""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+
+def intersect(a: Tuple[int, int], b: Tuple[int, int]) -> Optional[Tuple[int, int]]:
+    """Overlap of two half-open ``[lo, hi)`` ranges, or ``None`` if disjoint."""
+    lo, hi = max(a[0], b[0]), min(a[1], b[1])
+    return (lo, hi) if lo < hi else None
+
+
+def transfers_for_src(plan, src_rank):
+    """Transfers in ``plan`` originating from ``src_rank`` (any KV/Mamba
+    reshard transfer -- both expose a ``src_rank`` field)."""
+    return [t for t in plan if t.src_rank == src_rank]
+
+
+def transfers_for_dst(plan, dst_rank):
+    """Transfers in ``plan`` destined for ``dst_rank``."""
+    return [t for t in plan if t.dst_rank == dst_rank]
diff --git a/tests/unit_tests/inference/test_kv_reshard.py b/tests/unit_tests/inference/test_kv_reshard.py
new file mode 100644
index 00000000000..63b62bc0f0b
--- /dev/null
+++ b/tests/unit_tests/inference/test_kv_reshard.py
@@ -0,0 +1,191 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Correctness of hetero TP/PP/EP KV resharding (single process).
+
+We materialize a global KV tensor, split it into a *source* layout's
+shards, run the reshard plan to assemble a *destination* layout's
+shards, and assert each dst shard equals the direct split of the global
+KV. Sweeping many (Tp,Pp,Td,Pd) combos -- divisible, non-divisible,
+PP-changing, and EP-replicated -- exercises the range-intersection
+planner end to end without any distributed runtime.
+"""
+
+import pytest
+import torch
+
+from megatron.core.inference.disaggregation.kv_reshard import KVShardLayout, plan_kv_reshard
+from megatron.core.inference.disaggregation.utils import transfers_for_dst
+
+# global model
+L, Hh, BC, BS, HD = 12, 8, 2, 4, 5  # layers, kv-heads, block_count, block_size, head_dim
+
+
+def _global_kv():
+    # [2(K/V), L, BC, BS, H, HD] with unique values per (kv, layer, head)
+    g = torch.zeros(2, L, BC, BS, Hh, HD)
+    for kv in range(2):
+        for l in range(L):
+            for h in range(Hh):
+                g[kv, l, :, :, h, :] = (kv * 1_000_000) + l * 1000 + h
+    return g
+
+
+def _shard_of(global_kv, lay: KVShardLayout):
+    """The dst staging tensor a worker with layout `lay` should hold:
+    [BC, 2, local_layers, BS, local_heads, HD] (export's attn layout)."""
+    l0, l1 = lay.layer_range()
+    h0, h1 = lay.head_range()
+    # global_kv is [2, L, BC, BS, H, HD]; export layout is
+    # [BC, 2, layers, BS, heads, HD]
+    sub = global_kv[:, l0:l1, :, :, h0:h1, :]  # [2, ll, BC, BS, hh, HD]
+    return sub.permute(2, 0, 1, 3, 4, 5).contiguous()  # [BC,2,ll,BS,hh,HD]
+
+
+def _make_layouts(tp, pp, ep=1, etp=1):
+    outs = []
+    rank = 0
+    for p in range(pp):
+        for t in range(tp):
+            for e in range(ep):
+                for et in range(etp):
+                    outs.append(
+                        KVShardLayout(
+                            num_layers=L,
+                            num_heads=Hh,
+                            tp_size=tp,
+                            tp_rank=t,
+                            pp_size=pp,
+                            pp_rank=p,
+                            global_rank=rank,
+                            ep_size=ep,
+                            ep_rank=e,
+                            etp_size=etp,
+                            etp_rank=et,
+                        )
+                    )
+                    rank += 1
+    return outs
+
+
+def _run_reshard(src_layouts, dst_layouts):
+    g = _global_kv()
+    # src buffers = each src's correct shard of the global KV
+    src_buf = {s.global_rank: _shard_of(g, s) for s in src_layouts}
+    plan = plan_kv_reshard(src_layouts, dst_layouts)
+    by_rank = {s.global_rank: s for s in src_layouts}
+    out = {}
+    for d in dst_layouts:
+        dst = torch.full((BC, 2, d.local_num_layers(), BS, d.local_num_heads(), HD), -999.0)
+        for t in transfers_for_dst(plan, d.global_rank):
+            s = by_rank[t.src_rank]
+            block = src_buf[t.src_rank][:, :, t.src_layer_slice(s), :, t.src_head_slice(s), :]
+            dst[:, :, t.dst_layer_slice(d), :, t.dst_head_slice(d), :] = block
+        out[d.global_rank] = dst
+    return g, out
+
+
+@pytest.mark.parametrize(
+    "src,dst",
+    [
+        ((1, 1), (1, 1)),  # homogeneous
+        ((2, 1), (4, 1)),  # TP fan-out (divisible)
+        ((4, 1), (2, 1)),  # TP merge (divisible)
+        ((1, 2), (1, 3)),  # PP change (divisible both)
+        ((2, 2), (4, 3)),  # both change
+        ((2, 3), (4, 2)),  # TP + PP mixed
+    ],
+)
+def test_reshard_matches_direct_split(src, dst):
+    tp_s, pp_s = src
+    tp_d, pp_d = dst
+    # skip layouts that violate divisibility of the GLOBAL dims
+    if Hh % tp_s or Hh % tp_d or L % pp_s or L % pp_d:
+        pytest.skip("layout not divisible for this global model")
+    src_layouts = _make_layouts(tp_s, pp_s)
+    dst_layouts = _make_layouts(tp_d, pp_d)
+    g, out = _run_reshard(src_layouts, dst_layouts)
+    for d in dst_layouts:
+        expected = _shard_of(g, d)
+        got = out[d.global_rank]
+        assert torch.equal(got, expected), f"dst rank {d.global_rank} mismatch"
+        assert (got != -999.0).all(), "some dst entries never received"
+
+
+def _assert_one_source_per_shard(plan, src_layouts):
+    """Each attention shard (tp_rank, pp_rank) must be sourced by exactly
+    one rank -- no duplicate sends from EP/ETP replicas."""
+    src_by_rank = {s.global_rank: s for s in src_layouts}
+    shard_sources = {}
+    for t in plan:
+        s = src_by_rank[t.src_rank]
+        shard_sources.setdefault(s.kv_shard_key(), set()).add(t.src_rank)
+    for key, ranks in shard_sources.items():
+        assert len(ranks) == 1, f"shard {key} sourced by {ranks}"
+
+
+@pytest.mark.parametrize("ep,etp", [(2, 1), (1, 2), (2, 2)])
+def test_expert_replication_picks_single_source(ep, etp):
+    """EP- and/or ETP-replicated sources: each attention shard is sourced
+    once; every dst (any EP/ETP replica) still gets correct, complete data.
+    EP and ETP shard the expert FFN, not the KV, so they're pure replicas."""
+    src_layouts = _make_layouts(tp=2, pp=1, ep=ep, etp=etp)
+    dst_layouts = _make_layouts(tp=2, pp=1, ep=ep, etp=etp)
+    plan = plan_kv_reshard(src_layouts, dst_layouts)
+    _assert_one_source_per_shard(plan, src_layouts)
+    g, out = _run_reshard(src_layouts, dst_layouts)
+    for d in dst_layouts:
+        assert torch.equal(out[d.global_rank], _shard_of(g, d))
+
+
+def test_hetero_tp_with_expert_replication():
+    """Hetero attention TP merge (4->2) while sources are also ETP-replicated:
+    the reshard still merges heads correctly and dedupes the ETP replicas."""
+    src_layouts = _make_layouts(tp=4, pp=1, etp=2)  # 8 ranks, 4 attn shards x2
+    dst_layouts = _make_layouts(tp=2, pp=1)
+    plan = plan_kv_reshard(src_layouts, dst_layouts)
+    _assert_one_source_per_shard(plan, src_layouts)
+    g, out = _run_reshard(src_layouts, dst_layouts)
+    for d in dst_layouts:
+        assert torch.equal(out[d.global_rank], _shard_of(g, d))
+
+
+def test_one_prefill_to_multiple_decode_targets_of_different_parallelism():
+    """A single prefill source set reshards correctly to several decode
+    targets that each use a DIFFERENT (Tp,Pp) -- e.g. a heterogeneous
+    decode pool. Each target is an independent reshard (one plan call per
+    target replica); the planner imposes no shared parallelism across
+    targets."""
+    src_layouts = _make_layouts(tp=2, pp=2)  # prefill: TP2 x PP2
+    targets = [(4, 1), (2, 1), (1, 3), (4, 3)]  # decode replicas, all different
+    g = _global_kv()
+    for tp_d, pp_d in targets:
+        dst_layouts = _make_layouts(tp_d, pp_d)
+        _, out = _run_reshard(src_layouts, dst_layouts)
+        for d in dst_layouts:
+            assert torch.equal(
+                out[d.global_rank], _shard_of(g, d)
+            ), f"decode target TP{tp_d}xPP{pp_d} rank {d.global_rank} mismatch"
+
+
+def test_uneven_pp_attention_window():
+    """Attention layers split UNEVENLY across PP (hybrid-style) via explicit
+    (layer_start, num_local_layers); reshard to pp=1 still reconstructs the
+    global KV. The even-split default would map the wrong global layers here."""
+    src = [
+        KVShardLayout(L, Hh, 1, 0, 2, 0, 0, layer_start=0, num_local_layers=5),
+        KVShardLayout(L, Hh, 1, 0, 2, 1, 1, layer_start=5, num_local_layers=7),
+    ]
+    dst = [KVShardLayout(L, Hh, 1, 0, 1, 0, 2)]  # pp=1: all L layers on one rank
+    assert src[0].layer_range() == (0, 5) and src[1].layer_range() == (5, 12)
+    g, out = _run_reshard(src, dst)
+    for d in dst:
+        assert torch.equal(out[d.global_rank], _shard_of(g, d))
+
+
+def test_explicit_layer_window_is_all_or_nothing():
+    # Setting only one of (layer_start, num_local_layers) would silently fall
+    # back to the even-split count -- reject it.
+    with pytest.raises(ValueError):
+        KVShardLayout(L, Hh, 1, 0, 2, 0, 0, layer_start=0)
+    with pytest.raises(ValueError):
+        KVShardLayout(L, Hh, 1, 0, 2, 0, 0, num_local_layers=5)
diff --git a/tests/unit_tests/inference/test_mamba_reshard.py b/tests/unit_tests/inference/test_mamba_reshard.py
new file mode 100644
index 00000000000..4a197813ab9
--- /dev/null
+++ b/tests/unit_tests/inference/test_mamba_reshard.py
@@ -0,0 +1,185 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Hetero TP/PP reshard of Mamba conv/ssm state (pure, CPU).
+
+Builds a known global Mamba state, shards it to a source (tp,pp) the exact way
+mamba_mixer does ([x|B|C] conv bands + head-sharded ssm, layers split by PP),
+runs plan_mamba_reshard to a different destination (tp,pp), and asserts every
+destination rank ends up byte-identical to a direct shard of the global state.
+This validates the band/layer index math against the real sharding model
+without a hybrid checkpoint (the residual gap is a real-model functional run).
+"""
+
+import pytest
+import torch
+
+from megatron.core.inference.disaggregation.mamba_reshard import (
+    MambaShardLayout,
+    MambaStateDims,
+    plan_mamba_reshard,
+)
+
+
+def apply_conv_transfer(t, src_conv, dst_conv):
+    """Copy a conv sub-block in-memory (no transfer); conv is
+    ``(num_layers, conv_dim_local, d_conv)`` -- the band slices the channel axis."""
+    dst_conv[t.dst_layer, t.dst_lo : t.dst_hi, :] = src_conv[t.src_layer, t.src_lo : t.src_hi, :]
+
+
+def apply_ssm_transfer(t, src_ssm, dst_ssm):
+    """Copy an ssm sub-block in-memory; ssm is
+    ``(num_layers, nheads_local, headdim, d_state)`` -- the band slices heads."""
+    dst_ssm[t.dst_layer, t.dst_lo : t.dst_hi, :, :] = src_ssm[
+        t.src_layer, t.src_lo : t.src_hi, :, :
+    ]
+
+
+# Global model dims (chosen divisible by the tp values under test).
+NHEADS, HEADDIM, DSTATE, NGROUPS, DCONV = 8, 4, 2, 2, 3
+M = 4  # global Mamba layers
+D_INNER = NHEADS * HEADDIM  # 32
+G = NGROUPS * DSTATE  # 4  (B and C band global size)
+CONV_DIM = D_INNER + 2 * G  # 40
+
+
+def _global_state():
+    """Distinct value per (layer, channel, ...) so any mis-slice is caught."""
+    conv = torch.arange(M * CONV_DIM * DCONV, dtype=torch.float32).reshape(M, CONV_DIM, DCONV)
+    ssm = (
+        torch.arange(M * NHEADS * HEADDIM * DSTATE, dtype=torch.float32).reshape(
+            M, NHEADS, HEADDIM, DSTATE
+        )
+        + 10_000.0
+    )
+    return conv, ssm
+
+
+def _layouts(tp, pp):
+    """One MambaShardLayout per rank for a (tp, pp) instance; rank = p*tp + r.
+    PP splits the M layers evenly (contiguous per stage)."""
+    per = M // pp
+    out = {}
+    for p in range(pp):
+        for r in range(tp):
+            rank = p * tp + r
+            out[rank] = MambaShardLayout(
+                global_rank=rank,
+                tp_size=tp,
+                tp_rank=r,
+                layer_start=p * per,
+                num_layers=per,
+                dims=MambaStateDims(
+                    nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV
+                ),
+            )
+    return out
+
+
+def _shard(conv_g, ssm_g, lay: MambaShardLayout):
+    """Shard the global state to one rank exactly as mamba_mixer does."""
+    s, e = lay.layer_range()
+    r, tp = lay.tp_rank, lay.tp_size
+    di_l = D_INNER // tp
+    g_l = (NGROUPS // tp) * DSTATE
+    x = conv_g[s:e, 0:D_INNER][:, r * di_l : (r + 1) * di_l]
+    b = conv_g[s:e, D_INNER : D_INNER + G][:, r * g_l : (r + 1) * g_l]
+    c = conv_g[s:e, D_INNER + G : D_INNER + 2 * G][:, r * g_l : (r + 1) * g_l]
+    conv_l = torch.cat([x, b, c], dim=1).contiguous()
+    nh_l = NHEADS // tp
+    ssm_l = ssm_g[s:e, r * nh_l : (r + 1) * nh_l, :, :].contiguous()
+    return conv_l, ssm_l
+
+
+@pytest.mark.parametrize(
+    "src,dst",
+    [
+        ((2, 1), (1, 1)),  # TP2 -> TP1 (band merge)
+        ((1, 1), (2, 1)),  # TP1 -> TP2 (band split)
+        ((1, 2), (1, 1)),  # PP2 -> PP1 (layer merge)
+        ((1, 1), (1, 2)),  # PP1 -> PP2 (layer split)
+        ((2, 2), (1, 1)),  # both axes hetero
+        ((2, 1), (2, 1)),  # identity
+    ],
+)
+def test_mamba_reshard_reconstructs_destination(src, dst):
+    conv_g, ssm_g = _global_state()
+    src_lay, dst_lay = _layouts(*src), _layouts(*dst)
+
+    # Source per-rank tensors (as a prefill instance would hold them).
+    src_t = {rk: _shard(conv_g, ssm_g, lay) for rk, lay in src_lay.items()}
+    # Destination buffers, zero-filled at each rank's local shape.
+    dst_t = {}
+    for rk, lay in dst_lay.items():
+        dst_t[rk] = (
+            torch.zeros(lay.num_layers, lay.conv_dim_local, DCONV),
+            torch.zeros(lay.num_layers, lay.nheads_local, HEADDIM, DSTATE),
+        )
+
+    plan = plan_mamba_reshard(list(src_lay.values()), list(dst_lay.values()))
+    for t in plan:
+        if t.is_conv:
+            apply_conv_transfer(t, src_t[t.src_rank][0], dst_t[t.dst_rank][0])
+        else:
+            apply_ssm_transfer(t, src_t[t.src_rank][1], dst_t[t.dst_rank][1])
+
+    # Every destination rank must match a direct shard of the global state.
+    for rk, lay in dst_lay.items():
+        want_conv, want_ssm = _shard(conv_g, ssm_g, lay)
+        assert torch.equal(dst_t[rk][0], want_conv), f"conv mismatch at rank {rk} ({src}->{dst})"
+        assert torch.equal(dst_t[rk][1], want_ssm), f"ssm mismatch at rank {rk} ({src}->{dst})"
+
+
+def test_mamba_rejects_indivisible_groups():
+    """ngroups < tp_size would truncate the B/C bands to zero width; reject it
+    up front instead of silently dropping state."""
+    with pytest.raises(ValueError):
+        MambaShardLayout(
+            global_rank=0,
+            tp_size=4,
+            tp_rank=0,
+            layer_start=0,
+            num_layers=1,
+            dims=MambaStateDims(nheads=8, headdim=HEADDIM, d_state=DSTATE, ngroups=2, d_conv=DCONV),
+        )
+
+
+def test_mamba_dedupes_replica_sources():
+    """Two source ranks holding the same Mamba shard (same tp_rank+layer_start,
+    e.g. EP/DP replicas) are deduped: the shard is sourced from exactly one of
+    them (smallest global_rank), so no duplicate sends."""
+
+    def _lay(gr):
+        return MambaShardLayout(
+            global_rank=gr,
+            tp_size=1,
+            tp_rank=0,
+            layer_start=0,
+            num_layers=M,
+            dims=MambaStateDims(
+                nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV
+            ),
+        )
+
+    plan = plan_mamba_reshard([_lay(0), _lay(1)], [_lay(2)])
+    assert {t.src_rank for t in plan} == {0}  # only the smallest-rank replica sources
+
+
+def test_layout_wire_roundtrip():
+    """Layouts cross the coordinator as plain dicts (asdict) and are rebuilt via
+    MambaShardLayout(**dict); the nested dims dict must coerce back to
+    MambaStateDims so proxies (.headdim/.d_conv/...) keep working."""
+    import dataclasses
+
+    lay = MambaShardLayout(
+        global_rank=1,
+        tp_size=2,
+        tp_rank=1,
+        layer_start=0,
+        num_layers=M,
+        dims=MambaStateDims(
+            nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV
+        ),
+    )
+    rebuilt = MambaShardLayout(**dataclasses.asdict(lay))
+    assert rebuilt == lay
+    assert rebuilt.headdim == HEADDIM and rebuilt.d_conv == DCONV