diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 472917d56f0..6d17f89d47e 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3094,3 +3094,233 @@ def run(x, module=model):
         forward=case in ("bias", "unfused"),
         backward=case != "linear" and bool(targets),
     )
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("equal_splits", [False, True])
+def test_te_split_parameters_compile(compile_mode, dtype, equal_splits):
+    """Every split remains an autograd input across storage and parameter changes."""
+    import copy
+
+    torch._dynamo.reset()
+    counters.clear()
+    splits = ("q", "k", "v") if equal_splits else dict(q=64, k=32, v=32)
+    model = te.Linear(64, 192 if equal_splits else 128, parameters_split=splits, params_dtype=dtype)
+    model.k_weight.requires_grad_(False)
+    model.v_bias.requires_grad_(False)
+    reference = copy.deepcopy(model)
+    compiled = torch.compile(model, fullgraph=True, mode=compile_mode)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01)
+    for mutation in ("none", "optimizer", "data", "parameter", "assign", "dtype"):
+        for current in (model, reference):
+            if mutation == "data":
+                current.k_weight.data = current.k_weight.detach().clone() + 0.01
+                current.k_bias.data = current.k_bias.detach().clone() + 0.01
+            elif mutation == "parameter":
+                current.q_weight = torch.nn.Parameter(current.q_weight.detach().clone() + 0.01)
+            elif mutation == "assign":
+                state = {name: value.clone() for name, value in current.state_dict().items()}
+                current.load_state_dict(state, assign=True)
+            elif mutation == "dtype":
+                current.to(torch.float64).to(dtype)
+        if mutation == "optimizer":
+            optimizer.step()
+            ref_optimizer.step()
+        for _ in range(3):
+            torch.compiler.cudagraph_mark_step_begin()
+            inp = torch.randn(16, 64, device="cuda", dtype=dtype, requires_grad=True)
+            ref_inp = inp.detach().clone().requires_grad_()
+            model.zero_grad(set_to_none=True)
+            reference.zero_grad(set_to_none=True)
+            out = compiled(inp)
+            grad = torch.randn_like(out)
+            out.backward(grad)
+            actual = [out.detach().clone(), inp.grad.clone()]
+            actual.extend(None if p.grad is None else p.grad.clone() for p in model.parameters())
+            ref_out = reference(ref_inp)
+            ref_out.backward(grad)
+            expected = [ref_out, ref_inp.grad, *(p.grad for p in reference.parameters())]
+            for result, target in zip(actual, expected):
+                if target is None:
+                    assert result is None
+                else:
+                    torch.testing.assert_close(result, target, **dtype_tols(dtype))
+    if compile_mode == "reduce-overhead":
+        assert not counters["inductor"]["cudagraph_skips"]
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("fp8_recipe", [None, *_all_recipes], ids=recipe_id)
+@pytest.mark.parametrize("bias_mode", ["fused", "none", "returned"])
+def test_te_split_parameters_recipes(fp8_recipe, bias_mode):
+    autocast_recipe = fp8_recipe or recipe.Float8CurrentScaling()
+    options = dict(bias=bias_mode != "none", return_bias=bias_mode == "returned")
+    model = te.Linear(
+        128,
+        256,
+        parameters_split=dict(q=128, k=64, v=64),
+        params_dtype=torch.bfloat16,
+        **options,
+    )
+
+    def fn(inp):
+        with te.autocast(enabled=fp8_recipe is not None, recipe=autocast_recipe):
+            out = model(inp)
+        if bias_mode == "returned":
+            out, bias = out
+            out = out + bias
+        return out
+
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+    inp = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    grad = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16)
+    expected = fn(inp)
+    expected.backward(grad)
+    expected_grads = [t.grad.clone() for t in (inp, *model.parameters()) if t.requires_grad]
+    inp.grad = None
+    model.zero_grad(set_to_none=True)
+    actual = compiled(inp)
+    actual.backward(grad)
+    actual_grads = [t.grad for t in (inp, *model.parameters()) if t.requires_grad]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_grads, expected_grads, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("training", [False, True])
+def test_te_split_parameters_no_cat(training):
+    """Adjacent parameter storage does not launch a concatenation in either pass."""
+    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
+    compiled = torch.compile(model, fullgraph=True)
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16, requires_grad=training)
+
+    def run():
+        with torch.set_grad_enabled(training):
+            out = compiled(inp)
+            if training:
+                out.sum().backward()
+        model.zero_grad(set_to_none=True)
+        inp.grad = None
+
+    for _ in range(3):
+        run()
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof:
+        run()
+    assert "aten::cat" not in {event.key for event in prof.key_averages()}
+    model.k_weight.data = model.k_weight.detach().clone()
+    with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof:
+        run()
+    assert "aten::cat" in {event.key for event in prof.key_averages()}
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+def test_te_split_parameters_saved_versions():
+    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
+    compiled = torch.compile(model, fullgraph=True, backend="aot_eager")
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    out = compiled(inp)
+    with torch.no_grad():
+        model.v_weight.add_(0.01)
+    with pytest.raises(RuntimeError, match="modified by an inplace operation"):
+        out.sum().backward()
+
+
+@pytest.mark.parametrize(
+    "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
+)
+def test_concat_storage(layout):
+    from transformer_engine.pytorch.dynamo.concat import materialize_concat
+    from transformer_engine.pytorch.module._common import concat_input
+
+    storage = torch.arange(128, device="cuda").view(16, 8)
+    if layout == "offset":
+        storage = storage[2:10]
+    elif layout == "strided":
+        storage = storage[:, ::2]
+    elif layout == "singleton":
+        storage = storage[:, :1]
+    elif layout == "negative":
+        storage = torch._neg_view(storage)
+    elif layout == "conjugate":
+        storage = torch.complex(storage.float(), storage.float()).conj()
+    parts = list(storage.split([1, storage.shape[0] - 3, 2]))
+    if layout == "disjoint":
+        parts[1] = parts[1].clone()
+    tensor, deferred = concat_input(parts, defer=True)
+    assert tensor is None
+    assert isinstance(deferred, list)
+    result = materialize_concat(deferred)
+    torch.testing.assert_close(result, torch.cat(parts), rtol=0, atol=0)
+    if layout in ("adjacent", "offset"):
+        assert result.data_ptr() == parts[0].data_ptr()
+    else:
+        assert result.untyped_storage().data_ptr() != parts[0].untyped_storage().data_ptr()
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize(
+    "optional, contents",
+    [(False, "full"), (False, "empty"), (True, "full"), (True, "empty"), (True, "none")],
+)
+def test_tensor_list_custom_op(optional, contents):
+    from transformer_engine.pytorch.dynamo import register_custom_op
+
+    @dataclasses.dataclass
+    class Args:
+        inp: torch.Tensor
+        tensors: list[torch.Tensor] | None if optional else list[torch.Tensor]
+
+    op = register_custom_op(
+        op_name=f"test_tensor_list_inputs_{optional}_{contents}",
+        arg_type=Args,
+        impl=lambda args: [args.inp + (args.tensors is not None)]
+        + [tensor.square() for tensor in (args.tensors or ())],
+        fake_impl=lambda args: [args.inp, *(args.tensors or ())],
+    )
+    assert op is not None
+
+    def fn(inp, tensors):
+        return op(Args(inp, tensors))
+
+    tensors = [torch.randn(3), torch.randn(2, 4, dtype=torch.float64)]
+    if contents == "empty":
+        tensors = []
+    elif contents == "none":
+        tensors = None
+    inp = torch.randn(2)
+    expected = [inp + (contents != "none"), *[tensor.square() for tensor in (tensors or ())]]
+    torch.testing.assert_close(fn(inp, tensors), expected)
+    compiled = torch.compile(fn, fullgraph=True, backend="aot_eager")
+    torch.testing.assert_close(compiled(inp, tensors), expected)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("grad_target", ["input", "weight"])
+def test_te_split_parameters_grad_targets(grad_target):
+    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
+    for parameter in model.parameters():
+        parameter.requires_grad_(False)
+    if grad_target != "input":
+        getattr(model, f"v_{grad_target}").requires_grad_(True)
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    inp.requires_grad_(grad_target == "input")
+    compiled = torch.compile(model, fullgraph=True)
+    out = compiled(inp)
+    grad = torch.randn_like(out)
+    out.backward(grad)
+    actual = [out.detach().clone()]
+    tensors = [inp, *model.parameters()]
+    actual.extend(None if t.grad is None else t.grad.clone() for t in tensors)
+    for tensor in tensors:
+        tensor.grad = None
+    expected = model(inp)
+    expected.backward(grad)
+    for result, target in zip(actual, [expected, *(t.grad for t in tensors)]):
+        if target is None:
+            assert result is None
+        else:
+            torch.testing.assert_close(result, target, rtol=0, atol=0)
diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
new file mode 100644
index 00000000000..e6ac9882613
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Deferred row concatenation for tensors consumed inside a custom op."""
+
+from typing import List, Sequence, Tuple, Union
+
+import torch
+from torch._prims_common import make_contiguous_strides_for
+
+from .tensor_spec import TensorSpec
+from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
+
+
+def validate_concat(parts: List[torch.Tensor]) -> None:
+    """Require a nonempty list of non-quantized tensors."""
+    if not parts:
+        raise ValueError("Attempted to concatenate 0 tensors")
+    if any(
+        not isinstance(part, (torch.Tensor, TensorSpec))
+        or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
+        or (isinstance(part, TensorSpec) and part.is_quantized)
+        for part in parts
+    ):
+        raise TypeError("Deferred concatenation only supports non-quantized tensors")
+
+
+def concat_shape(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> Tuple[int, ...]:
+    """Shape after row concatenation."""
+    return (sum(t.shape[0] for t in parts), *parts[0].shape[1:])
+
+
+def concat_to_spec(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> TensorSpec:
+    """Describe the concatenation without accessing its storage."""
+    validate_concat(parts)
+    first = parts[0]
+    dtype = first.dtype
+    for tensor in parts[1:]:
+        dtype = torch.promote_types(dtype, tensor.dtype)
+    return TensorSpec(
+        shape=concat_shape(parts),
+        dtype=dtype,
+        device=first.device,
+        requires_grad=any(t.requires_grad for t in parts),
+    )
+
+
+def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor:
+    """Build a checked view, or copy parts whose storage is not adjacent."""
+    validate_concat(value)
+    if len(value) == 1:
+        return value[0]
+    first = value[0]
+    storage = first.untyped_storage()
+    offset = first.storage_offset()
+    for tensor in value:
+        if (
+            tensor.dtype != first.dtype
+            or tensor.device != first.device
+            or tensor.is_neg()
+            or tensor.is_conj()
+        ):
+            return torch.cat(value)
+        if (
+            tensor.shape[1:] != first.shape[1:]
+            or not tensor.is_contiguous()
+            or tensor.untyped_storage().data_ptr() != storage.data_ptr()
+            or tensor.storage_offset() != offset
+        ):
+            return torch.cat(value)
+        offset += tensor.numel()
+    if offset * first.element_size() > storage.nbytes():
+        return torch.cat(value)
+    shape = concat_shape(value)
+    return first.as_strided(shape, make_contiguous_strides_for(shape))
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 39d59dc2907..a3e555f5a62 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,6 +37,7 @@
   * ``TENSOR_OR_QUANTIZED`` -- a field that may be a plain tensor, a bare
     quantized storage, or ``None``: three slots (the tensor, its flat inner
     buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers.
+  * ``TENSOR_LIST`` -- a ``List[Tensor]`` field: one ``Tensor[]`` slot.
   * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes,
     quantizers -- value-opaque constants baked into the graph -- and nested
     collections of them), gathered into one shared ``OpaqueValueBundle`` slot.
@@ -396,6 +397,7 @@ class _FieldKind(Enum):
     """How one dataclass field crosses the custom-op boundary."""
 
     TENSOR = "tensor"  # one ``Tensor`` / ``Tensor?`` slot
+    TENSOR_LIST = "tensor_list"  # one ``Tensor[]`` / ``Tensor[]?`` slot
     TENSOR_OR_QUANTIZED = "tensor_or_quantized"  # 3 slots: tensor / inner / meta
     PROCESS_GROUP = "process_group"  # c10d group name inside the shared bundle
     SIMPLE = "simple"  # value carried verbatim inside the shared bundle
@@ -413,8 +415,7 @@ class _TensorOrQuantizedKind(Enum):
 _TQ_KIND_KEY = "__kind__"
 _SIMPLE_META_SLOT = "_simple_meta"
 
-# Matched by exact member set, so a bare quantized annotation or an accidental
-# extra union member is rejected rather than silently taken as tensor-or-quantized.
+# Tensor unions are matched by exact member sets.
 _TQ_MEMBERS = frozenset(get_args(TensorOrQuantized))
 
 
@@ -439,8 +440,13 @@ class _FieldPlan:
     slots: Tuple[_SlotSpec, ...]
 
 
+def _is_tensor_list_annot(annot: Any) -> bool:
+    """Whether an annotation is a list of tensors."""
+    return get_origin(annot) is list and get_args(annot) == (torch.Tensor,)
+
+
 def _is_tensor_storage_union(annot: Any) -> bool:
-    """Whether ``annot`` is exactly the tensor-or-quantized union."""
+    """Whether a tensor union can use the tensor / buffers / metadata slots."""
     if not _is_union(annot):
         return False
     members = frozenset(a for a in get_args(annot) if a is not type(None))
@@ -488,6 +494,9 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
         )
         return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots)
     stripped, is_optional = _strip_optional(annot)
+    if _is_tensor_list_annot(stripped):
+        slot = _SlotSpec(name, "Tensor[]?" if is_optional else "Tensor[]")
+        return _FieldPlan(name, _FieldKind.TENSOR_LIST, (slot,))
     if stripped is torch.Tensor:
         slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor")
         return _FieldPlan(name, _FieldKind.TENSOR, (slot,))
@@ -582,7 +591,7 @@ def tensor_field_names(self) -> Tuple[str, ...]:
         return tuple(
             f.name
             for f in self.fields
-            if f.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED)
+            if f.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_LIST, _FieldKind.TENSOR_OR_QUANTIZED)
         )
 
     def tensor_or_quantized_offsets(self) -> List[int]:
@@ -608,7 +617,11 @@ def resolve_grad_targets(self, input_tensors_for_grad: Sequence[str]) -> List[in
         index: Dict[str, int] = {}
         offset = 0
         for field in self.fields:
-            if field.kind in (_FieldKind.TENSOR, _FieldKind.TENSOR_OR_QUANTIZED):
+            if field.kind in (
+                _FieldKind.TENSOR,
+                _FieldKind.TENSOR_LIST,
+                _FieldKind.TENSOR_OR_QUANTIZED,
+            ):
                 # The gradient flows to the group's first slot -- for
                 # tensor-or-quantized that is the ``Tensor?`` slot, the one
                 # autograd sees the (subclass) tensor in.
@@ -631,7 +644,7 @@ def pack(self, obj: Any) -> Dict[str, Any]:
         for field in self.fields:
             value = getattr(obj, field.name, None)
             match field.kind:
-                case _FieldKind.TENSOR:
+                case _FieldKind.TENSOR | _FieldKind.TENSOR_LIST:
                     slots[field.slots[0].name] = value
                 case _FieldKind.TENSOR_OR_QUANTIZED:
                     _pack_tensor_or_quantized(field, value, slots)
@@ -665,7 +678,7 @@ def unpack(self, slots: Dict[str, Any]) -> Any:
         bundle = slots.get(_SIMPLE_META_SLOT)
         for field in self.fields:
             match field.kind:
-                case _FieldKind.TENSOR:
+                case _FieldKind.TENSOR | _FieldKind.TENSOR_LIST:
                     kwargs[field.name] = slots[field.slots[0].name]
                 case _FieldKind.TENSOR_OR_QUANTIZED:
                     kwargs[field.name] = _unpack_tensor_or_quantized(field, slots)
@@ -741,7 +754,9 @@ def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any:
     overrides: Dict[str, Any] = {}
     for name in tensor_field_names:
         value = getattr(obj, name, None)
-        if value is not None and not isinstance(value, TensorSpec):
+        if isinstance(value, list):
+            overrides[name] = [to_tensor_spec(tensor) for tensor in value]
+        elif value is not None and not isinstance(value, TensorSpec):
             overrides[name] = to_tensor_spec(value)
     if not overrides:
         return obj
@@ -1066,6 +1081,10 @@ def _autograd_backward(ctx, *grad_outputs):
         kwargs = bwd_plan.pack(bwd_obj)
         bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names]
         grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)]
+        if hasattr(bwd_obj, "postprocess_gradients"):
+            grads = bwd_obj.postprocess_gradients(grads)
+        if len(grads) != len(grad_targets):
+            raise RuntimeError("Expected one gradient per input_tensors_for_grad entry")
         ctx.backward_objects = None
         # One grad per input schema slot: default None, but a ``Tensor[]`` slot
         # (always recorded in ``fwd_tensor_list_lengths``) needs a
@@ -1074,6 +1093,12 @@ def _autograd_backward(ctx, *grad_outputs):
         for pos, length in ctx.fwd_tensor_list_lengths.items():
             out[pos] = [None] * length
         for pos, g in zip(grad_targets, grads):
+            if pos in ctx.fwd_tensor_list_lengths:
+                if g is not None:
+                    if not isinstance(g, list) or len(g) != ctx.fwd_tensor_list_lengths[pos]:
+                        raise RuntimeError("Expected one gradient per tensor-list element")
+                    out[pos] = g
+                continue
             if g is not None:
                 shape = ctx.grad_input_shapes.get(pos)
                 if shape is not None and g.shape != shape:
@@ -1344,9 +1369,10 @@ def register_custom_op_with_autograd(
     * ``setup_context(bwd_obj, fwd_args, user_outputs, ctx_attrs, saved)
       -> tensors_to_save`` -- populate ``bwd_obj`` from forward state; return the
       tensors to persist across the boundary.
-    * ``bwd_impl(bwd_args) -> grads`` -- exactly one grad per
-      ``input_tensors_for_grad`` entry, in that order (``None`` for a
-      non-differentiable input).
+    * ``bwd_impl(bwd_args) -> grads`` -- flat tensor grads (or ``None``).
+      Without a ``bwd_arg_type.postprocess_gradients(grads)`` hook, return one
+      grad per ``input_tensors_for_grad`` entry. With the hook, it maps the
+      flat results to those entries outside the op, including tensor-list grads.
     * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning
       :class:`TensorSpec` grads.
     * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if
@@ -1420,7 +1446,9 @@ def _register_custom_op_with_autograd_impl(
         tags=fwd_tags,
     )
     bwd_qualname = f"{_TE_OP_NAMESPACE}::{op_name}_backward_base"
-    num_grad_inputs = len(input_tensors_for_grad)
+    num_grad_inputs = (
+        None if hasattr(bwd_arg_type, "postprocess_gradients") else len(input_tensors_for_grad)
+    )
     bwd_op = _register_op(
         name=f"{op_name}_backward",
         arg_type=bwd_arg_type,
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 3d04882698a..62572a3a6f6 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -6,14 +6,16 @@
 
 import dataclasses
 import queue
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 
 import torch
 
 from .. import cpp_extensions as tex
 from ..constants import TE_DType
 from ..distributed import in_fp8_activation_recompute_phase
+from ..dynamo import TensorOrQuantized
 from ..export import is_in_onnx_export_mode
+from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
 from ..quantization import FP8GlobalStateManager
 from ..tensor.hybrid_tensor import HybridQuantizer
 from ..utils import get_default_init_method
@@ -239,6 +241,27 @@ def noop_cat(
     return _NoopCatFunc.apply(dim, *tensors)
 
 
+def concat_input(
+    tensors: Sequence[TensorOrQuantized], *, defer: bool = False
+) -> Tuple[Optional[TensorOrQuantized], Optional[List[torch.Tensor]]]:
+    """Return a parameter operand or its deferred parts."""
+    if len(tensors) == 1:
+        return tensors[0], None
+    if (
+        defer
+        and tensors
+        and all(
+            isinstance(t, torch.Tensor)
+            and not isinstance(t, (QuantizedTensor, QuantizedTensorStorage))
+            for t in tensors
+        )
+    ):
+        return None, list(tensors)
+    if torch.compiler.is_compiling():
+        return torch.cat(tensors), None
+    return noop_cat(tensors), None
+
+
 @dataclasses.dataclass
 class _ParameterInitMeta:
     """
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 9ead1e73651..5d150b8bac4 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -10,7 +10,7 @@
 import warnings
 from enum import Enum
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, Union
 from contextlib import contextmanager
 from types import MethodType
 
