From fdaea9e1f31542453f9d99e53587cf6376e6a803 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 5 Oct 2026 11:42:18 +0200
Subject: [PATCH 01/13] [PyTorch] Avoid concatenation copies for compiled split
 Linear parameters

Pass split Linear parameters through the custom-op boundary, reconstruct checked views at execution time and split full gradients outside backward. Preserve concatenation for disjoint storage and returned biases, handle promoted dtypes, and compute bias gradients when weights are frozen.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 214 ++++++++++++++++++
 .../pytorch/dynamo/concatenated_tensor.py     |  89 ++++++++
 .../pytorch/dynamo/custom_op.py               |  60 ++++-
 transformer_engine/pytorch/module/linear.py   |  59 +++--
 4 files changed, 396 insertions(+), 26 deletions(-)
 create mode 100644 transformer_engine/pytorch/dynamo/concatenated_tensor.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index eae6f0a8a24..195b9bc339a 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2720,3 +2720,217 @@ def test_te_ops_forward_kwargs_compile():
         for gain in (3.0, 5.0):
             _check_ops(compiled, model, x, dy, {"gain": gain})
     _assert_custom_ops(graphs[-1:], "_affineop", present=False)
+
+
+@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):
+    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=fp8_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_concatenated_tensor_storage(layout):
+    from transformer_engine.pytorch.dynamo.concatenated_tensor import ConcatenatedTensor
+
+    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()
+    result = ConcatenatedTensor(parts).materialize()
+    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("grad_target", ["input", "weight", "bias"])
+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)
+    if grad_target == "bias":
+        torch.testing.assert_close(model.v_bias.grad, grad[:, -32:].sum(0), rtol=0, atol=0)
+    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)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("bf16_part", ["q", "k"])
+def test_te_split_parameters_mixed_dtype_autocast(bf16_part):
+    model = te.Linear(64, 192, parameters_split=("q", "k", "v"), params_dtype=torch.float32)
+    name = f"{bf16_part}_weight"
+    setattr(model, name, torch.nn.Parameter(getattr(model, name).to(torch.bfloat16)))
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.float32, requires_grad=True)
+    grad = torch.randn(16, 192, device="cuda", dtype=torch.bfloat16)
+    torch._dynamo.reset()
+    compiled = torch.compile(model, fullgraph=True)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        expected = model(inp)
+    expected.backward(grad)
+    expected_grads = [tensor.grad.clone() for tensor in (inp, *model.parameters())]
+    inp.grad = None
+    model.zero_grad(set_to_none=True)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        actual = compiled(inp)
+    actual.backward(grad)
+    actual_grads = [tensor.grad for tensor in (inp, *model.parameters())]
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_grads, expected_grads, rtol=0, atol=0)
diff --git a/transformer_engine/pytorch/dynamo/concatenated_tensor.py b/transformer_engine/pytorch/dynamo/concatenated_tensor.py
new file mode 100644
index 00000000000..6f6c39ebc6f
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/concatenated_tensor.py
@@ -0,0 +1,89 @@
+# 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 dataclasses import dataclass
+import math
+from typing import List
+
+import torch
+from torch._prims_common import make_contiguous_strides_for
+
+from ..quantized_tensor import restore_from_func_ctx as restore_tensor_ctx
+
+
+@dataclass
+class ConcatenatedTensor:
+    """Keep every parameter visible to autograd until the opaque consumer runs."""
+
+    tensors: List[torch.Tensor]
+
+    @property
+    def shape(self):
+        """Combined shape along dimension zero."""
+        return (sum(t.shape[0] for t in self.tensors), *self.tensors[0].shape[1:])
+
+    @property
+    def dtype(self):
+        """Parameter dtype."""
+        dtype = self.tensors[0].dtype
+        for tensor in self.tensors[1:]:
+            dtype = torch.promote_types(dtype, tensor.dtype)
+        return dtype
+
+    @property
+    def device(self):
+        """Parameter device."""
+        return self.tensors[0].device
+
+    @property
+    def requires_grad(self):
+        """Whether any part requires a gradient."""
+        return any(t.requires_grad for t in self.tensors)
+
+    def numel(self):
+        """Combined element count."""
+        return math.prod(self.shape)
+
+    def materialize(self):
+        """Build a checked view, or copy parts whose storage is not adjacent."""
+        first = self.tensors[0]
+        storage = first.untyped_storage()
+        offset = first.storage_offset()
+        for tensor in self.tensors:
+            if tensor.dtype != first.dtype or tensor.device != first.device:
+                return torch.cat(self.tensors)
+            if (
+                tensor.shape[1:] != first.shape[1:]
+                or not tensor.is_contiguous()
+                or tensor.is_neg()
+                or tensor.is_conj()
+                or tensor.untyped_storage().data_ptr() != storage.data_ptr()
+                or tensor.storage_offset() != offset
+            ):
+                return torch.cat(self.tensors)
+            offset += tensor.numel()
+        if offset * first.element_size() > storage.nbytes():
+            return torch.cat(self.tensors)
+        return first.as_strided(self.shape, make_contiguous_strides_for(self.shape))
+
+
+def restore_from_func_ctx(ctx):
+    """Restore original split parameters saved by the custom-op framework."""
+    tensors = restore_tensor_ctx(ctx)
+    lengths = getattr(ctx, "concatenated_saved_lengths", None)
+    if lengths is None:
+        return tensors
+    restored = []
+    offset = 0
+    for length in lengths:
+        if length is None:
+            restored.append(tensors[offset])
+            offset += 1
+        else:
+            restored.append(ConcatenatedTensor(list(tensors[offset : offset + length])))
+            offset += length
+    ctx.concatenated_saved_lengths = None
+    return restored
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index a87374095a7..5225e293456 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,6 +37,10 @@
   * ``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.
+    Unions including ``ConcatenatedTensor`` also accept deferred row concatenations:
+    all parts cross in the buffer list, with full gradients split outside the op.
+    Saved parts use ``concatenated_tensor.restore_from_func_ctx`` in the backward
+    container; reconstruction of the full view happens only in the real kernel.
   * ``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.
@@ -106,6 +110,7 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
+from .concatenated_tensor import ConcatenatedTensor
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -408,13 +413,13 @@ class _TensorOrQuantizedKind(Enum):
     NONE = "none"
     TENSOR = "tensor"
     STORAGE = "storage"
+    CONCATENATED = "concatenated"
 
 
 _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))
 
 
@@ -437,14 +442,19 @@ class _FieldPlan:
     name: str
     kind: _FieldKind
     slots: Tuple[_SlotSpec, ...]
+    allows_concatenation: bool = False
 
 
 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))
-    return members == _TQ_MEMBERS
+    return members in (
+        _TQ_MEMBERS,
+        _TQ_MEMBERS | {ConcatenatedTensor},
+        frozenset((torch.Tensor, ConcatenatedTensor)),
+    )
 
 
 def _is_process_group_annot(annot: Any) -> bool:
@@ -486,7 +496,9 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__tensors", "Tensor[]"),
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
-        return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots)
+        return _FieldPlan(
+            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, ConcatenatedTensor in get_args(annot)
+        )
     stripped, is_optional = _strip_optional(annot)
     if stripped is torch.Tensor:
         slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor")
@@ -527,6 +539,12 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         slots[tensor_slot] = None
         slots[inner_slot] = []
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE})
+    elif isinstance(value, ConcatenatedTensor):
+        if not field.allows_concatenation:
+            raise TypeError(f"field {field.name!r} does not accept ConcatenatedTensor")
+        slots[tensor_slot] = None
+        slots[inner_slot] = value.tensors
+        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
         # ``Tensor?`` slot; subclass flattening (if any) is done by the
