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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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/26] [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(

From 303b09344d6fe3f1733c689343905cc05895b0bc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 15:20:12 +0100
Subject: [PATCH 14/26] [PyTorch] Materialize split LayerNormLinear parameters
 in eager forward

Share Linear's checked view/cat helper for split weights and fused biases.
Pass the parts as explicit autograd inputs, save full operands for backward
and return views of the full parameter gradients.

Keep backward computation and saved-tensor restoration unchanged. Compute
split shapes only when configuring weight quantization, and extend the
existing cat-count regression to eager LayerNormLinear.

Remove standalone storage-helper and tensor-list adapter tests at the
user's request.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  90 +--------------
 .../pytorch/module/layernorm_linear.py        | 103 +++++++++++++-----
 2 files changed, 81 insertions(+), 112 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 6ecd3ba55ed..2495383105c 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3366,10 +3366,13 @@ def fn(inp):
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize("training", [False, True])
-@pytest.mark.parametrize("execution_mode", ["eager", "compile"])
-def test_te_split_parameters_no_cat(training, execution_mode):
+@pytest.mark.parametrize(
+    "module, execution_mode",
+    [(te.Linear, "eager"), (te.Linear, "compile"), (te.LayerNormLinear, "eager")],
+)
+def test_te_split_parameters_no_cat(training, module, 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)
+    model = module(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
     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)
 
@@ -3405,87 +3408,6 @@ def test_te_split_parameters_saved_versions():
         out.sum().backward()
 
 
-@pytest.mark.parametrize(
-    "layout",
-    [
-        "adjacent",
-        "offset",
-        "single_row",
-        "disjoint",
-        "strided",
-        "singleton",
-        "negative",
-        "conjugate",
-    ],
-)
-def test_concat_storage(layout):
-    from transformer_engine.pytorch.dynamo.concat import materialize_concat
-    from transformer_engine.pytorch.module._common import concat_input
-
-    storage = torch.arange(128, device="cuda").view(16, 8)
-    if layout == "offset":
-        storage = storage[2:10]
-    elif layout == "strided":
-        storage = storage[:, ::2]
-    elif layout == "singleton":
-        storage = storage[:, :1]
-    elif layout == "negative":
-        storage = torch._neg_view(storage)
-    elif layout == "conjugate":
-        storage = torch.complex(storage.float(), storage.float()).conj()
-    parts = list(storage.split([1, storage.shape[0] - 3, 2]))
-    if layout == "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", "single_row"):
-        assert result.data_ptr() == parts[0].data_ptr()
-    else:
-        assert result.untyped_storage().data_ptr() != parts[0].untyped_storage().data_ptr()
-
-
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize(
-    "optional, contents",
-    [(False, "full"), (False, "empty"), (True, "full"), (True, "empty"), (True, "none")],
-)
-def test_tensor_list_custom_op(optional, contents):
-    from transformer_engine.pytorch.dynamo import register_custom_op
-
-    @dataclasses.dataclass
-    class Args:
-        inp: torch.Tensor
-        tensors: list[torch.Tensor] | None if optional else list[torch.Tensor]
-
-    op = register_custom_op(
-        op_name=f"test_tensor_list_inputs_{optional}_{contents}",
-        arg_type=Args,
-        impl=lambda args: [args.inp + (args.tensors is not None)]
-        + [tensor.square() for tensor in (args.tensors or ())],
-        fake_impl=lambda args: [args.inp, *(args.tensors or ())],
-    )
-    assert op is not None
-
-    def fn(inp, tensors):
-        return op(Args(inp, tensors))
-
-    tensors = [torch.randn(3), torch.randn(2, 4, dtype=torch.float64)]
-    if contents == "empty":
-        tensors = []
-    elif contents == "none":
-        tensors = None
-    inp = torch.randn(2)
-    expected = [inp + (contents != "none"), *[tensor.square() for tensor in (tensors or ())]]
-    torch.testing.assert_close(fn(inp, tensors), expected)
-    compiled = torch.compile(fn, fullgraph=True, backend="aot_eager")
-    torch.testing.assert_close(compiled(inp, tensors), expected)
-
-
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize("grad_target", ["input", "weight"])
 def test_te_split_parameters_grad_targets(grad_target):
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 03e6cf50de8..cf8c1d19a37 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -67,7 +67,7 @@
 from ._common import (
     apply_normalization,
     check_fp8_reduce_and_update,
-    noop_cat,
+    concat_input,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -79,6 +79,7 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
+from ..dynamo.concat import concat_shape, materialize_concat
 from ..dynamo import TensorOrQuantized
 from ...debug.pytorch.debug_state import TEDebugState
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
@@ -109,7 +110,7 @@ class LayerNormLinearFwdArgs:
     inp: torch.Tensor
     ln_weight: torch.Tensor
     ln_bias: Optional[torch.Tensor]
-    weight: TensorOrQuantized
+    weight: Optional[TensorOrQuantized]
     bias: Optional[torch.Tensor]
 
     # --- Non-differentiable cached tensors ---
@@ -181,6 +182,8 @@ class LayerNormLinearFwdArgs:
     # --- Misc ---
     cpu_offloading: bool
     is_grad_enabled: bool
+    weight_list: Optional[List[torch.Tensor]] = None
+    bias_list: Optional[List[torch.Tensor]] = None
 
     def any_requires_grad(self) -> bool:
         """Whether any differentiable input requires a gradient."""
@@ -315,8 +318,11 @@ def _layernorm_linear_forward_impl(
     inp = args.inp
     ln_weight = args.ln_weight
     ln_bias = args.ln_bias
-    weight = args.weight
-    bias = args.bias
+    weight = materialize_concat(args.weight_list) if args.weight_list else args.weight
+    bias = materialize_concat(args.bias_list) if args.bias_list else args.bias
+    if 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
@@ -1477,9 +1483,10 @@ def forward(
         inp: torch.Tensor,
         ln_weight: torch.Tensor,
         ln_bias: Optional[torch.Tensor],
-        weight: torch.Tensor,
+        weight: Optional[torch.Tensor],
         bias: Optional[torch.Tensor],
         fwd_args: LayerNormLinearFwdArgs,
+        *parameter_parts: torch.Tensor,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
         """Forward pass: compute the output and set up the autograd context.
 
@@ -1494,6 +1501,14 @@ def forward(
         fwd_args.ln_bias = ln_bias
         fwd_args.weight = weight
         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,
             ln_out_return,
@@ -1546,14 +1561,24 @@ def backward(
             nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
             FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
             nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
-        return (
-            dgrad,
-            dgamma,
-            dbeta,
-            wgrad,
-            grad_bias,
-            None,  # fwd_args
-        )
+        if len(ctx.needs_input_grad) > 6:
+            weight_sizes, bias_sizes = ctx.parameter_split_sizes
+            part_grads = []
+            for grad, sizes in ((wgrad, weight_sizes), (grad_bias, bias_sizes)):
+                if sizes:
+                    part_grads.extend(
+                        grad.split(sizes) if grad is not None else (None,) * len(sizes)
+                    )
+            return (
+                dgrad,
+                dgamma,
+                dbeta,
+                None if weight_sizes else wgrad,
+                None if bias_sizes else grad_bias,
+                None,
+                *part_grads,
+            )
+        return dgrad, dgamma, dbeta, wgrad, grad_bias, None
 
 
 class LayerNormLinear(TransformerEngineBaseModule):
@@ -2081,7 +2106,9 @@ def forward(
 
         try:
             # Get concatenated weight and bias tensors
-            weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+            weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors(
+                defer_concatenation=True
+            )
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2103,7 +2130,9 @@ 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_tensor,
+                    weight_shape=concat_shape(weight_list) if weight_list else None,
                 )
 
             cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight"
@@ -2148,6 +2177,8 @@ def forward(
                 ln_bias=self.layer_norm_bias,
                 weight=weight_tensor,
                 bias=linear_bias_tensor,
+                weight_list=weight_list,
+                bias_list=bias_list,
                 weight_workspace=weight_workspace,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
@@ -2157,9 +2188,15 @@ def forward(
                     if self.layer_norm_bias is not None
                     else False
                 ),
-                weight_requires_grad=weight_tensor.requires_grad,
+                weight_requires_grad=(
+                    any(t.requires_grad for t in weight_list)
+                    if weight_list
+                    else weight_tensor.requires_grad
+                ),
                 bias_requires_grad=(
-                    linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False
+                    any(t.requires_grad for t in bias_list)
+                    if bias_list
+                    else linear_bias_tensor is not None and linear_bias_tensor.requires_grad
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2216,7 +2253,18 @@ def forward(
                 is_grad_enabled=is_grad_enabled,
             )
 
-            if is_grad_enabled:
+            if is_grad_enabled and (weight_list or bias_list):
+                out, ln_out, new_weight_workspace = _LayerNormLinear.apply(
+                    inp,
+                    self.layer_norm_weight,
+                    self.layer_norm_bias,
+                    weight_tensor,
+                    linear_bias_tensor,
+                    fwd_args,
+                    *(weight_list or ()),
+                    *(bias_list or ()),
+                )
+            elif is_grad_enabled:
                 out, ln_out, new_weight_workspace = _LayerNormLinear.apply(
                     inp,
                     self.layer_norm_weight,
@@ -2300,16 +2348,15 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
             for name, q in zip(names, original_quantizers)
         )
 
-    def _get_weight_and_bias_tensors(self):
-        # 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):
+        weight, weight_list = concat_input(self._get_weight_tensors(), defer=defer_concatenation)
+        bias, bias_list = None, None
         if self.use_bias:
-            bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names])
-        else:
-            bias_tensor = getattr(self, self.bias_names[0])  # Unused
-        return weight_tensor, bias_tensor
+            bias, bias_list = 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, bias, weight_list, bias_list
 
     def onnx_forward(
         self,
@@ -2334,7 +2381,7 @@ def onnx_forward(
         ) = self._get_quantizers(fp8_output, False, is_grad_enabled)
         inp_dtype = inp.dtype
 
-        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+        weight_tensor, bias_tensor, _, _ = self._get_weight_and_bias_tensors()
         ln_out, ln_out_return = onnx_layernorm(
             inp,
             self.layer_norm_weight,

From 61e1854f31a91435f2f5d1c1c2b9ea56e64f2040 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 15:44:40 +0100
Subject: [PATCH 15/26] [PyTorch] Consolidate split parameter integration tests

Replace five split-parameter tests with one parameterized integration test.
Cover gradients, inference, recipes, bias modes, storage and parameter updates,
saved-version checks and concatenation counts in a shared execution flow.
Use ordinary torch.cat in the reference to verify deferred gradient routing.

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 2495383105c..c2ac0e683e8 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3271,166 +3271,164 @@ def run(x, module=model):
 
 
 @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
-
+@pytest.mark.parametrize(
+    "module, compile_mode, dtype, equal_splits, fp8_recipe, bias_mode, grads",
+    [
+        (te.Linear, mode, dtype, equal, None, "fused", "mixed")
+        for mode in ("default", "reduce-overhead")
+        for dtype in (torch.float16, torch.bfloat16, torch.float32)
+        for equal in (False, True)
+    ]
+    + [
+        (te.Linear, "default", torch.bfloat16, False, quant_recipe, bias, "all")
+        for quant_recipe in (None, *_all_recipes)
+        for bias in ("fused", "none", "returned")
+    ]
+    + [
+        (module, mode, torch.bfloat16, False, None, "fused", grads)
+        for module, mode, targets in (
+            (te.Linear, None, ("none", "mixed")),
+            (te.Linear, "default", ("none", "input", "weight")),
+            (te.LayerNormLinear, None, ("none", "mixed", "input", "weight")),
+        )
+        for grads in targets
+    ],
+)
+def test_te_split_parameters(
+    module, compile_mode, dtype, equal_splits, fp8_recipe, bias_mode, grads
+):
+    """Check outputs, gradients, storage changes, version checks and concatenation counts."""
     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)
+    training = grads != "none"
+    splits = ("q", "k", "v") if equal_splits else dict(q=128, k=64, v=64)
+    model = module(
+        128,
+        384 if equal_splits else 256,
+        parameters_split=splits,
+        params_dtype=dtype,
+        bias=bias_mode != "none",
+        return_bias=bias_mode == "returned",
+    )
+    for name, parameter in model.named_parameters():
+        parameter.requires_grad_(
+            grads == "all"
+            or (grads == "mixed" and name not in ("k_weight", "v_bias"))
+            or (grads == "weight" and name == "v_weight")
+        )
     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)
+    quant_recipe = copy.deepcopy(fp8_recipe)
+    if isinstance(quant_recipe, recipe.NVFP4BlockScaling):
+        quant_recipe.disable_stochastic_rounding = True
+
+    def run(inp, current=model):
+        if quant_recipe is None:
+            out = current(inp)
+        else:
+            with te.autocast(recipe=quant_recipe):
+                out = current(inp)
+        return out[0] + out[1] if bias_mode == "returned" else out
+
+    def reference_parameters(defer_concatenation=False):
+        weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
+        bias = (
+            torch.cat([getattr(reference, name) for name in reference.bias_names])
+            if reference.use_bias
+            else None
+        )
+        return weight, bias, None, None
+
+    reference._get_weight_and_bias_tensors = reference_parameters
+
+    if (
+        compile_mode == "default"
+        and dtype == torch.bfloat16
+        and not equal_splits
+        and fp8_recipe is None
+        and bias_mode == "fused"
+        and grads == "mixed"
+    ):
+        inp = torch.randn(16, 128, device="cuda", dtype=dtype, requires_grad=True)
+        out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
+        with torch.no_grad():
+            model.v_weight.add_(0.01)
+            reference.v_weight.add_(0.01)
+        with pytest.raises(RuntimeError, match="modified by an inplace operation"):
+            out.sum().backward()
+
+    actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
+    check_cats = (
+        compile_mode != "reduce-overhead"
+        and dtype == torch.bfloat16
+        and fp8_recipe is None
+        and bias_mode == "fused"
+    )
     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)
+                current.q_weight = torch.nn.Parameter(
+                    current.q_weight.detach().clone() + 0.01,
+                    requires_grad=current.q_weight.requires_grad,
+                )
             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":
+        if mutation == "optimizer" and training:
             optimizer.step()
             ref_optimizer.step()
-        for _ in range(3):
+        for iteration 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_()
+            inp = torch.randn(
+                16,
+                128,
+                device="cuda",
+                dtype=dtype,
+                requires_grad=grads in ("all", "mixed", "input"),
+            )
+            ref_inp = inp.detach().clone().requires_grad_(inp.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
+            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
+            rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
+            profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
+            with (
+                torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
+                if profile
+                else contextlib.nullcontext()
+            ) as prof:
+                with torch.set_grad_enabled(training):
+                    out = actual_forward(inp)
+                    if training:
+                        out.backward(grad)
+            if profile:
+                cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
+                expected_cats = (
+                    0
+                    if mutation in ("none", "optimizer")
+                    else (4 if compile_mode and training else 2)
+                )
+                assert cats == expected_cats
+            if rng_state is not None:
+                torch.cuda.set_rng_state(rng_state)
+            with torch.set_grad_enabled(training):
+                ref_out = run(ref_inp, reference)
+                if training:
+                    ref_out.backward(grad)
+            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
+            for actual, expected in zip(
+                (inp, *model.parameters()), (ref_inp, *reference.parameters())
+            ):
+                if not actual.requires_grad:
+                    assert actual.grad is None
                 else:
-                    torch.testing.assert_close(result, target, **dtype_tols(dtype))
+                    assert actual.grad is not None
+                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))
     if compile_mode == "reduce-overhead":
         assert not counters["inductor"]["cudagraph_skips"]
-
-
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("fp8_recipe", [None, *_all_recipes], ids=recipe_id)
-@pytest.mark.parametrize("bias_mode", ["fused", "none", "returned"])
-def test_te_split_parameters_recipes(fp8_recipe, bias_mode):
-    autocast_recipe = fp8_recipe or recipe.Float8CurrentScaling()
-    options = dict(bias=bias_mode != "none", return_bias=bias_mode == "returned")
-    model = te.Linear(
-        128,
-        256,
-        parameters_split=dict(q=128, k=64, v=64),
-        params_dtype=torch.bfloat16,
-        **options,
-    )
-
-    def fn(inp):
-        with te.autocast(enabled=fp8_recipe is not None, recipe=autocast_recipe):
-            out = model(inp)
-        if bias_mode == "returned":
-            out, bias = out
-            out = out + bias
-        return out
-
-    torch._dynamo.reset()
-    compiled = torch.compile(fn, fullgraph=True)
-    inp = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
-    grad = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16)
-    expected = fn(inp)
-    expected.backward(grad)
-    expected_grads = [t.grad.clone() for t in (inp, *model.parameters()) if t.requires_grad]
-    inp.grad = None
-    model.zero_grad(set_to_none=True)
-    actual = compiled(inp)
-    actual.backward(grad)
-    actual_grads = [t.grad for t in (inp, *model.parameters()) if t.requires_grad]
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-    torch.testing.assert_close(actual_grads, expected_grads, rtol=0, atol=0)
-
-
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("training", [False, True])
-@pytest.mark.parametrize(
-    "module, execution_mode",
-    [(te.Linear, "eager"), (te.Linear, "compile"), (te.LayerNormLinear, "eager")],
-)
-def test_te_split_parameters_no_cat(training, module, execution_mode):
-    """Adjacent parameter storage does not launch a concatenation in either pass."""
-    model = module(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
-    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 = run_model(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()
-    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")
-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.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("grad_target", ["input", "weight"])
-def test_te_split_parameters_grad_targets(grad_target):
-    model = te.Linear(64, 128, parameters_split=dict(q=64, k=32, v=32), params_dtype=torch.bfloat16)
-    for parameter in model.parameters():
-        parameter.requires_grad_(False)
-    if grad_target != "input":
-        getattr(model, f"v_{grad_target}").requires_grad_(True)
-    inp = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
-    inp.requires_grad_(grad_target == "input")
-    compiled = torch.compile(model, fullgraph=True)
-    out = compiled(inp)
-    grad = torch.randn_like(out)
-    out.backward(grad)
-    actual = [out.detach().clone()]
-    tensors = [inp, *model.parameters()]
-    actual.extend(None if t.grad is None else t.grad.clone() for t in tensors)
-    for tensor in tensors:
-        tensor.grad = None
-    expected = model(inp)
-    expected.backward(grad)
-    for result, target in zip(actual, [expected, *(t.grad for t in tensors)]):
-        if target is None:
-            assert result is None
-        else:
-            torch.testing.assert_close(result, target, rtol=0, atol=0)

From 576b505886bab169bad9648109dafa6ba04d4b89 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 16:21:22 +0100
Subject: [PATCH 16/26] [PyTorch] Avoid repeated split-parameter validation

Rely on concat_input to select nonempty ordinary tensor lists. Preserve
storage, layout, offset and bounds checks before reconstructing a view.

Extend the shared integration test with LayerNormLinear recipes and dtype/
split variants, plus mixed parameter dtypes under autocast. Keep dtype
promotion in fake concatenation to match torch.cat metadata.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py         | 34 +++++++++++++++------
 transformer_engine/pytorch/dynamo/concat.py | 16 ----------
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index c2ac0e683e8..b3083efad78 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3274,13 +3274,18 @@ def run(x, module=model):
 @pytest.mark.parametrize(
     "module, compile_mode, dtype, equal_splits, fp8_recipe, bias_mode, grads",
     [
-        (te.Linear, mode, dtype, equal, None, "fused", "mixed")
-        for mode in ("default", "reduce-overhead")
+        (module, mode, dtype, equal, None, "fused", "mixed")
+        for module, mode in (
+            (te.Linear, "default"),
+            (te.Linear, "reduce-overhead"),
+            (te.LayerNormLinear, None),
+        )
         for dtype in (torch.float16, torch.bfloat16, torch.float32)
         for equal in (False, True)
     ]
     + [
-        (te.Linear, "default", torch.bfloat16, False, quant_recipe, bias, "all")
+        (module, mode, torch.bfloat16, False, quant_recipe, bias, "all")
+        for module, mode in ((te.Linear, "default"), (te.LayerNormLinear, None))
         for quant_recipe in (None, *_all_recipes)
         for bias in ("fused", "none", "returned")
     ]
@@ -3289,7 +3294,7 @@ def run(x, module=model):
         for module, mode, targets in (
             (te.Linear, None, ("none", "mixed")),
             (te.Linear, "default", ("none", "input", "weight")),
-            (te.LayerNormLinear, None, ("none", "mixed", "input", "weight")),
+            (te.LayerNormLinear, None, ("none", "input", "weight")),
         )
         for grads in targets
     ],
@@ -3323,12 +3328,15 @@ def test_te_split_parameters(
     if isinstance(quant_recipe, recipe.NVFP4BlockScaling):
         quant_recipe.disable_stochastic_rounding = True
 
+    autocast_enabled = False
+
     def run(inp, current=model):
-        if quant_recipe is None:
-            out = current(inp)
-        else:
-            with te.autocast(recipe=quant_recipe):
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
+            if quant_recipe is None:
                 out = current(inp)
+            else:
+                with te.autocast(recipe=quant_recipe):
+                    out = current(inp)
         return out[0] + out[1] if bias_mode == "returned" else out
 
     def reference_parameters(defer_concatenation=False):
@@ -3365,7 +3373,11 @@ def reference_parameters(defer_concatenation=False):
         and fp8_recipe is None
         and bias_mode == "fused"
     )
-    for mutation in ("none", "optimizer", "data", "parameter", "assign", "dtype"):
+    mutations = ["none", "optimizer", "data", "parameter", "assign", "dtype"]
+    if dtype == torch.bfloat16 and fp8_recipe is None:
+        mutations.append("mixed_dtype")
+    for mutation in mutations:
+        autocast_enabled = mutation == "mixed_dtype"
         for current in (model, reference):
             if mutation == "data":
                 current.k_weight.data = current.k_weight.detach().clone() + 0.01
@@ -3380,6 +3392,10 @@ def reference_parameters(defer_concatenation=False):
                 current.load_state_dict(state, assign=True)
             elif mutation == "dtype":
                 current.to(torch.float64).to(dtype)
+            elif mutation == "mixed_dtype":
+                current.k_weight.data = current.k_weight.float()
+                if current.use_bias:
+                    current.k_bias.data = current.k_bias.float()
         if mutation == "optimizer" and training:
             optimizer.step()
             ref_optimizer.step()
diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
index 7c1f5217d91..01b7c559136 100644
--- a/transformer_engine/pytorch/dynamo/concat.py
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -10,20 +10,6 @@
 import torch
 
 from .tensor_spec import TensorSpec
-from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage
-
-
-def validate_concat(parts: List[torch.Tensor]) -> None:
-    """Require a nonempty list of non-quantized tensors."""
-    if not parts:
-        raise ValueError("Attempted to concatenate 0 tensors")
-    if any(
-        not isinstance(part, (torch.Tensor, TensorSpec))
-        or isinstance(part, (QuantizedTensor, QuantizedTensorStorage))
-        or (isinstance(part, TensorSpec) and part.is_quantized)
-        for part in parts
-    ):
-        raise TypeError("Deferred concatenation only supports non-quantized tensors")
 
 
 def concat_shape(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> Tuple[int, ...]:
@@ -33,7 +19,6 @@ def concat_shape(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> Tuple[int,
 
 def concat_to_spec(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> TensorSpec:
     """Describe the concatenation without accessing its storage."""
-    validate_concat(parts)
     first = parts[0]
     dtype = first.dtype
     for tensor in parts[1:]:
@@ -48,7 +33,6 @@ def concat_to_spec(parts: Sequence[Union[torch.Tensor, TensorSpec]]) -> TensorSp
 
 def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor:
     """Build a checked view, or copy parts whose storage is not adjacent."""
-    validate_concat(value)
     if len(value) == 1:
         return value[0]
     first = value[0]

From 06415ad12854f0788f3d675c728b087a1d7d5f80 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 17:15:08 +0100
Subject: [PATCH 17/26] Simplify deferred split parameter selection

Remove the deferred-concatenation switch from Linear and LayerNormLinear
parameter getters. Keep bias routing in the getters and explicitly
concatenate deferred operands in the ONNX inference paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py              |  2 +-
 .../pytorch/module/layernorm_linear.py           | 16 +++++++++-------
 transformer_engine/pytorch/module/linear.py      | 16 +++++++++-------
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index b3083efad78..c2fa518b311 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3339,7 +3339,7 @@ def run(inp, current=model):
                     out = current(inp)
         return out[0] + out[1] if bias_mode == "returned" else out
 
-    def reference_parameters(defer_concatenation=False):
+    def reference_parameters():
         weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
         bias = (
             torch.cat([getattr(reference, name) for name in reference.bias_names])
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index cf8c1d19a37..6f43f952c05 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -2106,9 +2106,7 @@ def forward(
 
         try:
             # Get concatenated weight and bias tensors
-            weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors(
-                defer_concatenation=True
-            )
+            weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors()
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2348,13 +2346,13 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
             for name, q in zip(names, original_quantizers)
         )
 
-    def _get_weight_and_bias_tensors(self, defer_concatenation=False):
-        weight, weight_list = concat_input(self._get_weight_tensors(), defer=defer_concatenation)
+    def _get_weight_and_bias_tensors(self):
+        weight, weight_list = concat_input(self._get_weight_tensors(), defer=True)
         bias, bias_list = None, None
         if self.use_bias:
             bias, bias_list = 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,
+                defer=self.apply_bias and not self.gemm_bias_unfused_add,
             )
         return weight, bias, weight_list, bias_list
 
@@ -2381,7 +2379,11 @@ def onnx_forward(
         ) = self._get_quantizers(fp8_output, False, is_grad_enabled)
         inp_dtype = inp.dtype
 
-        weight_tensor, bias_tensor, _, _ = self._get_weight_and_bias_tensors()
+        weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors()
+        if weight_list:
+            weight_tensor = torch.cat(weight_list)
+        if bias_list:
+            bias_tensor = torch.cat(bias_list)
         ln_out, ln_out_return = onnx_layernorm(
             inp,
             self.layer_norm_weight,
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 8abc383fb36..7518238115e 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -2511,9 +2511,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=True
-            )
+            weight, bias, weight_list, bias_list = self._get_weight_and_bias_tensors()
             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
@@ -2814,16 +2812,16 @@ def _forward_eager_fallback(
             fp8_grad=fp8_grad,
         )
 
-    def _get_weight_and_bias_tensors(self, defer_concatenation=False):
+    def _get_weight_and_bias_tensors(self):
         weights = self._get_weight_tensors()
-        weight, weight_list = concat_input(weights, defer=defer_concatenation)
+        weight, weight_list = concat_input(weights, defer=True)
         bias = None
         bias_list = None
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
             bias, bias_list = concat_input(
                 biases,
-                defer=defer_concatenation and self.apply_bias and not self.gemm_bias_unfused_add,
+                defer=self.apply_bias and not self.gemm_bias_unfused_add,
             )
         return weight, bias, weight_list, bias_list
 
@@ -2842,7 +2840,11 @@ 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, weight_list, bias_list = self._get_weight_and_bias_tensors()
+        if weight_list:
+            weight_tensor = torch.cat(weight_list)
+        if bias_list:
+            bias_tensor = torch.cat(bias_list)
         (
             input_quantizer,
             weight_quantizer,

From d6619876e2c82826939ba1b362a42ba1b663ce81 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 17:39:14 +0100
Subject: [PATCH 18/26] Clarify split parameter preparation names

Rename concat_input to prepare_parameter and drop the tensors suffix
from the Linear and LayerNormLinear getters, which can return split lists.
Use the same forward operand names in LayerNormLinear as in Linear.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  2 +-
 transformer_engine/pytorch/module/_common.py  |  2 +-
 .../pytorch/module/layernorm_linear.py        | 44 +++++++++----------
 transformer_engine/pytorch/module/linear.py   | 12 ++---
 4 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index c2fa518b311..85b2998a2bb 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3348,7 +3348,7 @@ def reference_parameters():
         )
         return weight, bias, None, None
 
-    reference._get_weight_and_bias_tensors = reference_parameters
+    reference._get_weight_and_bias = reference_parameters
 
     if (
         compile_mode == "default"
diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py
index 62572a3a6f6..6dbad7fad9b 100644
--- a/transformer_engine/pytorch/module/_common.py
+++ b/transformer_engine/pytorch/module/_common.py
@@ -241,7 +241,7 @@ def noop_cat(
     return _NoopCatFunc.apply(dim, *tensors)
 
 
-def concat_input(
+def prepare_parameter(
     tensors: Sequence[TensorOrQuantized], *, defer: bool = False
 ) -> Tuple[Optional[TensorOrQuantized], Optional[List[torch.Tensor]]]:
     """Return a parameter operand or its deferred parts."""
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 6f43f952c05..8b377a3c5a2 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -67,7 +67,7 @@
 from ._common import (
     apply_normalization,
     check_fp8_reduce_and_update,
-    concat_input,
+    prepare_parameter,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -2106,7 +2106,7 @@ def forward(
 
         try:
             # Get concatenated weight and bias tensors
-            weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors()
+            weight, bias, weight_list, bias_list = self._get_weight_and_bias()
 
             quantizers = (
                 self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled)
@@ -2129,7 +2129,7 @@ def forward(
             if weight_quantizer is not None and not debug:
                 weight_quantizer.optimize_for_gemm = self._enable_weight_preswizzle(
                     weight_quantizer,
-                    weight_tensor,
+                    weight,
                     weight_shape=concat_shape(weight_list) if weight_list else None,
                 )
 
@@ -2165,16 +2165,14 @@ def forward(
                 ub_bulk_wgrad = self.ub_bulk_wgrad
                 ub_bulk_dgrad = self.ub_bulk_dgrad
 
-            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
             fwd_args = LayerNormLinearFwdArgs(
                 # tensors
                 inp=inp,
                 ln_weight=self.layer_norm_weight,
                 ln_bias=self.layer_norm_bias,
-                weight=weight_tensor,
-                bias=linear_bias_tensor,
+                weight=weight,
+                bias=linear_bias,
                 weight_list=weight_list,
                 bias_list=bias_list,
                 weight_workspace=weight_workspace,
@@ -2189,12 +2187,12 @@ def forward(
                 weight_requires_grad=(
                     any(t.requires_grad for t in weight_list)
                     if weight_list
-                    else weight_tensor.requires_grad
+                    else weight.requires_grad
                 ),
                 bias_requires_grad=(
                     any(t.requires_grad for t in bias_list)
                     if bias_list
-                    else linear_bias_tensor is not None and linear_bias_tensor.requires_grad
+                    else linear_bias is not None and linear_bias.requires_grad
                 ),
                 # quantizers
                 input_quantizer=input_quantizer,
@@ -2256,8 +2254,8 @@ def forward(
                     inp,
                     self.layer_norm_weight,
                     self.layer_norm_bias,
-                    weight_tensor,
-                    linear_bias_tensor,
+                    weight,
+                    linear_bias,
                     fwd_args,
                     *(weight_list or ()),
                     *(bias_list or ()),
@@ -2267,8 +2265,8 @@ def forward(
                     inp,
                     self.layer_norm_weight,
                     self.layer_norm_bias,
-                    weight_tensor,
-                    linear_bias_tensor,
+                    weight,
+                    linear_bias,
                     fwd_args,
                 )
             else:
@@ -2277,8 +2275,8 @@ def forward(
                     inp,
                     self.layer_norm_weight,
                     self.layer_norm_bias,
-                    weight_tensor,
-                    linear_bias_tensor,
+                    weight,
+                    linear_bias,
                     fwd_args,
                 )
 
@@ -2291,12 +2289,12 @@ def forward(
             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:
             if self.return_layernorm_output:
-                return out, cast_if_needed(bias_tensor, self.activation_dtype), ln_out
-            return out, cast_if_needed(bias_tensor, self.activation_dtype)
+                return out, cast_if_needed(bias, self.activation_dtype), ln_out
+            return out, cast_if_needed(bias, self.activation_dtype)
         if self.return_layernorm_output:
             return out, ln_out
         return out
@@ -2346,11 +2344,11 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled):
             for name, q in zip(names, original_quantizers)
         )
 
-    def _get_weight_and_bias_tensors(self):
-        weight, weight_list = concat_input(self._get_weight_tensors(), defer=True)
+    def _get_weight_and_bias(self):
+        weight, weight_list = prepare_parameter(self._get_weight_tensors(), defer=True)
         bias, bias_list = None, None
         if self.use_bias:
-            bias, bias_list = concat_input(
+            bias, bias_list = prepare_parameter(
                 [getattr(self, name) for name in self.bias_names],
                 defer=self.apply_bias and not self.gemm_bias_unfused_add,
             )
@@ -2379,7 +2377,7 @@ def onnx_forward(
         ) = self._get_quantizers(fp8_output, False, is_grad_enabled)
         inp_dtype = inp.dtype
 
-        weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias_tensors()
+        weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias()
         if weight_list:
             weight_tensor = torch.cat(weight_list)
         if bias_list:
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 7518238115e..61d65134f34 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -34,7 +34,7 @@
 from ._common import (
     can_reconstruct_wgrad_input_from_original,
     check_fp8_reduce_and_update,
-    concat_input,
+    prepare_parameter,
     set_quantizer_amax_reduction_group,
     set_quantizer_usage_for_wgrad_all_gather,
     WeightGradStore,
@@ -2511,7 +2511,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()
+            weight, bias, weight_list, bias_list = self._get_weight_and_bias()
             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
@@ -2812,14 +2812,14 @@ def _forward_eager_fallback(
             fp8_grad=fp8_grad,
         )
 
-    def _get_weight_and_bias_tensors(self):
+    def _get_weight_and_bias(self):
         weights = self._get_weight_tensors()
-        weight, weight_list = concat_input(weights, defer=True)
+        weight, weight_list = prepare_parameter(weights, defer=True)
         bias = None
         bias_list = None
         if self.use_bias:
             biases = [getattr(self, name) for name in self.bias_names]
-            bias, bias_list = concat_input(
+            bias, bias_list = prepare_parameter(
                 biases,
                 defer=self.apply_bias and not self.gemm_bias_unfused_add,
             )
@@ -2840,7 +2840,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, weight_list, bias_list = self._get_weight_and_bias_tensors()
+        weight_tensor, bias_tensor, weight_list, bias_list = self._get_weight_and_bias()
         if weight_list:
             weight_tensor = torch.cat(weight_list)
         if bias_list:

From 15b9973ff0c9b7aa25ff6e5b88adeca13a1c787c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 17:47:31 +0100
Subject: [PATCH 19/26] Detect split backward inputs from saved metadata

Use the presence of parameter_split_sizes to select gradient splitting
in eager Linear and LayerNormLinear backward instead of counting the
fixed positional autograd inputs.

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

diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 8b377a3c5a2..ad270b6fc99 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -1561,7 +1561,7 @@ def backward(
             nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors")
             FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
             nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors")
-        if len(ctx.needs_input_grad) > 6:
+        if hasattr(ctx, "parameter_split_sizes"):
             weight_sizes, bias_sizes = ctx.parameter_split_sizes
             part_grads = []
             for grad, sizes in ((wgrad, weight_sizes), (grad_bias, bias_sizes)):
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 61d65134f34..3137eba8574 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -1948,7 +1948,7 @@ 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:
+        if hasattr(ctx, "parameter_split_sizes"):
             wgrad, dgrad, bgrad, _ = result
             weight_sizes, bias_sizes = ctx.parameter_split_sizes
             part_grads = []

From 63ccf35e020040ccb154f5bef10dcd17f131e71a Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 18:15:41 +0100
Subject: [PATCH 20/26] Reuse adapter metadata to restore split parameters

Restore parameter lists from named input-list lengths already tracked by
the custom-op adapter. Remove the duplicate lengths and the saved weight
alias flag from Linear backward arguments. Resolve a missing plain split
weight in the existing backward workspace reconstruction path.

Extend the split integration test with workspace caching, reconstruction
and high-precision/dequantized backward overrides.

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 85b2998a2bb..bae4a42775d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3289,6 +3289,19 @@ def run(x, module=model):
         for quant_recipe in (None, *_all_recipes)
         for bias in ("fused", "none", "returned")
     ]
+    + [
+        (
+            te.Linear,
+            "default",
+            torch.bfloat16,
+            False,
+            recipe.Float8CurrentScaling(backward_override=override),
+            "fused",
+            "all",
+        )
+        for override in ("high_precision", "dequantized")
+        if fp8_available
+    ]
     + [
         (module, mode, torch.bfloat16, False, None, "fused", grads)
         for module, mode, targets in (
@@ -3329,14 +3342,15 @@ def test_te_split_parameters(
         quant_recipe.disable_stochastic_rounding = True
 
     autocast_enabled = False
+    is_first_microbatch = None
 
     def run(inp, current=model):
         with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
             if quant_recipe is None:
-                out = current(inp)
+                out = current(inp, is_first_microbatch=is_first_microbatch)
             else:
                 with te.autocast(recipe=quant_recipe):
-                    out = current(inp)
+                    out = current(inp, is_first_microbatch=is_first_microbatch)
         return out[0] + out[1] if bias_mode == "returned" else out
 
     def reference_parameters():
@@ -3376,6 +3390,8 @@ def reference_parameters():
     mutations = ["none", "optimizer", "data", "parameter", "assign", "dtype"]
     if dtype == torch.bfloat16 and fp8_recipe is None:
         mutations.append("mixed_dtype")
+    if module is te.Linear and compile_mode and fp8_recipe is not None:
+        mutations.extend(("cache", "fsdp2"))
     for mutation in mutations:
         autocast_enabled = mutation == "mixed_dtype"
         for current in (model, reference):
@@ -3396,10 +3412,13 @@ def reference_parameters():
                 current.k_weight.data = current.k_weight.float()
                 if current.use_bias:
                     current.k_bias.data = current.k_bias.float()
+            elif mutation == "fsdp2":
+                current._is_fsdp2 = True
         if mutation == "optimizer" and training:
             optimizer.step()
             ref_optimizer.step()
         for iteration in range(3):
+            is_first_microbatch = iteration == 0 if mutation == "cache" else None
             torch.compiler.cudagraph_mark_step_begin()
             inp = torch.randn(
                 16,
@@ -3446,5 +3465,10 @@ def reference_parameters():
                 else:
                     assert actual.grad is not None
                     torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))
+            if mutation == "cache":
+                if iteration == 0:
+                    cached_workspace = model._fp8_workspaces["weight"]
+                else:
+                    assert model._fp8_workspaces["weight"] is cached_workspace
     if compile_mode == "reduce-overhead":
         assert not counters["inductor"]["cudagraph_skips"]
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 203547a9d5d..8d391941ec4 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1045,10 +1045,13 @@ 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)
+    fwd_slot_indices = {name: i for i, name in enumerate(fwd_plan.slot_names)}
 
     def _setup_context(ctx, inputs, output):
         ctx.fwd_tensor_list_lengths = {
-            i: len(value) for i, value in enumerate(inputs) if isinstance(value, list)
+            name: len(value)
+            for name, value in zip(fwd_plan.slot_names, inputs)
+            if isinstance(value, list)
         }
         fwd_obj = fwd_plan.unpack(dict(zip(fwd_plan.slot_names, inputs)))
         spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names())
@@ -1100,12 +1103,13 @@ def _autograd_backward(ctx, *grad_outputs):
         # (always recorded in ``fwd_tensor_list_lengths``) needs a
         # list-shaped no-grad of matching length.
         out: List[Any] = [None] * len(fwd_plan.slot_names)
-        for pos, length in ctx.fwd_tensor_list_lengths.items():
-            out[pos] = [None] * length
+        for name, length in ctx.fwd_tensor_list_lengths.items():
+            out[fwd_slot_indices[name]] = [None] * length
         for pos, g in zip(grad_targets, grads):
-            if pos in ctx.fwd_tensor_list_lengths:
+            name = fwd_plan.slot_names[pos]
+            if name in ctx.fwd_tensor_list_lengths:
                 if g is not None:
-                    if not isinstance(g, list) or len(g) != ctx.fwd_tensor_list_lengths[pos]:
+                    if not isinstance(g, list) or len(g) != ctx.fwd_tensor_list_lengths[name]:
                         raise RuntimeError("Expected one gradient per tensor-list element")
                     out[pos] = g
                 continue
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 3137eba8574..e1f84ff8266 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -246,8 +246,6 @@ class LinearBwdArgs:
     bias: Optional[torch.Tensor] = None
     weight_list: Optional[List[torch.Tensor]] = None
     bias_list: Optional[List[torch.Tensor]] = None
-    parameter_list_lengths: Tuple[int, int] = (0, 0)
-    weight_fp8_from_list: bool = False
 
     # --- Quantizers ---
     input_quantizer: Optional[Quantizer] = None
@@ -314,15 +312,13 @@ 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):
+        if len(saved) == 4:
             self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved
             return
         self.inputmat, self.weight_fp8, self.saved_weight, self.bias = saved[:4]
-        weight_count, bias_count = self.parameter_list_lengths
-        self.weight_list = list(saved[4 : 4 + weight_count]) if weight_count else None
-        self.bias_list = (
-            list(saved[4 + weight_count : 4 + weight_count + bias_count]) if bias_count else None
-        )
+        weight_count = ctx.fwd_tensor_list_lengths.get("weight_list", 0)
+        self.weight_list = saved[4 : 4 + weight_count] if weight_count else None
+        self.bias_list = saved[4 + weight_count :] or None
 
     def postprocess_gradients(self, grads):
         """Split gradients outside the op so the parts can share storage."""
@@ -1134,17 +1130,13 @@ def _linear_setup_ctx(
         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)
-    weight_list = fwd_args.weight_list or ()
-    bias_list = fwd_args.bias_list or ()
-    bwd_args.parameter_list_lengths = (len(weight_list), len(bias_list))
     return (
         saved_inputmat,
         wt_save,
         saved_weight,
         saved_bias,
-        *weight_list,
-        *bias_list,
+        *(fwd_args.weight_list or ()),
+        *(fwd_args.bias_list or ()),
     )
 
 
@@ -1159,7 +1151,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
     assert grad_output is not None
     inputmat = args.inputmat
     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
+    weight_fp8 = args.weight_fp8
     is_dist_weight = is_distributed_weight(saved_weight)
     bias = materialize_concat(args.bias_list) if args.bias_list else args.bias
     input_quantizer = args.input_quantizer
@@ -1402,6 +1394,8 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None
                 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)
+                elif args.weight_list:
+                    weight_fp8 = cast_if_needed(saved_weight, bwd_args.activation_dtype)
             elif (
                 is_dist_weight
                 and bwd_args.fp8

From 4fa64a8879f23e7aedbb0eca5609d372a04a7590 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 18:39:41 +0100
Subject: [PATCH 21/26] Cover split parameters in existing module tests

Move selected split scenarios into the Linear compile and LayerNormLinear
sanity tests, and add a split case to the existing microbatch cache test.
Share parameter-update assertions and remove the standalone split test.
Append complete scenarios without multiplying the existing parameter axes.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_sanity.py        | 130 +++++++++--
 tests/pytorch/test_torch_compile.py | 346 +++++++++-------------------
 tests/pytorch/utils.py              | 151 +++++++++++-
 3 files changed, 370 insertions(+), 257 deletions(-)

diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index 2627de023ac..b605f3a97b8 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -8,6 +8,7 @@
 import pytest
 import os
 import copy
+from itertools import product
 
 import transformer_engine
 import transformer_engine.pytorch as te
@@ -44,7 +45,12 @@
     layernorm_mlp,
     linear,
 )
-from utils import ModelConfig, recipe_id, skip_unsupported_backward_override
+from utils import (
+    ModelConfig,
+    recipe_id,
+    skip_unsupported_backward_override,
+    check_module_parameters,
+)
 
 # Only run FP8 tests on supported devices.
 fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
@@ -436,19 +442,100 @@ def test_sanity_normalization_amp(dtype, model, skip_wgrad, skip_dgrad, normaliz
     _test_sanity_normalization_amp(block, dtype, config, skip_wgrad, skip_dgrad)
 
 
-@pytest.mark.parametrize("dtype", param_types)
 @pytest.mark.parametrize(
-    "fp8_recipe",
-    fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []),
-    ids=recipe_id,
+    "dtype,fp8_recipe,backward_override,model,skip_wgrad,zero_centered_gamma,"
+    "skip_dgrad,normalization,microbatching,parameter_options",
+    [
+        (*case, None)
+        for case in product(
+            param_types,
+            fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []),
+            (None, "high_precision", "dequantized"),
+            ("small", "weird"),
+            all_boolean,
+            all_boolean,
+            all_boolean,
+            all_normalizations,
+            all_boolean,
+        )
+    ]
+    + [
+        pytest.param(
+            dtype,
+            None,
+            None,
+            "small",
+            False,
+            False,
+            False,
+            "LayerNorm",
+            False,
+            {"parameters_split": splits},
+            id=f"split-{dtype}",
+        )
+        for dtype, splits in (
+            (torch.float32, ("q", "k", "v")),
+            (torch.float16, dict(q=96, k=48, v=48)),
+            (torch.bfloat16, dict(q=96, k=48, v=48)),
+        )
+    ]
+    + [
+        pytest.param(
+            torch.bfloat16,
+            None,
+            None,
+            "small",
+            False,
+            False,
+            False,
+            "LayerNorm",
+            False,
+            {"parameters_split": ("q", "k", "v"), **options},
+            id=f"split-{name}",
+        )
+        for name, options in (
+            ("no-bias", {"bias": False}),
+            ("returned-bias", {"return_bias": True}),
+        )
+    ]
+    + [
+        pytest.param(
+            torch.bfloat16,
+            None,
+            None,
+            "small",
+            skip_wgrad,
+            False,
+            skip_dgrad,
+            "LayerNorm",
+            False,
+            {"parameters_split": ("q", "k", "v")},
+            id=f"split-{name}",
+        )
+        for name, skip_wgrad, skip_dgrad in (
+            ("inference", True, True),
+            ("input-grad", True, False),
+            ("weight-grad", False, True),
+        )
+    ]
+    + [
+        pytest.param(
+            torch.bfloat16,
+            q,
+            None,
+            "small",
+            False,
+            False,
+            False,
+            "LayerNorm",
+            False,
+            {"parameters_split": ("q", "k", "v")},
+            id=f"split-{recipe_id(q)}",
+        )
+        for q in fp8_recipes
+        if q is not None
+    ],
 )
-@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"])
-@pytest.mark.parametrize("model", ["small", "weird"])
-@pytest.mark.parametrize("skip_wgrad", all_boolean)
-@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
-@pytest.mark.parametrize("skip_dgrad", all_boolean)
-@pytest.mark.parametrize("normalization", all_normalizations)
-@pytest.mark.parametrize("microbatching", all_boolean)
 def test_sanity_layernorm_linear(
     dtype,
     fp8_recipe,
@@ -459,6 +546,7 @@ def test_sanity_layernorm_linear(
     skip_dgrad,
     normalization,
     microbatching,
+    parameter_options,
 ):
     config = model_configs[model]
 
@@ -484,8 +572,24 @@ def test_sanity_layernorm_linear(
         normalization=normalization,
         params_dtype=dtype,
         device="cuda",
+        **(parameter_options or {}),
     )
-    _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching)
+    if parameter_options:
+        check_storage = (
+            isinstance(parameter_options["parameters_split"], dict) and dtype == torch.bfloat16
+        )
+        grads = (
+            "none"
+            if skip_wgrad and skip_dgrad
+            else (
+                "input"
+                if skip_wgrad
+                else "weight" if skip_dgrad else "mixed" if check_storage else "all"
+            )
+        )
+        check_module_parameters(block, fp8_recipe, grads=grads, check_storage=check_storage)
+    else:
+        _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching)
 
 
 @pytest.mark.parametrize("dtype", param_types)
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index bae4a42775d..468fc310b99 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -71,7 +71,13 @@
 # Import from the local utils.py by explicit path: importing cutedsl makes a
 # top-level ``utils`` package visible that would shadow it.
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id
+from utils import (
+    ModelConfig,
+    dtype_tols,
+    get_available_attention_backends,
+    recipe_id,
+    check_module_parameters,
+)
 
 sys.path.pop(0)
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
@@ -1934,23 +1940,66 @@ def test_to_tensor_spec_quantized(factory, shape):
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("compile_mode", _compile_modes)
 @pytest.mark.parametrize(
-    "fp8_recipe",
-    [None, *_all_recipes],
-    ids=lambda r: "bf16" if r is None else type(r).__name__,
+    "fp8_recipe,compile_mode,dtype,parameters_split,bias_mode,grads",
+    [
+        (q, mode, torch.bfloat16, None, "fused", "all")
+        for q in (None, *_all_recipes)
+        for mode in _compile_modes
+    ]
+    + [
+        (None, "default", torch.float32, ("q", "k", "v"), "fused", "all"),
+        (None, "default", torch.float16, dict(q=128, k=64, v=64), "fused", "all"),
+        (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", "mixed"),
+        (None, "reduce-overhead", torch.bfloat16, ("q", "k", "v"), "fused", "all"),
+        (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "none", "all"),
+        (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "returned", "all"),
+    ]
+    + [
+        (None, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", grads)
+        for grads in ("none", "input", "weight")
+    ]
+    + [
+        (q, "default", torch.bfloat16, dict(q=128, k=64, v=64), "fused", "all")
+        for q in _all_recipes
+    ]
+    + [
+        (
+            recipe.Float8CurrentScaling(backward_override=override),
+            "default",
+            torch.bfloat16,
+            dict(q=128, k=64, v=64),
+            "fused",
+            "all",
+        )
+        for override in ("high_precision", "dequantized")
+        if fp8_available
+    ],
 )
-def test_te_linear_compiles(fp8_recipe, compile_mode):
-    """
-    torch.compile(fullgraph=True) of ``te.Linear`` under every built-in
-    recipe (plus the bf16-only baseline with no autocast), for both the default
-    backend and ``mode="reduce-overhead"`` (CUDA-graph trees).
-    """
-    dtype = torch.bfloat16
+def test_te_linear_compiles(fp8_recipe, compile_mode, dtype, parameters_split, bias_mode, grads):
+    """Check compiled Linear outputs and gradients, including split parameter updates."""
     device = "cuda"
-
-    # FP8 GEMMs require leading dimensions divisible by 16.
-    model = te.Linear(64, 32, params_dtype=dtype, device=device)
+    model = te.Linear(
+        128 if parameters_split else 64,
+        (256 if isinstance(parameters_split, dict) else 384) if parameters_split else 32,
+        params_dtype=dtype,
+        device=device,
+        parameters_split=parameters_split,
+        bias=bias_mode != "none",
+        return_bias=bias_mode == "returned",
+    )
+    if parameters_split:
+        with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
+            check_module_parameters(
+                model,
+                fp8_recipe,
+                compile_mode,
+                grads=grads,
+                check_storage=grads == "mixed",
+                recompute_workspace=isinstance(fp8_recipe, recipe.Float8CurrentScaling)
+                and fp8_recipe.backward_override is None,
+            )
+        return
 
     def fn(inp):
         if fp8_recipe is None:
@@ -1967,8 +2016,6 @@ def fn(inp):
         )
         model.zero_grad(set_to_none=True)
     compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
-
-    # Iterate a few times so reduce-overhead actually replays a captured graph.
     n_iters = 3 if compile_mode == "reduce-overhead" else 1
     with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
         for _ in range(n_iters):
@@ -2238,39 +2285,56 @@ def fn(inp):
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize("compile_mode", _compile_modes)
-@pytest.mark.parametrize("deferred_backward", [False, True])
 @pytest.mark.parametrize(
-    "fp8_recipe",
+    "fp8_recipe,compile_mode,deferred_backward,parameters_split",
     [
         pytest.param(
-            recipe.Float8CurrentScaling(),
-            marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
-        ),
-        pytest.param(
-            recipe.Float8BlockScaling(),
-            marks=pytest.mark.skipif(
-                not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling
+            q,
+            mode,
+            deferred,
+            None,
+            marks=pytest.mark.skipif(not available, reason=reason),
+            id=f"{recipe_id(q)}-{mode}-{deferred}",
+        )
+        for q, available, reason in (
+            (recipe.Float8CurrentScaling(), fp8_available, reason_for_no_fp8),
+            (
+                recipe.Float8BlockScaling(),
+                fp8_block_scaling_available,
+                reason_for_no_fp8_block_scaling,
             ),
-        ),
-        pytest.param(
-            recipe.MXFP8BlockScaling(),
-            marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8),
-        ),
+            (recipe.MXFP8BlockScaling(), mxfp8_available, reason_for_no_mxfp8),
+            (
+                recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
+                nvfp4_available,
+                reason_for_no_nvfp4,
+            ),
+        )
+        for mode in _compile_modes
+        for deferred in (False, True)
+    ]
+    + [
         pytest.param(
-            recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
-            marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
-        ),
+            recipe.Float8CurrentScaling(),
+            "default",
+            False,
+            dict(q=64, k=32, v=32),
+            marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
+            id="split-cache",
+        )
     ],
-    ids=recipe_id,
 )
 def test_te_linear_compile_is_first_microbatch(
-    fp8_recipe, compile_mode, deferred_backward, monkeypatch
+    fp8_recipe, compile_mode, deferred_backward, parameters_split, monkeypatch
 ):
     """Reuse cached weights across microbatches and refresh after optimizer updates."""
     dtype, device = torch.bfloat16, "cuda"
-    model = te.Linear(128, 128, params_dtype=dtype, device=device)
-    ref_model = te.Linear(128, 128, params_dtype=dtype, device=device)
+    model = te.Linear(
+        128, 128, params_dtype=dtype, device=device, parameters_split=parameters_split
+    )
+    ref_model = te.Linear(
+        128, 128, params_dtype=dtype, device=device, parameters_split=parameters_split
+    )
     ref_model.load_state_dict(model.state_dict())
 
     def fn(inp, is_first):
@@ -3268,207 +3332,3 @@ def run(x, module=model):
         forward=case in ("bias", "unfused"),
         backward=case != "linear" and bool(targets),
     )
-
-
-@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.parametrize(
-    "module, compile_mode, dtype, equal_splits, fp8_recipe, bias_mode, grads",
-    [
-        (module, mode, dtype, equal, None, "fused", "mixed")
-        for module, mode in (
-            (te.Linear, "default"),
-            (te.Linear, "reduce-overhead"),
-            (te.LayerNormLinear, None),
-        )
-        for dtype in (torch.float16, torch.bfloat16, torch.float32)
-        for equal in (False, True)
-    ]
-    + [
-        (module, mode, torch.bfloat16, False, quant_recipe, bias, "all")
-        for module, mode in ((te.Linear, "default"), (te.LayerNormLinear, None))
-        for quant_recipe in (None, *_all_recipes)
-        for bias in ("fused", "none", "returned")
-    ]
-    + [
-        (
-            te.Linear,
-            "default",
-            torch.bfloat16,
-            False,
-            recipe.Float8CurrentScaling(backward_override=override),
-            "fused",
-            "all",
-        )
-        for override in ("high_precision", "dequantized")
-        if fp8_available
-    ]
-    + [
-        (module, mode, torch.bfloat16, False, None, "fused", grads)
-        for module, mode, targets in (
-            (te.Linear, None, ("none", "mixed")),
-            (te.Linear, "default", ("none", "input", "weight")),
-            (te.LayerNormLinear, None, ("none", "input", "weight")),
-        )
-        for grads in targets
-    ],
-)
-def test_te_split_parameters(
-    module, compile_mode, dtype, equal_splits, fp8_recipe, bias_mode, grads
-):
-    """Check outputs, gradients, storage changes, version checks and concatenation counts."""
-    torch._dynamo.reset()
-    counters.clear()
-    training = grads != "none"
-    splits = ("q", "k", "v") if equal_splits else dict(q=128, k=64, v=64)
-    model = module(
-        128,
-        384 if equal_splits else 256,
-        parameters_split=splits,
-        params_dtype=dtype,
-        bias=bias_mode != "none",
-        return_bias=bias_mode == "returned",
-    )
-    for name, parameter in model.named_parameters():
-        parameter.requires_grad_(
-            grads == "all"
-            or (grads == "mixed" and name not in ("k_weight", "v_bias"))
-            or (grads == "weight" and name == "v_weight")
-        )
-    reference = copy.deepcopy(model)
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
-    ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01)
-    quant_recipe = copy.deepcopy(fp8_recipe)
-    if isinstance(quant_recipe, recipe.NVFP4BlockScaling):
-        quant_recipe.disable_stochastic_rounding = True
-
-    autocast_enabled = False
-    is_first_microbatch = None
-
-    def run(inp, current=model):
-        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
-            if quant_recipe is None:
-                out = current(inp, is_first_microbatch=is_first_microbatch)
-            else:
-                with te.autocast(recipe=quant_recipe):
-                    out = current(inp, is_first_microbatch=is_first_microbatch)
-        return out[0] + out[1] if bias_mode == "returned" else out
-
-    def reference_parameters():
-        weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
-        bias = (
-            torch.cat([getattr(reference, name) for name in reference.bias_names])
-            if reference.use_bias
-            else None
-        )
-        return weight, bias, None, None
-
-    reference._get_weight_and_bias = reference_parameters
-
-    if (
-        compile_mode == "default"
-        and dtype == torch.bfloat16
-        and not equal_splits
-        and fp8_recipe is None
-        and bias_mode == "fused"
-        and grads == "mixed"
-    ):
-        inp = torch.randn(16, 128, device="cuda", dtype=dtype, requires_grad=True)
-        out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
-        with torch.no_grad():
-            model.v_weight.add_(0.01)
-            reference.v_weight.add_(0.01)
-        with pytest.raises(RuntimeError, match="modified by an inplace operation"):
-            out.sum().backward()
-
-    actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
-    check_cats = (
-        compile_mode != "reduce-overhead"
-        and dtype == torch.bfloat16
-        and fp8_recipe is None
-        and bias_mode == "fused"
-    )
-    mutations = ["none", "optimizer", "data", "parameter", "assign", "dtype"]
-    if dtype == torch.bfloat16 and fp8_recipe is None:
-        mutations.append("mixed_dtype")
-    if module is te.Linear and compile_mode and fp8_recipe is not None:
-        mutations.extend(("cache", "fsdp2"))
-    for mutation in mutations:
-        autocast_enabled = mutation == "mixed_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,
-                    requires_grad=current.q_weight.requires_grad,
-                )
-            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)
-            elif mutation == "mixed_dtype":
-                current.k_weight.data = current.k_weight.float()
-                if current.use_bias:
-                    current.k_bias.data = current.k_bias.float()
-            elif mutation == "fsdp2":
-                current._is_fsdp2 = True
-        if mutation == "optimizer" and training:
-            optimizer.step()
-            ref_optimizer.step()
-        for iteration in range(3):
-            is_first_microbatch = iteration == 0 if mutation == "cache" else None
-            torch.compiler.cudagraph_mark_step_begin()
-            inp = torch.randn(
-                16,
-                128,
-                device="cuda",
-                dtype=dtype,
-                requires_grad=grads in ("all", "mixed", "input"),
-            )
-            ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad)
-            model.zero_grad(set_to_none=True)
-            reference.zero_grad(set_to_none=True)
-            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
-            rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
-            profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
-            with (
-                torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
-                if profile
-                else contextlib.nullcontext()
-            ) as prof:
-                with torch.set_grad_enabled(training):
-                    out = actual_forward(inp)
-                    if training:
-                        out.backward(grad)
-            if profile:
-                cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
-                expected_cats = (
-                    0
-                    if mutation in ("none", "optimizer")
-                    else (4 if compile_mode and training else 2)
-                )
-                assert cats == expected_cats
-            if rng_state is not None:
-                torch.cuda.set_rng_state(rng_state)
-            with torch.set_grad_enabled(training):
-                ref_out = run(ref_inp, reference)
-                if training:
-                    ref_out.backward(grad)
-            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
-            for actual, expected in zip(
-                (inp, *model.parameters()), (ref_inp, *reference.parameters())
-            ):
-                if not actual.requires_grad:
-                    assert actual.grad is None
-                else:
-                    assert actual.grad is not None
-                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))
-            if mutation == "cache":
-                if iteration == 0:
-                    cached_workspace = model._fp8_workspaces["weight"]
-                else:
-                    assert model._fp8_workspaces["weight"] is cached_workspace
-    if compile_mode == "reduce-overhead":
-        assert not counters["inductor"]["cudagraph_skips"]
diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py
index 2a848842b52..9fd6ef5dee7 100644
--- a/tests/pytorch/utils.py
+++ b/tests/pytorch/utils.py
@@ -5,11 +5,12 @@
 from __future__ import annotations
 
 import logging
+import copy
 import os
 import random
 import subprocess
 from collections.abc import Iterable
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from typing import Optional, Sequence, Tuple, Dict, Any, List
 from packaging.version import Version as PkgVersion
 
@@ -691,3 +692,151 @@ def verify_main_grad_accumulation(
                     "was not zeroed; downstream hooks reading .grad would "
                     "see stale bytes from the previous step."
                 )
+
+
+def check_module_parameters(
+    model,
+    fp8_recipe=None,
+    compile_mode=None,
+    *,
+    grads="all",
+    check_storage=False,
+    recompute_workspace=False,
+):
+    """Compare module outputs and gradients through parameter storage changes."""
+    torch._dynamo.reset()
+    training = grads != "none"
+    dtype = next(model.parameters()).dtype
+    in_features = model.in_features
+    bias_mode = "returned" if model.return_bias else "fused" if model.use_bias else "none"
+    for name, parameter in model.named_parameters():
+        parameter.requires_grad_(
+            grads == "all"
+            or (grads == "mixed" and name not in ("k_weight", "v_bias"))
+            or (grads == "weight" and name == "v_weight")
+        )
+    reference = copy.deepcopy(model)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01)
+    quant_recipe = copy.deepcopy(fp8_recipe)
+    if quant_recipe is not None and quant_recipe.nvfp4():
+        quant_recipe.disable_stochastic_rounding = True
+
+    autocast_enabled = False
+
+    def run(inp, current=model):
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
+            if quant_recipe is None:
+                out = current(inp)
+            else:
+                with transformer_engine.pytorch.autocast(recipe=quant_recipe):
+                    out = current(inp)
+        return out[0] + out[1] if bias_mode == "returned" else out
+
+    def reference_parameters():
+        weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
+        bias = (
+            torch.cat([getattr(reference, name) for name in reference.bias_names])
+            if reference.use_bias
+            else None
+        )
+        return weight, bias, None, None
+
+    reference._get_weight_and_bias = reference_parameters
+
+    if compile_mode == "default" and check_storage:
+        inp = torch.randn(16, in_features, device="cuda", dtype=dtype, requires_grad=True)
+        out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
+        with torch.no_grad():
+            model.v_weight.add_(0.01)
+            reference.v_weight.add_(0.01)
+        with pytest.raises(RuntimeError, match="modified by an inplace operation"):
+            out.sum().backward()
+
+    actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
+    check_cats = (
+        check_storage
+        and compile_mode != "reduce-overhead"
+        and dtype == torch.bfloat16
+        and fp8_recipe is None
+        and bias_mode == "fused"
+    )
+    mutations = (
+        ["none", "optimizer", "data", "parameter", "assign", "dtype", "mixed_dtype"]
+        if check_storage
+        else ["none"]
+    )
+    if recompute_workspace:
+        mutations.append("fsdp2")
+    for mutation in mutations:
+        autocast_enabled = mutation == "mixed_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,
+                    requires_grad=current.q_weight.requires_grad,
+                )
+            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)
+            elif mutation == "mixed_dtype":
+                current.k_weight.data = current.k_weight.float()
+                if current.use_bias:
+                    current.k_bias.data = current.k_bias.float()
+            elif mutation == "fsdp2":
+                current._is_fsdp2 = True
+        if mutation == "optimizer" and training:
+            optimizer.step()
+            ref_optimizer.step()
+        for iteration in range(3):
+            torch.compiler.cudagraph_mark_step_begin()
+            inp = torch.randn(
+                16,
+                in_features,
+                device="cuda",
+                dtype=dtype,
+                requires_grad=grads in ("all", "mixed", "input"),
+            )
+            ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad)
+            model.zero_grad(set_to_none=True)
+            reference.zero_grad(set_to_none=True)
+            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
+            rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
+            profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
+            with (
+                torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
+                if profile
+                else nullcontext()
+            ) as prof:
+                with torch.set_grad_enabled(training):
+                    out = actual_forward(inp)
+                    if training:
+                        out.backward(grad)
+            if profile:
+                cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
+                expected_cats = (
+                    0
+                    if mutation in ("none", "optimizer")
+                    else (4 if compile_mode and training else 2)
+                )
+                assert cats == expected_cats
+            if rng_state is not None:
+                torch.cuda.set_rng_state(rng_state)
+            with torch.set_grad_enabled(training):
+                ref_out = run(ref_inp, reference)
+                if training:
+                    ref_out.backward(grad)
+            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
+            for actual, expected in zip(
+                (inp, *model.parameters()), (ref_inp, *reference.parameters())
+            ):
+                if not actual.requires_grad:
+                    assert actual.grad is None
+                else:
+                    assert actual.grad is not None
+                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))

From 10105b32ab0706b409bba469316eefd4550fdd2b Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 19:17:00 +0100
Subject: [PATCH 22/26] Remove added split LayerNormLinear sanity scenarios

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_sanity.py | 130 ++++-------------------------------
 1 file changed, 13 insertions(+), 117 deletions(-)

diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index b605f3a97b8..2627de023ac 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -8,7 +8,6 @@
 import pytest
 import os
 import copy
-from itertools import product
 
 import transformer_engine
 import transformer_engine.pytorch as te
@@ -45,12 +44,7 @@
     layernorm_mlp,
     linear,
 )
-from utils import (
-    ModelConfig,
-    recipe_id,
-    skip_unsupported_backward_override,
-    check_module_parameters,
-)
+from utils import ModelConfig, recipe_id, skip_unsupported_backward_override
 
 # Only run FP8 tests on supported devices.
 fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
@@ -442,100 +436,19 @@ def test_sanity_normalization_amp(dtype, model, skip_wgrad, skip_dgrad, normaliz
     _test_sanity_normalization_amp(block, dtype, config, skip_wgrad, skip_dgrad)
 
 
+@pytest.mark.parametrize("dtype", param_types)
 @pytest.mark.parametrize(
-    "dtype,fp8_recipe,backward_override,model,skip_wgrad,zero_centered_gamma,"
-    "skip_dgrad,normalization,microbatching,parameter_options",
-    [
-        (*case, None)
-        for case in product(
-            param_types,
-            fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []),
-            (None, "high_precision", "dequantized"),
-            ("small", "weird"),
-            all_boolean,
-            all_boolean,
-            all_boolean,
-            all_normalizations,
-            all_boolean,
-        )
-    ]
-    + [
-        pytest.param(
-            dtype,
-            None,
-            None,
-            "small",
-            False,
-            False,
-            False,
-            "LayerNorm",
-            False,
-            {"parameters_split": splits},
-            id=f"split-{dtype}",
-        )
-        for dtype, splits in (
-            (torch.float32, ("q", "k", "v")),
-            (torch.float16, dict(q=96, k=48, v=48)),
-            (torch.bfloat16, dict(q=96, k=48, v=48)),
-        )
-    ]
-    + [
-        pytest.param(
-            torch.bfloat16,
-            None,
-            None,
-            "small",
-            False,
-            False,
-            False,
-            "LayerNorm",
-            False,
-            {"parameters_split": ("q", "k", "v"), **options},
-            id=f"split-{name}",
-        )
-        for name, options in (
-            ("no-bias", {"bias": False}),
-            ("returned-bias", {"return_bias": True}),
-        )
-    ]
-    + [
-        pytest.param(
-            torch.bfloat16,
-            None,
-            None,
-            "small",
-            skip_wgrad,
-            False,
-            skip_dgrad,
-            "LayerNorm",
-            False,
-            {"parameters_split": ("q", "k", "v")},
-            id=f"split-{name}",
-        )
-        for name, skip_wgrad, skip_dgrad in (
-            ("inference", True, True),
-            ("input-grad", True, False),
-            ("weight-grad", False, True),
-        )
-    ]
-    + [
-        pytest.param(
-            torch.bfloat16,
-            q,
-            None,
-            "small",
-            False,
-            False,
-            False,
-            "LayerNorm",
-            False,
-            {"parameters_split": ("q", "k", "v")},
-            id=f"split-{recipe_id(q)}",
-        )
-        for q in fp8_recipes
-        if q is not None
-    ],
+    "fp8_recipe",
+    fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []),
+    ids=recipe_id,
 )
+@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"])
+@pytest.mark.parametrize("model", ["small", "weird"])
+@pytest.mark.parametrize("skip_wgrad", all_boolean)
+@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
+@pytest.mark.parametrize("skip_dgrad", all_boolean)
+@pytest.mark.parametrize("normalization", all_normalizations)
+@pytest.mark.parametrize("microbatching", all_boolean)
 def test_sanity_layernorm_linear(
     dtype,
     fp8_recipe,
@@ -546,7 +459,6 @@ def test_sanity_layernorm_linear(
     skip_dgrad,
     normalization,
     microbatching,
-    parameter_options,
 ):
     config = model_configs[model]
 
@@ -572,24 +484,8 @@ def test_sanity_layernorm_linear(
         normalization=normalization,
         params_dtype=dtype,
         device="cuda",
-        **(parameter_options or {}),
     )
-    if parameter_options:
-        check_storage = (
-            isinstance(parameter_options["parameters_split"], dict) and dtype == torch.bfloat16
-        )
-        grads = (
-            "none"
-            if skip_wgrad and skip_dgrad
-            else (
-                "input"
-                if skip_wgrad
-                else "weight" if skip_dgrad else "mixed" if check_storage else "all"
-            )
-        )
-        check_module_parameters(block, fp8_recipe, grads=grads, check_storage=check_storage)
-    else:
-        _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching)
+    _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching)
 
 
 @pytest.mark.parametrize("dtype", param_types)

From b04c1e340e8a02e0b4f2a605b6943bcdd576ce3e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 28 Oct 2026 19:27:00 +0100
Subject: [PATCH 23/26] Keep module output and gradient checks local to compile
 tests

Move the parameter-update helper from the shared test utilities into
test_torch_compile.py and name it _check_module_outputs_and_grads.
Restore the shared utility module to its upstream contents.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 158 ++++++++++++++++++++++++++--
 tests/pytorch/utils.py              | 151 +-------------------------
 2 files changed, 151 insertions(+), 158 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 468fc310b99..875fac662e0 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -71,13 +71,7 @@
 # Import from the local utils.py by explicit path: importing cutedsl makes a
 # top-level ``utils`` package visible that would shadow it.
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from utils import (
-    ModelConfig,
-    dtype_tols,
-    get_available_attention_backends,
-    recipe_id,
-    check_module_parameters,
-)
+from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id
 
 sys.path.pop(0)
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
@@ -1939,6 +1933,154 @@ def test_to_tensor_spec_quantized(factory, shape):
 # ---------------------------------------------------------------------------
 
 
+def _check_module_outputs_and_grads(
+    model,
+    fp8_recipe=None,
+    compile_mode=None,
+    *,
+    grads="all",
+    check_storage=False,
+    recompute_workspace=False,
+):
+    """Compare module outputs and gradients through parameter storage changes."""
+    torch._dynamo.reset()
+    training = grads != "none"
+    dtype = next(model.parameters()).dtype
+    in_features = model.in_features
+    bias_mode = "returned" if model.return_bias else "fused" if model.use_bias else "none"
+    for name, parameter in model.named_parameters():
+        parameter.requires_grad_(
+            grads == "all"
+            or (grads == "mixed" and name not in ("k_weight", "v_bias"))
+            or (grads == "weight" and name == "v_weight")
+        )
+    reference = copy.deepcopy(model)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01)
+    quant_recipe = copy.deepcopy(fp8_recipe)
+    if quant_recipe is not None and quant_recipe.nvfp4():
+        quant_recipe.disable_stochastic_rounding = True
+
+    autocast_enabled = False
+
+    def run(inp, current=model):
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
+            if quant_recipe is None:
+                out = current(inp)
+            else:
+                with te.autocast(recipe=quant_recipe):
+                    out = current(inp)
+        return out[0] + out[1] if bias_mode == "returned" else out
+
+    def reference_parameters():
+        weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
+        bias = (
+            torch.cat([getattr(reference, name) for name in reference.bias_names])
+            if reference.use_bias
+            else None
+        )
+        return weight, bias, None, None
+
+    reference._get_weight_and_bias = reference_parameters
+
+    if compile_mode == "default" and check_storage:
+        inp = torch.randn(16, in_features, device="cuda", dtype=dtype, requires_grad=True)
+        out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
+        with torch.no_grad():
+            model.v_weight.add_(0.01)
+            reference.v_weight.add_(0.01)
+        with pytest.raises(RuntimeError, match="modified by an inplace operation"):
+            out.sum().backward()
+
+    actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
+    check_cats = (
+        check_storage
+        and compile_mode != "reduce-overhead"
+        and dtype == torch.bfloat16
+        and fp8_recipe is None
+        and bias_mode == "fused"
+    )
+    mutations = (
+        ["none", "optimizer", "data", "parameter", "assign", "dtype", "mixed_dtype"]
+        if check_storage
+        else ["none"]
+    )
+    if recompute_workspace:
+        mutations.append("fsdp2")
+    for mutation in mutations:
+        autocast_enabled = mutation == "mixed_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,
+                    requires_grad=current.q_weight.requires_grad,
+                )
+            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)
+            elif mutation == "mixed_dtype":
+                current.k_weight.data = current.k_weight.float()
+                if current.use_bias:
+                    current.k_bias.data = current.k_bias.float()
+            elif mutation == "fsdp2":
+                current._is_fsdp2 = True
+        if mutation == "optimizer" and training:
+            optimizer.step()
+            ref_optimizer.step()
+        for iteration in range(3):
+            torch.compiler.cudagraph_mark_step_begin()
+            inp = torch.randn(
+                16,
+                in_features,
+                device="cuda",
+                dtype=dtype,
+                requires_grad=grads in ("all", "mixed", "input"),
+            )
+            ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad)
+            model.zero_grad(set_to_none=True)
+            reference.zero_grad(set_to_none=True)
+            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
+            rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
+            profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
+            with (
+                torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
+                if profile
+                else contextlib.nullcontext()
+            ) as prof:
+                with torch.set_grad_enabled(training):
+                    out = actual_forward(inp)
+                    if training:
+                        out.backward(grad)
+            if profile:
+                cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
+                expected_cats = (
+                    0
+                    if mutation in ("none", "optimizer")
+                    else (4 if compile_mode and training else 2)
+                )
+                assert cats == expected_cats
+            if rng_state is not None:
+                torch.cuda.set_rng_state(rng_state)
+            with torch.set_grad_enabled(training):
+                ref_out = run(ref_inp, reference)
+                if training:
+                    ref_out.backward(grad)
+            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
+            for actual, expected in zip(
+                (inp, *model.parameters()), (ref_inp, *reference.parameters())
+            ):
+                if not actual.requires_grad:
+                    assert actual.grad is None
+                else:
+                    assert actual.grad is not None
+                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))
+
+
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
 @pytest.mark.parametrize(
     "fp8_recipe,compile_mode,dtype,parameters_split,bias_mode,grads",
@@ -1990,7 +2132,7 @@ def test_te_linear_compiles(fp8_recipe, compile_mode, dtype, parameters_split, b
     )
     if parameters_split:
         with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
-            check_module_parameters(
+            _check_module_outputs_and_grads(
                 model,
                 fp8_recipe,
                 compile_mode,
diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py
index 9fd6ef5dee7..2a848842b52 100644
--- a/tests/pytorch/utils.py
+++ b/tests/pytorch/utils.py
@@ -5,12 +5,11 @@
 from __future__ import annotations
 
 import logging
-import copy
 import os
 import random
 import subprocess
 from collections.abc import Iterable
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from typing import Optional, Sequence, Tuple, Dict, Any, List
 from packaging.version import Version as PkgVersion
 
@@ -692,151 +691,3 @@ def verify_main_grad_accumulation(
                     "was not zeroed; downstream hooks reading .grad would "
                     "see stale bytes from the previous step."
                 )
-
-
-def check_module_parameters(
-    model,
-    fp8_recipe=None,
-    compile_mode=None,
-    *,
-    grads="all",
-    check_storage=False,
-    recompute_workspace=False,
-):
-    """Compare module outputs and gradients through parameter storage changes."""
-    torch._dynamo.reset()
-    training = grads != "none"
-    dtype = next(model.parameters()).dtype
-    in_features = model.in_features
-    bias_mode = "returned" if model.return_bias else "fused" if model.use_bias else "none"
-    for name, parameter in model.named_parameters():
-        parameter.requires_grad_(
-            grads == "all"
-            or (grads == "mixed" and name not in ("k_weight", "v_bias"))
-            or (grads == "weight" and name == "v_weight")
-        )
-    reference = copy.deepcopy(model)
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
-    ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01)
-    quant_recipe = copy.deepcopy(fp8_recipe)
-    if quant_recipe is not None and quant_recipe.nvfp4():
-        quant_recipe.disable_stochastic_rounding = True
-
-    autocast_enabled = False
-
-    def run(inp, current=model):
-        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast_enabled):
-            if quant_recipe is None:
-                out = current(inp)
-            else:
-                with transformer_engine.pytorch.autocast(recipe=quant_recipe):
-                    out = current(inp)
-        return out[0] + out[1] if bias_mode == "returned" else out
-
-    def reference_parameters():
-        weight = torch.cat([getattr(reference, name) for name in reference.weight_names])
-        bias = (
-            torch.cat([getattr(reference, name) for name in reference.bias_names])
-            if reference.use_bias
-            else None
-        )
-        return weight, bias, None, None
-
-    reference._get_weight_and_bias = reference_parameters
-
-    if compile_mode == "default" and check_storage:
-        inp = torch.randn(16, in_features, device="cuda", dtype=dtype, requires_grad=True)
-        out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
-        with torch.no_grad():
-            model.v_weight.add_(0.01)
-            reference.v_weight.add_(0.01)
-        with pytest.raises(RuntimeError, match="modified by an inplace operation"):
-            out.sum().backward()
-
-    actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
-    check_cats = (
-        check_storage
-        and compile_mode != "reduce-overhead"
-        and dtype == torch.bfloat16
-        and fp8_recipe is None
-        and bias_mode == "fused"
-    )
-    mutations = (
-        ["none", "optimizer", "data", "parameter", "assign", "dtype", "mixed_dtype"]
-        if check_storage
-        else ["none"]
-    )
-    if recompute_workspace:
-        mutations.append("fsdp2")
-    for mutation in mutations:
-        autocast_enabled = mutation == "mixed_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,
-                    requires_grad=current.q_weight.requires_grad,
-                )
-            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)
-            elif mutation == "mixed_dtype":
-                current.k_weight.data = current.k_weight.float()
-                if current.use_bias:
-                    current.k_bias.data = current.k_bias.float()
-            elif mutation == "fsdp2":
-                current._is_fsdp2 = True
-        if mutation == "optimizer" and training:
-            optimizer.step()
-            ref_optimizer.step()
-        for iteration in range(3):
-            torch.compiler.cudagraph_mark_step_begin()
-            inp = torch.randn(
-                16,
-                in_features,
-                device="cuda",
-                dtype=dtype,
-                requires_grad=grads in ("all", "mixed", "input"),
-            )
-            ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad)
-            model.zero_grad(set_to_none=True)
-            reference.zero_grad(set_to_none=True)
-            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
-            rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
-            profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
-            with (
-                torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
-                if profile
-                else nullcontext()
-            ) as prof:
-                with torch.set_grad_enabled(training):
-                    out = actual_forward(inp)
-                    if training:
-                        out.backward(grad)
-            if profile:
-                cats = sum(event.count for event in prof.key_averages() if event.key == "aten::cat")
-                expected_cats = (
-                    0
-                    if mutation in ("none", "optimizer")
-                    else (4 if compile_mode and training else 2)
-                )
-                assert cats == expected_cats
-            if rng_state is not None:
-                torch.cuda.set_rng_state(rng_state)
-            with torch.set_grad_enabled(training):
-                ref_out = run(ref_inp, reference)
-                if training:
-                    ref_out.backward(grad)
-            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
-            for actual, expected in zip(
-                (inp, *model.parameters()), (ref_inp, *reference.parameters())
-            ):
-                if not actual.requires_grad:
-                    assert actual.grad is None
-                else:
-                    assert actual.grad is not None
-                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))

From 4d29b8e25ce7aea92c79d006190ab02e23f5aea5 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 29 Oct 2026 11:58:00 +0100
Subject: [PATCH 24/26] Use one execution path for compiled Linear tests

Run split and unsplit scenarios through the same local output and gradient
checker. Keep CUDA graph warmup and exact unsplit comparisons while
removing the duplicated test body and its early return.

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 875fac662e0..06270955a81 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -1947,6 +1947,11 @@ def _check_module_outputs_and_grads(
     training = grads != "none"
     dtype = next(model.parameters()).dtype
     in_features = model.in_features
+    tolerances = (
+        dict(atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+        if len(model.weight_names) == 1
+        else dtype_tols(dtype)
+    )
     bias_mode = "returned" if model.return_bias else "fused" if model.use_bias else "none"
     for name, parameter in model.named_parameters():
         parameter.requires_grad_(
@@ -1984,7 +1989,7 @@ def reference_parameters():
     reference._get_weight_and_bias = reference_parameters
 
     if compile_mode == "default" and check_storage:
-        inp = torch.randn(16, in_features, device="cuda", dtype=dtype, requires_grad=True)
+        inp = torch.randn(32, in_features, device="cuda", dtype=dtype, requires_grad=True)
         out = torch.compile(model, fullgraph=True, backend="aot_eager")(inp)
         with torch.no_grad():
             model.v_weight.add_(0.01)
@@ -1992,6 +1997,19 @@ def reference_parameters():
         with pytest.raises(RuntimeError, match="modified by an inplace operation"):
             out.sum().backward()
 
+    if compile_mode == "reduce-overhead":
+        _cudagraph_warmup(
+            run,
+            torch.randn(
+                32,
+                in_features,
+                device="cuda",
+                dtype=dtype,
+                requires_grad=grads in ("all", "mixed", "input"),
+            ),
+            backward=training,
+        )
+        model.zero_grad(set_to_none=True)
     actual_forward = torch.compile(run, fullgraph=True, mode=compile_mode) if compile_mode else run
     check_cats = (
         check_storage
@@ -2035,7 +2053,7 @@ def reference_parameters():
         for iteration in range(3):
             torch.compiler.cudagraph_mark_step_begin()
             inp = torch.randn(
-                16,
+                32,
                 in_features,
                 device="cuda",
                 dtype=dtype,
@@ -2044,7 +2062,7 @@ def reference_parameters():
             ref_inp = inp.detach().clone().requires_grad_(inp.requires_grad)
             model.zero_grad(set_to_none=True)
             reference.zero_grad(set_to_none=True)
-            grad = torch.randn(16, model.out_features, device="cuda", dtype=dtype)
+            grad = torch.randn(32, model.out_features, device="cuda", dtype=dtype)
             rng_state = torch.cuda.get_rng_state() if quant_recipe is not None else None
             profile = check_cats and iteration == 2 and mutation in ("none", "optimizer", "data")
             with (
@@ -2070,7 +2088,7 @@ def reference_parameters():
                 ref_out = run(ref_inp, reference)
                 if training:
                     ref_out.backward(grad)
-            torch.testing.assert_close(out, ref_out, **dtype_tols(dtype))
+            torch.testing.assert_close(out, ref_out, **tolerances)
             for actual, expected in zip(
                 (inp, *model.parameters()), (ref_inp, *reference.parameters())
             ):
@@ -2078,7 +2096,7 @@ def reference_parameters():
                     assert actual.grad is None
                 else:
                     assert actual.grad is not None
-                    torch.testing.assert_close(actual.grad, expected.grad, **dtype_tols(dtype))
+                    torch.testing.assert_close(actual.grad, expected.grad, **tolerances)
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
@@ -2130,39 +2148,17 @@ def test_te_linear_compiles(fp8_recipe, compile_mode, dtype, parameters_split, b
         bias=bias_mode != "none",
         return_bias=bias_mode == "returned",
     )
-    if parameters_split:
-        with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
-            _check_module_outputs_and_grads(
-                model,
-                fp8_recipe,
-                compile_mode,
-                grads=grads,
-                check_storage=grads == "mixed",
-                recompute_workspace=isinstance(fp8_recipe, recipe.Float8CurrentScaling)
-                and fp8_recipe.backward_override is None,
-            )
-        return
-
-    def fn(inp):
-        if fp8_recipe is None:
-            return model(inp)
-        with te.autocast(recipe=fp8_recipe):
-            return model(inp)
-
-    torch._dynamo.reset()
-    if compile_mode == "reduce-overhead":
-        _cudagraph_warmup(
-            fn,
-            torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True),
-            backward=True,
-        )
-        model.zero_grad(set_to_none=True)
-    compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
-    n_iters = 3 if compile_mode == "reduce-overhead" else 1
     with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
-        for _ in range(n_iters):
-            base = torch.randn(32, 64, dtype=dtype, device=device)
-            _assert_close_eager_compiled(fn, compiled, model, base)
+        _check_module_outputs_and_grads(
+            model,
+            fp8_recipe,
+            compile_mode,
+            grads=grads,
+            check_storage=grads == "mixed",
+            recompute_workspace=bool(parameters_split)
+            and isinstance(fp8_recipe, recipe.Float8CurrentScaling)
+            and fp8_recipe.backward_override is None,
+        )
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")

From 905f9545ad09a8145da2770d8031fbd19606a76e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 29 Oct 2026 12:33:19 +0100
Subject: [PATCH 25/26] Preserve safe concatenation and no-bias LayerNormLinear
 behavior

Fall back to cat when the first split is empty, since its trailing
strides need not describe the concatenated tensor. Restore the existing
empty bias tensor for LayerNormLinear without bias.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/concat.py           | 2 ++
 transformer_engine/pytorch/module/layernorm_linear.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/transformer_engine/pytorch/dynamo/concat.py b/transformer_engine/pytorch/dynamo/concat.py
index 01b7c559136..b0f1a315aa8 100644
--- a/transformer_engine/pytorch/dynamo/concat.py
+++ b/transformer_engine/pytorch/dynamo/concat.py
@@ -36,6 +36,8 @@ def materialize_concat(value: List[torch.Tensor]) -> torch.Tensor:
     if len(value) == 1:
         return value[0]
     first = value[0]
+    if first.shape[0] == 0:
+        return torch.cat(value)
     storage = first.untyped_storage()
     offset = first.storage_offset()
     for tensor in value:
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index ad270b6fc99..9e81f30f928 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -2352,6 +2352,8 @@ def _get_weight_and_bias(self):
                 [getattr(self, name) for name in self.bias_names],
                 defer=self.apply_bias and not self.gemm_bias_unfused_add,
             )
+        else:
+            bias = getattr(self, self.bias_names[0])
         return weight, bias, weight_list, bias_list
 
     def onnx_forward(

From bc6faaf53acbda2569015510bd8d675e53dd3fb1 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 29 Oct 2026 12:53:43 +0100
Subject: [PATCH 26/26] Test eager LayerNormLinear gradients with shared and
 separate split storage

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 86 ++++++++++++++++++++++++++++------
 1 file changed, 71 insertions(+), 15 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index fcae3dda604..5fec1860c39 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -1475,7 +1475,9 @@ def test_mha_accuracy(dtype, bs, model, mask_type):
             assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])
 
 
-def _test_granular_accuracy(block, bs, dtype, config, delay_wgrad_compute=False, recipe=None):
+def _test_granular_accuracy(
+    block, bs, dtype, config, delay_wgrad_compute=False, recipe=None, grad_output=None
+):
     reset_rng_states()
     fp8 = recipe is not None
     if fp8:
@@ -1493,8 +1495,11 @@ def _test_granular_accuracy(block, bs, dtype, config, delay_wgrad_compute=False,
         out = block(inp_hidden_states)
         if isinstance(out, (List, Tuple)):
             out = out[0]
-    loss = out.sum()
-    loss.backward()
+    if grad_output is None:
+        loss = out.sum()
+        loss.backward()
+    else:
+        out.backward(grad_output)
     if delay_wgrad_compute:
         block.backward_dw()
 
@@ -1897,15 +1902,35 @@ def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma):
         assert_allclose(te_output, torch_output, atol[dtype], rtol[dtype])
 
 
-@pytest.mark.parametrize("dtype", param_types)
-@pytest.mark.parametrize("bs", batch_sizes)
-@pytest.mark.parametrize("model", ["small"])
-@pytest.mark.parametrize("normalization", all_normalizations)
-@pytest.mark.parametrize("zero_centered_gamma", all_boolean)
-@pytest.mark.parametrize("return_bias", all_boolean)
-@pytest.mark.parametrize("bias", all_boolean)
+@pytest.mark.parametrize(
+    "dtype,bs,model,normalization,zero_centered_gamma,return_bias,bias,split_storage",
+    [
+        (dtype, bs, "small", normalization, zero_centered_gamma, return_bias, bias, None)
+        for dtype in param_types
+        for bs in batch_sizes
+        for normalization in all_normalizations
+        for zero_centered_gamma in all_boolean
+        for return_bias in all_boolean
+        for bias in all_boolean
+    ]
+    + [
+        pytest.param(
+            torch.bfloat16,
+            1,
+            "small",
+            "LayerNorm",
+            False,
+            False,
+            True,
+            storage,
+            id=f"split-{storage}-storage",
+            marks=pytest.mark.skipif(not is_bf16_available(), reason="BF16 is not supported"),
+        )
+        for storage in ("shared", "separate")
+    ],
+)
 def test_layernorm_linear_accuracy(
-    dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias
+    dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias, split_storage
 ):
     config = model_configs[model]
 
@@ -1920,8 +1945,25 @@ def test_layernorm_linear_accuracy(
         return_bias=return_bias,
         bias=bias,
         device="cuda",
+        parameters_split=(
+            dict(q=2 * config.hidden_size, k=config.hidden_size, v=config.hidden_size)
+            if split_storage
+            else None
+        ),
     )
 
+    te_module = te_ln_linear.te_module
+    weights = [getattr(te_module, name) for name in te_module.weight_names]
+    biases = [getattr(te_module, name) for name in te_module.bias_names] if bias else []
+    if split_storage:
+        for parts in (weights, biases):
+            if split_storage == "separate":
+                for part in parts:
+                    part.data = part.detach().clone()
+            assert len({part.untyped_storage().data_ptr() for part in parts}) == (
+                len(parts) if split_storage == "separate" else 1
+            )
+
     torch_ln_linear = (
         TorchLayerNormLinear(
             config.hidden_size,
@@ -1944,12 +1986,26 @@ def test_layernorm_linear_accuracy(
             torch_ln_linear.layernorm.bias = Parameter(
                 te_ln_linear.te_module.layer_norm_bias.clone()
             )
-        torch_ln_linear.linear.weight = Parameter(te_ln_linear.te_module.weight.clone())
+        torch_ln_linear.linear.weight = Parameter(torch.cat(weights))
         if bias:
-            torch_ln_linear.linear.bias = Parameter(te_ln_linear.te_module.bias.clone())
+            torch_ln_linear.linear.bias = Parameter(torch.cat(biases))
 
-    te_outputs = _test_granular_accuracy(te_ln_linear, bs, dtype, config)
-    torch_outputs = _test_granular_accuracy(torch_ln_linear, bs, dtype, config)
+    grad_output = (
+        torch.randn(config.max_seqlen_q, bs, 4 * config.hidden_size, dtype=dtype, device="cuda")
+        if split_storage
+        else None
+    )
+    te_outputs = _test_granular_accuracy(te_ln_linear, bs, dtype, config, grad_output=grad_output)
+    torch_outputs = _test_granular_accuracy(
+        torch_ln_linear, bs, dtype, config, grad_output=grad_output
+    )
+    if split_storage:
+        torch_outputs = [
+            *torch_outputs[:-2],
+            *torch_outputs[-2].split(te_module.parameter_split_sizes),
+            *torch_outputs[-1].split(te_module.parameter_split_sizes),
+        ]
+    assert len(te_outputs) == len(torch_outputs)
 
     atol = {
         torch.float32: 2.5e-4,