@@ -1322,7 +1322,9 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
     def _enable_weight_preswizzle(
         self,
         quantizer: Quantizer,
-        weight: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        *,
+        weight_shape: Optional[Sequence[int]] = None,
     ) -> bool:
         """Whether to fuse scale-factor swizzling into weight quantization.
 
@@ -1339,7 +1341,8 @@ def _enable_weight_preswizzle(
         if isinstance(quantizer, MXFP8Quantizer):
             return True
         if isinstance(quantizer, NVFP4Quantizer):
-            rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+            shape = weight.shape if weight_shape is None else weight_shape
+            rows, cols = math.prod(shape[:-1]), shape[-1]
             arch_supported = get_device_compute_capability() >= (10, 0)
             if quantizer.with_rht:
                 return arch_supported and rows % 64 == 0 and cols % 128 == 0
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 7f00e19644a..93cc83525e2 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -32,7 +32,7 @@
 from ._common import (
     can_reconstruct_wgrad_input_from_original,
     check_fp8_reduce_and_update,
-    noop_cat,
+    concat_input,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -84,6 +84,7 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
+from ..dynamo.concat import concat_shape, concat_to_spec, materialize_concat
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -110,7 +111,7 @@ class LinearFwdArgs:
     """Single-argument bag for the forward path of :class:`_Linear`."""
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
-    weight: TensorOrQuantized
+    weight: Optional[TensorOrQuantized]
     inp: torch.Tensor
     bias: Optional[torch.Tensor]
 
@@ -178,6 +179,8 @@ class LinearFwdArgs:
     # --- Misc ---
     cpu_offloading: bool
     is_grad_enabled: bool
+    weight_list: Optional[List[torch.Tensor]] = None
+    bias_list: Optional[List[torch.Tensor]] = None
 
     def compile_unsupported_reason(self) -> Optional[str]:
         """Reason this config can't use the torch.compile custom-op path (else None)."""