@@ -555,6 +573,8 @@ def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any
         return None
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
+    if kind == _TensorOrQuantizedKind.CONCATENATED:
+        return ConcatenatedTensor(slots[inner_slot])
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -984,8 +1004,14 @@ def _register_base_op(
     ``pack_result``.
     """
 
+    concatenated_fields = tuple(f.name for f in plan.fields if f.allows_concatenation)
+
     def _impl(*flat: Any) -> List[torch.Tensor]:
         obj = plan.unpack(dict(zip(plan.slot_names, flat)))
+        for name in concatenated_fields:
+            value = getattr(obj, name)
+            if isinstance(value, ConcatenatedTensor):
+                setattr(obj, name, value.materialize())
         return pack_result(impl(obj))
 
     def _fake(*flat: Any) -> List[torch.Tensor]:
@@ -1039,7 +1065,16 @@ def _setup_context(ctx, inputs, output):
             out_plan.ctx_attrs,
             tuple(saved_list),
         )
-        tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ()))
+        saved = []
+        ctx.concatenated_saved_lengths = []
+        for value in tensors_to_save_from_setup or ():
+            if isinstance(value, ConcatenatedTensor):
+                ctx.concatenated_saved_lengths.append(len(value.tensors))
+                saved.extend(value.tensors)
+            else:
+                ctx.concatenated_saved_lengths.append(None)
+                saved.append(value)
+        tensors_to_save, tensor_objects = prepare_for_saving(*saved)
         ctx.tensor_objects = tensor_objects
         ctx.save_for_backward(*tensors_to_save)
         ctx.backward_objects = bwd_obj
@@ -1047,6 +1082,13 @@ def _setup_context(ctx, inputs, output):
         # Input shapes for the grad slots (SymInt-safe on ctx): a bwd impl may
         # rederive shapes lossily (e.g. rank-1 inputs come back rank-2), so the
         # returned grads are viewed back to the true input shapes below.
+        ctx.concatenated_grad_shapes = {
+            pos: [tensor.shape for tensor in inputs[pos + 1]]
+            for pos in grad_targets
+            if pos in fwd_plan.tensor_or_quantized_offsets()
+            and inputs[pos] is None
+            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.CONCATENATED
+        }
         ctx.grad_input_shapes = {
             pos: inputs[pos].shape for pos in grad_targets if isinstance(inputs[pos], torch.Tensor)
         }
@@ -1073,12 +1115,18 @@ 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.concatenated_grad_shapes:
+                shapes = ctx.concatenated_grad_shapes[pos]
+                if g is not None:
+                    out[pos + 1] = list(torch.split(g, [shape[0] for shape in shapes], dim=0))
+                continue
             if g is not None:
                 shape = ctx.grad_input_shapes.get(pos)
                 if shape is not None and g.shape != shape:
                     g = g.view(shape)
             out[pos] = g
         ctx.grad_input_shapes = None
+        ctx.concatenated_grad_shapes = None
         return tuple(out)
 
     fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 7f00e19644a..acde6334999 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -33,6 +33,7 @@
     can_reconstruct_wgrad_input_from_original,
     check_fp8_reduce_and_update,
     noop_cat,
+    sum_bias_grad,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -82,8 +83,8 @@
     QuantizedTensorStorage,
     Quantizer,
     prepare_for_saving,
-    restore_from_func_ctx,
 )
+from ..dynamo.concatenated_tensor import ConcatenatedTensor, restore_from_func_ctx
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -110,9 +111,9 @@ class LinearFwdArgs:
     """Single-argument bag for the forward path of :class:`_Linear`."""
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
-    weight: TensorOrQuantized
+    weight: Union[TensorOrQuantized, ConcatenatedTensor]
     inp: torch.Tensor
-    bias: Optional[torch.Tensor]
+    bias: Optional[Union[torch.Tensor, ConcatenatedTensor]]
 
     # --- Non-differentiable cached tensors ---
     # TensorOrQuantized so a cached quantized workspace can cross the op boundary.
@@ -238,9 +239,9 @@ class LinearBwdArgs:
     # --- Saved / restored tensors (populated at backward entry) ---
     grad_output: Optional[torch.Tensor] = None
     inputmat: Optional[TensorOrQuantized] = None
-    weight_fp8: Optional[TensorOrQuantized] = None
-    saved_weight: Optional[TensorOrQuantized] = None
-    bias: Optional[torch.Tensor] = None
+    weight_fp8: Optional[Union[TensorOrQuantized, ConcatenatedTensor]] = None
+    saved_weight: Optional[Union[TensorOrQuantized, ConcatenatedTensor]] = None
+    bias: Optional[Union[torch.Tensor, ConcatenatedTensor]] = None
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -1263,6 +1264,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             grad_output_quantizer,
         )
         nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
+        if bwd_args.use_bias and grad_bias is None and not bwd_args.requires_wgrad:
+            grad_bias = sum_bias_grad(grad_output)
 
         # --------------------------------------------------
         # Grad output tensor is ready for computing grad input...
@@ -1775,11 +1778,7 @@ def _linear_backward_fake(
         )
 
     grad_bias = None
-    # FP8 backward computes bgrad in grad_output_preprocess whenever bias is
-    # used; in high precision it is fused into the wgrad GEMM, so it only
-    # exists when wgrad runs.
-    fp8_bwd = args.fp8 and args.backward_override is None
-    if args.use_bias and (args.requires_wgrad or fp8_bwd):
+    if args.use_bias:
         grad_bias = TensorSpec(
             shape=(out_features,), dtype=out_dtype, device=args.grad_output.device
         )
@@ -1931,7 +1930,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
@@ -2399,7 +2400,9 @@ 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_tensor, bias_tensor = self._get_weight_and_bias_tensors(
+                defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
+            )
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2546,6 +2549,12 @@ def forward(
                         msg=f"te.Linear falling back to eager: {fallback_reason}"
                     )
                     use_compiled_op = False
+                    weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+                    linear_bias_tensor = (
+                        bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None
+                    )
+                    fwd_args.weight = weight_tensor
+                    fwd_args.bias = linear_bias_tensor
 
             if use_compiled_op:
                 check_gemm_dims(inp, weight_tensor, self.fp8)
@@ -2693,14 +2702,24 @@ 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_tensor = (
+            ConcatenatedTensor(weights)
+            if defer_concatenation and len(weights) > 1
+            else noop_cat(weights)
+        )
+        bias_tensor = None
         if self.use_bias:
-            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
-        else:
-            bias_tensor = None
+            biases = [getattr(self, name) for name in self.bias_names]
+            bias_tensor = (
+                ConcatenatedTensor(biases)
+                if defer_concatenation
+                and len(biases) > 1
+                and self.apply_bias
+                and not self.gemm_bias_unfused_add
+                else noop_cat(biases)
+            )
         return weight_tensor, bias_tensor
 
     def onnx_forward(

From 289cff12aedad7d8ce09442ca2366b16b42795b3 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 5 Oct 2026 12:40:05 +0200
Subject: [PATCH 02/13] [PyTorch] Make deferred parameter concatenation
 explicit

Replace ConcatenatedTensor with DeferredCat(parts), expose metadata through to_spec and keep materialization inside the custom op. Preserve deferred backward operands for non-FP8 training, reject quantized parts, and make dimension checks consume explicit metadata.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 44 +++++++++-
 .../pytorch/dynamo/custom_op.py               | 32 ++++----
 ...concatenated_tensor.py => deferred_cat.py} | 80 ++++++++++---------
 transformer_engine/pytorch/module/base.py     |  5 +-
 transformer_engine/pytorch/module/linear.py   | 39 +++++----
 transformer_engine/pytorch/utils.py           |  8 +-
 6 files changed, 132 insertions(+), 76 deletions(-)
 rename transformer_engine/pytorch/dynamo/{concatenated_tensor.py => deferred_cat.py} (50%)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 195b9bc339a..1f8081d9f6f 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2857,8 +2857,8 @@ def test_te_split_parameters_saved_versions():
 @pytest.mark.parametrize(
     "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
 )
-def test_concatenated_tensor_storage(layout):
-    from transformer_engine.pytorch.dynamo.concatenated_tensor import ConcatenatedTensor
+def test_deferred_cat_storage(layout):
+    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
 
     storage = torch.arange(128, device="cuda").view(16, 8)
     if layout == "offset":
@@ -2874,7 +2874,7 @@ def test_concatenated_tensor_storage(layout):
     parts = list(storage.split([1, storage.shape[0] - 3, 2]))
     if layout == "disjoint":
         parts[1] = parts[1].clone()
-    result = ConcatenatedTensor(parts).materialize()
+    result = DeferredCat(parts).materialize()
     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()
@@ -2882,6 +2882,44 @@ def test_concatenated_tensor_storage(layout):
         assert result.untyped_storage().data_ptr() != parts[0].untyped_storage().data_ptr()
 
 
+@pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"])
+@pytest.mark.parametrize("mixed_dtype", [False, True])
+def test_deferred_cat_spec(fake, mixed_dtype):
+    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
+
+    with FakeTensorMode() if fake else contextlib.nullcontext():
+        parts = [
+            torch.empty(3, 8, device="cuda", dtype=torch.bfloat16),
+            torch.empty(
+                5,
+                8,
+                device="cuda",
+                dtype=torch.float32 if mixed_dtype else torch.bfloat16,
+                requires_grad=True,
+            ),
+        ]
+        expected = torch.cat(parts)
+        spec = DeferredCat(parts).to_spec()
+        assert spec.shape == tuple(expected.shape)
+        assert spec.dtype == expected.dtype
+        assert spec.device == expected.device
+        assert spec.requires_grad == expected.requires_grad
+        assert not spec.is_quantized
+
+
+@pytest.mark.parametrize("internal", [False, True])
+def test_deferred_cat_rejects_quantized_parts(internal):
+    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
+
+    quantizer = _current_scaling()
+    quantizer.internal = internal
+    tensor = TensorSpec(
+        shape=(4, 8), dtype=torch.bfloat16, quantizer=quantizer, device=torch.device("cpu")
+    ).create_tensor()
+    with pytest.raises(TypeError, match="non-quantized tensors"):
+        DeferredCat([tensor, tensor])
+
+
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize("grad_target", ["input", "weight", "bias"])
 def test_te_split_parameters_grad_targets(grad_target):
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 5225e293456..e3cddbd6d65 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,9 +37,9 @@
   * ``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.
-    Unions including ``ConcatenatedTensor`` also accept deferred row concatenations:
+    Unions including ``DeferredCat`` also accept deferred row concatenations:
     all parts cross in the buffer list, with full gradients split outside the op.
-    Saved parts use ``concatenated_tensor.restore_from_func_ctx`` in the backward
+    Saved parts use ``deferred_cat.restore_from_func_ctx`` in the backward
     container; reconstruction of the full view happens only in the real kernel.
   * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes,
     quantizers -- value-opaque constants baked into the graph -- and nested
@@ -110,7 +110,7 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .concatenated_tensor import ConcatenatedTensor
+from .deferred_cat import DeferredCat
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -452,8 +452,8 @@ def _is_tensor_storage_union(annot: Any) -> bool:
     members = frozenset(a for a in get_args(annot) if a is not type(None))
     return members in (
         _TQ_MEMBERS,
-        _TQ_MEMBERS | {ConcatenatedTensor},
-        frozenset((torch.Tensor, ConcatenatedTensor)),
+        _TQ_MEMBERS | {DeferredCat},
+        frozenset((torch.Tensor, DeferredCat)),
     )
 
 
@@ -497,7 +497,7 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
         return _FieldPlan(
-            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, ConcatenatedTensor in get_args(annot)
+            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, DeferredCat in get_args(annot)
         )
     stripped, is_optional = _strip_optional(annot)
     if stripped is torch.Tensor:
@@ -539,11 +539,11 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         slots[tensor_slot] = None
         slots[inner_slot] = []
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE})
-    elif isinstance(value, ConcatenatedTensor):
+    elif isinstance(value, DeferredCat):
         if not field.allows_concatenation:
-            raise TypeError(f"field {field.name!r} does not accept ConcatenatedTensor")
+            raise TypeError(f"field {field.name!r} does not accept DeferredCat")
         slots[tensor_slot] = None
-        slots[inner_slot] = value.tensors
+        slots[inner_slot] = value.parts
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
@@ -574,7 +574,7 @@ def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
     if kind == _TensorOrQuantizedKind.CONCATENATED:
-        return ConcatenatedTensor(slots[inner_slot])
+        return DeferredCat(slots[inner_slot])
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -761,7 +761,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, DeferredCat):
+            overrides[name] = value.to_spec()
+        elif value is not None and not isinstance(value, TensorSpec):
             overrides[name] = to_tensor_spec(value)
     if not overrides:
         return obj
@@ -1010,7 +1012,7 @@ def _impl(*flat: Any) -> List[torch.Tensor]:
         obj = plan.unpack(dict(zip(plan.slot_names, flat)))
         for name in concatenated_fields:
             value = getattr(obj, name)
-            if isinstance(value, ConcatenatedTensor):
+            if isinstance(value, DeferredCat):
                 setattr(obj, name, value.materialize())
         return pack_result(impl(obj))
 
@@ -1068,9 +1070,9 @@ def _setup_context(ctx, inputs, output):
         saved = []
         ctx.concatenated_saved_lengths = []
         for value in tensors_to_save_from_setup or ():
-            if isinstance(value, ConcatenatedTensor):
-                ctx.concatenated_saved_lengths.append(len(value.tensors))
-                saved.extend(value.tensors)
+            if isinstance(value, DeferredCat):
+                ctx.concatenated_saved_lengths.append(len(value.parts))
+                saved.extend(value.parts)
             else:
                 ctx.concatenated_saved_lengths.append(None)
                 saved.append(value)
diff --git a/transformer_engine/pytorch/dynamo/concatenated_tensor.py b/transformer_engine/pytorch/dynamo/deferred_cat.py
similarity index 50%
rename from transformer_engine/pytorch/dynamo/concatenated_tensor.py
rename to transformer_engine/pytorch/dynamo/deferred_cat.py
index 6f6c39ebc6f..cae9527807f 100644
--- a/transformer_engine/pytorch/dynamo/concatenated_tensor.py
+++ b/transformer_engine/pytorch/dynamo/deferred_cat.py
@@ -5,69 +5,73 @@
 """Deferred row concatenation for tensors consumed inside a custom op."""
 
 from dataclasses import dataclass
-import math
 from typing import List
 
 import torch
 from torch._prims_common import make_contiguous_strides_for
 
-from ..quantized_tensor import restore_from_func_ctx as restore_tensor_ctx
+from .tensor_spec import TensorSpec
+from ..quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    restore_from_func_ctx as restore_tensor_ctx,
+)
 
 
 @dataclass
-class ConcatenatedTensor:
+class DeferredCat:
     """Keep every parameter visible to autograd until the opaque consumer runs."""
 
-    tensors: List[torch.Tensor]
+    parts: List[torch.Tensor]
 
-    @property
-    def shape(self):
-        """Combined shape along dimension zero."""
-        return (sum(t.shape[0] for t in self.tensors), *self.tensors[0].shape[1:])
+    def __post_init__(self):
+        if not self.parts:
+            raise ValueError("DeferredCat requires at least one tensor")
+        if any(
+            not isinstance(part, torch.Tensor)
+            or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
+            for part in self.parts
+        ):
+            raise TypeError("DeferredCat only supports non-quantized tensors")
 
-    @property
-    def dtype(self):
-        """Parameter dtype."""
-        dtype = self.tensors[0].dtype
-        for tensor in self.tensors[1:]:
+    def to_spec(self) -> TensorSpec:
+        """Describe the concatenation without accessing its storage."""
+        first = self.parts[0]
+        dtype = first.dtype
+        for tensor in self.parts[1:]:
             dtype = torch.promote_types(dtype, tensor.dtype)
-        return dtype
-
-    @property
-    def device(self):
-        """Parameter device."""
-        return self.tensors[0].device
-
-    @property
-    def requires_grad(self):
-        """Whether any part requires a gradient."""
-        return any(t.requires_grad for t in self.tensors)
-
-    def numel(self):
-        """Combined element count."""
-        return math.prod(self.shape)
+        return TensorSpec(
+            shape=(sum(t.shape[0] for t in self.parts), *first.shape[1:]),
+            dtype=dtype,
+            device=first.device,
+            requires_grad=any(t.requires_grad for t in self.parts),
+        )
 
     def materialize(self):
         """Build a checked view, or copy parts whose storage is not adjacent."""
-        first = self.tensors[0]
+        first = self.parts[0]
         storage = first.untyped_storage()
         offset = first.storage_offset()
-        for tensor in self.tensors:
-            if tensor.dtype != first.dtype or tensor.device != first.device:
-                return torch.cat(self.tensors)
+        for tensor in self.parts:
             if (
-                tensor.shape[1:] != first.shape[1:]
-                or not tensor.is_contiguous()
+                tensor.dtype != first.dtype
+                or tensor.device != first.device
                 or tensor.is_neg()
                 or tensor.is_conj()
+            ):
+                return torch.cat(self.parts)
+            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(self.tensors)
+                return torch.cat(self.parts)
             offset += tensor.numel()
         if offset * first.element_size() > storage.nbytes():
-            return torch.cat(self.tensors)
-        return first.as_strided(self.shape, make_contiguous_strides_for(self.shape))
+            return torch.cat(self.parts)
+        shape = (sum(t.shape[0] for t in self.parts), *first.shape[1:])
+        return first.as_strided(shape, make_contiguous_strides_for(shape))
 
 
 def restore_from_func_ctx(ctx):
@@ -83,7 +87,7 @@ def restore_from_func_ctx(ctx):
             restored.append(tensors[offset])
             offset += 1
         else:
-            restored.append(ConcatenatedTensor(list(tensors[offset : offset + length])))
+            restored.append(DeferredCat(list(tensors[offset : offset + length])))
             offset += length
     ctx.concatenated_saved_lengths = None
     return restored
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 9ead1e73651..8909720d025 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -46,6 +46,7 @@
     _fsdp_gather_tensors,
 )
 from ..constants import dist_group_type
+from ..dynamo.tensor_spec import TensorSpec
 from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
@@ -1322,7 +1323,7 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
     def _enable_weight_preswizzle(
         self,
         quantizer: Quantizer,
-        weight: torch.Tensor,
+        weight: Union[torch.Tensor, TensorSpec],
     ) -> bool:
         """Whether to fuse scale-factor swizzling into weight quantization.
 
@@ -1339,7 +1340,7 @@ def _enable_weight_preswizzle(
         if isinstance(quantizer, MXFP8Quantizer):
             return True
         if isinstance(quantizer, NVFP4Quantizer):
-            rows, cols = weight.numel() // weight.shape[-1], weight.shape[-1]
+            rows, cols = math.prod(weight.shape[:-1]), weight.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 acde6334999..e9b3f1d7cdc 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -84,7 +84,7 @@
     Quantizer,
     prepare_for_saving,
 )
-from ..dynamo.concatenated_tensor import ConcatenatedTensor, restore_from_func_ctx
+from ..dynamo.deferred_cat import DeferredCat, restore_from_func_ctx
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -111,9 +111,9 @@ class LinearFwdArgs:
     """Single-argument bag for the forward path of :class:`_Linear`."""
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
-    weight: Union[TensorOrQuantized, ConcatenatedTensor]
+    weight: Union[TensorOrQuantized, DeferredCat]
     inp: torch.Tensor
-    bias: Optional[Union[torch.Tensor, ConcatenatedTensor]]
+    bias: Optional[Union[torch.Tensor, DeferredCat]]
 
     # --- Non-differentiable cached tensors ---
     # TensorOrQuantized so a cached quantized workspace can cross the op boundary.
@@ -239,9 +239,9 @@ class LinearBwdArgs:
     # --- Saved / restored tensors (populated at backward entry) ---
     grad_output: Optional[torch.Tensor] = None
     inputmat: Optional[TensorOrQuantized] = None
-    weight_fp8: Optional[Union[TensorOrQuantized, ConcatenatedTensor]] = None
-    saved_weight: Optional[Union[TensorOrQuantized, ConcatenatedTensor]] = None
-    bias: Optional[Union[torch.Tensor, ConcatenatedTensor]] = None
+    weight_fp8: Optional[Union[TensorOrQuantized, DeferredCat]] = None
+    saved_weight: Optional[Union[TensorOrQuantized, DeferredCat]] = None
+    bias: Optional[Union[torch.Tensor, DeferredCat]] = None
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -2403,6 +2403,12 @@ def forward(
             weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
                 defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
+            weight_meta = (
+                weight_tensor.to_spec() if isinstance(weight_tensor, DeferredCat) else weight_tensor
+            )
+            bias_meta = (
+                bias_tensor.to_spec() if isinstance(bias_tensor, DeferredCat) else bias_tensor
+            )
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2424,7 +2430,7 @@ 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_meta
                 )
 
             use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
@@ -2454,7 +2460,7 @@ def forward(
                 backward_override = None
             custom = is_custom(input_quantizer) or is_custom(weight_quantizer)
             backward_input_needs_gather = (
-                weight_tensor.requires_grad
+                weight_meta.requires_grad
                 and self.parallel_mode == "column"
                 and self.sequence_parallel
             )
@@ -2487,9 +2493,9 @@ def forward(
                 weight_workspace=weight_workspace,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
-                weight_requires_grad=weight_tensor.requires_grad,
+                weight_requires_grad=weight_meta.requires_grad,
                 bias_requires_grad=(
-                    linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False
+                    bias_meta.requires_grad if linear_bias_tensor is not None else False
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2557,7 +2563,7 @@ def forward(
                     fwd_args.bias = linear_bias_tensor
 
             if use_compiled_op:
-                check_gemm_dims(inp, weight_tensor, self.fp8)
+                check_gemm_dims(inp.shape, weight_meta.shape, self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
             else:
                 out, new_weight_workspace = _linear_eager(
@@ -2705,19 +2711,24 @@ def _forward_eager_fallback(
     def _get_weight_and_bias_tensors(self, defer_concatenation=False):
         weights = self._get_weight_tensors()
         weight_tensor = (
-            ConcatenatedTensor(weights)
-            if defer_concatenation and len(weights) > 1
+            DeferredCat(weights)
+            if defer_concatenation
+            and len(weights) > 1
+            and all(not isinstance(w, (QuantizedTensor, QuantizedTensorStorage)) for w in weights)
             else noop_cat(weights)
         )
         bias_tensor = None
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
             bias_tensor = (
-                ConcatenatedTensor(biases)
+                DeferredCat(biases)
                 if defer_concatenation
                 and len(biases) > 1
                 and self.apply_bias
                 and not self.gemm_bias_unfused_add
+                and all(
+                    not isinstance(b, (QuantizedTensor, QuantizedTensorStorage)) for b in biases
+                )
                 else noop_cat(biases)
             )
         return weight_tensor, bias_tensor
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index ca77f7b8bb3..53e27ae6f6b 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -716,21 +716,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"

From 43a33bb348549658408a145ecd2c85c8836651ee Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 5 Oct 2026 16:23:23 +0200
Subject: [PATCH 03/13] [PyTorch] Simplify split Linear parameter handling

Keep singleton operands unwrapped and represent split parameters with immutable ParameterParts through the ConcatInput alias. Centralize eager and compiled input preparation, let the custom-op adapter restore saved parts, and select the execution path before concatenating parameters. Name backward weights by their role and update the existing MLA caller.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  87 +++++-
 .../pytorch/attention/fused_mla_q_uproj.py    |   4 +-
 .../pytorch/dynamo/custom_op.py               | 116 ++++----
 .../{deferred_cat.py => parameter_parts.py}   |  33 +--
 transformer_engine/pytorch/module/_common.py  |  26 +-
 transformer_engine/pytorch/module/linear.py   | 252 ++++++------------
 6 files changed, 270 insertions(+), 248 deletions(-)
 rename transformer_engine/pytorch/dynamo/{deferred_cat.py => parameter_parts.py} (72%)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 1f8081d9f6f..ba0432b68cb 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2048,7 +2048,7 @@ def fn(inp):
             )
 
 
-# Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a
+# Configs rejected by Linear._compile_eager_fallback_reason() that a
 # single-GPU unit test can construct. Distributed-only reasons (fsdp_group,
 # DistributedWeight) and CPU offloading need machinery this file doesn't have;
 # delayed scaling is a hard error (check_recipe_support), tested separately.
@@ -2782,6 +2782,7 @@ def test_te_split_parameters_compile(compile_mode, dtype, equal_splits):
 @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,
@@ -2792,7 +2793,7 @@ def test_te_split_parameters_recipes(fp8_recipe, bias_mode):
     )
 
     def fn(inp):
-        with te.autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe):
+        with te.autocast(enabled=fp8_recipe is not None, recipe=autocast_recipe):
             out = model(inp)
         if bias_mode == "returned":
             out, bias = out
@@ -2854,11 +2855,73 @@ def test_te_split_parameters_saved_versions():
         out.sum().backward()
 
 
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+def test_te_split_parameters_saved_hooks():
+    import copy
+
+    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
+    reference = copy.deepcopy(model)
+    compiled = torch.compile(model, fullgraph=True, backend="aot_eager")
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    ref_inp = inp.detach().clone().requires_grad_()
+    with torch.autograd.graph.saved_tensors_hooks(lambda t: t.clone(), lambda t: t):
+        out = compiled(inp)
+        out.sum().backward()
+    expected = reference(ref_inp)
+    expected.sum().backward()
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)
+    torch.testing.assert_close(inp.grad, ref_inp.grad, rtol=0, atol=0)
+    for actual, target in zip(model.parameters(), reference.parameters()):
+        torch.testing.assert_close(actual.grad, target.grad, rtol=0, atol=0)
+
+
+def test_te_split_parameters_fallback_check(monkeypatch):
+    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32))
+    inp = torch.empty(16, 64, device="cuda")
+
+    def unexpected_materialization(*args, **kwargs):
+        raise AssertionError("Fallback eligibility must not concatenate parameters")
+
+    monkeypatch.setattr(model, "_get_weight_and_bias_tensors", unexpected_materialization)
+    assert model._compile_eager_fallback_reason(inp, None, False, False, True, False) is None
+    assert (
+        model._compile_eager_fallback_reason(inp, None, True, False, True, False)
+        == "differentiable fp8_output=True"
+    )
+
+
+@pytest.mark.parametrize("defer", [False, True])
+def test_concat_input_single_tensor(defer):
+    from transformer_engine.pytorch.module._common import concat_input
+
+    tensor = torch.nn.Parameter(torch.randn(4, 8, device="cuda"))
+    assert concat_input([tensor], defer=defer) is tensor
+
+
+@pytest.mark.parametrize("disjoint", [False, True])
+def test_concat_input_eager_gradients(disjoint):
+    from transformer_engine.pytorch.module._common import concat_input
+
+    parts = [torch.nn.Parameter(t) for t in torch.randn(8, 8, device="cuda").split((3, 5))]
+    if disjoint:
+        parts[1] = torch.nn.Parameter(parts[1].detach().clone())
+    result = concat_input(parts)
+    expected = torch.cat(parts)
+    grad = torch.randn_like(result)
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+    torch.testing.assert_close(
+        torch.autograd.grad(result, parts, grad),
+        torch.autograd.grad(expected, parts, grad),
+        rtol=0,
+        atol=0,
+    )
+
+
 @pytest.mark.parametrize(
     "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
 )
-def test_deferred_cat_storage(layout):
-    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
+def test_parameter_parts_storage(layout):
+    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
 
     storage = torch.arange(128, device="cuda").view(16, 8)
     if layout == "offset":
@@ -2874,7 +2937,7 @@ def test_deferred_cat_storage(layout):
     parts = list(storage.split([1, storage.shape[0] - 3, 2]))
     if layout == "disjoint":
         parts[1] = parts[1].clone()
-    result = DeferredCat(parts).materialize()
+    result = ParameterParts(parts).materialize()
     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()
@@ -2884,8 +2947,8 @@ def test_deferred_cat_storage(layout):
 
 @pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"])
 @pytest.mark.parametrize("mixed_dtype", [False, True])
-def test_deferred_cat_spec(fake, mixed_dtype):
-    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
+def test_parameter_parts_spec(fake, mixed_dtype):
+    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
 
     with FakeTensorMode() if fake else contextlib.nullcontext():
         parts = [
@@ -2899,7 +2962,7 @@ def test_deferred_cat_spec(fake, mixed_dtype):
             ),
         ]
         expected = torch.cat(parts)
-        spec = DeferredCat(parts).to_spec()
+        spec = ParameterParts(parts).to_spec()
         assert spec.shape == tuple(expected.shape)
         assert spec.dtype == expected.dtype
         assert spec.device == expected.device
@@ -2908,8 +2971,9 @@ def test_deferred_cat_spec(fake, mixed_dtype):
 
 
 @pytest.mark.parametrize("internal", [False, True])
-def test_deferred_cat_rejects_quantized_parts(internal):
-    from transformer_engine.pytorch.dynamo.deferred_cat import DeferredCat
+def test_parameter_parts_rejects_quantized_parts(internal):
+    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
+    from transformer_engine.pytorch.module._common import concat_input
 
     quantizer = _current_scaling()
     quantizer.internal = internal
@@ -2917,7 +2981,8 @@ def test_deferred_cat_rejects_quantized_parts(internal):
         shape=(4, 8), dtype=torch.bfloat16, quantizer=quantizer, device=torch.device("cpu")
     ).create_tensor()
     with pytest.raises(TypeError, match="non-quantized tensors"):
-        DeferredCat([tensor, tensor])
+        ParameterParts([tensor, tensor])
+    assert concat_input([tensor], defer=True) is tensor
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
index b15e53e6ee3..913aea2ccae 100644
--- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
+++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
@@ -299,8 +299,8 @@ def backward_linear(
         bwd_args = LinearBwdArgs(
             grad_output=grad_output,
             inputmat=x_saved,
-            weight_fp8=w_q,
-            saved_weight=w_q,
+            weight_for_dgrad=w_q,
+            original_weight=w_q,
             grad_output_quantizer=grad_output_quantizer,
             inp_shape=x_saved.shape,
             activation_dtype=act_dtype,
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index e3cddbd6d65..503e3b8d4c4 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,10 +37,9 @@
   * ``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.
-    Unions including ``DeferredCat`` also accept deferred row concatenations:
+    Unions including ``ParameterParts`` also accept deferred row concatenations:
     all parts cross in the buffer list, with full gradients split outside the op.
-    Saved parts use ``deferred_cat.restore_from_func_ctx`` in the backward
-    container; reconstruction of the full view happens only in the real kernel.
+    The adapter saves/restores the parts; only the real kernel materializes them.
   * ``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.
@@ -72,7 +71,7 @@
   * on ``backward()`` the incoming flat grads are sliced per user output from the
     stashed plan (a ``grad_outputs`` field on the backward args receives the
     whole tuple; otherwise ``grad_output`` receives the first output's grad),
-    the container's optional ``setup_saved_tensors`` hook restores the saved
+    the container's optional ``set_saved_tensors`` hook receives the restored
     tensors, then the *backward op* runs the real ``bwd_impl`` and returns the
     flat grads (``bwd_fake_impl`` is its data-free fake).
 
@@ -110,13 +109,14 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .deferred_cat import DeferredCat
+from .parameter_parts import ParameterParts
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
     Quantizer,
     _quantized_tensor_passthrough_ops,
     prepare_for_saving,
+    restore_from_func_ctx,
 )
 from ..utils import record_compile_disabled
 
@@ -413,7 +413,7 @@ class _TensorOrQuantizedKind(Enum):
     NONE = "none"
     TENSOR = "tensor"
     STORAGE = "storage"
-    CONCATENATED = "concatenated"
+    PARTS = "parts"
 
 
 _TQ_KIND_KEY = "__kind__"
@@ -452,8 +452,8 @@ def _is_tensor_storage_union(annot: Any) -> bool:
     members = frozenset(a for a in get_args(annot) if a is not type(None))
     return members in (
         _TQ_MEMBERS,
-        _TQ_MEMBERS | {DeferredCat},
-        frozenset((torch.Tensor, DeferredCat)),
+        _TQ_MEMBERS | {ParameterParts},
+        frozenset((torch.Tensor, ParameterParts)),
     )
 
 
@@ -497,7 +497,7 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
         return _FieldPlan(
-            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, DeferredCat in get_args(annot)
+            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, ParameterParts in get_args(annot)
         )
     stripped, is_optional = _strip_optional(annot)
     if stripped is torch.Tensor:
@@ -539,12 +539,12 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         slots[tensor_slot] = None
         slots[inner_slot] = []
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE})
-    elif isinstance(value, DeferredCat):
+    elif isinstance(value, ParameterParts):
         if not field.allows_concatenation:
-            raise TypeError(f"field {field.name!r} does not accept DeferredCat")
+            raise TypeError(f"field {field.name!r} does not accept ParameterParts")
         slots[tensor_slot] = None
-        slots[inner_slot] = value.parts
-        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
+        slots[inner_slot] = list(value.parts)
+        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.PARTS})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
         # ``Tensor?`` slot; subclass flattening (if any) is done by the
@@ -564,7 +564,9 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         )
 
 
-def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any:
+def _unpack_tensor_or_quantized(
+    field: _FieldPlan, slots: Dict[str, Any], *, materialize_parts: bool = False
+) -> Any:
     """Inverse of :func:`_pack_tensor_or_quantized`."""
     tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots)
     meta = slots[meta_slot]
@@ -573,8 +575,9 @@ def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any
         return None
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
-    if kind == _TensorOrQuantizedKind.CONCATENATED:
-        return DeferredCat(slots[inner_slot])
+    if kind == _TensorOrQuantizedKind.PARTS:
+        parts = ParameterParts(tuple(slots[inner_slot]))
+        return parts.materialize() if materialize_parts else parts
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -676,7 +679,7 @@ def pack(self, obj: Any) -> Dict[str, Any]:
             slots[_SIMPLE_META_SLOT] = OpaqueValueBundle(simple)
         return slots
 
-    def unpack(self, slots: Dict[str, Any]) -> Any:
+    def unpack(self, slots: Dict[str, Any], *, materialize_parts: bool = False) -> Any:
         """Rebuild a fresh ``arg_type`` instance from the op's flat slot dict.
 
         Inverse of :meth:`pack`.
@@ -688,7 +691,9 @@ def unpack(self, slots: Dict[str, Any]) -> Any:
                 case _FieldKind.TENSOR:
                     kwargs[field.name] = slots[field.slots[0].name]
                 case _FieldKind.TENSOR_OR_QUANTIZED:
-                    kwargs[field.name] = _unpack_tensor_or_quantized(field, slots)
+                    kwargs[field.name] = _unpack_tensor_or_quantized(
+                        field, slots, materialize_parts=materialize_parts
+                    )
                 case _FieldKind.PROCESS_GROUP:
                     if bundle is not None:
                         name = bundle[field.name]
@@ -761,7 +766,7 @@ 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 isinstance(value, DeferredCat):
+        if isinstance(value, ParameterParts):
             overrides[name] = value.to_spec()
         elif value is not None and not isinstance(value, TensorSpec):
             overrides[name] = to_tensor_spec(value)
@@ -1006,14 +1011,8 @@ def _register_base_op(
     ``pack_result``.
     """
 
-    concatenated_fields = tuple(f.name for f in plan.fields if f.allows_concatenation)
-
     def _impl(*flat: Any) -> List[torch.Tensor]:
-        obj = plan.unpack(dict(zip(plan.slot_names, flat)))
-        for name in concatenated_fields:
-            value = getattr(obj, name)
-            if isinstance(value, DeferredCat):
-                setattr(obj, name, value.materialize())
+        obj = plan.unpack(dict(zip(plan.slot_names, flat)), materialize_parts=True)
         return pack_result(impl(obj))
 
     def _fake(*flat: Any) -> List[torch.Tensor]:
@@ -1029,6 +1028,32 @@ def _fake(*flat: Any) -> List[torch.Tensor]:
     return op
 
 
+def _flatten_saved_parameter_parts(ctx, tensors):
+    lengths = [len(t.parts) if isinstance(t, ParameterParts) else None for t in tensors]
+    ctx.parameter_parts_lengths = lengths if any(n is not None for n in lengths) else None
+    if ctx.parameter_parts_lengths is None:
+        return tensors
+    return [part for t in tensors for part in (t.parts if isinstance(t, ParameterParts) else (t,))]
+
+
+def _restore_saved_tensors(ctx):
+    tensors = restore_from_func_ctx(ctx)
+    lengths = getattr(ctx, "parameter_parts_lengths", None)
+    if lengths is None:
+        return tensors
+    restored = []
+    offset = 0
+    for length in lengths:
+        if length is None:
+            restored.append(tensors[offset])
+            offset += 1
+        else:
+            restored.append(ParameterParts(tuple(tensors[offset : offset + length])))
+            offset += length
+    ctx.parameter_parts_lengths = None
+    return restored
+
+
 def _register_autograd_for_op(
     *,
     fwd_op: Any,
@@ -1047,6 +1072,9 @@ def _register_autograd_for_op(
     the plan on ``ctx`` so backward can slice its grads per user output.
     """
     bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields)
+    supports_parts = any(f.allows_concatenation for f in fwd_plan.fields)
+    tensor_storage_offsets = frozenset(fwd_plan.tensor_or_quantized_offsets())
+    set_saved_tensors = getattr(bwd_plan.arg_type, "set_saved_tensors", None)
 
     def _setup_context(ctx, inputs, output):
         ctx.fwd_tensor_list_lengths = {
@@ -1067,15 +1095,9 @@ def _setup_context(ctx, inputs, output):
             out_plan.ctx_attrs,
             tuple(saved_list),
         )
-        saved = []
-        ctx.concatenated_saved_lengths = []
-        for value in tensors_to_save_from_setup or ():
-            if isinstance(value, DeferredCat):
-                ctx.concatenated_saved_lengths.append(len(value.parts))
-                saved.extend(value.parts)
-            else:
-                ctx.concatenated_saved_lengths.append(None)
-                saved.append(value)
+        saved = tensors_to_save_from_setup or ()
+        if supports_parts:
+            saved = _flatten_saved_parameter_parts(ctx, saved)
         tensors_to_save, tensor_objects = prepare_for_saving(*saved)
         ctx.tensor_objects = tensor_objects
         ctx.save_for_backward(*tensors_to_save)
@@ -1084,12 +1106,12 @@ def _setup_context(ctx, inputs, output):
         # Input shapes for the grad slots (SymInt-safe on ctx): a bwd impl may
         # rederive shapes lossily (e.g. rank-1 inputs come back rank-2), so the
         # returned grads are viewed back to the true input shapes below.
-        ctx.concatenated_grad_shapes = {
+        ctx.parameter_parts_grad_shapes = {
             pos: [tensor.shape for tensor in inputs[pos + 1]]
             for pos in grad_targets
-            if pos in fwd_plan.tensor_or_quantized_offsets()
+            if pos in tensor_storage_offsets
             and inputs[pos] is None
-            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.CONCATENATED
+            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.PARTS
         }
         ctx.grad_input_shapes = {
             pos: inputs[pos].shape for pos in grad_targets if isinstance(inputs[pos], torch.Tensor)
@@ -1097,7 +1119,9 @@ def _setup_context(ctx, inputs, output):
 
     def _autograd_backward(ctx, *grad_outputs):
         bwd_obj = ctx.backward_objects
-        if hasattr(bwd_obj, "setup_saved_tensors"):
+        if set_saved_tensors is not None:
+            set_saved_tensors(bwd_obj, _restore_saved_tensors(ctx))
+        elif hasattr(bwd_obj, "setup_saved_tensors"):
             bwd_obj.setup_saved_tensors(ctx)
         ctx.tensor_objects = None
         user_grads = _slice_user_grads(ctx.output_ranges, grad_outputs[0])
@@ -1117,8 +1141,8 @@ 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.concatenated_grad_shapes:
-                shapes = ctx.concatenated_grad_shapes[pos]
+            if pos in ctx.parameter_parts_grad_shapes:
+                shapes = ctx.parameter_parts_grad_shapes[pos]
                 if g is not None:
                     out[pos + 1] = list(torch.split(g, [shape[0] for shape in shapes], dim=0))
                 continue
@@ -1128,7 +1152,7 @@ def _autograd_backward(ctx, *grad_outputs):
                     g = g.view(shape)
             out[pos] = g
         ctx.grad_input_shapes = None
-        ctx.concatenated_grad_shapes = None
+        ctx.parameter_parts_grad_shapes = None
         return tuple(out)
 
     fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context)
@@ -1387,15 +1411,17 @@ def register_custom_op_with_autograd(
       non-differentiable input).
     * ``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
-      absent.
+    * ``bwd_arg_type.set_saved_tensors(tensors)`` -- optional hook receiving
+      restored tensors, storage objects or parameter parts in forward save order.
+      Otherwise, the legacy ``setup_saved_tensors(ctx)`` hook is used if present.
 
     How the backward container is populated: ``setup_context`` fills the
     ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from
     forward state and returns the tensors to persist; the framework saves them
     via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework
     restores them into the container's tensor fields through the
-    ``setup_saved_tensors`` hook and sets the incoming gradient directly --
+    ``set_saved_tensors`` (or legacy ``setup_saved_tensors``) hook and sets the
+    incoming gradient directly --
     into a ``grad_outputs`` field (tuple, one grad per user output) if
     ``bwd_arg_type`` declares one, else into ``grad_output`` (the first user
     output's grad) -- so ``bwd_impl`` receives a fully-populated
diff --git a/transformer_engine/pytorch/dynamo/deferred_cat.py b/transformer_engine/pytorch/dynamo/parameter_parts.py
similarity index 72%
rename from transformer_engine/pytorch/dynamo/deferred_cat.py
rename to transformer_engine/pytorch/dynamo/parameter_parts.py
index cae9527807f..b314b2efb33 100644
--- a/transformer_engine/pytorch/dynamo/deferred_cat.py
+++ b/transformer_engine/pytorch/dynamo/parameter_parts.py
@@ -5,7 +5,7 @@
 """Deferred row concatenation for tensors consumed inside a custom op."""
 
 from dataclasses import dataclass
-from typing import List
+from typing import Tuple, TypeVar, Union
 
 import torch
 from torch._prims_common import make_contiguous_strides_for
@@ -14,25 +14,25 @@
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
-    restore_from_func_ctx as restore_tensor_ctx,
 )
 
 
-@dataclass
-class DeferredCat:
+@dataclass(frozen=True, slots=True)
+class ParameterParts:
     """Keep every parameter visible to autograd until the opaque consumer runs."""
 
-    parts: List[torch.Tensor]
+    parts: Tuple[torch.Tensor, ...]
 
     def __post_init__(self):
+        object.__setattr__(self, "parts", tuple(self.parts))
         if not self.parts:
-            raise ValueError("DeferredCat requires at least one tensor")
+            raise ValueError("ParameterParts requires at least one tensor")
         if any(
             not isinstance(part, torch.Tensor)
             or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
             for part in self.parts
         ):
-            raise TypeError("DeferredCat only supports non-quantized tensors")
+            raise TypeError("ParameterParts only supports non-quantized tensors")
 
     def to_spec(self) -> TensorSpec:
         """Describe the concatenation without accessing its storage."""
@@ -74,20 +74,5 @@ def materialize(self):
         return first.as_strided(shape, make_contiguous_strides_for(shape))
 
 
-def restore_from_func_ctx(ctx):
-    """Restore original split parameters saved by the custom-op framework."""
-    tensors = restore_tensor_ctx(ctx)
-    lengths = getattr(ctx, "concatenated_saved_lengths", None)
-    if lengths is None:
-        return tensors
-    restored = []
-    offset = 0
-    for length in lengths:
-        if length is None:
-            restored.append(tensors[offset])
-            offset += 1
-        else:
-            restored.append(DeferredCat(list(tensors[offset : offset + length])))
-            offset += length
-    ctx.concatenated_saved_lengths = None
-    return restored
+_TensorT = TypeVar("_TensorT")
+ConcatInput = Union[_TensorT, ParameterParts]
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 3d04882698a..7785eda7eb4 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -6,14 +6,17 @@
 
 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 ..dynamo.parameter_parts import ConcatInput, ParameterParts
 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 +242,27 @@ def noop_cat(
     return _NoopCatFunc.apply(dim, *tensors)
 
 
+def concat_input(
+    tensors: Sequence[TensorOrQuantized], *, defer: bool = False
+) -> ConcatInput[TensorOrQuantized]:
+    """Prepare one parameter operand for eager or a compiled consumer."""
+    if len(tensors) == 1:
+        return tensors[0]
+    if (
+        defer
+        and tensors
+        and all(
+            isinstance(t, torch.Tensor)
+            and not isinstance(t, (QuantizedTensor, QuantizedTensorStorage))
+            for t in tensors
+        )
+    ):
+        return ParameterParts(tuple(tensors))
+    if torch.compiler.is_compiling():
+        return torch.cat(tensors)
+    return noop_cat(tensors)
+
+
 @dataclasses.dataclass
 class _ParameterInitMeta:
     """
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index e9b3f1d7cdc..a9691acaee6 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,
     sum_bias_grad,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
@@ -83,8 +83,9 @@
     QuantizedTensorStorage,
     Quantizer,
     prepare_for_saving,
+    restore_from_func_ctx,
 )
-from ..dynamo.deferred_cat import DeferredCat, restore_from_func_ctx
+from ..dynamo.parameter_parts import ConcatInput, ParameterParts
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -111,9 +112,9 @@ class LinearFwdArgs:
     """Single-argument bag for the forward path of :class:`_Linear`."""
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
-    weight: Union[TensorOrQuantized, DeferredCat]
+    weight: ConcatInput[TensorOrQuantized]
     inp: torch.Tensor
-    bias: Optional[Union[torch.Tensor, DeferredCat]]
+    bias: Optional[ConcatInput[torch.Tensor]]
 
     # --- Non-differentiable cached tensors ---
     # TensorOrQuantized so a cached quantized workspace can cross the op boundary.
@@ -180,57 +181,6 @@ class LinearFwdArgs:
     cpu_offloading: bool
     is_grad_enabled: bool
 
-    def compile_unsupported_reason(self) -> Optional[str]:
-        """Reason this config can't use the torch.compile custom-op path (else None)."""
-        if self.debug:
-            return "debug instrumentation (nvidia-dlfw-inspect)"
-        if is_distributed_weight(self.weight):
-            return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
-        if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)):
-            return "a quantized input tensor"
-        if self.fsdp_group is not None:
-            return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
-        if (
-            self.fp8_output
-            and self.is_grad_enabled
-            and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad)
-        ):
-            return "differentiable fp8_output=True"
-        if self.cpu_offloading:
-            return "CPU activation offloading"
-        if self.wgrad_store is not None:
-            # Non-None only when delayed wgrad compute is on (see Linear.forward).
-            return "delayed wgrad compute (wgrad_store)"
-        if (
-            self.grad_input_quantizer is not None
-            and self.is_grad_enabled
-            and self.input_requires_grad
-            and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad)
-        ):
-            # A quantized dgrad can't cross the op boundary: grads are packed
-            # one plain Tensor[] slot each (_pack_bwd_result).
-            return "a quantized input grad (fp8_grad=True)"
-        if self.cache_weight and self.fp8:
-            # The cached workspace is updated in place on the first microbatch,
-            # which the functional op (mutates_args=()) can't express. Without
-            # FP8 no workspace exists, so is_first_microbatch is inert.
-            return "FP8 weight caching (is_first_microbatch)"
-        if self.fuse_wgrad_accumulation:
-            return "fuse_wgrad_accumulation (main_grad)"
-        for quantizer in (
-            self.input_quantizer,
-            self.weight_quantizer,
-            self.output_quantizer,
-            self.grad_input_quantizer,
-            self.grad_weight_quantizer,
-            self.grad_output_quantizer,
-        ):
-            # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe
-            # quantizers are not value-opaque and can't cross the custom-op boundary.
-            if quantizer is not None and not is_value_opaque_quantizer(quantizer):
-                return "a quantizer not registered as a torch.compile value-opaque type"
-        return None
-
 
 @dataclass(slots=True)
 class LinearBwdArgs:
@@ -239,9 +189,9 @@ class LinearBwdArgs:
     # --- Saved / restored tensors (populated at backward entry) ---
     grad_output: Optional[torch.Tensor] = None
     inputmat: Optional[TensorOrQuantized] = None
-    weight_fp8: Optional[Union[TensorOrQuantized, DeferredCat]] = None
-    saved_weight: Optional[Union[TensorOrQuantized, DeferredCat]] = None
-    bias: Optional[Union[torch.Tensor, DeferredCat]] = None
+    weight_for_dgrad: Optional[ConcatInput[TensorOrQuantized]] = None
+    original_weight: Optional[ConcatInput[TensorOrQuantized]] = None
+    bias: Optional[ConcatInput[torch.Tensor]] = None
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -305,16 +255,9 @@ class LinearBwdArgs:
     # --- Per-backward scratch state (populated inside _linear_backward_impl) ---
     ub_obj_gradout: Optional[Any] = None
 
-    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
+    def set_saved_tensors(self, tensors) -> None:
+        """Assign restored operands in their forward save order."""
+        self.inputmat, self.weight_for_dgrad, self.original_weight, self.bias = tensors
 
 
 def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int:
@@ -745,7 +688,7 @@ def _linear_forward_impl(
         saved_tensor_aliases = (
             "inp" if saved_inputmat is inp else None,
             wt_alias,
-            "weight",  # ``saved_weight`` slot is always the weight parameter
+            "weight",  # ``original_weight`` slot is always the weight parameter
             "bias" if bias is not None else None,
         )
         tensors_to_save_from_forward = (
@@ -911,7 +854,7 @@ def _linear_forward_fake(
 
     # ------------------------------------------------------
     # Backward state -- saved-tensor layout
-    # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing.
+    # (saved_inputmat, wt_save, original_weight, bias) with name-based aliasing.
     # ------------------------------------------------------
     tensors_to_save_from_forward = None
     ctx_attrs = None
@@ -965,7 +908,7 @@ def _linear_forward_fake(
                 shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
             )
 
-        # Slot 2 -- ``saved_weight`` (always aliased to ``weight``).
+        # Slot 2 -- ``original_weight`` (always aliased to ``weight``).
         # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent).
         saved_tensor_aliases = (
             inputmat_alias,
@@ -1085,8 +1028,8 @@ def _linear_setup_ctx(
         bwd_args.grad_weight_quantizer = None
         bwd_args.grad_output_quantizer = None
 
-    saved_inputmat, wt_save, saved_weight, saved_bias = tensors_to_save_from_forward
-    inputmat_alias, wt_save_alias, saved_weight_alias, bias_alias = ctx_attrs[
+    saved_inputmat, wt_save, original_weight, saved_bias = tensors_to_save_from_forward
+    inputmat_alias, wt_save_alias, original_weight_alias, bias_alias = ctx_attrs[
         "saved_tensor_aliases"
     ]
     bwd_args.owns_input = inputmat_alias != "inp"
@@ -1098,26 +1041,26 @@ def _linear_setup_ctx(
         wt_save = fwd_outputs[1]
     elif wt_save_alias == "weight_workspace":
         wt_save = fwd_args.weight_workspace
-    if saved_weight_alias == "weight":
-        saved_weight = weight
+    if original_weight_alias == "weight":
+        original_weight = weight
     if bias_alias == "bias":
         saved_bias = bias
-    return (saved_inputmat, wt_save, saved_weight, saved_bias)
+    return (saved_inputmat, wt_save, original_weight, saved_bias)
 
 
 def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]:
     """Backward implementation for the linear layer.
 
     Caller must have populated ``args.grad_output`` and run
-    ``args.setup_saved_tensors(ctx)`` before invocation.
+    ``args.set_saved_tensors(tensors)`` before invocation.
     """
     bwd_args = args
     grad_output = args.grad_output
     assert grad_output is not None
     inputmat = args.inputmat
-    weight_fp8 = args.weight_fp8
-    saved_weight = args.saved_weight
-    is_dist_weight = is_distributed_weight(saved_weight)
+    weight_for_dgrad = args.weight_for_dgrad
+    original_weight = args.original_weight
+    is_dist_weight = is_distributed_weight(original_weight)
     bias = args.bias
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
@@ -1166,20 +1109,20 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
                 origin_weight_python_object.main_grad = main_grad
 
         # Gather intermediate/activation tensors if needed
-        # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already
+        # NOTE: weight_for_dgrad = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already
         #       shards/unshards the base weights so we don't do it ourselves
         nvtx_range_push(f"{nvtx_label}.fsdp_gather")
         _fsdp_gather_tensors(
             bwd_args.fsdp_group,
             bwd_args.fsdp_shapes,
             inputmat,
-            weight_fp8,
+            weight_for_dgrad,
         )
         nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
 
         # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes).
         if bwd_args.inp_shape is None:
-            in_features = saved_weight.shape[-1]
+            in_features = original_weight.shape[-1]
             inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args)
             bwd_args.inp_shape = (
                 torch.Size([in_features])
@@ -1341,32 +1284,32 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
         # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when
         # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd.
         if is_dist_weight:
-            weight_fp8 = materialize_weight_for_backward(saved_weight)[0]
+            weight_for_dgrad = materialize_weight_for_backward(original_weight)[0]
 
         if bwd_args.requires_dgrad:
 
             # FSDP2: Re-create workspace from all-gathered weight when
             # workspace was not saved. (Issue #2681)
-            # Use saved_weight (the original weight parameter) since
-            # weight_fp8 is only set when workspace was saved.
-            if weight_fp8 is None:
-                if isinstance(saved_weight, QuantizedTensorStorage):
+            # Use original_weight (the original weight parameter) since
+            # weight_for_dgrad is only set when workspace was saved.
+            if weight_for_dgrad is None:
+                if isinstance(original_weight, QuantizedTensorStorage):
                     # saved weight is already set to right usages by
                     # fsdp2 quantized-tensor hooks when workspace was not saved.
-                    weight_fp8 = saved_weight
+                    weight_for_dgrad = original_weight
                 elif bwd_args.weight_quantizer is not None:
                     bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True)
-                    weight_fp8 = bwd_args.weight_quantizer(saved_weight)
+                    weight_for_dgrad = bwd_args.weight_quantizer(original_weight)
             elif (
                 is_dist_weight
                 and bwd_args.fp8
                 and bwd_args.weight_quantizer is not None
-                and not isinstance(weight_fp8, QuantizedTensorStorage)
+                and not isinstance(weight_for_dgrad, QuantizedTensorStorage)
             ):
                 # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer
                 # so the dgrad operand isn't cast by the delayed recipe.
                 bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True)
-                weight_fp8 = bwd_args.weight_quantizer(weight_fp8)
+                weight_for_dgrad = bwd_args.weight_quantizer(weight_for_dgrad)
 
             # Make sure required data is available
             if isinstance(grad_output, QuantizedTensorStorage):
@@ -1374,9 +1317,9 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             if (
                 bwd_args.fp8
                 and weight_quantizer is not None
-                and isinstance(weight_fp8, QuantizedTensorStorage)
+                and isinstance(weight_for_dgrad, QuantizedTensorStorage)
             ):
-                weight_fp8.update_usage(columnwise_usage=True)
+                weight_for_dgrad.update_usage(columnwise_usage=True)
 
             # Choose whether to use GEMM kernel with split accumulator
             use_split_accumulator = bwd_args.dgrad_use_split_accumulator
@@ -1403,18 +1346,18 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             # Note: dx = dy * w
 
             nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
-            weight_for_dgrad = weight_fp8
+            gemm_weight = weight_for_dgrad
             if bwd_args.backward_override == "dequantized":
-                if isinstance(weight_for_dgrad, QuantizedTensorStorage):
-                    weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype)
+                if isinstance(gemm_weight, QuantizedTensorStorage):
+                    gemm_weight = gemm_weight.dequantize(dtype=bwd_args.activation_dtype)
                 else:
-                    weight_for_dgrad = cast_if_needed(weight_for_dgrad, bwd_args.activation_dtype)
+                    gemm_weight = cast_if_needed(gemm_weight, bwd_args.activation_dtype)
             elif bwd_args.backward_override == "high_precision":
-                weight_for_dgrad = saved_weight
-                if isinstance(weight_for_dgrad, QuantizedTensorStorage):
-                    weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype)
+                gemm_weight = original_weight
+                if isinstance(gemm_weight, QuantizedTensorStorage):
+                    gemm_weight = gemm_weight.dequantize(dtype=bwd_args.activation_dtype)
             gemm_out, *_, reduce_scatter_out = general_gemm(
-                weight_for_dgrad,
+                gemm_weight,
                 grad_output,
                 layout="NN",
                 grad=True,
@@ -1434,8 +1377,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             # and 2d block-scaled weights in TE managed memory. So we need to clear
             # it here.
             # (Issues #2681, #2717)
-            if bwd_args.is_fsdp2 and isinstance(weight_fp8, QuantizedTensorStorage):
-                clear_columnwise_cache(weight_fp8)
+            if bwd_args.is_fsdp2 and isinstance(weight_for_dgrad, QuantizedTensorStorage):
+                clear_columnwise_cache(weight_for_dgrad)
 
             # Prepare grad input tensor
             # Note: Perform tensor-parallel communication
@@ -1644,7 +1587,7 @@ def wgrad_gemm(
                 # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad
                 # (async; overlap with the next layer's bwd via the cascade).
                 if is_dist_weight:
-                    wgrad = finalize_weight_grads(saved_weight, [wgrad])[0]
+                    wgrad = finalize_weight_grads(original_weight, [wgrad])[0]
 
                 # Update grad bias if needed
                 if grad_bias is None:
@@ -1716,7 +1659,7 @@ def wgrad_gemm(
 
     # Scatter fp8 weight buffers
     if bwd_args.fp8 and not bwd_args.is_weight_param_quantized:
-        _fsdp_scatter_tensors(bwd_args.fsdp_group, weight_fp8)
+        _fsdp_scatter_tensors(bwd_args.fsdp_group, weight_for_dgrad)
     return (
         wgrad,
         dgrad if bwd_args.requires_dgrad else None,
@@ -1738,7 +1681,7 @@ def _linear_backward_fake(
             "(fsdp_group is not None); use FSDP2 or MCore FSDP."
         )
 
-    weight = args.saved_weight
+    weight = args.original_weight
     out_dtype = args.activation_dtype
     out_features, in_features = weight.shape
 
@@ -1866,7 +1809,7 @@ def backward(
         """Backward pass: compute gradients and reduce FP8 scaling factors."""
         bwd_args: LinearBwdArgs = ctx.backward_objects
         bwd_args.grad_output = grad_output
-        bwd_args.setup_saved_tensors(ctx)
+        bwd_args.set_saved_tensors(restore_from_func_ctx(ctx))
         nvtx_label = "transformer_engine._Linear.backward"
         if bwd_args.ub_name is not None:
             nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}"
@@ -2388,7 +2331,10 @@ 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:
+        use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
+        if _linear_op is None and torch.compiler.is_compiling():
+            warn_if_compile_disabled()
+        if use_compiled_op:
             reason = self._compile_eager_fallback_reason(
                 inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug
             )
@@ -2400,16 +2346,6 @@ def forward(
 
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
-            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
-                defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
-            )
-            weight_meta = (
-                weight_tensor.to_spec() if isinstance(weight_tensor, DeferredCat) else weight_tensor
-            )
-            bias_meta = (
-                bias_tensor.to_spec() if isinstance(bias_tensor, DeferredCat) else bias_tensor
-            )
-
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
                 if not debug
@@ -2420,6 +2356,26 @@ def forward(
                     debug = False
                     quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
 
+            if use_compiled_op and any(
+                q is not None and not is_value_opaque_quantizer(q) for q in quantizers
+            ):
+                reason = "a quantizer not registered as a torch.compile value-opaque type"
+                warn_compile_eager_fallback(reason)
+                torch._dynamo.graph_break(msg=f"te.Linear falling back to eager: {reason}")
+                use_compiled_op = False
+
+            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
+                defer_concatenation=use_compiled_op
+            )
+            weight_meta = (
+                weight_tensor.to_spec()
+                if isinstance(weight_tensor, ParameterParts)
+                else weight_tensor
+            )
+            bias_meta = (
+                bias_tensor.to_spec() if isinstance(bias_tensor, ParameterParts) else bias_tensor
+            )
+
             (
                 input_quantizer,
                 weight_quantizer,
@@ -2433,9 +2389,6 @@ def forward(
                     weight_quantizer, weight_meta
                 )
 
-            use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
-            if _linear_op is None and torch.compiler.is_compiling():
-                warn_if_compile_disabled()
             if use_compiled_op:
                 # Process groups cross the op boundary separately from quantizers.
                 for quantizer in (input_quantizer, grad_output_quantizer):
@@ -2546,22 +2499,6 @@ def forward(
                 is_grad_enabled=is_grad_enabled,
             )
 
-            if use_compiled_op:
-                # Safety net for quantizer-dependent conditions only.
-                fallback_reason = fwd_args.compile_unsupported_reason()
-                if fallback_reason is not None:
-                    warn_compile_eager_fallback(fallback_reason)
-                    torch._dynamo.graph_break(
-                        msg=f"te.Linear falling back to eager: {fallback_reason}"
-                    )
-                    use_compiled_op = False
-                    weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
-                    linear_bias_tensor = (
-                        bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None
-                    )
-                    fwd_args.weight = weight_tensor
-                    fwd_args.bias = linear_bias_tensor
-
             if use_compiled_op:
                 check_gemm_dims(inp.shape, weight_meta.shape, self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
@@ -2654,12 +2591,10 @@ def _compile_eager_fallback_reason(
         is_grad_enabled: bool,
         debug: bool,
     ) -> Optional[str]:
-        """Why this call can't use the compiled op (else None), decided before
-        prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
+        """Check configuration before prepare_forward, without concatenating parameters."""
         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"
@@ -2667,8 +2602,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"
@@ -2709,29 +2646,14 @@ def _forward_eager_fallback(
         )
 
     def _get_weight_and_bias_tensors(self, defer_concatenation=False):
-        weights = self._get_weight_tensors()
-        weight_tensor = (
-            DeferredCat(weights)
-            if defer_concatenation
-            and len(weights) > 1
-            and all(not isinstance(w, (QuantizedTensor, QuantizedTensorStorage)) for w in weights)
-            else noop_cat(weights)
-        )
-        bias_tensor = None
+        weight = concat_input(self._get_weight_tensors(), defer=defer_concatenation)
+        bias = None
         if self.use_bias:
-            biases = [getattr(self, name) for name in self.bias_names]
-            bias_tensor = (
-                DeferredCat(biases)
-                if defer_concatenation
-                and len(biases) > 1
-                and self.apply_bias
-                and not self.gemm_bias_unfused_add
-                and all(
-                    not isinstance(b, (QuantizedTensor, QuantizedTensorStorage)) for b in biases
-                )
-                else noop_cat(biases)
+            bias = concat_input(
+                [getattr(self, name) for name in self.bias_names],
+                defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
             )
-        return weight_tensor, bias_tensor
+        return weight, bias
 
     def onnx_forward(
         self,

From 5d5c1802c10e281a2b64b1be42a65e4631f7cb77 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 5 Oct 2026 16:56:29 +0200
Subject: [PATCH 04/13] [PyTorch] Limit the split-parameter API refactor

Restore compile_unsupported_reason, setup_saved_tensors(ctx), the original backward operand names, and the existing adapter contract. Remove the unrelated MLA changes. Retain the parameter-parts alias and helper, split-parameter correctness fixes, and focused regression coverage.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |   2 +-
 .../pytorch/attention/fused_mla_q_uproj.py    |   4 +-
 .../pytorch/dynamo/custom_op.py               |  98 +++-----
 .../pytorch/dynamo/parameter_parts.py         |  20 ++
 transformer_engine/pytorch/module/linear.py   | 220 ++++++++++++------
 5 files changed, 203 insertions(+), 141 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ba0432b68cb..19eb4e245ee 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2048,7 +2048,7 @@ def fn(inp):
             )
 
 
-# Configs rejected by Linear._compile_eager_fallback_reason() that a
+# Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a
 # single-GPU unit test can construct. Distributed-only reasons (fsdp_group,
 # DistributedWeight) and CPU offloading need machinery this file doesn't have;
 # delayed scaling is a hard error (check_recipe_support), tested separately.
diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
index 913aea2ccae..b15e53e6ee3 100644
--- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
+++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
@@ -299,8 +299,8 @@ def backward_linear(
         bwd_args = LinearBwdArgs(
             grad_output=grad_output,
             inputmat=x_saved,
-            weight_for_dgrad=w_q,
-            original_weight=w_q,
+            weight_fp8=w_q,
+            saved_weight=w_q,
             grad_output_quantizer=grad_output_quantizer,
             inp_shape=x_saved.shape,
             activation_dtype=act_dtype,
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 503e3b8d4c4..a81ebccae81 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -39,7 +39,8 @@
     buffers, and a ``__kind__`` tag) so a quantized tensor crosses as its buffers.
     Unions including ``ParameterParts`` also accept deferred row concatenations:
     all parts cross in the buffer list, with full gradients split outside the op.
-    The adapter saves/restores the parts; only the real kernel materializes them.
+    Saved parts use ``parameter_parts.restore_from_func_ctx`` in the backward
+    container; reconstruction of the full view happens only in the real kernel.
   * ``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.
@@ -71,7 +72,7 @@
   * on ``backward()`` the incoming flat grads are sliced per user output from the
     stashed plan (a ``grad_outputs`` field on the backward args receives the
     whole tuple; otherwise ``grad_output`` receives the first output's grad),
-    the container's optional ``set_saved_tensors`` hook receives the restored
+    the container's optional ``setup_saved_tensors`` hook restores the saved
     tensors, then the *backward op* runs the real ``bwd_impl`` and returns the
     flat grads (``bwd_fake_impl`` is its data-free fake).
 
@@ -116,7 +117,6 @@
     Quantizer,
     _quantized_tensor_passthrough_ops,
     prepare_for_saving,
-    restore_from_func_ctx,
 )
 from ..utils import record_compile_disabled
 
@@ -413,7 +413,7 @@ class _TensorOrQuantizedKind(Enum):
     NONE = "none"
     TENSOR = "tensor"
     STORAGE = "storage"
-    PARTS = "parts"
+    CONCATENATED = "concatenated"
 
 
 _TQ_KIND_KEY = "__kind__"
@@ -544,7 +544,7 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
             raise TypeError(f"field {field.name!r} does not accept ParameterParts")
         slots[tensor_slot] = None
         slots[inner_slot] = list(value.parts)
-        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.PARTS})
+        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
         # ``Tensor?`` slot; subclass flattening (if any) is done by the
@@ -564,9 +564,7 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         )
 
 
-def _unpack_tensor_or_quantized(
-    field: _FieldPlan, slots: Dict[str, Any], *, materialize_parts: bool = False
-) -> Any:
+def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any:
     """Inverse of :func:`_pack_tensor_or_quantized`."""
     tensor_slot, inner_slot, meta_slot = (s.name for s in field.slots)
     meta = slots[meta_slot]
@@ -575,9 +573,8 @@ def _unpack_tensor_or_quantized(
         return None
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
-    if kind == _TensorOrQuantizedKind.PARTS:
-        parts = ParameterParts(tuple(slots[inner_slot]))
-        return parts.materialize() if materialize_parts else parts
+    if kind == _TensorOrQuantizedKind.CONCATENATED:
+        return ParameterParts(tuple(slots[inner_slot]))
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -679,7 +676,7 @@ def pack(self, obj: Any) -> Dict[str, Any]:
             slots[_SIMPLE_META_SLOT] = OpaqueValueBundle(simple)
         return slots
 
-    def unpack(self, slots: Dict[str, Any], *, materialize_parts: bool = False) -> Any:
+    def unpack(self, slots: Dict[str, Any]) -> Any:
         """Rebuild a fresh ``arg_type`` instance from the op's flat slot dict.
 
         Inverse of :meth:`pack`.
@@ -691,9 +688,7 @@ def unpack(self, slots: Dict[str, Any], *, materialize_parts: bool = False) -> A
                 case _FieldKind.TENSOR:
                     kwargs[field.name] = slots[field.slots[0].name]
                 case _FieldKind.TENSOR_OR_QUANTIZED:
-                    kwargs[field.name] = _unpack_tensor_or_quantized(
-                        field, slots, materialize_parts=materialize_parts
-                    )
+                    kwargs[field.name] = _unpack_tensor_or_quantized(field, slots)
                 case _FieldKind.PROCESS_GROUP:
                     if bundle is not None:
                         name = bundle[field.name]
@@ -1011,8 +1006,14 @@ def _register_base_op(
     ``pack_result``.
     """
 
+    concatenated_fields = tuple(f.name for f in plan.fields if f.allows_concatenation)
+
     def _impl(*flat: Any) -> List[torch.Tensor]:
-        obj = plan.unpack(dict(zip(plan.slot_names, flat)), materialize_parts=True)
+        obj = plan.unpack(dict(zip(plan.slot_names, flat)))
+        for name in concatenated_fields:
+            value = getattr(obj, name)
+            if isinstance(value, ParameterParts):
+                setattr(obj, name, value.materialize())
         return pack_result(impl(obj))
 
     def _fake(*flat: Any) -> List[torch.Tensor]:
@@ -1028,32 +1029,6 @@ def _fake(*flat: Any) -> List[torch.Tensor]:
     return op
 
 
-def _flatten_saved_parameter_parts(ctx, tensors):
-    lengths = [len(t.parts) if isinstance(t, ParameterParts) else None for t in tensors]
-    ctx.parameter_parts_lengths = lengths if any(n is not None for n in lengths) else None
-    if ctx.parameter_parts_lengths is None:
-        return tensors
-    return [part for t in tensors for part in (t.parts if isinstance(t, ParameterParts) else (t,))]
-
-
-def _restore_saved_tensors(ctx):
-    tensors = restore_from_func_ctx(ctx)
-    lengths = getattr(ctx, "parameter_parts_lengths", None)
-    if lengths is None:
-        return tensors
-    restored = []
-    offset = 0
-    for length in lengths:
-        if length is None:
-            restored.append(tensors[offset])
-            offset += 1
-        else:
-            restored.append(ParameterParts(tuple(tensors[offset : offset + length])))
-            offset += length
-    ctx.parameter_parts_lengths = None
-    return restored
-
-
 def _register_autograd_for_op(
     *,
     fwd_op: Any,
@@ -1072,9 +1047,6 @@ def _register_autograd_for_op(
     the plan on ``ctx`` so backward can slice its grads per user output.
     """
     bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields)
-    supports_parts = any(f.allows_concatenation for f in fwd_plan.fields)
-    tensor_storage_offsets = frozenset(fwd_plan.tensor_or_quantized_offsets())
-    set_saved_tensors = getattr(bwd_plan.arg_type, "set_saved_tensors", None)
 
     def _setup_context(ctx, inputs, output):
         ctx.fwd_tensor_list_lengths = {
@@ -1095,9 +1067,15 @@ def _setup_context(ctx, inputs, output):
             out_plan.ctx_attrs,
             tuple(saved_list),
         )
-        saved = tensors_to_save_from_setup or ()
-        if supports_parts:
-            saved = _flatten_saved_parameter_parts(ctx, saved)
+        saved = []
+        ctx.concatenated_saved_lengths = []
+        for value in tensors_to_save_from_setup or ():
+            if isinstance(value, ParameterParts):
+                ctx.concatenated_saved_lengths.append(len(value.parts))
+                saved.extend(value.parts)
+            else:
+                ctx.concatenated_saved_lengths.append(None)
+                saved.append(value)
         tensors_to_save, tensor_objects = prepare_for_saving(*saved)
         ctx.tensor_objects = tensor_objects
         ctx.save_for_backward(*tensors_to_save)
@@ -1106,12 +1084,12 @@ def _setup_context(ctx, inputs, output):
         # Input shapes for the grad slots (SymInt-safe on ctx): a bwd impl may
         # rederive shapes lossily (e.g. rank-1 inputs come back rank-2), so the
         # returned grads are viewed back to the true input shapes below.
-        ctx.parameter_parts_grad_shapes = {
+        ctx.concatenated_grad_shapes = {
             pos: [tensor.shape for tensor in inputs[pos + 1]]
             for pos in grad_targets
-            if pos in tensor_storage_offsets
+            if pos in fwd_plan.tensor_or_quantized_offsets()
             and inputs[pos] is None
-            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.PARTS
+            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.CONCATENATED
         }
         ctx.grad_input_shapes = {
             pos: inputs[pos].shape for pos in grad_targets if isinstance(inputs[pos], torch.Tensor)
@@ -1119,9 +1097,7 @@ def _setup_context(ctx, inputs, output):
 
     def _autograd_backward(ctx, *grad_outputs):
         bwd_obj = ctx.backward_objects
-        if set_saved_tensors is not None:
-            set_saved_tensors(bwd_obj, _restore_saved_tensors(ctx))
-        elif hasattr(bwd_obj, "setup_saved_tensors"):
+        if hasattr(bwd_obj, "setup_saved_tensors"):
             bwd_obj.setup_saved_tensors(ctx)
         ctx.tensor_objects = None
         user_grads = _slice_user_grads(ctx.output_ranges, grad_outputs[0])
@@ -1141,8 +1117,8 @@ 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.parameter_parts_grad_shapes:
-                shapes = ctx.parameter_parts_grad_shapes[pos]
+            if pos in ctx.concatenated_grad_shapes:
+                shapes = ctx.concatenated_grad_shapes[pos]
                 if g is not None:
                     out[pos + 1] = list(torch.split(g, [shape[0] for shape in shapes], dim=0))
                 continue
@@ -1152,7 +1128,7 @@ def _autograd_backward(ctx, *grad_outputs):
                     g = g.view(shape)
             out[pos] = g
         ctx.grad_input_shapes = None
-        ctx.parameter_parts_grad_shapes = None
+        ctx.concatenated_grad_shapes = None
         return tuple(out)
 
     fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context)
@@ -1411,17 +1387,15 @@ def register_custom_op_with_autograd(
       non-differentiable input).
     * ``bwd_fake_impl(bwd_args)`` -- data-free twin of ``bwd_impl`` returning
       :class:`TensorSpec` grads.
-    * ``bwd_arg_type.set_saved_tensors(tensors)`` -- optional hook receiving
-      restored tensors, storage objects or parameter parts in forward save order.
-      Otherwise, the legacy ``setup_saved_tensors(ctx)`` hook is used if present.
+    * ``bwd_arg_type.setup_saved_tensors(ctx)`` -- optional hook; skipped if
+      absent.
 
     How the backward container is populated: ``setup_context`` fills the
     ``bwd_arg_type`` instance's non-tensor fields (quantizers, config) from
     forward state and returns the tensors to persist; the framework saves them
     via ``ctx.save_for_backward``. Before ``bwd_impl`` runs, the framework
     restores them into the container's tensor fields through the
-    ``set_saved_tensors`` (or legacy ``setup_saved_tensors``) hook and sets the
-    incoming gradient directly --
+    ``setup_saved_tensors`` hook and sets the incoming gradient directly --
     into a ``grad_outputs`` field (tuple, one grad per user output) if
     ``bwd_arg_type`` declares one, else into ``grad_output`` (the first user
     output's grad) -- so ``bwd_impl`` receives a fully-populated
diff --git a/transformer_engine/pytorch/dynamo/parameter_parts.py b/transformer_engine/pytorch/dynamo/parameter_parts.py
index b314b2efb33..c0ad545d783 100644
--- a/transformer_engine/pytorch/dynamo/parameter_parts.py
+++ b/transformer_engine/pytorch/dynamo/parameter_parts.py
@@ -14,6 +14,7 @@
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
+    restore_from_func_ctx as restore_tensor_ctx,
 )
 
 
@@ -76,3 +77,22 @@ def materialize(self):
 
 _TensorT = TypeVar("_TensorT")
 ConcatInput = Union[_TensorT, ParameterParts]
+
+
+def restore_from_func_ctx(ctx):
+    """Restore original split parameters saved by the custom-op framework."""
+    tensors = restore_tensor_ctx(ctx)
+    lengths = getattr(ctx, "concatenated_saved_lengths", None)
+    if lengths is None:
+        return tensors
+    restored = []
+    offset = 0
+    for length in lengths:
+        if length is None:
+            restored.append(tensors[offset])
+            offset += 1
+        else:
+            restored.append(ParameterParts(tuple(tensors[offset : offset + length])))
+            offset += length
+    ctx.concatenated_saved_lengths = None
+    return restored
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index a9691acaee6..8f66102c40c 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -83,9 +83,8 @@
     QuantizedTensorStorage,
     Quantizer,
     prepare_for_saving,
-    restore_from_func_ctx,
 )
-from ..dynamo.parameter_parts import ConcatInput, ParameterParts
+from ..dynamo.parameter_parts import ConcatInput, ParameterParts, restore_from_func_ctx
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -181,6 +180,57 @@ class LinearFwdArgs:
     cpu_offloading: bool
     is_grad_enabled: bool
 
+    def compile_unsupported_reason(self) -> Optional[str]:
+        """Reason this config can't use the torch.compile custom-op path (else None)."""
+        if self.debug:
+            return "debug instrumentation (nvidia-dlfw-inspect)"
+        if is_distributed_weight(self.weight):
+            return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
+        if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)):
+            return "a quantized input tensor"
+        if self.fsdp_group is not None:
+            return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
+        if (
+            self.fp8_output
+            and self.is_grad_enabled
+            and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad)
+        ):
+            return "differentiable fp8_output=True"
+        if self.cpu_offloading:
+            return "CPU activation offloading"
+        if self.wgrad_store is not None:
+            # Non-None only when delayed wgrad compute is on (see Linear.forward).
+            return "delayed wgrad compute (wgrad_store)"
+        if (
+            self.grad_input_quantizer is not None
+            and self.is_grad_enabled
+            and self.input_requires_grad
+            and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad)
+        ):
+            # A quantized dgrad can't cross the op boundary: grads are packed
+            # one plain Tensor[] slot each (_pack_bwd_result).
+            return "a quantized input grad (fp8_grad=True)"
+        if self.cache_weight and self.fp8:
+            # The cached workspace is updated in place on the first microbatch,
+            # which the functional op (mutates_args=()) can't express. Without
+            # FP8 no workspace exists, so is_first_microbatch is inert.
+            return "FP8 weight caching (is_first_microbatch)"
+        if self.fuse_wgrad_accumulation:
+            return "fuse_wgrad_accumulation (main_grad)"
+        for quantizer in (
+            self.input_quantizer,
+            self.weight_quantizer,
+            self.output_quantizer,
+            self.grad_input_quantizer,
+            self.grad_weight_quantizer,
+            self.grad_output_quantizer,
+        ):
+            # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe
+            # quantizers are not value-opaque and can't cross the custom-op boundary.
+            if quantizer is not None and not is_value_opaque_quantizer(quantizer):
+                return "a quantizer not registered as a torch.compile value-opaque type"
+        return None
+
 
 @dataclass(slots=True)
 class LinearBwdArgs:
@@ -189,8 +239,8 @@ class LinearBwdArgs:
     # --- Saved / restored tensors (populated at backward entry) ---
     grad_output: Optional[torch.Tensor] = None
     inputmat: Optional[TensorOrQuantized] = None
-    weight_for_dgrad: Optional[ConcatInput[TensorOrQuantized]] = None
-    original_weight: Optional[ConcatInput[TensorOrQuantized]] = None
+    weight_fp8: Optional[ConcatInput[TensorOrQuantized]] = None
+    saved_weight: Optional[ConcatInput[TensorOrQuantized]] = None
     bias: Optional[ConcatInput[torch.Tensor]] = None
 
     # --- Quantizers ---
@@ -255,9 +305,16 @@ class LinearBwdArgs:
     # --- Per-backward scratch state (populated inside _linear_backward_impl) ---
     ub_obj_gradout: Optional[Any] = None
 
-    def set_saved_tensors(self, tensors) -> None:
-        """Assign restored operands in their forward save order."""
-        self.inputmat, self.weight_for_dgrad, self.original_weight, self.bias = tensors
+    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
 
 
 def _out_leading_from_inp(leading: int, args: Union[LinearFwdArgs, LinearBwdArgs]) -> int:
@@ -688,7 +745,7 @@ def _linear_forward_impl(
         saved_tensor_aliases = (
             "inp" if saved_inputmat is inp else None,
             wt_alias,
-            "weight",  # ``original_weight`` slot is always the weight parameter
+            "weight",  # ``saved_weight`` slot is always the weight parameter
             "bias" if bias is not None else None,
         )
         tensors_to_save_from_forward = (
@@ -854,7 +911,7 @@ def _linear_forward_fake(
 
     # ------------------------------------------------------
     # Backward state -- saved-tensor layout
-    # (saved_inputmat, wt_save, original_weight, bias) with name-based aliasing.
+    # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing.
     # ------------------------------------------------------
     tensors_to_save_from_forward = None
     ctx_attrs = None
@@ -908,7 +965,7 @@ def _linear_forward_fake(
                 shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
             )
 
-        # Slot 2 -- ``original_weight`` (always aliased to ``weight``).
+        # Slot 2 -- ``saved_weight`` (always aliased to ``weight``).
         # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent).
         saved_tensor_aliases = (
             inputmat_alias,
@@ -1028,8 +1085,8 @@ def _linear_setup_ctx(
         bwd_args.grad_weight_quantizer = None
         bwd_args.grad_output_quantizer = None
 
-    saved_inputmat, wt_save, original_weight, saved_bias = tensors_to_save_from_forward
-    inputmat_alias, wt_save_alias, original_weight_alias, bias_alias = ctx_attrs[
+    saved_inputmat, wt_save, saved_weight, saved_bias = tensors_to_save_from_forward
+    inputmat_alias, wt_save_alias, saved_weight_alias, bias_alias = ctx_attrs[
         "saved_tensor_aliases"
     ]
     bwd_args.owns_input = inputmat_alias != "inp"
@@ -1041,26 +1098,26 @@ def _linear_setup_ctx(
         wt_save = fwd_outputs[1]
     elif wt_save_alias == "weight_workspace":
         wt_save = fwd_args.weight_workspace
-    if original_weight_alias == "weight":
-        original_weight = weight
+    if saved_weight_alias == "weight":
+        saved_weight = weight
     if bias_alias == "bias":
         saved_bias = bias
-    return (saved_inputmat, wt_save, original_weight, saved_bias)
+    return (saved_inputmat, wt_save, saved_weight, saved_bias)
 
 
 def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]:
     """Backward implementation for the linear layer.
 
     Caller must have populated ``args.grad_output`` and run
-    ``args.set_saved_tensors(tensors)`` before invocation.
+    ``args.setup_saved_tensors(ctx)`` before invocation.
     """
     bwd_args = args
     grad_output = args.grad_output
     assert grad_output is not None
     inputmat = args.inputmat
-    weight_for_dgrad = args.weight_for_dgrad
-    original_weight = args.original_weight
-    is_dist_weight = is_distributed_weight(original_weight)
+    weight_fp8 = args.weight_fp8
+    saved_weight = args.saved_weight
+    is_dist_weight = is_distributed_weight(saved_weight)
     bias = args.bias
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
@@ -1109,20 +1166,20 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
                 origin_weight_python_object.main_grad = main_grad
 
         # Gather intermediate/activation tensors if needed
-        # NOTE: weight_for_dgrad = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already
+        # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already
         #       shards/unshards the base weights so we don't do it ourselves
         nvtx_range_push(f"{nvtx_label}.fsdp_gather")
         _fsdp_gather_tensors(
             bwd_args.fsdp_group,
             bwd_args.fsdp_shapes,
             inputmat,
-            weight_for_dgrad,
+            weight_fp8,
         )
         nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
 
         # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes).
         if bwd_args.inp_shape is None:
-            in_features = original_weight.shape[-1]
+            in_features = saved_weight.shape[-1]
             inp_leading = _inp_leading_from_out(grad_output.shape[0], bwd_args)
             bwd_args.inp_shape = (
                 torch.Size([in_features])
@@ -1284,32 +1341,32 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
         # Distributed weight (e.g. GTP): re-gather the sharded weight; runs even when
         # requires_dgrad=False so the prev_w prefetch is issued for the next layer's bwd.
         if is_dist_weight:
-            weight_for_dgrad = materialize_weight_for_backward(original_weight)[0]
+            weight_fp8 = materialize_weight_for_backward(saved_weight)[0]
 
         if bwd_args.requires_dgrad:
 
             # FSDP2: Re-create workspace from all-gathered weight when
             # workspace was not saved. (Issue #2681)
-            # Use original_weight (the original weight parameter) since
-            # weight_for_dgrad is only set when workspace was saved.
-            if weight_for_dgrad is None:
-                if isinstance(original_weight, QuantizedTensorStorage):
+            # Use saved_weight (the original weight parameter) since
+            # weight_fp8 is only set when workspace was saved.
+            if weight_fp8 is None:
+                if isinstance(saved_weight, QuantizedTensorStorage):
                     # saved weight is already set to right usages by
                     # fsdp2 quantized-tensor hooks when workspace was not saved.
-                    weight_for_dgrad = original_weight
+                    weight_fp8 = saved_weight
                 elif bwd_args.weight_quantizer is not None:
                     bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True)
-                    weight_for_dgrad = bwd_args.weight_quantizer(original_weight)
+                    weight_fp8 = bwd_args.weight_quantizer(saved_weight)
             elif (
                 is_dist_weight
                 and bwd_args.fp8
                 and bwd_args.weight_quantizer is not None
-                and not isinstance(weight_for_dgrad, QuantizedTensorStorage)
+                and not isinstance(weight_fp8, QuantizedTensorStorage)
             ):
                 # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer
                 # so the dgrad operand isn't cast by the delayed recipe.
                 bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True)
-                weight_for_dgrad = bwd_args.weight_quantizer(weight_for_dgrad)
+                weight_fp8 = bwd_args.weight_quantizer(weight_fp8)
 
             # Make sure required data is available
             if isinstance(grad_output, QuantizedTensorStorage):
@@ -1317,9 +1374,9 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             if (
                 bwd_args.fp8
                 and weight_quantizer is not None
-                and isinstance(weight_for_dgrad, QuantizedTensorStorage)
+                and isinstance(weight_fp8, QuantizedTensorStorage)
             ):
-                weight_for_dgrad.update_usage(columnwise_usage=True)
+                weight_fp8.update_usage(columnwise_usage=True)
 
             # Choose whether to use GEMM kernel with split accumulator
             use_split_accumulator = bwd_args.dgrad_use_split_accumulator
@@ -1346,18 +1403,18 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             # Note: dx = dy * w
 
             nvtx_range_push(f"{nvtx_label}.dgrad_gemm")
-            gemm_weight = weight_for_dgrad
+            weight_for_dgrad = weight_fp8
             if bwd_args.backward_override == "dequantized":
-                if isinstance(gemm_weight, QuantizedTensorStorage):
-                    gemm_weight = gemm_weight.dequantize(dtype=bwd_args.activation_dtype)
+                if isinstance(weight_for_dgrad, QuantizedTensorStorage):
+                    weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype)
                 else:
-                    gemm_weight = cast_if_needed(gemm_weight, bwd_args.activation_dtype)
+                    weight_for_dgrad = cast_if_needed(weight_for_dgrad, bwd_args.activation_dtype)
             elif bwd_args.backward_override == "high_precision":
-                gemm_weight = original_weight
-                if isinstance(gemm_weight, QuantizedTensorStorage):
-                    gemm_weight = gemm_weight.dequantize(dtype=bwd_args.activation_dtype)
+                weight_for_dgrad = saved_weight
+                if isinstance(weight_for_dgrad, QuantizedTensorStorage):
+                    weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype)
             gemm_out, *_, reduce_scatter_out = general_gemm(
-                gemm_weight,
+                weight_for_dgrad,
                 grad_output,
                 layout="NN",
                 grad=True,
@@ -1377,8 +1434,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             # and 2d block-scaled weights in TE managed memory. So we need to clear
             # it here.
             # (Issues #2681, #2717)
-            if bwd_args.is_fsdp2 and isinstance(weight_for_dgrad, QuantizedTensorStorage):
-                clear_columnwise_cache(weight_for_dgrad)
+            if bwd_args.is_fsdp2 and isinstance(weight_fp8, QuantizedTensorStorage):
+                clear_columnwise_cache(weight_fp8)
 
             # Prepare grad input tensor
             # Note: Perform tensor-parallel communication
@@ -1587,7 +1644,7 @@ def wgrad_gemm(
                 # Distributed weight (e.g. GTP): reduce-scatter the freshly computed wgrad
                 # (async; overlap with the next layer's bwd via the cascade).
                 if is_dist_weight:
-                    wgrad = finalize_weight_grads(original_weight, [wgrad])[0]
+                    wgrad = finalize_weight_grads(saved_weight, [wgrad])[0]
 
                 # Update grad bias if needed
                 if grad_bias is None:
@@ -1659,7 +1716,7 @@ def wgrad_gemm(
 
     # Scatter fp8 weight buffers
     if bwd_args.fp8 and not bwd_args.is_weight_param_quantized:
-        _fsdp_scatter_tensors(bwd_args.fsdp_group, weight_for_dgrad)
+        _fsdp_scatter_tensors(bwd_args.fsdp_group, weight_fp8)
     return (
         wgrad,
         dgrad if bwd_args.requires_dgrad else None,
@@ -1681,7 +1738,7 @@ def _linear_backward_fake(
             "(fsdp_group is not None); use FSDP2 or MCore FSDP."
         )
 
-    weight = args.original_weight
+    weight = args.saved_weight
     out_dtype = args.activation_dtype
     out_features, in_features = weight.shape
 
@@ -1809,7 +1866,7 @@ def backward(
         """Backward pass: compute gradients and reduce FP8 scaling factors."""
         bwd_args: LinearBwdArgs = ctx.backward_objects
         bwd_args.grad_output = grad_output
-        bwd_args.set_saved_tensors(restore_from_func_ctx(ctx))
+        bwd_args.setup_saved_tensors(ctx)
         nvtx_label = "transformer_engine._Linear.backward"
         if bwd_args.ub_name is not None:
             nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}"
@@ -2331,10 +2388,7 @@ def forward(
             if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()):
                 fp8_grad = True
 
-        use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
-        if _linear_op is None and torch.compiler.is_compiling():
-            warn_if_compile_disabled()
-        if use_compiled_op:
+        if torch.compiler.is_compiling() and _linear_op is not None:
             reason = self._compile_eager_fallback_reason(
                 inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug
             )
@@ -2346,26 +2400,8 @@ def forward(
 
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
-            quantizers = (
-                self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
-                if not debug
-                else self._get_debug_quantizers(fp8_output, fp8_grad, is_grad_enabled)
-            )
-            if debug:
-                if self.no_debug_features_active(quantizers):
-                    debug = False
-                    quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
-
-            if use_compiled_op and any(
-                q is not None and not is_value_opaque_quantizer(q) for q in quantizers
-            ):
-                reason = "a quantizer not registered as a torch.compile value-opaque type"
-                warn_compile_eager_fallback(reason)
-                torch._dynamo.graph_break(msg=f"te.Linear falling back to eager: {reason}")
-                use_compiled_op = False
-
             weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
-                defer_concatenation=use_compiled_op
+                defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
             weight_meta = (
                 weight_tensor.to_spec()
@@ -2376,6 +2412,16 @@ def forward(
                 bias_tensor.to_spec() if isinstance(bias_tensor, ParameterParts) else bias_tensor
             )
 
+            quantizers = (
+                self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
+                if not debug
+                else self._get_debug_quantizers(fp8_output, fp8_grad, is_grad_enabled)
+            )
+            if debug:
+                if self.no_debug_features_active(quantizers):
+                    debug = False
+                    quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
+
             (
                 input_quantizer,
                 weight_quantizer,
@@ -2389,6 +2435,9 @@ def forward(
                     weight_quantizer, weight_meta
                 )
 
+            use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
+            if _linear_op is None and torch.compiler.is_compiling():
+                warn_if_compile_disabled()
             if use_compiled_op:
                 # Process groups cross the op boundary separately from quantizers.
                 for quantizer in (input_quantizer, grad_output_quantizer):
@@ -2499,6 +2548,22 @@ def forward(
                 is_grad_enabled=is_grad_enabled,
             )
 
+            if use_compiled_op:
+                # Safety net for quantizer-dependent conditions only.
+                fallback_reason = fwd_args.compile_unsupported_reason()
+                if fallback_reason is not None:
+                    warn_compile_eager_fallback(fallback_reason)
+                    torch._dynamo.graph_break(
+                        msg=f"te.Linear falling back to eager: {fallback_reason}"
+                    )
+                    use_compiled_op = False
+                    weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+                    linear_bias_tensor = (
+                        bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None
+                    )
+                    fwd_args.weight = weight_tensor
+                    fwd_args.bias = linear_bias_tensor
+
             if use_compiled_op:
                 check_gemm_dims(inp.shape, weight_meta.shape, self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
@@ -2591,7 +2656,8 @@ def _compile_eager_fallback_reason(
         is_grad_enabled: bool,
         debug: bool,
     ) -> Optional[str]:
-        """Check configuration before prepare_forward, without concatenating parameters."""
+        """Why this call can't use the compiled op (else None), decided before
+        prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
         if any(is_distributed_weight(getattr(self, name)) for name in self.weight_names):
@@ -2646,14 +2712,16 @@ def _forward_eager_fallback(
         )
 
     def _get_weight_and_bias_tensors(self, defer_concatenation=False):
-        weight = concat_input(self._get_weight_tensors(), defer=defer_concatenation)
-        bias = None
+        weights = self._get_weight_tensors()
+        weight_tensor = concat_input(weights, defer=defer_concatenation)
+        bias_tensor = None
         if self.use_bias:
-            bias = concat_input(
-                [getattr(self, name) for name in self.bias_names],
+            biases = [getattr(self, name) for name in self.bias_names]
+            bias_tensor = concat_input(
+                biases,
                 defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
             )
-        return weight, bias
+        return weight_tensor, bias_tensor
 
     def onnx_forward(
         self,

From 73d58af26c7464c73b308785f822456ef58a5bfc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 12 Oct 2026 16:24:51 +0200
Subject: [PATCH 05/13] [PyTorch] Keep the frozen-weight bias fix out of split
 parameter support

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py         | 4 +---
 transformer_engine/pytorch/module/linear.py | 9 +++++----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 225c1e8d03d..835159e8ab2 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3360,7 +3360,7 @@ def test_parameter_parts_rejects_quantized_parts(internal):
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("grad_target", ["input", "weight", "bias"])
+@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():
@@ -3373,8 +3373,6 @@ def test_te_split_parameters_grad_targets(grad_target):
     out = compiled(inp)
     grad = torch.randn_like(out)
     out.backward(grad)
-    if grad_target == "bias":
-        torch.testing.assert_close(model.v_bias.grad, grad[:, -32:].sum(0), rtol=0, atol=0)
     actual = [out.detach().clone()]
     tensors = [inp, *model.parameters()]
     actual.extend(None if t.grad is None else t.grad.clone() for t in tensors)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 8f66102c40c..543235a54ec 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -33,7 +33,6 @@
     can_reconstruct_wgrad_input_from_original,
     check_fp8_reduce_and_update,
     concat_input,
-    sum_bias_grad,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -1264,8 +1263,6 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
             grad_output_quantizer,
         )
         nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess")
-        if bwd_args.use_bias and grad_bias is None and not bwd_args.requires_wgrad:
-            grad_bias = sum_bias_grad(grad_output)
 
         # --------------------------------------------------
         # Grad output tensor is ready for computing grad input...
@@ -1778,7 +1775,11 @@ def _linear_backward_fake(
         )
 
     grad_bias = None
-    if args.use_bias:
+    # FP8 backward computes bgrad in grad_output_preprocess whenever bias is
+    # used; in high precision it is fused into the wgrad GEMM, so it only
+    # exists when wgrad runs.
+    fp8_bwd = args.fp8 and args.backward_override is None
+    if args.use_bias and (args.requires_wgrad or fp8_bwd):
         grad_bias = TensorSpec(
             shape=(out_features,), dtype=out_dtype, device=args.grad_output.device
         )

From 122a2c946e4408aa00315b9a3774f2110ad86c41 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 14 Oct 2026 14:54:47 +0200
Subject: [PATCH 06/13] [PyTorch] Read split parameter shape and gradient flags
 directly

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/dynamo/parameter_parts.py         | 16 ++++++++++---
 transformer_engine/pytorch/module/base.py     |  4 ++--
 transformer_engine/pytorch/module/linear.py   | 23 ++++++-------------
 3 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/parameter_parts.py b/transformer_engine/pytorch/dynamo/parameter_parts.py
index c0ad545d783..7e84124fd5d 100644
--- a/transformer_engine/pytorch/dynamo/parameter_parts.py
+++ b/transformer_engine/pytorch/dynamo/parameter_parts.py
@@ -35,6 +35,16 @@ def __post_init__(self):
         ):
             raise TypeError("ParameterParts only supports non-quantized tensors")
 
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        """Shape after row concatenation."""
+        return (sum(t.shape[0] for t in self.parts), *self.parts[0].shape[1:])
+
+    @property
+    def requires_grad(self) -> bool:
+        """Whether any part requires gradients."""
+        return any(t.requires_grad for t in self.parts)
+
     def to_spec(self) -> TensorSpec:
         """Describe the concatenation without accessing its storage."""
         first = self.parts[0]
@@ -42,10 +52,10 @@ def to_spec(self) -> TensorSpec:
         for tensor in self.parts[1:]:
             dtype = torch.promote_types(dtype, tensor.dtype)
         return TensorSpec(
-            shape=(sum(t.shape[0] for t in self.parts), *first.shape[1:]),
+            shape=self.shape,
             dtype=dtype,
             device=first.device,
-            requires_grad=any(t.requires_grad for t in self.parts),
+            requires_grad=self.requires_grad,
         )
 
     def materialize(self):
@@ -71,7 +81,7 @@ def materialize(self):
             offset += tensor.numel()
         if offset * first.element_size() > storage.nbytes():
             return torch.cat(self.parts)
-        shape = (sum(t.shape[0] for t in self.parts), *first.shape[1:])
+        shape = self.shape
         return first.as_strided(shape, make_contiguous_strides_for(shape))
 
 
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 8909720d025..6fc8175d821 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -46,7 +46,7 @@
     _fsdp_gather_tensors,
 )
 from ..constants import dist_group_type
-from ..dynamo.tensor_spec import TensorSpec
+from ..dynamo.parameter_parts import ConcatInput
 from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
@@ -1323,7 +1323,7 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
     def _enable_weight_preswizzle(
         self,
         quantizer: Quantizer,
-        weight: Union[torch.Tensor, TensorSpec],
+        weight: ConcatInput[torch.Tensor],
     ) -> bool:
         """Whether to fuse scale-factor swizzling into weight quantization.
 
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 543235a54ec..585f2509b0b 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -83,7 +83,7 @@
     Quantizer,
     prepare_for_saving,
 )
-from ..dynamo.parameter_parts import ConcatInput, ParameterParts, restore_from_func_ctx
+from ..dynamo.parameter_parts import ConcatInput, restore_from_func_ctx
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -2404,14 +2404,7 @@ def forward(
             weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
                 defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
-            weight_meta = (
-                weight_tensor.to_spec()
-                if isinstance(weight_tensor, ParameterParts)
-                else weight_tensor
-            )
-            bias_meta = (
-                bias_tensor.to_spec() if isinstance(bias_tensor, ParameterParts) else bias_tensor
-            )
+            weight_requires_grad = weight_tensor.requires_grad
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2433,7 +2426,7 @@ def forward(
             ) = quantizers
             if weight_quantizer is not None and not debug:
                 weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
-                    weight_quantizer, weight_meta
+                    weight_quantizer, weight_tensor
                 )
 
             use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
@@ -2463,9 +2456,7 @@ def forward(
                 backward_override = None
             custom = is_custom(input_quantizer) or is_custom(weight_quantizer)
             backward_input_needs_gather = (
-                weight_meta.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:
@@ -2496,9 +2487,9 @@ def forward(
                 weight_workspace=weight_workspace,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
-                weight_requires_grad=weight_meta.requires_grad,
+                weight_requires_grad=weight_requires_grad,
                 bias_requires_grad=(
-                    bias_meta.requires_grad if linear_bias_tensor is not None else False
+                    linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2566,7 +2557,7 @@ def forward(
                     fwd_args.bias = linear_bias_tensor
 
             if use_compiled_op:
-                check_gemm_dims(inp.shape, weight_meta.shape, self.fp8)
+                check_gemm_dims(inp.shape, weight_tensor.shape, self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
             else:
                 out, new_weight_workspace = _linear_eager(

From 156576ecee92e40b339eccd4b58909c6a6d16086 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 14 Oct 2026 15:25:50 +0200
Subject: [PATCH 07/13] [PyTorch] Pass split parameters as tensor lists

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  24 ++--
 transformer_engine/pytorch/dynamo/concat.py   | 112 ++++++++++++++++++
 .../pytorch/dynamo/custom_op.py               |  46 +++----
 .../pytorch/dynamo/parameter_parts.py         | 108 -----------------
 transformer_engine/pytorch/module/_common.py  |   4 +-
 transformer_engine/pytorch/module/base.py     |   5 +-
 transformer_engine/pytorch/module/linear.py   |  15 ++-
 7 files changed, 168 insertions(+), 146 deletions(-)
 create mode 100644 transformer_engine/pytorch/dynamo/concat.py
 delete mode 100644 transformer_engine/pytorch/dynamo/parameter_parts.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 835159e8ab2..d5bebb66666 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3294,8 +3294,9 @@ def test_concat_input_eager_gradients(disjoint):
 @pytest.mark.parametrize(
     "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
 )
-def test_parameter_parts_storage(layout):
-    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
+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":
@@ -3311,7 +3312,9 @@ def test_parameter_parts_storage(layout):
     parts = list(storage.split([1, storage.shape[0] - 3, 2]))
     if layout == "disjoint":
         parts[1] = parts[1].clone()
-    result = ParameterParts(parts).materialize()
+    deferred = concat_input(parts, defer=True)
+    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()
@@ -3321,8 +3324,8 @@ def test_parameter_parts_storage(layout):
 
 @pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"])
 @pytest.mark.parametrize("mixed_dtype", [False, True])
-def test_parameter_parts_spec(fake, mixed_dtype):
-    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
+def test_concat_spec(fake, mixed_dtype):
+    from transformer_engine.pytorch.dynamo.concat import concat_to_spec
 
     with FakeTensorMode() if fake else contextlib.nullcontext():
         parts = [
@@ -3336,7 +3339,7 @@ def test_parameter_parts_spec(fake, mixed_dtype):
             ),
         ]
         expected = torch.cat(parts)
-        spec = ParameterParts(parts).to_spec()
+        spec = concat_to_spec(parts)
         assert spec.shape == tuple(expected.shape)
         assert spec.dtype == expected.dtype
         assert spec.device == expected.device
@@ -3345,8 +3348,8 @@ def test_parameter_parts_spec(fake, mixed_dtype):
 
 
 @pytest.mark.parametrize("internal", [False, True])
-def test_parameter_parts_rejects_quantized_parts(internal):
-    from transformer_engine.pytorch.dynamo.parameter_parts import ParameterParts
+def test_concat_rejects_quantized_parts(internal):
+    from transformer_engine.pytorch.dynamo.concat import concat_to_spec, materialize_concat
     from transformer_engine.pytorch.module._common import concat_input
 
     quantizer = _current_scaling()
@@ -3354,8 +3357,9 @@ def test_parameter_parts_rejects_quantized_parts(internal):
     tensor = TensorSpec(
         shape=(4, 8), dtype=torch.bfloat16, quantizer=quantizer, device=torch.device("cpu")
     ).create_tensor()
-    with pytest.raises(TypeError, match="non-quantized tensors"):
-        ParameterParts([tensor, tensor])
+    for concatenate in (concat_to_spec, materialize_concat):
+        with pytest.raises(TypeError, match="non-quantized tensors"):
+            concatenate([tensor, tensor])
     assert concat_input([tensor], defer=True) is tensor
 
 
diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
new file mode 100644
index 00000000000..596e6f3957c
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -0,0 +1,112 @@
+# 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, Tuple, TypeVar, Union
+
+import torch
+from torch._prims_common import make_contiguous_strides_for
+
+from .tensor_spec import TensorSpec
+from ..quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    restore_from_func_ctx as restore_tensor_ctx,
+)
+
+_TensorT = TypeVar("_TensorT")
+ConcatInput = Union[_TensorT, List[torch.Tensor]]
+
+
+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)
+        or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
+        for part in parts
+    ):
+        raise TypeError("Deferred concatenation only supports non-quantized tensors")
+
+
+def concat_shape(value: ConcatInput[torch.Tensor]) -> Tuple[int, ...]:
+    """Shape after row concatenation."""
+    if isinstance(value, list):
+        return (sum(t.shape[0] for t in value), *value[0].shape[1:])
+    return value.shape
+
+
+def concat_requires_grad(value: ConcatInput[torch.Tensor]) -> bool:
+    """Whether any input requires gradients."""
+    if isinstance(value, list):
+        return any(t.requires_grad for t in value)
+    return value.requires_grad
+
+
+def concat_to_spec(parts: List[torch.Tensor]) -> 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=concat_requires_grad(parts),
+    )
+
+
+def materialize_concat(value: ConcatInput[torch.Tensor]) -> torch.Tensor:
+    """Build a checked view, or copy parts whose storage is not adjacent."""
+    if not isinstance(value, list):
+        return value
+    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))
+
+
+def restore_from_func_ctx(ctx):
+    """Restore original split parameters saved by the custom-op framework."""
+    tensors = restore_tensor_ctx(ctx)
+    lengths = getattr(ctx, "concatenated_saved_lengths", None)
+    if lengths is None:
+        return tensors
+    restored = []
+    offset = 0
+    for length in lengths:
+        if length is None:
+            restored.append(tensors[offset])
+            offset += 1
+        else:
+            restored.append(list(tensors[offset : offset + length]))
+            offset += length
+    ctx.concatenated_saved_lengths = None
+    return restored
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 047297c96ea..49bc6aed8d5 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,9 +37,9 @@
   * ``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.
-    Unions including ``ParameterParts`` also accept deferred row concatenations:
+    Unions including ``List[Tensor]`` also accept deferred row concatenations:
     all parts cross in the buffer list, with full gradients split outside the op.
-    Saved parts use ``parameter_parts.restore_from_func_ctx`` in the backward
+    Saved parts use ``concat.restore_from_func_ctx`` in the backward
     container; reconstruction of the full view happens only in the real kernel.
   * ``SIMPLE`` -- every remaining simple value (scalars, enums, sizes,
     quantizers -- value-opaque constants baked into the graph -- and nested
@@ -110,7 +110,7 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .parameter_parts import ParameterParts
+from .concat import concat_to_spec, materialize_concat, validate_concat
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -445,16 +445,18 @@ class _FieldPlan:
     allows_concatenation: bool = False
 
 
+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 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))
-    return members in (
-        _TQ_MEMBERS,
-        _TQ_MEMBERS | {ParameterParts},
-        frozenset((torch.Tensor, ParameterParts)),
-    )
+    tensors = frozenset(a for a in members if not _is_tensor_list_annot(a))
+    return tensors == _TQ_MEMBERS or (tensors == frozenset((torch.Tensor,)) and tensors != members)
 
 
 def _is_process_group_annot(annot: Any) -> bool:
@@ -497,7 +499,10 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
         return _FieldPlan(
-            name, _FieldKind.TENSOR_OR_QUANTIZED, slots, ParameterParts in get_args(annot)
+            name,
+            _FieldKind.TENSOR_OR_QUANTIZED,
+            slots,
+            any(_is_tensor_list_annot(a) for a in get_args(annot)),
         )
     stripped, is_optional = _strip_optional(annot)
     if stripped is torch.Tensor:
@@ -539,11 +544,12 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         slots[tensor_slot] = None
         slots[inner_slot] = []
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE})
-    elif isinstance(value, ParameterParts):
+    elif isinstance(value, list):
         if not field.allows_concatenation:
-            raise TypeError(f"field {field.name!r} does not accept ParameterParts")
+            raise TypeError(f"field {field.name!r} does not accept a tensor list")
+        validate_concat(value)
         slots[tensor_slot] = None
-        slots[inner_slot] = list(value.parts)
+        slots[inner_slot] = value
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
@@ -574,7 +580,7 @@ def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
     if kind == _TensorOrQuantizedKind.CONCATENATED:
-        return ParameterParts(tuple(slots[inner_slot]))
+        return slots[inner_slot]
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -761,8 +767,8 @@ 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 isinstance(value, ParameterParts):
-            overrides[name] = value.to_spec()
+        if isinstance(value, list):
+            overrides[name] = concat_to_spec(value)
         elif value is not None and not isinstance(value, TensorSpec):
             overrides[name] = to_tensor_spec(value)
     if not overrides:
@@ -1013,8 +1019,8 @@ def _impl(*flat: Any) -> List[torch.Tensor]:
         obj = plan.unpack(dict(zip(plan.slot_names, flat)))
         for name in concatenated_fields:
             value = getattr(obj, name)
-            if isinstance(value, ParameterParts):
-                setattr(obj, name, value.materialize())
+            if isinstance(value, list):
+                setattr(obj, name, materialize_concat(value))
         return pack_result(impl(obj))
 
     def _fake(*flat: Any) -> List[torch.Tensor]:
@@ -1071,9 +1077,9 @@ def _setup_context(ctx, inputs, output):
         saved = []
         ctx.concatenated_saved_lengths = []
         for value in tensors_to_save_from_setup or ():
-            if isinstance(value, ParameterParts):
-                ctx.concatenated_saved_lengths.append(len(value.parts))
-                saved.extend(value.parts)
+            if isinstance(value, list):
+                ctx.concatenated_saved_lengths.append(len(value))
+                saved.extend(value)
             else:
                 ctx.concatenated_saved_lengths.append(None)
                 saved.append(value)
diff --git a/transformer_engine/pytorch/dynamo/parameter_parts.py b/transformer_engine/pytorch/dynamo/parameter_parts.py
deleted file mode 100644
index 7e84124fd5d..00000000000
--- a/transformer_engine/pytorch/dynamo/parameter_parts.py
+++ /dev/null
@@ -1,108 +0,0 @@
-# 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 dataclasses import dataclass
-from typing import Tuple, TypeVar, Union
-
-import torch
-from torch._prims_common import make_contiguous_strides_for
-
-from .tensor_spec import TensorSpec
-from ..quantized_tensor import (
-    QuantizedTensor,
-    QuantizedTensorStorage,
-    restore_from_func_ctx as restore_tensor_ctx,
-)
-
-
-@dataclass(frozen=True, slots=True)
-class ParameterParts:
-    """Keep every parameter visible to autograd until the opaque consumer runs."""
-
-    parts: Tuple[torch.Tensor, ...]
-
-    def __post_init__(self):
-        object.__setattr__(self, "parts", tuple(self.parts))
-        if not self.parts:
-            raise ValueError("ParameterParts requires at least one tensor")
-        if any(
-            not isinstance(part, torch.Tensor)
-            or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
-            for part in self.parts
-        ):
-            raise TypeError("ParameterParts only supports non-quantized tensors")
-
-    @property
-    def shape(self) -> Tuple[int, ...]:
-        """Shape after row concatenation."""
-        return (sum(t.shape[0] for t in self.parts), *self.parts[0].shape[1:])
-
-    @property
-    def requires_grad(self) -> bool:
-        """Whether any part requires gradients."""
-        return any(t.requires_grad for t in self.parts)
-
-    def to_spec(self) -> TensorSpec:
-        """Describe the concatenation without accessing its storage."""
-        first = self.parts[0]
-        dtype = first.dtype
-        for tensor in self.parts[1:]:
-            dtype = torch.promote_types(dtype, tensor.dtype)
-        return TensorSpec(
-            shape=self.shape,
-            dtype=dtype,
-            device=first.device,
-            requires_grad=self.requires_grad,
-        )
-
-    def materialize(self):
-        """Build a checked view, or copy parts whose storage is not adjacent."""
-        first = self.parts[0]
-        storage = first.untyped_storage()
-        offset = first.storage_offset()
-        for tensor in self.parts:
-            if (
-                tensor.dtype != first.dtype
-                or tensor.device != first.device
-                or tensor.is_neg()
-                or tensor.is_conj()
-            ):
-                return torch.cat(self.parts)
-            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(self.parts)
-            offset += tensor.numel()
-        if offset * first.element_size() > storage.nbytes():
-            return torch.cat(self.parts)
-        shape = self.shape
-        return first.as_strided(shape, make_contiguous_strides_for(shape))
-
-
-_TensorT = TypeVar("_TensorT")
-ConcatInput = Union[_TensorT, ParameterParts]
-
-
-def restore_from_func_ctx(ctx):
-    """Restore original split parameters saved by the custom-op framework."""
-    tensors = restore_tensor_ctx(ctx)
-    lengths = getattr(ctx, "concatenated_saved_lengths", None)
-    if lengths is None:
-        return tensors
-    restored = []
-    offset = 0
-    for length in lengths:
-        if length is None:
-            restored.append(tensors[offset])
-            offset += 1
-        else:
-            restored.append(ParameterParts(tuple(tensors[offset : offset + length])))
-            offset += length
-    ctx.concatenated_saved_lengths = None
-    return restored
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 7785eda7eb4..9b42f39ae94 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -14,7 +14,7 @@
 from ..constants import TE_DType
 from ..distributed import in_fp8_activation_recompute_phase
 from ..dynamo import TensorOrQuantized
-from ..dynamo.parameter_parts import ConcatInput, ParameterParts
+from ..dynamo.concat import ConcatInput
 from ..export import is_in_onnx_export_mode
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
 from ..quantization import FP8GlobalStateManager
@@ -257,7 +257,7 @@ def concat_input(
             for t in tensors
         )
     ):
-        return ParameterParts(tuple(tensors))
+        return list(tensors)
     if torch.compiler.is_compiling():
         return torch.cat(tensors)
     return noop_cat(tensors)
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 6fc8175d821..975c9fcddc7 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -46,7 +46,7 @@
     _fsdp_gather_tensors,
 )
 from ..constants import dist_group_type
-from ..dynamo.parameter_parts import ConcatInput
+from ..dynamo.concat import ConcatInput, concat_shape
 from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
@@ -1340,7 +1340,8 @@ def _enable_weight_preswizzle(
         if isinstance(quantizer, MXFP8Quantizer):
             return True
         if isinstance(quantizer, NVFP4Quantizer):
-            rows, cols = math.prod(weight.shape[:-1]), weight.shape[-1]
+            shape = concat_shape(weight)
+            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 585f2509b0b..5446b255f87 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -83,7 +83,12 @@
     Quantizer,
     prepare_for_saving,
 )
-from ..dynamo.parameter_parts import ConcatInput, restore_from_func_ctx
+from ..dynamo.concat import (
+    ConcatInput,
+    concat_shape,
+    concat_requires_grad,
+    restore_from_func_ctx,
+)
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -2404,7 +2409,7 @@ def forward(
             weight_tensor, bias_tensor = self._get_weight_and_bias_tensors(
                 defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
-            weight_requires_grad = weight_tensor.requires_grad
+            weight_requires_grad = concat_requires_grad(weight_tensor)
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2489,7 +2494,9 @@ def forward(
                 input_requires_grad=inp.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
+                    concat_requires_grad(linear_bias_tensor)
+                    if linear_bias_tensor is not None
+                    else False
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2557,7 +2564,7 @@ def forward(
                     fwd_args.bias = linear_bias_tensor
 
             if use_compiled_op:
-                check_gemm_dims(inp.shape, weight_tensor.shape, self.fp8)
+                check_gemm_dims(inp.shape, concat_shape(weight_tensor), self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
             else:
                 out, new_weight_workspace = _linear_eager(

From dd01b3ab647ea0c539c158168243bd10cc0c2ff4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 14 Oct 2026 15:47:17 +0200
Subject: [PATCH 08/13] [PyTorch] Use representation-neutral names for Linear
 operands

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/module/linear.py | 44 ++++++++++-----------
 1 file changed, 20 insertions(+), 24 deletions(-)

diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 5446b255f87..4be269b962d 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -2406,10 +2406,10 @@ 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 = self._get_weight_and_bias_tensors(
                 defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
-            weight_requires_grad = concat_requires_grad(weight_tensor)
+            weight_requires_grad = concat_requires_grad(weight)
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2431,7 +2431,7 @@ 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
                 )
 
             use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
@@ -2479,24 +2479,20 @@ 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_workspace=weight_workspace,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
                 weight_requires_grad=weight_requires_grad,
                 bias_requires_grad=(
-                    concat_requires_grad(linear_bias_tensor)
-                    if linear_bias_tensor is not None
-                    else False
+                    concat_requires_grad(linear_bias) if linear_bias is not None else False
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2556,19 +2552,19 @@ def forward(
                         msg=f"te.Linear falling back to eager: {fallback_reason}"
                     )
                     use_compiled_op = False
-                    weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
-                    linear_bias_tensor = (
-                        bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None
+                    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_tensor
-                    fwd_args.bias = linear_bias_tensor
+                    fwd_args.weight = weight
+                    fwd_args.bias = linear_bias
 
             if use_compiled_op:
-                check_gemm_dims(inp.shape, concat_shape(weight_tensor), self.fp8)
+                check_gemm_dims(inp.shape, concat_shape(weight), 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:
@@ -2579,10 +2575,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):
@@ -2712,15 +2708,15 @@ def _forward_eager_fallback(
 
     def _get_weight_and_bias_tensors(self, defer_concatenation=False):
         weights = self._get_weight_tensors()
-        weight_tensor = concat_input(weights, defer=defer_concatenation)
-        bias_tensor = None
+        weight = concat_input(weights, defer=defer_concatenation)
+        bias = None
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
-            bias_tensor = concat_input(
+            bias = concat_input(
                 biases,
                 defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
             )
-        return weight_tensor, bias_tensor
+        return weight, bias
 
     def onnx_forward(
         self,

From 88b47c3e9862f9fb5cd8f1b804a253c3027375f8 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 16 Oct 2026 14:49:22 +0200
Subject: [PATCH 09/13] [PyTorch] Pass split Linear weights in separate
 tensor-list fields

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  61 ++++++++-
 transformer_engine/pytorch/dynamo/concat.py   |  54 ++------
 .../pytorch/dynamo/custom_op.py               |  91 +++++---------
 transformer_engine/pytorch/module/_common.py  |  13 +-
 transformer_engine/pytorch/module/base.py     |   9 +-
 transformer_engine/pytorch/module/linear.py   | 117 ++++++++++++------
 6 files changed, 183 insertions(+), 162 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d5bebb66666..a0261854e9d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3269,7 +3269,9 @@ def test_concat_input_single_tensor(defer):
     from transformer_engine.pytorch.module._common import concat_input
 
     tensor = torch.nn.Parameter(torch.randn(4, 8, device="cuda"))
-    assert concat_input([tensor], defer=defer) is tensor
+    result, parts = concat_input([tensor], defer=defer)
+    assert result is tensor
+    assert not parts
 
 
 @pytest.mark.parametrize("disjoint", [False, True])
@@ -3279,7 +3281,8 @@ def test_concat_input_eager_gradients(disjoint):
     parts = [torch.nn.Parameter(t) for t in torch.randn(8, 8, device="cuda").split((3, 5))]
     if disjoint:
         parts[1] = torch.nn.Parameter(parts[1].detach().clone())
-    result = concat_input(parts)
+    result, deferred = concat_input(parts)
+    assert not deferred
     expected = torch.cat(parts)
     grad = torch.randn_like(result)
     torch.testing.assert_close(result, expected, rtol=0, atol=0)
@@ -3312,7 +3315,8 @@ def test_concat_storage(layout):
     parts = list(storage.split([1, storage.shape[0] - 3, 2]))
     if layout == "disjoint":
         parts[1] = parts[1].clone()
-    deferred = concat_input(parts, defer=True)
+    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)
@@ -3339,7 +3343,7 @@ def test_concat_spec(fake, mixed_dtype):
             ),
         ]
         expected = torch.cat(parts)
-        spec = concat_to_spec(parts)
+        spec = concat_to_spec([to_tensor_spec(t) for t in parts] if fake else parts)
         assert spec.shape == tuple(expected.shape)
         assert spec.dtype == expected.dtype
         assert spec.device == expected.device
@@ -3347,6 +3351,32 @@ def test_concat_spec(fake, mixed_dtype):
         assert not spec.is_quantized
 
 
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+def test_tensor_list_custom_op():
+    from transformer_engine.pytorch.dynamo import register_custom_op
+
+    @dataclasses.dataclass
+    class Args:
+        tensors: list[torch.Tensor]
+
+    op = register_custom_op(
+        op_name="test_tensor_list_inputs",
+        arg_type=Args,
+        impl=lambda args: [tensor.square() for tensor in args.tensors],
+        fake_impl=lambda args: args.tensors,
+    )
+    assert op is not None
+
+    def fn(tensors):
+        return op(Args(tensors))
+
+    tensors = [torch.randn(3), torch.randn(2, 4, dtype=torch.float64)]
+    expected = [tensor.square() for tensor in tensors]
+    torch.testing.assert_close(fn(tensors), expected)
+    compiled = torch.compile(fn, fullgraph=True, backend="aot_eager")
+    torch.testing.assert_close(compiled(tensors), expected)
+
+
 @pytest.mark.parametrize("internal", [False, True])
 def test_concat_rejects_quantized_parts(internal):
     from transformer_engine.pytorch.dynamo.concat import concat_to_spec, materialize_concat
@@ -3360,7 +3390,9 @@ def test_concat_rejects_quantized_parts(internal):
     for concatenate in (concat_to_spec, materialize_concat):
         with pytest.raises(TypeError, match="non-quantized tensors"):
             concatenate([tensor, tensor])
-    assert concat_input([tensor], defer=True) is tensor
+    result, parts = concat_input([tensor], defer=True)
+    assert result is tensor
+    assert not parts
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
@@ -3413,3 +3445,22 @@ def test_te_split_parameters_mixed_dtype_autocast(bf16_part):
     actual_grads = [tensor.grad for tensor in (inp, *model.parameters())]
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
     torch.testing.assert_close(actual_grads, expected_grads, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("split", [False, True])
+@pytest.mark.parametrize("bias_mode", ["fused", "none", "returned"])
+def test_te_split_parameters_onnx_path(split, bias_mode):
+    model = te.Linear(
+        64,
+        128,
+        parameters_split=dict(q=64, k=32, v=32) if split else None,
+        bias=bias_mode != "none",
+        return_bias=bias_mode == "returned",
+        params_dtype=torch.bfloat16,
+    )
+    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        expected = model(inp)
+        with te.onnx_export(enabled=True):
+            actual = model(inp)
+    torch.testing.assert_close(actual, expected, **dtype_tols(torch.bfloat16))
diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
index 596e6f3957c..e6ac9882613 100644
--- a/transformer_engine/pytorch/dynamo/concat.py
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -4,20 +4,13 @@
 
 """Deferred row concatenation for tensors consumed inside a custom op."""
 
-from typing import List, Tuple, TypeVar, Union
+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,
-    restore_from_func_ctx as restore_tensor_ctx,
-)
-
-_TensorT = TypeVar("_TensorT")
-ConcatInput = Union[_TensorT, List[torch.Tensor]]
+from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
 
 
 def validate_concat(parts: List[torch.Tensor]) -> None:
@@ -25,28 +18,20 @@ def validate_concat(parts: List[torch.Tensor]) -> None:
     if not parts:
         raise ValueError("Attempted to concatenate 0 tensors")
     if any(
-        not isinstance(part, torch.Tensor)
+        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(value: ConcatInput[torch.Tensor]) -> Tuple[int, ...]:
+def concat_shape(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> Tuple[int, ...]:
     """Shape after row concatenation."""
-    if isinstance(value, list):
-        return (sum(t.shape[0] for t in value), *value[0].shape[1:])
-    return value.shape
-
-
-def concat_requires_grad(value: ConcatInput[torch.Tensor]) -> bool:
-    """Whether any input requires gradients."""
-    if isinstance(value, list):
-        return any(t.requires_grad for t in value)
-    return value.requires_grad
+    return (sum(t.shape[0] for t in parts), *parts[0].shape[1:])
 
 
-def concat_to_spec(parts: List[torch.Tensor]) -> TensorSpec:
+def concat_to_spec(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> TensorSpec:
     """Describe the concatenation without accessing its storage."""
     validate_concat(parts)
     first = parts[0]
@@ -57,14 +42,12 @@ def concat_to_spec(parts: List[torch.Tensor]) -> TensorSpec:
         shape=concat_shape(parts),
         dtype=dtype,
         device=first.device,
-        requires_grad=concat_requires_grad(parts),
+        requires_grad=any(t.requires_grad for t in parts),
     )
 
 
-def materialize_concat(value: ConcatInput[torch.Tensor]) -> torch.Tensor:
+def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor:
     """Build a checked view, or copy parts whose storage is not adjacent."""
-    if not isinstance(value, list):
-        return value
     validate_concat(value)
     if len(value) == 1:
         return value[0]
@@ -91,22 +74,3 @@ def materialize_concat(value: ConcatInput[torch.Tensor]) -> torch.Tensor:
         return torch.cat(value)
     shape = concat_shape(value)
     return first.as_strided(shape, make_contiguous_strides_for(shape))
-
-
-def restore_from_func_ctx(ctx):
-    """Restore original split parameters saved by the custom-op framework."""
-    tensors = restore_tensor_ctx(ctx)
-    lengths = getattr(ctx, "concatenated_saved_lengths", None)
-    if lengths is None:
-        return tensors
-    restored = []
-    offset = 0
-    for length in lengths:
-        if length is None:
-            restored.append(tensors[offset])
-            offset += 1
-        else:
-            restored.append(list(tensors[offset : offset + length]))
-            offset += length
-    ctx.concatenated_saved_lengths = None
-    return restored
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 49bc6aed8d5..d380305c2e7 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -37,10 +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.
-    Unions including ``List[Tensor]`` also accept deferred row concatenations:
-    all parts cross in the buffer list, with full gradients split outside the op.
-    Saved parts use ``concat.restore_from_func_ctx`` in the backward
-    container; reconstruction of the full view happens only in the real kernel.
+  * ``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.
@@ -110,7 +107,6 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .concat import concat_to_spec, materialize_concat, validate_concat
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -401,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[]`` 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,7 +410,6 @@ class _TensorOrQuantizedKind(Enum):
     NONE = "none"
     TENSOR = "tensor"
     STORAGE = "storage"
-    CONCATENATED = "concatenated"
 
 
 _TQ_KIND_KEY = "__kind__"
@@ -442,7 +438,6 @@ class _FieldPlan:
     name: str
     kind: _FieldKind
     slots: Tuple[_SlotSpec, ...]
-    allows_concatenation: bool = False
 
 
 def _is_tensor_list_annot(annot: Any) -> bool:
@@ -455,8 +450,7 @@ def _is_tensor_storage_union(annot: Any) -> bool:
     if not _is_union(annot):
         return False
     members = frozenset(a for a in get_args(annot) if a is not type(None))
-    tensors = frozenset(a for a in members if not _is_tensor_list_annot(a))
-    return tensors == _TQ_MEMBERS or (tensors == frozenset((torch.Tensor,)) and tensors != members)
+    return members == _TQ_MEMBERS
 
 
 def _is_process_group_annot(annot: Any) -> bool:
@@ -498,12 +492,9 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__tensors", "Tensor[]"),
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
-        return _FieldPlan(
-            name,
-            _FieldKind.TENSOR_OR_QUANTIZED,
-            slots,
-            any(_is_tensor_list_annot(a) for a in get_args(annot)),
-        )
+        return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots)
+    if _is_tensor_list_annot(annot):
+        return _FieldPlan(name, _FieldKind.TENSOR_LIST, (_SlotSpec(name, "Tensor[]"),))
     stripped, is_optional = _strip_optional(annot)
     if stripped is torch.Tensor:
         slot = _SlotSpec(name, "Tensor?" if is_optional else "Tensor")
@@ -544,13 +535,6 @@ def _pack_tensor_or_quantized(field: _FieldPlan, value: Any, slots: Dict[str, An
         slots[tensor_slot] = None
         slots[inner_slot] = []
         slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.NONE})
-    elif isinstance(value, list):
-        if not field.allows_concatenation:
-            raise TypeError(f"field {field.name!r} does not accept a tensor list")
-        validate_concat(value)
-        slots[tensor_slot] = None
-        slots[inner_slot] = value
-        slots[meta_slot] = OpaqueValueBundle({_TQ_KIND_KEY: _TensorOrQuantizedKind.CONCATENATED})
     elif isinstance(value, torch.Tensor):
         # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
         # ``Tensor?`` slot; subclass flattening (if any) is done by the
@@ -579,8 +563,6 @@ def _unpack_tensor_or_quantized(field: _FieldPlan, slots: Dict[str, Any]) -> Any
         return None
     if kind == _TensorOrQuantizedKind.TENSOR:
         return slots[tensor_slot]
-    if kind == _TensorOrQuantizedKind.CONCATENATED:
-        return slots[inner_slot]
     return _storage_unflatten(meta, slots[inner_slot])
 
 
@@ -608,7 +590,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]:
@@ -634,7 +616,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.
@@ -657,7 +643,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)
@@ -691,7 +677,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)
@@ -768,7 +754,7 @@ def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any:
     for name in tensor_field_names:
         value = getattr(obj, name, None)
         if isinstance(value, list):
-            overrides[name] = concat_to_spec(value)
+            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:
@@ -1013,14 +999,8 @@ def _register_base_op(
     ``pack_result``.
     """
 
-    concatenated_fields = tuple(f.name for f in plan.fields if f.allows_concatenation)
-
     def _impl(*flat: Any) -> List[torch.Tensor]:
         obj = plan.unpack(dict(zip(plan.slot_names, flat)))
-        for name in concatenated_fields:
-            value = getattr(obj, name)
-            if isinstance(value, list):
-                setattr(obj, name, materialize_concat(value))
         return pack_result(impl(obj))
 
     def _fake(*flat: Any) -> List[torch.Tensor]:
@@ -1074,16 +1054,7 @@ def _setup_context(ctx, inputs, output):
             out_plan.ctx_attrs,
             tuple(saved_list),
         )
-        saved = []
-        ctx.concatenated_saved_lengths = []
-        for value in tensors_to_save_from_setup or ():
-            if isinstance(value, list):
-                ctx.concatenated_saved_lengths.append(len(value))
-                saved.extend(value)
-            else:
-                ctx.concatenated_saved_lengths.append(None)
-                saved.append(value)
-        tensors_to_save, tensor_objects = prepare_for_saving(*saved)
+        tensors_to_save, tensor_objects = prepare_for_saving(*(tensors_to_save_from_setup or ()))
         ctx.tensor_objects = tensor_objects
         ctx.save_for_backward(*tensors_to_save)
         ctx.backward_objects = bwd_obj
@@ -1091,13 +1062,6 @@ def _setup_context(ctx, inputs, output):
         # Input shapes for the grad slots (SymInt-safe on ctx): a bwd impl may
         # rederive shapes lossily (e.g. rank-1 inputs come back rank-2), so the
         # returned grads are viewed back to the true input shapes below.
-        ctx.concatenated_grad_shapes = {
-            pos: [tensor.shape for tensor in inputs[pos + 1]]
-            for pos in grad_targets
-            if pos in fwd_plan.tensor_or_quantized_offsets()
-            and inputs[pos] is None
-            and inputs[pos + 2][_TQ_KIND_KEY] == _TensorOrQuantizedKind.CONCATENATED
-        }
         ctx.grad_input_shapes = {
             pos: inputs[pos].shape for pos in grad_targets if isinstance(inputs[pos], torch.Tensor)
         }
@@ -1116,6 +1080,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
@@ -1124,10 +1092,11 @@ 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.concatenated_grad_shapes:
-                shapes = ctx.concatenated_grad_shapes[pos]
+            if pos in ctx.fwd_tensor_list_lengths:
                 if g is not None:
-                    out[pos + 1] = list(torch.split(g, [shape[0] for shape in shapes], dim=0))
+                    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)
@@ -1135,7 +1104,6 @@ def _autograd_backward(ctx, *grad_outputs):
                     g = g.view(shape)
             out[pos] = g
         ctx.grad_input_shapes = None
-        ctx.concatenated_grad_shapes = None
         return tuple(out)
 
     fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context)
@@ -1400,9 +1368,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
@@ -1476,7 +1445,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 9b42f39ae94..4189768a308 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -14,7 +14,6 @@
 from ..constants import TE_DType
 from ..distributed import in_fp8_activation_recompute_phase
 from ..dynamo import TensorOrQuantized
-from ..dynamo.concat import ConcatInput
 from ..export import is_in_onnx_export_mode
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
 from ..quantization import FP8GlobalStateManager
@@ -244,10 +243,10 @@ def noop_cat(
 
 def concat_input(
     tensors: Sequence[TensorOrQuantized], *, defer: bool = False
-) -> ConcatInput[TensorOrQuantized]:
-    """Prepare one parameter operand for eager or a compiled consumer."""
+) -> Tuple[Optional[TensorOrQuantized], List[torch.Tensor]]:
+    """Return a parameter operand or its deferred parts."""
     if len(tensors) == 1:
-        return tensors[0]
+        return tensors[0], []
     if (
         defer
         and tensors
@@ -257,10 +256,10 @@ def concat_input(
             for t in tensors
         )
     ):
-        return list(tensors)
+        return None, list(tensors)
     if torch.compiler.is_compiling():
-        return torch.cat(tensors)
-    return noop_cat(tensors)
+        return torch.cat(tensors), []
+    return noop_cat(tensors), []
 
 
 @dataclasses.dataclass
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index 975c9fcddc7..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
 
@@ -46,7 +46,6 @@
     _fsdp_gather_tensors,
 )
 from ..constants import dist_group_type
-from ..dynamo.concat import ConcatInput, concat_shape
 from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
@@ -1323,7 +1322,9 @@ def _get_weight_quantizers(self) -> List[Quantizer]:
     def _enable_weight_preswizzle(
         self,
         quantizer: Quantizer,
-        weight: ConcatInput[torch.Tensor],
+        weight: Optional[torch.Tensor],
+        *,
+        weight_shape: Optional[Sequence[int]] = None,
     ) -> bool:
         """Whether to fuse scale-factor swizzling into weight quantization.
 
@@ -1340,7 +1341,7 @@ def _enable_weight_preswizzle(
         if isinstance(quantizer, MXFP8Quantizer):
             return True
         if isinstance(quantizer, NVFP4Quantizer):
-            shape = concat_shape(weight)
+            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:
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 4be269b962d..9fdc55fb5ee 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -4,7 +4,7 @@
 
 """Linear API"""
 
-from dataclasses import dataclass, replace as dataclass_replace
+from dataclasses import dataclass, field, replace as dataclass_replace
 from typing import Any, Callable, Dict, Optional, Tuple, Union, List
 import warnings
 import weakref
@@ -82,13 +82,9 @@
     QuantizedTensorStorage,
     Quantizer,
     prepare_for_saving,
-)
-from ..dynamo.concat import (
-    ConcatInput,
-    concat_shape,
-    concat_requires_grad,
     restore_from_func_ctx,
 )
+from ..dynamo.concat import concat_shape, concat_to_spec, materialize_concat
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
@@ -115,9 +111,9 @@ class LinearFwdArgs:
     """Single-argument bag for the forward path of :class:`_Linear`."""
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
-    weight: ConcatInput[TensorOrQuantized]
+    weight: Optional[TensorOrQuantized]
     inp: torch.Tensor
-    bias: Optional[ConcatInput[torch.Tensor]]
+    bias: Optional[torch.Tensor]
 
     # --- Non-differentiable cached tensors ---
     # TensorOrQuantized so a cached quantized workspace can cross the op boundary.
@@ -183,6 +179,8 @@ class LinearFwdArgs:
     # --- Misc ---
     cpu_offloading: bool
     is_grad_enabled: bool
+    weight_list: List[torch.Tensor] = field(default_factory=list)
+    bias_list: List[torch.Tensor] = field(default_factory=list)
 
     def compile_unsupported_reason(self) -> Optional[str]:
         """Reason this config can't use the torch.compile custom-op path (else None)."""
@@ -243,9 +241,13 @@ class LinearBwdArgs:
     # --- Saved / restored tensors (populated at backward entry) ---
     grad_output: Optional[torch.Tensor] = None
     inputmat: Optional[TensorOrQuantized] = None
-    weight_fp8: Optional[ConcatInput[TensorOrQuantized]] = None
-    saved_weight: Optional[ConcatInput[TensorOrQuantized]] = None
-    bias: Optional[ConcatInput[torch.Tensor]] = None
+    weight_fp8: Optional[TensorOrQuantized] = None
+    saved_weight: Optional[TensorOrQuantized] = None
+    bias: Optional[torch.Tensor] = None
+    weight_list: List[torch.Tensor] = field(default_factory=list)
+    bias_list: List[torch.Tensor] = field(default_factory=list)
+    parameter_list_lengths: Tuple[int, int] = (0, 0)
+    weight_fp8_from_list: bool = False
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -311,14 +313,28 @@ 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)
+        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])
+        self.bias_list = list(saved[4 + weight_count : 4 + weight_count + bias_count])
+
+    def postprocess_gradients(self, grads):
+        """Split gradients outside the op so the parts can share storage."""
+        wgrad, dgrad, bgrad = grads
+        weight_grads = [None] * len(self.weight_list)
+        bias_grads = [None] * len(self.bias_list)
+        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:
@@ -366,10 +382,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
@@ -779,9 +795,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
@@ -1023,7 +1039,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)
@@ -1106,7 +1122,16 @@ def _linear_setup_ctx(
         saved_weight = weight
     if bias_alias == "bias":
         saved_bias = bias
-    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)
+    bwd_args.parameter_list_lengths = (len(fwd_args.weight_list), len(fwd_args.bias_list))
+    return (
+        saved_inputmat,
+        wt_save,
+        saved_weight,
+        saved_bias,
+        *fwd_args.weight_list,
+        *fwd_args.bias_list,
+    )
 
 
 def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]:
@@ -1119,10 +1144,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
@@ -1740,7 +1765,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
 
@@ -1795,7 +1820,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,
@@ -2406,10 +2431,13 @@ def forward(
 
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
-            weight, bias = self._get_weight_and_bias_tensors(
+            weight, bias, weight_list, bias_list = self._get_weight_and_bias_tensors(
                 defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
             )
-            weight_requires_grad = concat_requires_grad(weight)
+            weight_shape = concat_shape(weight_list) if weight_list else weight.shape
+            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)
@@ -2431,7 +2459,7 @@ def forward(
             ) = quantizers
             if weight_quantizer is not None and not debug:
                 weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
-                    weight_quantizer, weight
+                    weight_quantizer, weight, weight_shape=weight_shape
                 )
 
             use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
@@ -2487,12 +2515,16 @@ def forward(
                 weight=weight,
                 inp=inp,
                 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_requires_grad,
                 bias_requires_grad=(
-                    concat_requires_grad(linear_bias) if linear_bias 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,
@@ -2552,15 +2584,17 @@ def forward(
                         msg=f"te.Linear falling back to eager: {fallback_reason}"
                     )
                     use_compiled_op = False
-                    weight, bias = self._get_weight_and_bias_tensors()
+                    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 = []
+                    fwd_args.bias_list = []
 
             if use_compiled_op:
-                check_gemm_dims(inp.shape, concat_shape(weight), self.fp8)
+                check_gemm_dims(inp.shape, weight_shape, self.fp8)
                 out, new_weight_workspace = _linear_op(fwd_args)
             else:
                 out, new_weight_workspace = _linear_eager(
@@ -2708,15 +2742,16 @@ def _forward_eager_fallback(
 
     def _get_weight_and_bias_tensors(self, defer_concatenation=False):
         weights = self._get_weight_tensors()
-        weight = concat_input(weights, defer=defer_concatenation)
+        weight, weight_list = concat_input(weights, defer=defer_concatenation)
         bias = None
+        bias_list = []
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
-            bias = concat_input(
+            bias, bias_list = concat_input(
                 biases,
                 defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
             )
-        return weight, bias
+        return weight, bias, weight_list, bias_list
 
     def onnx_forward(
         self,
@@ -2733,7 +2768,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,

From bcf13f6519178aaf728267e890f2e8ad79ff70ee Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 16 Oct 2026 15:54:35 +0200
Subject: [PATCH 10/13] [PyTorch] Skip empty split-parameter bookkeeping in
 eager Linear

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/module/linear.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 9fdc55fb5ee..00a379b6391 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -314,6 +314,9 @@ class LinearBwdArgs:
     def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None:
         """Pull saved tensors from ``ctx`` into the fields backward consumes."""
         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])
@@ -1122,6 +1125,8 @@ def _linear_setup_ctx(
         saved_weight = weight
     if bias_alias == "bias":
         saved_bias = 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)
     bwd_args.parameter_list_lengths = (len(fwd_args.weight_list), len(fwd_args.bias_list))
     return (
@@ -2419,7 +2424,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
             )
@@ -2432,9 +2439,9 @@ def forward(
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
             weight, bias, weight_list, bias_list = self._get_weight_and_bias_tensors(
-                defer_concatenation=torch.compiler.is_compiling() and _linear_op is not None
+                defer_concatenation=use_compiled_op
             )
-            weight_shape = concat_shape(weight_list) if weight_list else weight.shape
+            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
             )
@@ -2462,8 +2469,7 @@ def forward(
                     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.
@@ -2594,7 +2600,9 @@ def forward(
                     fwd_args.bias_list = []
 
             if use_compiled_op:
-                check_gemm_dims(inp.shape, weight_shape, 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(

From 623baf85f4e2f95c1ff0234fc160b10377959788 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 16 Oct 2026 16:23:38 +0200
Subject: [PATCH 11/13] [PyTorch] Make deferred Linear parameter lists optional

Use None when Linear has no deferred parameter parts, including eager
forward, saved backward arguments and compile fallbacks. This avoids
allocating empty weight and bias lists on the ordinary eager path.

Transport Optional[List[Tensor]] through Tensor[]? custom-op inputs,
preserving None without allocating replacement lists. Keep split-gradient
handling outside the backward op and let the adapter supply absent list
gradients when needed.

Cover ordinary and optional lists, empty lists and None in the adapter
test. Validate the full compile suite (384 passed, 62 skipped, one existing
XPASS), 55 focused cases, 30 numerical Linear cases, pre-commit and Python
L0 lint.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 37 ++++++++++++-------
 .../pytorch/dynamo/custom_op.py               |  7 ++--
 transformer_engine/pytorch/module/_common.py  |  8 ++--
 transformer_engine/pytorch/module/linear.py   | 34 +++++++++--------
 4 files changed, 51 insertions(+), 35 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index a0261854e9d..06ddb4153ee 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3271,7 +3271,7 @@ def test_concat_input_single_tensor(defer):
     tensor = torch.nn.Parameter(torch.randn(4, 8, device="cuda"))
     result, parts = concat_input([tensor], defer=defer)
     assert result is tensor
-    assert not parts
+    assert parts is None
 
 
 @pytest.mark.parametrize("disjoint", [False, True])
@@ -3282,7 +3282,7 @@ def test_concat_input_eager_gradients(disjoint):
     if disjoint:
         parts[1] = torch.nn.Parameter(parts[1].detach().clone())
     result, deferred = concat_input(parts)
-    assert not deferred
+    assert deferred is None
     expected = torch.cat(parts)
     grad = torch.randn_like(result)
     torch.testing.assert_close(result, expected, rtol=0, atol=0)
@@ -3352,29 +3352,40 @@ def test_concat_spec(fake, mixed_dtype):
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-def test_tensor_list_custom_op():
+@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:
-        tensors: list[torch.Tensor]
+        inp: torch.Tensor
+        tensors: list[torch.Tensor] | None if optional else list[torch.Tensor]
 
     op = register_custom_op(
-        op_name="test_tensor_list_inputs",
+        op_name=f"test_tensor_list_inputs_{optional}_{contents}",
         arg_type=Args,
-        impl=lambda args: [tensor.square() for tensor in args.tensors],
-        fake_impl=lambda args: args.tensors,
+        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(tensors):
-        return op(Args(tensors))
+    def fn(inp, tensors):
+        return op(Args(inp, tensors))
 
     tensors = [torch.randn(3), torch.randn(2, 4, dtype=torch.float64)]
-    expected = [tensor.square() for tensor in tensors]
-    torch.testing.assert_close(fn(tensors), expected)
+    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(tensors), expected)
+    torch.testing.assert_close(compiled(inp, tensors), expected)
 
 
 @pytest.mark.parametrize("internal", [False, True])
@@ -3392,7 +3403,7 @@ def test_concat_rejects_quantized_parts(internal):
             concatenate([tensor, tensor])
     result, parts = concat_input([tensor], defer=True)
     assert result is tensor
-    assert not parts
+    assert parts is None
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index d380305c2e7..a3e555f5a62 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -397,7 +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[]`` 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
@@ -493,9 +493,10 @@ def _parse_field(name: str, annot: Any) -> _FieldPlan:
             _SlotSpec(name + "__meta", _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
         )
         return _FieldPlan(name, _FieldKind.TENSOR_OR_QUANTIZED, slots)
-    if _is_tensor_list_annot(annot):
-        return _FieldPlan(name, _FieldKind.TENSOR_LIST, (_SlotSpec(name, "Tensor[]"),))
     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,))
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 4189768a308..62572a3a6f6 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -243,10 +243,10 @@ def noop_cat(
 
 def concat_input(
     tensors: Sequence[TensorOrQuantized], *, defer: bool = False
-) -> Tuple[Optional[TensorOrQuantized], List[torch.Tensor]]:
+) -> Tuple[Optional[TensorOrQuantized], Optional[List[torch.Tensor]]]:
     """Return a parameter operand or its deferred parts."""
     if len(tensors) == 1:
-        return tensors[0], []
+        return tensors[0], None
     if (
         defer
         and tensors
@@ -258,8 +258,8 @@ def concat_input(
     ):
         return None, list(tensors)
     if torch.compiler.is_compiling():
-        return torch.cat(tensors), []
-    return noop_cat(tensors), []
+        return torch.cat(tensors), None
+    return noop_cat(tensors), None
 
 
 @dataclasses.dataclass
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 00a379b6391..93cc83525e2 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -4,7 +4,7 @@
 
 """Linear API"""
 
-from dataclasses import dataclass, field, replace as dataclass_replace
+from dataclasses import dataclass, replace as dataclass_replace
 from typing import Any, Callable, Dict, Optional, Tuple, Union, List
 import warnings
 import weakref
@@ -179,8 +179,8 @@ class LinearFwdArgs:
     # --- Misc ---
     cpu_offloading: bool
     is_grad_enabled: bool
-    weight_list: List[torch.Tensor] = field(default_factory=list)
-    bias_list: List[torch.Tensor] = field(default_factory=list)
+    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)."""
@@ -244,8 +244,8 @@ class LinearBwdArgs:
     weight_fp8: Optional[TensorOrQuantized] = None
     saved_weight: Optional[TensorOrQuantized] = None
     bias: Optional[torch.Tensor] = None
-    weight_list: List[torch.Tensor] = field(default_factory=list)
-    bias_list: List[torch.Tensor] = field(default_factory=list)
+    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
 
@@ -319,14 +319,16 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None:
             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])
-        self.bias_list = list(saved[4 + weight_count : 4 + weight_count + bias_count])
+        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] * len(self.weight_list)
-        bias_grads = [None] * len(self.bias_list)
+        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:
@@ -1128,14 +1130,16 @@ def _linear_setup_ctx(
     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)
-    bwd_args.parameter_list_lengths = (len(fwd_args.weight_list), len(fwd_args.bias_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,
-        *fwd_args.weight_list,
-        *fwd_args.bias_list,
+        *weight_list,
+        *bias_list,
     )
 
 
@@ -2596,8 +2600,8 @@ def forward(
                     )
                     fwd_args.weight = weight
                     fwd_args.bias = linear_bias
-                    fwd_args.weight_list = []
-                    fwd_args.bias_list = []
+                    fwd_args.weight_list = None
+                    fwd_args.bias_list = None
 
             if use_compiled_op:
                 check_gemm_dims(
@@ -2752,7 +2756,7 @@ 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 = []
+        bias_list = None
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
             bias, bias_list = concat_input(

From 7b280931785f122791ead2a93212b88faa846500 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 16 Oct 2026 17:07:58 +0200
Subject: [PATCH 12/13] [PyTorch] Trim split parameter tests

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 151 ----------------------------
 1 file changed, 151 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 06ddb4153ee..6d17f89d47e 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3229,71 +3229,6 @@ def test_te_split_parameters_saved_versions():
         out.sum().backward()
 
 
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-def test_te_split_parameters_saved_hooks():
-    import copy
-
-    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
-    reference = copy.deepcopy(model)
-    compiled = torch.compile(model, fullgraph=True, backend="aot_eager")
-    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    ref_inp = inp.detach().clone().requires_grad_()
-    with torch.autograd.graph.saved_tensors_hooks(lambda t: t.clone(), lambda t: t):
-        out = compiled(inp)
-        out.sum().backward()
-    expected = reference(ref_inp)
-    expected.sum().backward()
-    torch.testing.assert_close(out, expected, rtol=0, atol=0)
-    torch.testing.assert_close(inp.grad, ref_inp.grad, rtol=0, atol=0)
-    for actual, target in zip(model.parameters(), reference.parameters()):
-        torch.testing.assert_close(actual.grad, target.grad, rtol=0, atol=0)
-
-
-def test_te_split_parameters_fallback_check(monkeypatch):
-    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32))
-    inp = torch.empty(16, 64, device="cuda")
-
-    def unexpected_materialization(*args, **kwargs):
-        raise AssertionError("Fallback eligibility must not concatenate parameters")
-
-    monkeypatch.setattr(model, "_get_weight_and_bias_tensors", unexpected_materialization)
-    assert model._compile_eager_fallback_reason(inp, None, False, False, True, False) is None
-    assert (
-        model._compile_eager_fallback_reason(inp, None, True, False, True, False)
-        == "differentiable fp8_output=True"
-    )
-
-
-@pytest.mark.parametrize("defer", [False, True])
-def test_concat_input_single_tensor(defer):
-    from transformer_engine.pytorch.module._common import concat_input
-
-    tensor = torch.nn.Parameter(torch.randn(4, 8, device="cuda"))
-    result, parts = concat_input([tensor], defer=defer)
-    assert result is tensor
-    assert parts is None
-
-
-@pytest.mark.parametrize("disjoint", [False, True])
-def test_concat_input_eager_gradients(disjoint):
-    from transformer_engine.pytorch.module._common import concat_input
-
-    parts = [torch.nn.Parameter(t) for t in torch.randn(8, 8, device="cuda").split((3, 5))]
-    if disjoint:
-        parts[1] = torch.nn.Parameter(parts[1].detach().clone())
-    result, deferred = concat_input(parts)
-    assert deferred is None
-    expected = torch.cat(parts)
-    grad = torch.randn_like(result)
-    torch.testing.assert_close(result, expected, rtol=0, atol=0)
-    torch.testing.assert_close(
-        torch.autograd.grad(result, parts, grad),
-        torch.autograd.grad(expected, parts, grad),
-        rtol=0,
-        atol=0,
-    )
-
-
 @pytest.mark.parametrize(
     "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
 )
@@ -3326,31 +3261,6 @@ def test_concat_storage(layout):
         assert result.untyped_storage().data_ptr() != parts[0].untyped_storage().data_ptr()
 
 
-@pytest.mark.parametrize("fake", [False, True], ids=["eager", "fake"])
-@pytest.mark.parametrize("mixed_dtype", [False, True])
-def test_concat_spec(fake, mixed_dtype):
-    from transformer_engine.pytorch.dynamo.concat import concat_to_spec
-
-    with FakeTensorMode() if fake else contextlib.nullcontext():
-        parts = [
-            torch.empty(3, 8, device="cuda", dtype=torch.bfloat16),
-            torch.empty(
-                5,
-                8,
-                device="cuda",
-                dtype=torch.float32 if mixed_dtype else torch.bfloat16,
-                requires_grad=True,
-            ),
-        ]
-        expected = torch.cat(parts)
-        spec = concat_to_spec([to_tensor_spec(t) for t in parts] if fake else parts)
-        assert spec.shape == tuple(expected.shape)
-        assert spec.dtype == expected.dtype
-        assert spec.device == expected.device
-        assert spec.requires_grad == expected.requires_grad
-        assert not spec.is_quantized
-
-
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize(
     "optional, contents",
@@ -3388,24 +3298,6 @@ def fn(inp, tensors):
     torch.testing.assert_close(compiled(inp, tensors), expected)
 
 
-@pytest.mark.parametrize("internal", [False, True])
-def test_concat_rejects_quantized_parts(internal):
-    from transformer_engine.pytorch.dynamo.concat import concat_to_spec, materialize_concat
-    from transformer_engine.pytorch.module._common import concat_input
-
-    quantizer = _current_scaling()
-    quantizer.internal = internal
-    tensor = TensorSpec(
-        shape=(4, 8), dtype=torch.bfloat16, quantizer=quantizer, device=torch.device("cpu")
-    ).create_tensor()
-    for concatenate in (concat_to_spec, materialize_concat):
-        with pytest.raises(TypeError, match="non-quantized tensors"):
-            concatenate([tensor, tensor])
-    result, parts = concat_input([tensor], defer=True)
-    assert result is tensor
-    assert parts is None
-
-
 @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):
@@ -3432,46 +3324,3 @@ def test_te_split_parameters_grad_targets(grad_target):
             assert result is None
         else:
             torch.testing.assert_close(result, target, rtol=0, atol=0)
-
-
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("bf16_part", ["q", "k"])
-def test_te_split_parameters_mixed_dtype_autocast(bf16_part):
-    model = te.Linear(64, 192, parameters_split=("q", "k", "v"), params_dtype=torch.float32)
-    name = f"{bf16_part}_weight"
-    setattr(model, name, torch.nn.Parameter(getattr(model, name).to(torch.bfloat16)))
-    inp = torch.randn(16, 64, device="cuda", dtype=torch.float32, requires_grad=True)
-    grad = torch.randn(16, 192, device="cuda", dtype=torch.bfloat16)
-    torch._dynamo.reset()
-    compiled = torch.compile(model, fullgraph=True)
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        expected = model(inp)
-    expected.backward(grad)
-    expected_grads = [tensor.grad.clone() for tensor in (inp, *model.parameters())]
-    inp.grad = None
-    model.zero_grad(set_to_none=True)
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        actual = compiled(inp)
-    actual.backward(grad)
-    actual_grads = [tensor.grad for tensor in (inp, *model.parameters())]
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-    torch.testing.assert_close(actual_grads, expected_grads, rtol=0, atol=0)
-
-
-@pytest.mark.parametrize("split", [False, True])
-@pytest.mark.parametrize("bias_mode", ["fused", "none", "returned"])
-def test_te_split_parameters_onnx_path(split, bias_mode):
-    model = te.Linear(
-        64,
-        128,
-        parameters_split=dict(q=64, k=32, v=32) if split else None,
-        bias=bias_mode != "none",
-        return_bias=bias_mode == "returned",
-        params_dtype=torch.bfloat16,
-    )
-    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        expected = model(inp)
-        with te.onnx_export(enabled=True):
-            actual = model(inp)
-    torch.testing.assert_close(actual, expected, **dtype_tols(torch.bfloat16))

From 40974d38e4ab782c28a2f4aca3abb5cd193be6ab Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 21 Oct 2026 08:18:11 -0700
Subject: [PATCH 13/13] [PyTorch] Materialize split Linear parameters inside
 eager forward

Use the same checked view or concatenation in eager and compiled Linear.
Pass individual parts as explicit autograd inputs and save the materialized
operands for eager backward, which returns views of the full gradients.
Preserve the single-parameter and external-bias paths.

Derive the view strides directly from contiguous parts to avoid symbolic
stride-helper overhead. Extend the existing storage and cat-count tests.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py         | 28 +++++++---
 transformer_engine/pytorch/dynamo/concat.py |  4 +-
 transformer_engine/pytorch/module/linear.py | 60 +++++++++++++++------
 3 files changed, 68 insertions(+), 24 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 6d17f89d47e..f8415ed7d80 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3192,15 +3192,16 @@ def fn(inp):
 
 @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):
+@pytest.mark.parametrize("execution_mode", ["eager", "compile"])
+def test_te_split_parameters_no_cat(training, execution_mode):
     """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)
+    run_model = torch.compile(model, fullgraph=True) if execution_mode == "compile" else model
     inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16, requires_grad=training)
 
     def run():
         with torch.set_grad_enabled(training):
-            out = compiled(inp)
+            out = run_model(inp)
             if training:
                 out.sum().backward()
         model.zero_grad(set_to_none=True)
@@ -3214,7 +3215,8 @@ def run():
     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()}
+    cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
+    assert cats == (2 if training and execution_mode == "compile" else 1)
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
@@ -3230,7 +3232,17 @@ def test_te_split_parameters_saved_versions():
 
 
 @pytest.mark.parametrize(
-    "layout", ["adjacent", "offset", "disjoint", "strided", "singleton", "negative", "conjugate"]
+    "layout",
+    [
+        "adjacent",
+        "offset",
+        "single_row",
+        "disjoint",
+        "strided",
+        "singleton",
+        "negative",
+        "conjugate",
+    ],
 )
 def test_concat_storage(layout):
     from transformer_engine.pytorch.dynamo.concat import materialize_concat
@@ -3248,14 +3260,16 @@ def test_concat_storage(layout):
     elif layout == "conjugate":
         storage = torch.complex(storage.float(), storage.float()).conj()
     parts = list(storage.split([1, storage.shape[0] - 3, 2]))
-    if layout == "disjoint":
+    if layout == "single_row":
+        parts[0] = parts[0].as_strided(parts[0].shape, (0, *parts[0].stride()[1:]))
+    elif 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"):
+    if layout in ("adjacent", "offset", "single_row"):
         assert result.data_ptr() == parts[0].data_ptr()
     else:
         assert result.untyped_storage().data_ptr() != parts[0].untyped_storage().data_ptr()
diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
index e6ac9882613..7c1f5217d91 100644
--- a/transformer_engine/pytorch/dynamo/concat.py
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -4,10 +4,10 @@
 
 """Deferred row concatenation for tensors consumed inside a custom op."""
 
+import math
 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
@@ -73,4 +73,4 @@ def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor:
     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))
+    return first.as_strided(shape, (math.prod(shape[1:]), *first.stride()[1:]))
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 93cc83525e2..cc2c436e00f 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -378,6 +378,8 @@ def _fake_workspace_valid(workspace: TensorSpec, quantizer: Optional[Quantizer])
 
 def _linear_forward_impl(
     args: LinearFwdArgs,
+    *,
+    save_materialized_parameters: bool = False,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], Optional[Dict]]:
     """Forward implementation for the linear layer.
 
@@ -391,6 +393,9 @@ def _linear_forward_impl(
     is_dist_weight = is_distributed_weight(weight)
     inp = args.inp
     bias = materialize_concat(args.bias_list) if args.bias_list else args.bias
+    if save_materialized_parameters and (args.weight_list or args.bias_list):
+        args.weight, args.bias = weight, bias
+        args.weight_list = args.bias_list = None
     input_quantizer = args.input_quantizer
     weight_quantizer = args.weight_quantizer
     output_quantizer = args.output_quantizer
@@ -1848,16 +1853,16 @@ class _Linear(torch.autograd.Function):
     @staticmethod
     def forward(
         ctx,
-        weight: torch.Tensor,
+        weight: Optional[torch.Tensor],
         inp: torch.Tensor,
         bias: Optional[torch.Tensor],
         fwd_args: LinearFwdArgs,
+        *parameter_parts: torch.Tensor,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Forward pass: compute linear output and set up autograd context.
 
-        ``weight``, ``inp`` and ``bias`` are positional Tensor arguments so
-        autograd tracks them; they are immediately re-attached to ``fwd_args``
-        so every downstream helper can be invoked with a single argument.
+        Tensor operands, including split parameter parts, are explicit
+        autograd inputs and are re-attached to ``fwd_args``.
 
         ``weight_workspace`` is intentionally NOT a positional input: it is a
         non-differentiable cached tensor passed in via
@@ -1867,12 +1872,20 @@ def forward(
         fwd_args.weight = weight
         fwd_args.inp = inp
         fwd_args.bias = bias
+        if parameter_parts:
+            weight_count = len(fwd_args.weight_list or ())
+            fwd_args.weight_list = list(parameter_parts[:weight_count]) if weight_count else None
+            fwd_args.bias_list = list(parameter_parts[weight_count:]) or None
+            ctx.parameter_split_sizes = (
+                tuple(t.shape[0] for t in fwd_args.weight_list or ()),
+                tuple(t.shape[0] for t in fwd_args.bias_list or ()),
+            )
         (
             out,
             new_weight_workspace,
             tensors_to_save_from_forward,
             ctx_attrs,
-        ) = _linear_forward_impl(fwd_args)
+        ) = _linear_forward_impl(fwd_args, save_materialized_parameters=True)
         if ctx is not None:
             bwd_args = LinearBwdArgs()
             tensors_to_save_from_setup = _linear_setup_ctx(
@@ -1911,6 +1924,22 @@ def backward(
         if bwd_args.ub_name is not None:
             nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}"
         result = _linear_backward_impl(bwd_args) + (None,)  # fwd_args grad slot
+        if len(ctx.needs_input_grad) > 4:
+            wgrad, dgrad, bgrad, _ = result
+            weight_sizes, bias_sizes = ctx.parameter_split_sizes
+            part_grads = []
+            for grad, sizes in ((wgrad, weight_sizes), (bgrad, bias_sizes)):
+                if sizes:
+                    part_grads.extend(
+                        grad.split(sizes) if grad is not None else (None,) * len(sizes)
+                    )
+            result = (
+                None if weight_sizes else wgrad,
+                dgrad,
+                None if bias_sizes else bgrad,
+                None,
+                *part_grads,
+            )
         reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors
         # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs,
         # main_grad closure) so they don't outlive backward via ctx under retain_graph.
@@ -1925,7 +1954,7 @@ def backward(
 
 @no_torch_dynamo()
 def _linear_eager(
-    weight_tensor: torch.Tensor,
+    weight_tensor: Optional[torch.Tensor],
     inp: torch.Tensor,
     bias: Optional[torch.Tensor],
     fwd_args: LinearFwdArgs,
@@ -1933,6 +1962,15 @@ def _linear_eager(
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     """Run ``_Linear`` eagerly, bypassing Dynamo."""
     if is_grad_enabled:
+        if fwd_args.weight_list or fwd_args.bias_list:
+            return _Linear.apply(
+                weight_tensor,
+                inp,
+                bias,
+                fwd_args,
+                *(fwd_args.weight_list or ()),
+                *(fwd_args.bias_list or ()),
+            )
         return _Linear.apply(weight_tensor, inp, bias, fwd_args)
     return _Linear.forward(None, weight_tensor, inp, bias, fwd_args)
 
@@ -2443,7 +2481,7 @@ def forward(
         inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor))
         try:
             weight, bias, weight_list, bias_list = self._get_weight_and_bias_tensors(
-                defer_concatenation=use_compiled_op
+                defer_concatenation=True
             )
             weight_shape = concat_shape(weight_list) if weight_list else None
             weight_requires_grad = (
@@ -2594,14 +2632,6 @@ 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(