@@ -241,6 +244,10 @@ class LinearBwdArgs:
     weight_fp8: Optional[TensorOrQuantized] = None
     saved_weight: Optional[TensorOrQuantized] = None
     bias: Optional[torch.Tensor] = None
+    weight_list: Optional[List[torch.Tensor]] = None
+    bias_list: Optional[List[torch.Tensor]] = None
+    parameter_list_lengths: Tuple[int, int] = (0, 0)
+    weight_fp8_from_list: bool = False
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -306,14 +313,33 @@ class LinearBwdArgs:
 
     def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None:
         """Pull saved tensors from ``ctx`` into the fields backward consumes."""
-        (
-            self.inputmat,
-            self.weight_fp8,
-            self.saved_weight,
-            self.bias,
-        ) = restore_from_func_ctx(
-            ctx
-        )  # pylint: disable=unbalanced-tuple-unpacking
+        saved = restore_from_func_ctx(ctx)
+        if self.parameter_list_lengths == (0, 0):
+            self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved
+            return
+        self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved[:4]
+        weight_count, bias_count = self.parameter_list_lengths
+        self.weight_list = list(saved[4 : 4 + weight_count]) if weight_count else None
+        self.bias_list = (
+            list(saved[4 + weight_count : 4 + weight_count + bias_count]) if bias_count else None
+        )
+
+    def postprocess_gradients(self, grads):
+        """Split gradients outside the op so the parts can share storage."""
+        wgrad, dgrad, bgrad = grads
+        weight_grads = None
+        bias_grads = None
+        if self.weight_list and wgrad is not None:
+            weight_grads = list(wgrad.split([t.shape[0] for t in self.weight_list]))
+        if self.bias_list and bgrad is not None:
+            bias_grads = list(bgrad.split([t.shape[0] for t in self.bias_list]))
+        return (
+            None if self.weight_list else wgrad,
+            dgrad,
+            None if self.bias_list else bgrad,
+            weight_grads,
+            bias_grads,
+        )
 
 
 def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int:
@@ -361,10 +387,10 @@ def _linear_forward_impl(
     cache). The last two are ``None`` when gradients are disabled.
     """
 
-    weight = args.weight
-    is_dist_weight = is_distributed_weight(args.weight)
+    weight = materialize_concat(args.weight_list) if args.weight_list else args.weight
+    is_dist_weight = is_distributed_weight(weight)
     inp = args.inp
-    bias = args.bias
+    bias = materialize_concat(args.bias_list) if args.bias_list else args.bias
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
     output_quantizer = args.output_quantizer
@@ -774,9 +800,9 @@ def _linear_forward_fake(
             "(fsdp_group is not None); use FSDP2 or MCore FSDP."
         )
 
-    weight = args.weight
+    weight = concat_to_spec(args.weight_list) if args.weight_list else args.weight
     inp = args.inp
-    bias = args.bias
+    bias = concat_to_spec(args.bias_list) if args.bias_list else args.bias
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
     output_quantizer = args.output_quantizer
@@ -1018,7 +1044,7 @@ def _linear_setup_ctx(
     bwd_args.grad_output_quantizer = fwd_args.grad_output_quantizer
 
     # Differentiability summary
-    bwd_args.use_bias = bias is not None
+    bwd_args.use_bias = bias is not None or bool(fwd_args.bias_list)
     bwd_args.requires_dgrad = fwd_args.input_requires_grad
     bwd_args.requires_wgrad = fwd_args.weight_requires_grad
     # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True)
@@ -1101,7 +1127,20 @@ def _linear_setup_ctx(
         saved_weight = weight
     if bias_alias == "bias":
         saved_bias = bias
-    return (saved_inputmat, wt_save, saved_weight, saved_bias)
+    if not fwd_args.weight_list and not fwd_args.bias_list:
+        return saved_inputmat, wt_save, saved_weight, saved_bias
+    bwd_args.weight_fp8_from_list = wt_save_alias == "weight" and bool(fwd_args.weight_list)
+    weight_list = fwd_args.weight_list or ()
+    bias_list = fwd_args.bias_list or ()
+    bwd_args.parameter_list_lengths = (len(weight_list), len(bias_list))
+    return (
+        saved_inputmat,
+        wt_save,
+        saved_weight,
+        saved_bias,
+        *weight_list,
+        *bias_list,
+    )
 
 
 def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]:
@@ -1114,10 +1153,10 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
     grad_output = args.grad_output
     assert grad_output is not None
     inputmat = args.inputmat
-    weight_fp8 = args.weight_fp8
-    saved_weight = args.saved_weight
+    saved_weight = materialize_concat(args.weight_list) if args.weight_list else args.saved_weight
+    weight_fp8 = saved_weight if args.weight_fp8_from_list else args.weight_fp8
     is_dist_weight = is_distributed_weight(saved_weight)
-    bias = args.bias
+    bias = materialize_concat(args.bias_list) if args.bias_list else args.bias
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
     grad_input_quantizer = args.grad_input_quantizer
@@ -1735,7 +1774,7 @@ def _linear_backward_fake(
             "(fsdp_group is not None); use FSDP2 or MCore FSDP."
         )
 
-    weight = args.saved_weight
+    weight = concat_to_spec(args.weight_list) if args.weight_list else args.saved_weight
     out_dtype = args.activation_dtype
     out_features, in_features = weight.shape
 
@@ -1790,7 +1829,7 @@ def _linear_backward_fake(
 # Custom op used under ``torch.compile``.
 _linear_op = register_custom_op_with_autograd(
     op_name="linear",
-    input_tensors_for_grad=["weight", "inp", "bias"],
+    input_tensors_for_grad=["weight", "inp", "bias", "weight_list", "bias_list"],
     fwd_arg_type=LinearFwdArgs,
     fwd_impl=_linear_forward_impl,
     fwd_fake_impl=_linear_forward_fake,
@@ -1931,7 +1970,9 @@ class Linear(TransformerEngineBaseModule):
                       (preferably an OrderedDict) is provided, the keys are used as names and
                       values as split sizes along dim 0. The resulting parameters will have
                       names that end in ``_weight`` or ``_bias``, so trailing underscores are
-                      stripped from any provided names.
+                      stripped from any provided names. Under ``torch.compile``, adjacent
+                      parts sharing storage are consumed without a concatenation copy.
+                      Disjoint parts and returned split biases still require concatenation.
     device : Union[torch.device, str], default = "cuda"
           The device on which the parameters of the model will be allocated. It is the user's
           responsibility to ensure all parameters are moved to the GPU before running the
@@ -2387,7 +2428,9 @@ def forward(
             if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()):
                 fp8_grad = True
 
-        if torch.compiler.is_compiling() and _linear_op is not None:
+        is_compiling = torch.compiler.is_compiling()
+        use_compiled_op = is_compiling and _linear_op is not None
+        if use_compiled_op:
             reason = self._compile_eager_fallback_reason(
                 inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug
             )
@@ -2399,7 +2442,13 @@ def forward(
 
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
-            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+            weight, bias, weight_list, bias_list = self._get_weight_and_bias_tensors(
+                defer_concatenation=use_compiled_op
+            )
+            weight_shape = concat_shape(weight_list) if weight_list else None
+            weight_requires_grad = (
+                any(t.requires_grad for t in weight_list) if weight_list else weight.requires_grad
+            )
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2421,11 +2470,10 @@ def forward(
             ) = quantizers
             if weight_quantizer is not None and not debug:
                 weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
-                    weight_quantizer, weight_tensor
+                    weight_quantizer, weight, weight_shape=weight_shape
                 )
 
-            use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
-            if _linear_op is None and torch.compiler.is_compiling():
+            if _linear_op is None and is_compiling:
                 warn_if_compile_disabled()
             if use_compiled_op:
                 # Process groups cross the op boundary separately from quantizers.
@@ -2451,9 +2499,7 @@ def forward(
                 backward_override = None
             custom = is_custom(input_quantizer) or is_custom(weight_quantizer)
             backward_input_needs_gather = (
-                weight_tensor.requires_grad
-                and self.parallel_mode == "column"
-                and self.sequence_parallel
+                weight_requires_grad and self.parallel_mode == "column" and self.sequence_parallel
             )
 
             if debug:
@@ -2471,22 +2517,24 @@ def forward(
                 ub_bulk_dgrad = self.ub_bulk_dgrad
                 ub_bulk_wgrad = self.ub_bulk_wgrad
 
-            linear_bias_tensor = (
-                bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None
-            )
+            linear_bias = bias if (self.apply_bias and not self.gemm_bias_unfused_add) else None
             wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None
 
             fwd_args = LinearFwdArgs(
                 # tensors
-                weight=weight_tensor,
+                weight=weight,
                 inp=inp,
-                bias=linear_bias_tensor,
+                bias=linear_bias,
+                weight_list=weight_list,
+                bias_list=bias_list,
                 weight_workspace=weight_workspace,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
-                weight_requires_grad=weight_tensor.requires_grad,
+                weight_requires_grad=weight_requires_grad,
                 bias_requires_grad=(
-                    linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False
+                    any(t.requires_grad for t in bias_list)
+                    if bias_list
+                    else linear_bias is not None and linear_bias.requires_grad
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2546,13 +2594,23 @@ def forward(
                         msg=f"te.Linear falling back to eager: {fallback_reason}"
                     )
                     use_compiled_op = False
+                    weight, bias, _, _ = self._get_weight_and_bias_tensors()
+                    linear_bias = (
+                        bias if self.apply_bias and not self.gemm_bias_unfused_add else None
+                    )
+                    fwd_args.weight = weight
+                    fwd_args.bias = linear_bias
+                    fwd_args.weight_list = None
+                    fwd_args.bias_list = None
 
             if use_compiled_op:
-                check_gemm_dims(inp, weight_tensor, self.fp8)
+                check_gemm_dims(
+                    inp.shape, weight_shape if weight_shape is not None else weight.shape, self.fp8
+                )
                 out, new_weight_workspace = _linear_op(fwd_args)
             else:
                 out, new_weight_workspace = _linear_eager(
-                    weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled
+                    weight, inp, linear_bias, fwd_args, is_grad_enabled
                 )
 
             if new_weight_workspace is not None and cache_name is not None:
@@ -2563,10 +2621,10 @@ def forward(
         finally:
             self.end_forward()
         if self.gemm_bias_unfused_add:
-            out = out + cast_if_needed(bias_tensor, self.activation_dtype)
+            out = out + cast_if_needed(bias, self.activation_dtype)
 
         if self.return_bias:
-            return out, cast_if_needed(bias_tensor, self.activation_dtype)
+            return out, cast_if_needed(bias, self.activation_dtype)
         return out
 
     def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
@@ -2643,8 +2701,7 @@ def _compile_eager_fallback_reason(
         prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
-        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
-        if is_distributed_weight(weight_tensor):
+        if any(is_distributed_weight(getattr(self, name)) for name in self.weight_names):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
         if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)):
             return "a quantized input tensor"
@@ -2652,8 +2709,10 @@ def _compile_eager_fallback_reason(
             return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
         any_requires_grad = (
             inp.requires_grad
-            or weight_tensor.requires_grad
-            or (bias_tensor is not None and bias_tensor.requires_grad)
+            or any(getattr(self, name).requires_grad for name in self.weight_names)
+            or (
+                self.use_bias and any(getattr(self, name).requires_grad for name in self.bias_names)
+            )
         )
         if fp8_output and is_grad_enabled and any_requires_grad:
             return "differentiable fp8_output=True"
@@ -2693,15 +2752,18 @@ def _forward_eager_fallback(
             fp8_grad=fp8_grad,
         )
 
-    def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        # Get concatenated weight and bias tensors
-        unfused_weights = self._get_weight_tensors()
-        weight_tensor = noop_cat(unfused_weights)
+    def _get_weight_and_bias_tensors(self, defer_concatenation=False):
+        weights = self._get_weight_tensors()
+        weight, weight_list = concat_input(weights, defer=defer_concatenation)
+        bias = None
+        bias_list = None
         if self.use_bias:
-            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
-        else:
-            bias_tensor = None
-        return weight_tensor, bias_tensor
+            biases = [getattr(self, name) for name in self.bias_names]
+            bias, bias_list = concat_input(
+                biases,
+                defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
+            )
+        return weight, bias, weight_list, bias_list
 
     def onnx_forward(
         self,
@@ -2718,7 +2780,7 @@ def onnx_forward(
 
         assert_warmed_up(self)
         assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export."
-        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+        weight_tensor, bias_tensor, _, _ = self._get_weight_and_bias_tensors()
         (
             input_quantizer,
             weight_quantizer,
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index aa21ed06e05..d7fcf2c6dbf 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -718,21 +718,21 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None:
             )
 
 
-def check_gemm_dims(inp: torch.Tensor, weight: torch.Tensor, fp8: bool) -> None:
+def check_gemm_dims(inp_shape: Sequence[int], weight_shape: Sequence[int], fp8: bool) -> None:
     """Emit the TN GEMM (``y = x @ w^T``) dim constraints as ``torch._check``
     guards at trace time. torch.compile path only; eager validation lives in
     the op impl. Messages are constant: Dynamo forbids tensor closures here.
     """
     # pylint: disable=protected-access
     torch._check(
-        inp.shape[-1] == weight.shape[-1],
+        inp_shape[-1] == weight_shape[-1],
         lambda: "GEMM not possible: input last dim must equal in_features",
     )
     if not fp8:
         return
-    for tensor, name in ((inp, "input"), (weight, "weight")):
+    for shape, name in ((inp_shape, "input"), (weight_shape, "weight")):
         torch._check(
-            math.prod(tensor.shape[:-1]) % 8 == 0 and tensor.shape[-1] % 16 == 0,
+            math.prod(shape[:-1]) % 8 == 0 and shape[-1] % 16 == 0,
             lambda n=name: (
                 f"FP8 execution requires the {n}'s product of all dimensions except the"
                 " last to be divisible by 8 and its last dimension to be divisible by 16"