[core] Support tensor parallelism for model inference (CUDA, Neuron) (#13718)

* draft:add neuron as a legit backend

* feat: neuron-specific changes in the pipeline

* tests: eager tests

* draft: start with tp for flux2

* fix: style

* fix:apr_02 beta

* feat:add wan

* fix:pixart

* fix: rewrite flux swiglu activation to avoid gather op in neuron IR

* test: pixart compile mode on neuron

* cleanup & fix style

* merge: another change

* review: cleanup+suggestions

* fix: CIs style

* tests: add test units for tp

* fix: in case of text-encoder(s) on CPU

* review:cleanup+add test

* fix: style

* doc: remove it for now

* Add from_single_file support for SkyReelsV2 and ChronoEdit transformers (#13946)

SkyReels-V2 and ChronoEdit are both built on Wan, and their transformers have
the same keys as WanTransformer3DModel, so they reuse
convert_wan_transformer_to_diffusers (like WanVACE / WanAnimate). This lets the
community GGUF builds load directly.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* multi-GPU VAE Fix for Cosmos 3 (#13924)

fix(cosmos3): pin VAE latent norm buffers to encode output device

Under sharded placement (device_map="balanced"), vae.encode() runs on the
VAE's own device while the mean/inv_std buffers were pinned to x.device,
causing a cross-device RuntimeError. Compute raw_mu first, then pin the
normalization buffers to its device so all tensors share one device.

Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* docs: fix repeated word typo in set_timesteps docstring (#13876)

* docs: fix repeated word typo in set_timesteps docstring

Removed the duplicate word "schedule" from the docstring for the sigmas argument in EulerDiscreteScheduler.set_timesteps.

* Update scheduling_euler_discrete.py

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* clean some stuff to simplify code.

* clean more to remove permutation related shenanigans.

* revert: put torch.chunk back

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Address Sayak's TP review nits (docstrings + stale comments)

- _modeling_parallel: fix stale comment claiming attention processors read
  tp_degree at runtime (no longer true after the PackedColwise/Rowwise
  refactor; processors are TP-agnostic via head_dim).
- tensor_parallel: single backticks instead of double in docstrings, add
  input/output examples to _blocks_to_block_sizes and _resolve_tp_plan, and
  drop the stale "applies the Flux2 fused-weight permutations" line from
  apply_tensor_parallel (permuters were removed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix stale permuter reference in Flux2ParallelSelfAttnProcessor comment

The proportional QKV/MLP slice is now guaranteed by PackedColwiseParallel
block-by-block sharding, not the removed fused-weight permuter.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Derive TP backend inside apply_tensor_parallel from the mesh device type

Drops the `backend` parameter: apply_tensor_parallel already holds the TP
mesh, so it derives neuron-vs-default from `config._mesh.device_type` itself.
The caller no longer computes it. Documents why torch_device can't be used
(Neuron reports as "cpu" via torch's accelerator API). Removes the now-unused
is_torch_neuronx_available import from modeling_utils.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Simplify TP backend derivation to the mesh device type alone

Drop the redundant is_torch_neuronx_available() guard: a DeviceMesh("neuron", ...)
only exists when running on Neuron, so tp_mesh.device_type is the single source
of truth.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove incorrect "Neuron reports as cpu" rationale from TP backend docstring

Verified on the current native torch_neuronx drop that Neuron does surface as
the torch accelerator (torch._C._get_accelerator().type == "neuron",
get_device() == "neuron"); the old "reports as cpu" claim was stale. The
mesh-device-type derivation is unchanged — it's the device being sharded onto.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop stale "XLA" labels from Neuron device/runtime comments

Neuron is a native torch backend now (not XLA-routed): remove the "(XLA)"
parenthetical from the randn-on-CPU note in torch_utils and the "XLA runtime"
mention from the Neuron TP test docstring. Behavior unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* make style: reflow tensor_parallel docstrings

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Revert "revert: put torch.chunk back"

This reverts commit 7ea75f771a.

* fix: sharding for neuron + validate flux 1

* fix: change for check_repository_consistency

* feat: support qwen image as well

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: address agent review

* review: apply suggestions

* Update docs/source/en/training/distributed_inference.md

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* review: change example tp size to 4

* review: add flux tp test

* review: address comments

* review: improve the doc

* review:apply qwen rotary_emb per device

* review:remove redundant device def after PR #14383

* review: doc title suggestion

* review: restore copy comment

* review: remove cp divisibility explanation from doc on parallelism guide

* fix style

---------

Co-authored-by: Hz_Zhang <47402297+HaozheZhang6@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Ramkumar R <ramkumarashvanth09@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

docs/source/en/api/parallel.md

@@ -22,3 +22,9 @@ Parallelism strategies help speed up diffusion transformers by distributing comp
 [[autodoc]] ContextParallelConfig
 
 [[autodoc]] hooks.apply_context_parallel
+
+## TensorParallelConfig
+
+[[autodoc]] TensorParallelConfig
+
+[[autodoc]] hooks.apply_tensor_parallel

docs/source/en/training/distributed_inference.md

@@ -431,3 +431,133 @@ pipeline = DiffusionPipeline.from_pretrained(
     CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
 ).to(device)
 ```
+
+## Tensor parallelism
+
+[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) shards the weight matrices of a model across devices. Each device holds a column-wise (`"colwise"`) or row-wise (`"rowwise"`) slice of each layer, computes a partial result, and an `AllReduce`/`AllGather` at the layer boundary reconstructs the full output. Unlike context parallelism, it reduces the per-device *weight* memory, which is useful for models that do not fit on a single device.
+
+Pass a [`TensorParallelConfig`] to [`~ModelMixin.enable_parallelism`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style).
+
+```py
+import torch
+from torch import distributed as dist
+from diffusers import DiffusionPipeline, TensorParallelConfig
+
+def setup_distributed():
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl")
+    rank = dist.get_rank()
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+    return device
+
+def main():
+    device = setup_distributed()
+    world_size = dist.get_world_size()
+
+    pipeline = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16
+    )  # weights stay on CPU
+
+    # Shard the transformer first, then move only each rank's slice onto the accelerator.
+    pipeline.transformer.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
+    pipeline.transformer.to(device)
+
+    # Move the remaining, non-sharded components onto the accelerator individually.
+    pipeline.text_encoder.to(device)
+    pipeline.vae.to(device)
+
+    generator = torch.Generator().manual_seed(42)
+    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
+    if dist.get_rank() == 0:
+        image.save("output.png")
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+if __name__ == "__main__":
+    main()
+```
+
+```shell
+torchrun --nproc-per-node 4 tensor_parallel_flux.py
+```
+
+`tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.
+
+### Writing a tensor parallelism plan
+
+Tensor parallelism only works on models that define a `_tp_plan`, a flat class attribute mapping module-name globs to a sharding style. Writing one is mostly a matter of pairing each projection that *expands* the hidden dimension with the projection that *contracts* it back.
+
+Each key may contain **at most one `*`**, and the prefix before it must resolve to an [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) so a single entry covers every block. A key without a `*` applies to the model itself. Paths are relative to the model.
+
+#### Colwise and rowwise
+
+| Style | Shards | Each rank | Use for |
+|---|---|---|---|
+| `"colwise"` | output features (`weight` dim 0) | computes a slice of the output | `to_q`, `to_k`, `to_v`, FFN in-projection |
+| `"rowwise"` | input features (`weight` dim 1) | computes a partial sum | `to_out.0`, FFN out-projection |
+
+Always pair them in that order. A `"colwise"` projection leaves its output sharded, the following `"rowwise"` projection consumes that shard directly, and a single `AllReduce` at the block boundary reconstructs the result. Sharding the pair any other way forces a gather in the middle and communicates far more.
+
+For attention this means each rank owns a subset of heads, which is why `tp_degree` must divide the head count. Encoder-stream duplicates (`add_q_proj`, `to_add_out`, `ff_context`) follow the same pattern as their image-stream counterparts.
+
+#### Fused projections
+
+When one `Linear` packs several logical tensors along the dimension being sharded, plain `"colwise"`/`"rowwise"` slices straight across the concatenation and misaligns the pieces. Use `PackedColwiseParallel`/`PackedRowwiseParallel` instead, which shard each packed block independently.
+
+```py
+# in src/diffusers/models/transformers/your_model.py
+from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
+```
+
+`blocks` is a list of proportional integers whose sum divides the packed dimension — `[1, 1]` for a SwiGLU gate+up projection of equal halves, or `[1, 1, 1, 3, 3]` for a fused Q+K+V+gate+up projection with `mlp_ratio=3`.
+
+```py
+"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
+```
+
+When the block sizes are only known from the config, omit the argument and store the absolute sizes on the `Linear` during `__init__` instead, as `_tp_packed_col_blocks` or `_tp_packed_row_blocks`.
+
+```py
+# in the attention module's __init__
+self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
+```
+
+```py
+# in the model's _tp_plan
+"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
+```
+
+Both forms appear in one plan in [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py): the double-stream SwiGLU pair passes `blocks` literally, while the fused single-stream QKV+MLP projection stores its sizes on the `Linear` in `Flux2ParallelSelfAttention.__init__` because they depend on `mlp_ratio` and `mlp_mult_factor`.
+
+#### What to leave out
+
+Anything absent from the plan stays replicated on every rank, which is the right choice for normalization layers, AdaLN modulation (`img_mod`/`txt_mod`), patch and text embeddings, and the final `norm_out`/`proj_out`. These are small, so sharding them saves little memory while adding communication.
+
+#### Constraints and verification
+
+- `tp_degree` must divide `config.num_attention_heads`. This is validated in [`~ModelMixin.enable_parallelism`].
+- Every packed block must *individually* be divisible by `tp_degree`, not just their sum.
+
+Validate a new plan numerically rather than by eye: generate with a fixed seed on a single device, then again under tensor parallelism, and compare the outputs. A misplaced `"colwise"`/`"rowwise"` usually still runs and produces a plausible but wrong image.
+
+> [!TIP]
+> Start from an existing plan for a similar architecture. [`QwenImageTransformer2DModel`] is fully unfused and every entry is plain `"colwise"`/`"rowwise"`, [`FluxTransformer2DModel`] adds a single packed row-wise projection, and [`Flux2Transformer2DModel`] covers both packed styles — see [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py), [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py), and [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) respectively.
+
+## Choosing a strategy
+
+The strategies above solve different problems, and the useful question is not which is fastest in the abstract but what you are running out of.
+
+| Strategy | Splits | Reduces | Latency for one prompt | Best when |
+|---|---|---|---|---|
+| [Accelerate](#accelerate) / [DDP](#pytorch-distributed) | prompts across replicas | nothing — each device holds a full copy | unchanged | the model already fits and you have many prompts |
+| [`device_map`](#device_map) | components across devices | weight memory | slightly worse | the model doesn't fit and the interconnect is slow |
+| [Context parallelism](#context-parallelism) | the input sequence | activation memory | lower | sequences are long — high resolution or video |
+| [Tensor parallelism](#tensor-parallelism) | weight matrices | weight memory | lower | one component's weights don't fit and the interconnect is fast |
+
+Some practical guidance:
+
+- **Throughput on many prompts, model already fits.** Use data parallelism. It is the only strategy here that scales throughput linearly without touching the model, and it leaves single-prompt latency alone.
+- **A single component's weights don't fit.** Reach for tensor parallelism first, since it lowers both memory and latency. It communicates at every block boundary, so it wants a fast interconnect like NVLink; over PCIe that per-layer traffic can outweigh the compute it saves, and `device_map` becomes the better choice. `device_map` also handles the case where the components are individually fine but collectively too large.
+- **Activations, not weights, are the problem.** This is the long-sequence regime — large images, many frames — and context parallelism is the direct answer. Pick a backend from the [Ulysses/Ring benchmarks](#ulysses-attention) above.
+- **Both weights and sequence are too large.** Combine tensor and context parallelism. [`TensorParallelConfig`] accepts a `mesh` argument so both can share one device mesh. This combination is not yet validated, so treat it as experimental.

src/diffusers/__init__.py

@@ -347,6 +347,7 @@ else:
             "StableCascadeUNet",
             "T2IAdapter",
             "T5FilmDecoder",
+            "TensorParallelConfig",
             "Transformer2DModel",
             "TransformerTemporalModel",
             "UNet1DModel",
@@ -1216,6 +1217,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             StableAudioDiTModel,
             T2IAdapter,
             T5FilmDecoder,
+            TensorParallelConfig,
             Transformer2DModel,
             TransformerTemporalModel,
             UNet1DModel,

src/diffusers/hooks/__init__.py

@@ -27,4 +27,5 @@ if is_torch_available():
     from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
     from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
     from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
+    from .tensor_parallel import apply_tensor_parallel
     from .text_kv_cache import TextKVCacheConfig, apply_text_kv_cache

src/diffusers/hooks/tensor_parallel.py

@@ -0,0 +1,273 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from ..models._modeling_parallel import TensorParallelConfig
+from ..utils import get_logger
+
+
+logger = get_logger(__name__)  # pylint: disable=invalid-name
+
+_SUPPORTED_TP_DEVICES = ("cuda", "neuron")
+
+
+class PackedColwiseParallel:
+    """Column-wise sharding for fused projections with heterogeneous block structure.
+
+    `blocks` is a list of proportional integers whose sum divides the weight's row count. For example, `[1, 1]` for a
+    SwiGLU gate+linear projection (two equal halves) or `[1, 1, 1, 3, 3]` for a Q+K+V+gate+linear projection with
+    `mlp_ratio=3`. If `blocks` is `None`, the Linear module must carry a `_tp_packed_col_blocks` attribute set during
+    model `__init__`.
+    """
+
+    def __init__(self, blocks: "list[int] | None" = None):
+        self.blocks = blocks
+
+
+class PackedRowwiseParallel:
+    """Row-wise sharding for fused projections with heterogeneous block structure.
+
+    `blocks` describes the input-column partition of the fused Linear (e.g. `[1, 3]` when the input concatenates an
+    attention projection and an MLP projection with `mlp_ratio=3`). If `blocks` is `None`, the module must carry a
+    `_tp_packed_row_blocks` attribute.
+    """
+
+    def __init__(self, blocks: "list[int] | None" = None):
+        self.blocks = blocks
+
+
+def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]":
+    """Convert proportional block counts to absolute sizes.
+
+    `blocks` is a list of positive integers interpreted as proportional weights. Their sum must divide `total_size`
+    evenly. Returns a list of absolute sizes that sum to `total_size`. For example, `_blocks_to_block_sizes(1152, [1,
+    1, 1, 3, 3])` returns `[128, 128, 128, 384, 384]`.
+    """
+    total = sum(blocks)
+    if total_size % total != 0:
+        raise ValueError(
+            f"Cannot split {total_size} into proportional blocks {blocks}: "
+            f"sum({blocks})={total} does not divide {total_size}."
+        )
+    unit = total_size // total
+    return [b * unit for b in blocks]
+
+
+def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list:
+    """Group a flat `_tp_plan` into per-block `(submodule, {relative_path: style})` plans.
+
+    Each glob is split at its single `*`; the prefix must resolve to a `ModuleList` and the suffix is the per-element
+    key. Grouping by block lets the caller issue one `parallelize_module` call per block, which `RowwiseParallel` needs
+    to attach its input redistribution at the block boundary.
+
+    Example: when `transformer_blocks` is a `ModuleList` of length 2, the input `{"transformer_blocks.*.ff.linear_out":
+    "rowwise"}` returns `[(transformer_blocks[0], {"ff.linear_out": "rowwise"}), (transformer_blocks[1],
+    {"ff.linear_out": "rowwise"})]`.
+    """
+    grouped: dict[int, tuple] = {}
+    order: list[int] = []
+
+    for pattern, style in tp_plan.items():
+        if pattern.count("*") > 1:
+            raise ValueError(f"Wildcard '*' can only be used once in a `_tp_plan` key, got '{pattern}'.")
+
+        if "*" in pattern:
+            prefix, _, suffix = pattern.partition("*")
+            container = model
+            for atom in prefix.strip(".").split("."):
+                container = getattr(container, atom)
+            if not isinstance(container, torch.nn.ModuleList):
+                raise ValueError(
+                    f"`_tp_plan` wildcard '{pattern}' must expand over a `ModuleList`, but "
+                    f"'{prefix.strip('.')}' resolved to '{container.__class__.__name__}'."
+                )
+            relative, blocks = suffix.strip("."), list(container)
+        else:
+            relative, blocks = pattern, [model]
+
+        for block in blocks:
+            key = id(block)
+            if key not in grouped:
+                grouped[key] = (block, {})
+                order.append(key)
+            grouped[key][1][relative] = style
+
+    return [grouped[key] for key in order]
+
+
+def _styles(relative_plan: dict) -> dict:
+    """Map a `{relative_path: style}` plan to `parallelize_module` style instances.
+
+    Values may be plain strings (`"colwise"` / `"rowwise"`) or `PackedColwiseParallel` / `PackedRowwiseParallel` marker
+    instances. Returns `{relative_path: ColwiseParallel() | RowwiseParallel() | <packed impl>}`, each subclassed to
+    reject a sharded dim that is not divisible by the TP degree.
+    """
+    import torch.nn as nn
+    from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    def _make_packed_col(marker: PackedColwiseParallel) -> ColwiseParallel:
+        _blocks = marker.blocks
+
+        class _PackedColwiseImpl(ColwiseParallel):
+            def _partition_linear_fn(self, name, module, device_mesh):
+                blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_col_blocks")
+                rank = device_mesh.get_local_rank()
+                tp_size = device_mesh.size()
+                # Both weight (`[out, in]`) and bias (`[out]`) are sharded row-wise (dim 0) with the same per-block
+                # slicing so each rank's bias rows line up with its weight rows for the packed layout.
+                for param_name, param in module.named_parameters():
+                    full = distribute_tensor(
+                        param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank
+                    ).to_local()
+                    block_sizes = _blocks_to_block_sizes(full.shape[0], blocks)
+                    parts, offset = [], 0
+                    for bs in block_sizes:
+                        if bs % tp_size != 0:
+                            raise ValueError(
+                                f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
+                                f"{bs} is not divisible by {tp_size}."
+                            )
+                        chunk = bs // tp_size
+                        parts.append(full[offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
+                        offset += bs
+                    local = torch.cat(parts, dim=0)
+                    dist_param = nn.Parameter(
+                        DTensor.from_local(local, device_mesh, [Shard(0)], run_check=False),
+                        requires_grad=param.requires_grad,
+                    )
+                    module.register_parameter(param_name, dist_param)
+
+        return _PackedColwiseImpl()
+
+    def _make_packed_row(marker: PackedRowwiseParallel) -> RowwiseParallel:
+        _blocks = marker.blocks
+
+        class _PackedRowwiseImpl(RowwiseParallel):
+            def _partition_linear_fn(self, name, module, device_mesh):
+                blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_row_blocks")
+                rank = device_mesh.get_local_rank()
+                tp_size = device_mesh.size()
+                for param_name, param in module.named_parameters():
+                    if param_name == "weight":
+                        full = distribute_tensor(
+                            param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank
+                        ).to_local()
+                        block_sizes = _blocks_to_block_sizes(full.shape[1], blocks)
+                        parts, offset = [], 0
+                        for bs in block_sizes:
+                            if bs % tp_size != 0:
+                                raise ValueError(
+                                    f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
+                                    f"{bs} is not divisible by {tp_size}."
+                                )
+                            chunk = bs // tp_size
+                            parts.append(full[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
+                            offset += bs
+                        local = torch.cat(parts, dim=1)
+                        dist_param = nn.Parameter(
+                            DTensor.from_local(local, device_mesh, [Shard(1)], run_check=False),
+                            requires_grad=param.requires_grad,
+                        )
+                    else:
+                        dist_param = nn.Parameter(
+                            distribute_tensor(param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank),
+                            requires_grad=param.requires_grad,
+                        )
+                    module.register_parameter(param_name, dist_param)
+
+        return _PackedRowwiseImpl()
+
+    # `distribute_tensor` accepts an indivisible shard dim and just gives the trailing ranks a smaller (or empty)
+    # slice, so an uneven split does not raise here — it surfaces much later as a shape or numerics error, because
+    # the attention head split and the paired colwise/rowwise Linear both assume equal shards. Reject it up front,
+    # matching what the packed styles above and the Neuron pre-shard path already do.
+    def _make_checked_col(path: str) -> ColwiseParallel:
+        class _CheckedColwiseImpl(ColwiseParallel):
+            def _partition_linear_fn(self, name, module, device_mesh):
+                tp_size = device_mesh.size()
+                out_features = module.weight.shape[0]
+                if out_features % tp_size != 0:
+                    raise ValueError(
+                        f"Cannot colwise-shard '{path}' weight rows ({out_features}) across {tp_size} "
+                        f"tensor-parallel ranks: not divisible by {tp_size}."
+                    )
+                super()._partition_linear_fn(name, module, device_mesh)
+
+        return _CheckedColwiseImpl()
+
+    def _make_checked_row(path: str) -> RowwiseParallel:
+        class _CheckedRowwiseImpl(RowwiseParallel):
+            def _partition_linear_fn(self, name, module, device_mesh):
+                tp_size = device_mesh.size()
+                in_features = module.weight.shape[1]
+                if in_features % tp_size != 0:
+                    raise ValueError(
+                        f"Cannot rowwise-shard '{path}' weight columns ({in_features}) across {tp_size} "
+                        f"tensor-parallel ranks: not divisible by {tp_size}."
+                    )
+                super()._partition_linear_fn(name, module, device_mesh)
+
+        return _CheckedRowwiseImpl()
+
+    resolved = {}
+    for path, style in relative_plan.items():
+        if style == "colwise":
+            resolved[path] = _make_checked_col(path)
+        elif style == "rowwise":
+            resolved[path] = _make_checked_row(path)
+        elif isinstance(style, PackedColwiseParallel):
+            resolved[path] = _make_packed_col(style)
+        elif isinstance(style, PackedRowwiseParallel):
+            resolved[path] = _make_packed_row(style)
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+    return resolved
+
+
+def apply_tensor_parallel(
+    model: torch.nn.Module,
+    config: TensorParallelConfig,
+    tp_plan: dict,
+) -> None:
+    """Apply tensor parallel on a model from its flat `_tp_plan`."""
+    tp_mesh = config._mesh
+    if tp_mesh is None:
+        raise ValueError("`config._mesh` is None. Call `config.setup(rank, world_size, device)` before applying TP.")
+
+    if tp_mesh.device_type not in _SUPPORTED_TP_DEVICES:
+        raise ValueError(
+            f"Tensor parallelism is not supported on device type '{tp_mesh.device_type}'. Supported device types are "
+            f"{list(_SUPPORTED_TP_DEVICES)}. The device type comes from the `mesh` passed to `TensorParallelConfig`, "
+            f"or from the active accelerator when the mesh is built from `tp_degree`."
+        )
+
+    backend = "neuron" if tp_mesh.device_type == "neuron" else "default"
+    groups = _resolve_tp_plan(model, tp_plan)
+    logger.debug(f"Applying tensor parallel (backend={backend}) over {len(groups)} module group(s) on mesh {tp_mesh}.")
+
+    if backend == "neuron":
+        from .tensor_parallel_neuron import _apply_tp_neuron
+
+        _apply_tp_neuron(model, tp_mesh, groups)
+        return
+
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    for submodule, relative_plan in groups:
+        parallelize_module(submodule, tp_mesh, _styles(relative_plan))

src/diffusers/hooks/tensor_parallel_neuron.py

@@ -0,0 +1,178 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Neuron backend for tensor parallelism, dispatched from `apply_tensor_parallel` when the TP mesh is on Neuron.
+
+The difference from the generic path is a workaround for a Neuron NRT bug: consecutive `reduce_scatter` collectives for
+large weight tensors (≥ 5120×5120) can fail when all layers are distributed in a single `parallelize_module` call. The
+fix is to pre-shard each weight locally on CPU via `DTensor.from_local` *before* calling `parallelize_module`; the
+latter then sees already-placed DTensors, skips the collective for weights, but still registers the required
+input/output hooks for the forward pass.
+"""
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+
+def _neuron_styles(relative_plan: dict) -> dict:
+    """Map a `{relative_path: style}` plan to no-op-partition styles for Neuron.
+
+    Weights (and biases) are pre-sharded in `_pre_shard_and_tp`, so `parallelize_module` runs only to register the
+    forward hooks; `_partition_linear_fn` must not re-partition. Packed and plain styles share hook behavior, so both
+    collapse onto the two no-op styles.
+    """
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
+
+    class _NeuronColwise(ColwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(0) via DTensor.from_local; parallelize_module runs only for the hooks
+
+    class _NeuronRowwise(RowwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(1) via DTensor.from_local; parallelize_module runs only for the hooks
+
+    resolved = {}
+    for path, style in relative_plan.items():
+        if style == "colwise" or isinstance(style, PackedColwiseParallel):
+            resolved[path] = _NeuronColwise()
+        elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
+            resolved[path] = _NeuronRowwise()
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+    return resolved
+
+
+def _pre_shard_and_tp(
+    module: nn.Module,
+    tp_mesh: "torch.distributed.device_mesh.DeviceMesh",
+    original_plan: dict,
+    rank: int,
+    tp_size: int,
+) -> None:
+    """Pre-shard Linear weights via `DTensor.from_local`, then call `parallelize_module`.
+
+    Workaround for a Neuron NRT bug where consecutive `reduce_scatter` calls for large weight tensors (≥ 5120×5120)
+    fail when all layers are distributed in a single `parallelize_module` call. Pre-sharding each weight on CPU means
+    it is already an on-device DTensor when `parallelize_module` runs (via `_neuron_styles`), so the collective is
+    skipped while the forward hooks are still registered.
+    """
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel, _blocks_to_block_sizes
+
+    device = torch.neuron.current_device()
+
+    for path, orig_style in original_plan.items():
+        # Resolve nested attribute path (e.g. "attn.to_q" or "attn.to_out.0")
+        submod = module
+        for part in path.split("."):
+            submod = getattr(submod, part)
+
+        if not hasattr(submod, "weight"):
+            raise ValueError(f"`_tp_plan` entry '{path}' does not resolve to a module with a `weight` parameter.")
+
+        w = submod.weight.data  # CPU at this point
+        b = submod.bias.data if submod.bias is not None else None
+        if isinstance(orig_style, PackedColwiseParallel):
+            blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_col_blocks")
+            block_sizes = _blocks_to_block_sizes(w.shape[0], blocks)
+            parts, bias_parts, offset = [], [], 0
+            for bs in block_sizes:
+                if bs % tp_size != 0:
+                    raise ValueError(
+                        f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
+                        f"{bs} is not divisible by {tp_size}."
+                    )
+                chunk = bs // tp_size
+                sl = slice(offset + rank * chunk, offset + (rank + 1) * chunk)
+                parts.append(w[sl, :].contiguous())
+                if b is not None:
+                    bias_parts.append(b[sl].contiguous())
+                offset += bs
+            shard = torch.cat(parts, dim=0).to(device)
+            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(0)]))
+            if b is not None:
+                bias_shard = torch.cat(bias_parts, dim=0).to(device)
+                submod.bias = nn.Parameter(DTensor.from_local(bias_shard, tp_mesh, [Shard(0)]))
+        elif isinstance(orig_style, PackedRowwiseParallel):
+            blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_row_blocks")
+            block_sizes = _blocks_to_block_sizes(w.shape[1], blocks)
+            parts, offset = [], 0
+            for bs in block_sizes:
+                if bs % tp_size != 0:
+                    raise ValueError(
+                        f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
+                        f"{bs} is not divisible by {tp_size}."
+                    )
+                chunk = bs // tp_size
+                parts.append(w[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
+                offset += bs
+            shard = torch.cat(parts, dim=1).to(device)
+            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)]))
+            if b is not None:  # rowwise bias is added post-reduction → keep it replicated
+                submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()]))
+        elif orig_style == "colwise":
+            if w.shape[0] % tp_size != 0:
+                raise ValueError(
+                    f"Cannot colwise-shard '{path}' weight rows ({w.shape[0]}) across {tp_size} "
+                    f"tensor-parallel ranks: not divisible by {tp_size}."
+                )
+            rows = w.shape[0] // tp_size
+            sl = slice(rank * rows, (rank + 1) * rows)
+            submod.weight = nn.Parameter(DTensor.from_local(w[sl, :].contiguous().to(device), tp_mesh, [Shard(0)]))
+            if b is not None:
+                submod.bias = nn.Parameter(DTensor.from_local(b[sl].contiguous().to(device), tp_mesh, [Shard(0)]))
+        elif orig_style == "rowwise":
+            if w.shape[1] % tp_size != 0:
+                raise ValueError(
+                    f"Cannot rowwise-shard '{path}' weight columns ({w.shape[1]}) across {tp_size} "
+                    f"tensor-parallel ranks: not divisible by {tp_size}."
+                )
+            cols = w.shape[1] // tp_size
+            shard = w[:, rank * cols : (rank + 1) * cols].contiguous().to(device)
+            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)]))
+            if b is not None:  # rowwise bias is added post-reduction → keep it replicated
+                submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()]))
+
+    # parallelize_module is now a no-op for weight distribution (already DTensors)
+    # but still registers the input/output hooks required for the forward pass.
+    parallelize_module(module, tp_mesh, _neuron_styles(original_plan))
+
+
+def _apply_tp_neuron(
+    model: nn.Module,
+    tp_mesh: "torch.distributed.device_mesh.DeviceMesh",
+    groups: list,
+) -> None:
+    """Apply tensor parallelism on Neuron from resolved `_tp_plan` groups.
+
+    `groups` is produced by `diffusers.hooks.tensor_parallel._resolve_tp_plan` — the same source of truth used by the
+    generic path, so the two backends shard identical layers. For each `(block, relative_plan)` group this pre-shards
+    the weights via `DTensor.from_local` (Neuron NRT consecutive-reduce-scatter workaround), then calls
+    `parallelize_module` to register the forward hooks.
+
+    Model weights must be on CPU when this is called.
+    """
+    rank = dist.get_rank()
+    tp_size = tp_mesh.size()
+
+    for block, relative_plan in groups:
+        _pre_shard_and_tp(block, tp_mesh, relative_plan, rank, tp_size)

src/diffusers/models/__init__.py

@@ -24,7 +24,7 @@ from ..utils import (
 _import_structure = {}
 
 if is_torch_available():
-    _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig"]
+    _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig", "TensorParallelConfig"]
     _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
     _import_structure["attention_dispatch"] = ["AttentionBackendName", "attention_backend"]
     _import_structure["auto_model"] = ["AutoModel"]
@@ -169,7 +169,7 @@ if is_torch_available():
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     if is_torch_available():
-        from ._modeling_parallel import ContextParallelConfig, ParallelConfig
+        from ._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig
         from .adapter import MultiAdapter, T2IAdapter
         from .attention_dispatch import AttentionBackendName, attention_backend
         from .auto_model import AutoModel

src/diffusers/models/_modeling_parallel.py

@@ -35,7 +35,6 @@ logger = get_logger(__name__)  # pylint: disable=invalid-name
 # - Unified Attention
 # - More dispatcher attention backends
 # - CFG/Data Parallel
-# - Tensor Parallel
 
 
 @dataclass
@@ -154,6 +153,46 @@ class ContextParallelConfig:
         self._ulysses_local_rank = self._ulysses_mesh.get_local_rank()
 
 
+@dataclass
+class TensorParallelConfig:
+    """
+    Configuration for tensor parallelism.
+
+    Tensor parallelism shards weight matrices (column-wise and row-wise) across devices. Each device computes a partial
+    result; an AllReduce/AllGather at layer boundaries reconstructs the full output. Uses
+    `torch.distributed.tensor.parallelize_module` with `ColwiseParallel` / `RowwiseParallel` sharding styles. Supported
+    device types are `"cuda"` and `"neuron"`.
+
+    Args:
+        tp_degree (`int`, defaults to `1`):
+            Number of devices to shard across. Must be a divisor of the number of attention heads (and FFN hidden
+            dimensions) of the model being parallelised.
+        mesh (`torch.distributed.device_mesh.DeviceMesh`, *optional*):
+            A custom device mesh to use. If provided, `tp_degree` is inferred from `mesh.size()` and the argument is
+            ignored. Useful when combining TP with other parallelism strategies (e.g. CP) that share the same mesh.
+    """
+
+    tp_degree: int = 1
+    mesh: torch.distributed.device_mesh.DeviceMesh | None = None
+
+    _rank: int = None
+    _world_size: int = None
+    _device: torch.device = None
+    _mesh: torch.distributed.device_mesh.DeviceMesh = None
+    _tp_degree: int = None
+
+    def __post_init__(self):
+        if self.tp_degree < 1:
+            raise ValueError("`tp_degree` must be >= 1.")
+
+    def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.distributed.device_mesh.DeviceMesh):
+        self._rank = rank
+        self._world_size = world_size
+        self._device = device
+        self._mesh = mesh
+        self._tp_degree = mesh.size()
+
+
 @dataclass
 class ParallelConfig:
     """
@@ -162,15 +201,25 @@ class ParallelConfig:
     Args:
         context_parallel_config (`ContextParallelConfig`, *optional*):
             Configuration for context parallelism.
+        tensor_parallel_config (`TensorParallelConfig`, *optional*):
+            Configuration for tensor parallelism.
     """
 
     context_parallel_config: ContextParallelConfig | None = None
+    tensor_parallel_config: TensorParallelConfig | None = None
 
     _rank: int = None
     _world_size: int = None
     _device: torch.device = None
     _mesh: torch.distributed.device_mesh.DeviceMesh = None
 
+    def __post_init__(self):
+        if self.context_parallel_config is not None and self.tensor_parallel_config is not None:
+            raise ValueError(
+                "Combining context parallelism and tensor parallelism in a single `ParallelConfig` is not supported. "
+                "Please specify only one of `context_parallel_config` or `tensor_parallel_config`."
+            )
+
     def setup(
         self,
         rank: int,
@@ -185,6 +234,8 @@ class ParallelConfig:
         self._mesh = mesh
         if self.context_parallel_config is not None:
             self.context_parallel_config.setup(rank, world_size, device, mesh)
+        if self.tensor_parallel_config is not None:
+            self.tensor_parallel_config.setup(rank, world_size, device, mesh)
 
 
 @dataclass(frozen=True)

src/diffusers/models/modeling_utils.py

@@ -63,7 +63,12 @@ from ..utils import (
 from ..utils.distributed_utils import is_torch_dist_rank_zero
 from ..utils.hub_utils import PushToHubMixin, load_or_create_model_card, populate_model_card
 from ..utils.torch_utils import empty_device_cache
-from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig
+from ._modeling_parallel import (
+    ContextParallelConfig,
+    ContextParallelModelPlan,
+    ParallelConfig,
+    TensorParallelConfig,
+)
 from .model_loading_utils import (
     _caching_allocator_warmup,
     _determine_device_map,
@@ -254,6 +259,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
     _repeated_blocks = []
     _parallel_config = None
     _cp_plan = None
+    _tp_plan = None
     _skip_keys = None
 
     def __init__(self):
@@ -1601,7 +1607,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
     def enable_parallelism(
         self,
         *,
-        config: ParallelConfig | ContextParallelConfig,
+        config: ParallelConfig | ContextParallelConfig | TensorParallelConfig,
         cp_plan: dict[str, ContextParallelModelPlan] | None = None,
     ):
         logger.warning(
@@ -1620,6 +1626,8 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
 
         if isinstance(config, ContextParallelConfig):
             config = ParallelConfig(context_parallel_config=config)
+        elif isinstance(config, TensorParallelConfig):
+            config = ParallelConfig(tensor_parallel_config=config)
 
         rank = torch.distributed.get_rank()
         world_size = torch.distributed.get_world_size()
@@ -1665,19 +1673,30 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
                 mesh_shape=cp_config.mesh_shape,
                 mesh_dim_names=cp_config.mesh_dim_names,
             )
+        elif config.tensor_parallel_config is not None:
+            tp_config = config.tensor_parallel_config
+            mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
+                device_type=device_type,
+                mesh_shape=(tp_config.tp_degree,),
+                mesh_dim_names=("tp",),
+            )
 
+        # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`.
         config.setup(rank, world_size, device, mesh=mesh)
         self._parallel_config = config
 
-        for module in self.modules():
-            if not isinstance(module, attention_classes):
-                continue
-            processor = module.processor
-            if processor is None or not hasattr(processor, "_parallel_config"):
-                continue
-            processor._parallel_config = config
-
+        # Only context parallelism needs the config inside attention: it replaces the attention computation itself
+        # (Ulysses all-to-all / ring). Tensor parallelism only shards `Linear` weights, so each rank runs the ordinary
+        # attention op over its own heads and the processors must stay unaware of it.
         if config.context_parallel_config is not None:
+            for module in self.modules():
+                if not isinstance(module, attention_classes):
+                    continue
+                processor = module.processor
+                if processor is None or not hasattr(processor, "_parallel_config"):
+                    continue
+                processor._parallel_config = config
+
             if cp_plan is None and self._cp_plan is None:
                 raise ValueError(
                     "`cp_plan` must be provided either as an argument or set in the model's `_cp_plan` attribute."
@@ -1685,6 +1704,21 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
             cp_plan = cp_plan if cp_plan is not None else self._cp_plan
             apply_context_parallel(self, config.context_parallel_config, cp_plan)
 
+        if config.tensor_parallel_config is not None:
+            if self._tp_plan is None:
+                raise ValueError(
+                    "`_tp_plan` must be set on the model class to use tensor parallelism. "
+                    f"'{self.__class__.__name__}' does not define one."
+                )
+            tp_degree = config.tensor_parallel_config._tp_degree
+            num_heads = getattr(self.config, "num_attention_heads", None)
+            if num_heads is not None and num_heads % tp_degree != 0:
+                raise ValueError(f"`tp_degree` ({tp_degree}) must divide the number of attention heads ({num_heads}).")
+
+            from ..hooks.tensor_parallel import apply_tensor_parallel
+
+            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
+
     @classmethod
     def _load_pretrained_model(
         cls,

src/diffusers/models/transformers/transformer_bria_fibo.py

@@ -64,7 +64,6 @@ def _get_qkv_projections(attn: "BriaFiboAttention", hidden_states, encoder_hidde
     return _get_projections(attn, hidden_states, encoder_hidden_states)
 
 
-# Copied from diffusers.models.transformers.transformer_flux.FluxAttnProcessor with FluxAttnProcessor->BriaFiboAttnProcessor, FluxAttention->BriaFiboAttention
 class BriaFiboAttnProcessor:
     _attention_backend = None
     _parallel_config = None

src/diffusers/models/transformers/transformer_flux.py

@@ -21,6 +21,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
+from ...hooks.tensor_parallel import PackedRowwiseParallel
 from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import apply_lora_scale, logging
 from ...utils.torch_utils import maybe_adjust_dtype_for_device, maybe_allow_in_graph
@@ -92,17 +93,19 @@ class FluxAttnProcessor:
             attn, hidden_states, encoder_hidden_states
         )
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank
+        # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic.
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
 
         if attn.added_kv_proj_dim is not None:
-            encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim))
+            encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim))
+            encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim))
 
             encoder_query = attn.norm_added_q(encoder_query)
             encoder_key = attn.norm_added_k(encoder_key)
@@ -362,6 +365,9 @@ class FluxSingleTransformerBlock(nn.Module):
         self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
         self.act_mlp = nn.GELU(approximate="tanh")
         self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+        # proj_out input concatenates the attention output (`dim`) and the MLP branch (`mlp_hidden_dim`); each is
+        # independently column-sharded, so PackedRowwiseParallel splits its input columns along these block sizes.
+        self.proj_out._tp_packed_row_blocks = [dim, self.mlp_hidden_dim]
 
         self.attn = FluxAttention(
             query_dim=dim,
@@ -573,6 +579,35 @@ class FluxTransformer2DModel(
         },
         "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
     }
+    # Tensor-parallel plan: how each block's fused/unfused Linears shard across the TP mesh. Flux1 is unfused, so
+    # nearly everything is plain "colwise"/"rowwise" (torch's ColwiseParallel / RowwiseParallel). The one exception is
+    # the single-stream `proj_out`, whose input is `cat([attn_output, mlp_hidden_states])` — two independently
+    # column-sharded streams — so it needs PackedRowwiseParallel with block sizes stored on the Linear at init.
+    # AdaLN modulation (norm1/norm1_context, the single-block norm, norm_out) and the QK-norms stay replicated
+    # (intentionally absent here): modulation indexes the full hidden dim, and QK-norm applies over head_dim after
+    # the heads are already split, so every rank needs the whole vector.
+    _tp_plan = {
+        # double-stream (MMDiT) blocks
+        "transformer_blocks.*.attn.to_q": "colwise",
+        "transformer_blocks.*.attn.to_k": "colwise",
+        "transformer_blocks.*.attn.to_v": "colwise",
+        "transformer_blocks.*.attn.to_out.0": "rowwise",
+        "transformer_blocks.*.attn.add_q_proj": "colwise",
+        "transformer_blocks.*.attn.add_k_proj": "colwise",
+        "transformer_blocks.*.attn.add_v_proj": "colwise",
+        "transformer_blocks.*.attn.to_add_out": "rowwise",
+        "transformer_blocks.*.ff.net.0.proj": "colwise",
+        "transformer_blocks.*.ff.net.2": "rowwise",
+        "transformer_blocks.*.ff_context.net.0.proj": "colwise",
+        "transformer_blocks.*.ff_context.net.2": "rowwise",
+        # single-stream (parallel attention + MLP) blocks
+        "single_transformer_blocks.*.attn.to_q": "colwise",
+        "single_transformer_blocks.*.attn.to_k": "colwise",
+        "single_transformer_blocks.*.attn.to_v": "colwise",
+        "single_transformer_blocks.*.proj_mlp": "colwise",
+        # input = cat([attn_output (inner_dim), mlp_hidden_states (mlp_hidden_dim)]); blocks set in __init__
+        "single_transformer_blocks.*.proj_out": PackedRowwiseParallel(),
+    }
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_flux2.py

@@ -21,6 +21,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
+from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
 from ...loaders import FluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import BaseOutput, apply_lora_scale, logging
 from ...utils.torch_utils import maybe_adjust_dtype_for_device
@@ -343,17 +344,19 @@ class Flux2AttnProcessor:
             attn, hidden_states, encoder_hidden_states
         )
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank
+        # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic.
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
 
         if attn.added_kv_proj_dim is not None:
-            encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim))
+            encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim))
+            encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim))
 
             encoder_query = attn.norm_added_q(encoder_query)
             encoder_key = attn.norm_added_k(encoder_key)
@@ -423,17 +426,17 @@ class Flux2KVAttnProcessor:
             attn, hidden_states, encoder_hidden_states
         )
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
 
         if attn.added_kv_proj_dim is not None:
-            encoder_query = encoder_query.unflatten(-1, (attn.heads, -1))
-            encoder_key = encoder_key.unflatten(-1, (attn.heads, -1))
-            encoder_value = encoder_value.unflatten(-1, (attn.heads, -1))
+            encoder_query = encoder_query.unflatten(-1, (-1, attn.head_dim))
+            encoder_key = encoder_key.unflatten(-1, (-1, attn.head_dim))
+            encoder_value = encoder_value.unflatten(-1, (-1, attn.head_dim))
 
             encoder_query = attn.norm_added_q(encoder_query)
             encoder_key = attn.norm_added_k(encoder_key)
@@ -583,16 +586,21 @@ class Flux2ParallelSelfAttnProcessor:
     ) -> torch.Tensor:
         # Parallel in (QKV + MLP in) projection
         hidden_states = attn.to_qkv_mlp_proj(hidden_states)
-        qkv, mlp_hidden_states = torch.split(
-            hidden_states, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1
-        )
+
+        # Split the fused output into its QKV and MLP halves by their global ratio, so the same code path works
+        # whether or not `to_qkv_mlp_proj` is column-sharded (PackedColwiseParallel shards each block, so every
+        # rank's contiguous slice keeps the QKV and MLP chunks proportional). No tensor-parallel state is read here.
+        qkv_dim = 3 * attn.inner_dim
+        mlp_dim = attn.mlp_hidden_dim * attn.mlp_mult_factor
+        local_qkv = hidden_states.shape[-1] * qkv_dim // (qkv_dim + mlp_dim)
+        qkv, mlp_hidden_states = torch.split(hidden_states, [local_qkv, hidden_states.shape[-1] - local_qkv], dim=-1)
 
         # Handle the attention logic
         query, key, value = qkv.chunk(3, dim=-1)
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
@@ -651,15 +659,21 @@ class Flux2KVParallelSelfAttnProcessor:
     ) -> torch.Tensor:
         # Parallel in (QKV + MLP in) projection
         hidden_states_proj = attn.to_qkv_mlp_proj(hidden_states)
+
+        # Split by the global QKV:MLP ratio so the path is identical with or without column sharding (see
+        # `Flux2ParallelSelfAttnProcessor`). No tensor-parallel state is read here.
+        qkv_dim = 3 * attn.inner_dim
+        mlp_dim = attn.mlp_hidden_dim * attn.mlp_mult_factor
+        local_qkv = hidden_states_proj.shape[-1] * qkv_dim // (qkv_dim + mlp_dim)
         qkv, mlp_hidden_states = torch.split(
-            hidden_states_proj, [3 * attn.inner_dim, attn.mlp_hidden_dim * attn.mlp_mult_factor], dim=-1
+            hidden_states_proj, [local_qkv, hidden_states_proj.shape[-1] - local_qkv], dim=-1
         )
 
         query, key, value = qkv.chunk(3, dim=-1)
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
@@ -754,6 +768,10 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
         self.to_qkv_mlp_proj = torch.nn.Linear(
             self.query_dim, self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias
         )
+        # Block structure for tensor-parallel packed sharding: Q, K, V, gate, linear
+        self.to_qkv_mlp_proj._tp_packed_col_blocks = [self.inner_dim] * 3 + [
+            self.mlp_hidden_dim
+        ] * self.mlp_mult_factor
         self.mlp_act_fn = Flux2SwiGLU()
 
         # QK Norm
@@ -762,6 +780,8 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
 
         # Fused attention output projection + MLP output projection
         self.to_out = torch.nn.Linear(self.inner_dim + self.mlp_hidden_dim, self.out_dim, bias=out_bias)
+        # Block structure for tensor-parallel packed sharding: attn columns, mlp columns
+        self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
 
         if processor is None:
             processor = self._default_processor_cls()
@@ -1090,6 +1110,34 @@ class Flux2Transformer2DModel(
         "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
     }
 
+    # Tensor-parallel sharding plan: a flat mapping of module-name globs (relative to the model) to a
+    # parallel style.  Plain strings ("colwise" / "rowwise") map to torch's ColwiseParallel /
+    # RowwiseParallel.  Fused projections use PackedColwiseParallel / PackedRowwiseParallel, which
+    # take the per-block structure from either an explicit `blocks` argument or a
+    # `_tp_packed_col_blocks` / `_tp_packed_row_blocks` attribute stored on the Linear at init.
+    # AdaLN modulation (double_stream_modulation_img/_txt, single_stream_modulation, norm_out) and the QK-norms stay
+    # replicated (intentionally absent here): modulation indexes the full hidden dim, and QK-norm applies over
+    # head_dim after the heads are already split, so every rank needs the whole vector.
+    _tp_plan = {
+        # double-stream (cross-attention + FFN) blocks
+        "transformer_blocks.*.attn.to_q": "colwise",
+        "transformer_blocks.*.attn.to_k": "colwise",
+        "transformer_blocks.*.attn.to_v": "colwise",
+        "transformer_blocks.*.attn.to_out.0": "rowwise",
+        "transformer_blocks.*.attn.add_q_proj": "colwise",
+        "transformer_blocks.*.attn.add_k_proj": "colwise",
+        "transformer_blocks.*.attn.add_v_proj": "colwise",
+        "transformer_blocks.*.attn.to_add_out": "rowwise",
+        # SwiGLU gate+linear are equal halves → blocks=[1,1]
+        "transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
+        "transformer_blocks.*.ff.linear_out": "rowwise",
+        "transformer_blocks.*.ff_context.linear_in": PackedColwiseParallel([1, 1]),
+        "transformer_blocks.*.ff_context.linear_out": "rowwise",
+        # single-stream: block structure stored on the Linear by Flux2ParallelSelfAttention.__init__
+        "single_transformer_blocks.*.attn.to_qkv_mlp_proj": PackedColwiseParallel(),
+        "single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
+    }
+
     @register_to_config
     def __init__(
         self,

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import math
 from math import prod
 from typing import Any
@@ -141,6 +142,38 @@ def apply_rotary_emb_qwen(
         return x_out.type_as(x)
 
 
+def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """
+    Apply rotary embeddings to `x` using real-valued cos/sin, for backends without a complex dtype.
+
+    Numerically equivalent to `apply_rotary_emb_qwen(..., use_real=False)`, which multiplies `x` by a complex
+    exponential. Neuron has no complex tensor support, so the rotation angles are carried as reals and cos/sin are
+    taken here instead.
+
+    Args:
+        x (`torch.Tensor`): Query or key tensor to rotate, shape `[B, S, H, D]`.
+        freqs (`torch.Tensor`): Rotation angles, shape `[S, D // 2]`.
+
+    Returns:
+        `torch.Tensor`: `x` with rotary embeddings applied.
+    """
+    # Adjacent feature pairs (2k, 2k+1) share angle k, so each angle is repeated twice along the last dim; unsqueeze
+    # the head axis so the freqs broadcast over heads (this is what keeps it tensor-parallel-agnostic).
+    cos = torch.cos(freqs).repeat_interleave(2, dim=-1).unsqueeze(1)  # [S, 1, D]
+    sin = torch.sin(freqs).repeat_interleave(2, dim=-1).unsqueeze(1)  # [S, 1, D]
+    x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)  # [B, S, H, D]
+    return (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+
+# RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron cannot
+# represent. Callers select by `device.type` and fall back to the default for any backend not listed here.
+ROPE_PER_DEVICE = {
+    "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False),
+    "neuron": apply_rotary_emb_qwen_neuron,
+}
+
+
 def compute_text_seq_len_from_mask(
     encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: torch.Tensor | None
 ) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
@@ -220,7 +253,6 @@ class QwenEmbedRope(nn.Module):
             dim=1,
         )
 
-        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
         self.scale_rope = scale_rope
 
     def rope_params(self, index, dim, theta=10000):
@@ -236,6 +268,11 @@ class QwenEmbedRope(nn.Module):
     @lru_cache_unless_export(maxsize=None)
     def _get_device_freqs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
         """Return pos_freqs and neg_freqs on the given device."""
+        if device is not None and device.type == "neuron":
+            # Neuron has no complex dtype, so send the rotation angles instead and let
+            # `apply_rotary_emb_qwen_neuron` take cos/sin on device. `torch.angle` runs on CPU while the freqs are
+            # still complex; wrapping into (-pi, pi] is harmless because only cos/sin of the angle are used.
+            return torch.angle(self.pos_freqs).to(device), torch.angle(self.neg_freqs).to(device)
         return self.pos_freqs.to(device), self.neg_freqs.to(device)
 
     def forward(
@@ -360,6 +397,11 @@ class QwenEmbedLayer3DRope(nn.Module):
     @lru_cache_unless_export(maxsize=None)
     def _get_device_freqs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
         """Return pos_freqs and neg_freqs on the given device."""
+        if device is not None and device.type == "neuron":
+            # Neuron has no complex dtype, so send the rotation angles instead and let
+            # `apply_rotary_emb_qwen_neuron` take cos/sin on device. `torch.angle` runs on CPU while the freqs are
+            # still complex; wrapping into (-pi, pi] is harmless because only cos/sin of the angle are used.
+            return torch.angle(self.pos_freqs).to(device), torch.angle(self.neg_freqs).to(device)
         return self.pos_freqs.to(device), self.neg_freqs.to(device)
 
     def forward(
@@ -521,14 +563,17 @@ class QwenDoubleStreamAttnProcessor2_0:
         txt_key = attn.add_k_proj(encoder_hidden_states)
         txt_value = attn.add_v_proj(encoder_hidden_states)
 
-        # Reshape for multi-head attention
-        img_query = img_query.unflatten(-1, (attn.heads, -1))
-        img_key = img_key.unflatten(-1, (attn.heads, -1))
-        img_value = img_value.unflatten(-1, (attn.heads, -1))
+        # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank
+        # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic.
+        # The shared Attention has no `head_dim` attribute; derive it from the unsharded `inner_dim`/`heads`.
+        head_dim = attn.inner_dim // attn.heads
+        img_query = img_query.unflatten(-1, (-1, head_dim))
+        img_key = img_key.unflatten(-1, (-1, head_dim))
+        img_value = img_value.unflatten(-1, (-1, head_dim))
 
-        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
-        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
-        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+        txt_query = txt_query.unflatten(-1, (-1, head_dim))
+        txt_key = txt_key.unflatten(-1, (-1, head_dim))
+        txt_value = txt_value.unflatten(-1, (-1, head_dim))
 
         # Apply QK normalization
         if attn.norm_q is not None:
@@ -543,10 +588,11 @@ class QwenDoubleStreamAttnProcessor2_0:
         # Apply RoPE
         if image_rotary_emb is not None:
             img_freqs, txt_freqs = image_rotary_emb
-            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
-            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
-            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
-            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
+            apply_rope = ROPE_PER_DEVICE.get(img_query.device.type, ROPE_PER_DEVICE["cuda"])
+            img_query = apply_rope(img_query, img_freqs)
+            img_key = apply_rope(img_key, img_freqs)
+            txt_query = apply_rope(txt_query, txt_freqs)
+            txt_key = apply_rope(txt_key, txt_freqs)
 
         # Concatenate for joint attention
         # Order: [text, image]
@@ -790,6 +836,25 @@ class QwenImageTransformer2DModel(
         },
         "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
     }
+    # Tensor-parallel plan: how each block's fused/unfused Linears shard across the TP mesh. Qwen-Image is fully
+    # unfused (separate Q/K/V and separate img/txt MLPs), so every entry is plain "colwise"/"rowwise" (torch's
+    # ColwiseParallel / RowwiseParallel) and no packed sharding is needed. AdaLN modulation (img_mod/txt_mod) and
+    # norm_out stay replicated (intentionally absent here).
+    _tp_plan = {
+        # double-stream (MMDiT) blocks
+        "transformer_blocks.*.attn.to_q": "colwise",
+        "transformer_blocks.*.attn.to_k": "colwise",
+        "transformer_blocks.*.attn.to_v": "colwise",
+        "transformer_blocks.*.attn.to_out.0": "rowwise",
+        "transformer_blocks.*.attn.add_q_proj": "colwise",
+        "transformer_blocks.*.attn.add_k_proj": "colwise",
+        "transformer_blocks.*.attn.add_v_proj": "colwise",
+        "transformer_blocks.*.attn.to_add_out": "rowwise",
+        "transformer_blocks.*.img_mlp.net.0.proj": "colwise",
+        "transformer_blocks.*.img_mlp.net.2": "rowwise",
+        "transformer_blocks.*.txt_mlp.net.0.proj": "colwise",
+        "transformer_blocks.*.txt_mlp.net.2": "rowwise",
+    }
 
     @register_to_config
     def __init__(

src/diffusers/utils/dummy_pt_objects.py

@@ -2235,6 +2235,21 @@ class T5FilmDecoder(metaclass=DummyObject):
         requires_backends(cls, ["torch"])
 
 
+class TensorParallelConfig(metaclass=DummyObject):
+    _backends = ["torch"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+
 class Transformer2DModel(metaclass=DummyObject):
     _backends = ["torch"]
 

src/diffusers/utils/torch_utils.py

@@ -200,7 +200,7 @@ def randn_tensor(
     layout = layout or torch.strided
     device = device or torch.device("cpu")
 
-    # Neuron (XLA) does not support creating random tensors directly on device; always use CPU
+    # Neuron does not support creating random tensors directly on device; always use CPU
     if device.type == "neuron":
         rand_device = torch.device("cpu")
 

tests/conftest.py

@@ -49,6 +49,7 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "modelopt: marks tests for NVIDIA ModelOpt quantization functionality")
     config.addinivalue_line("markers", "sdnq: marks tests for SDNQ quantization functionality")
     config.addinivalue_line("markers", "context_parallel: marks tests for context parallel inference functionality")
+    config.addinivalue_line("markers", "tensor_parallel: marks tests for tensor parallel inference functionality")
     config.addinivalue_line("markers", "slow: mark test as slow")
     config.addinivalue_line("markers", "nightly: mark test as nightly")
 

tests/models/testing_utils/__init__.py

@@ -17,7 +17,11 @@ from .compile import TorchCompileTesterMixin
 from .ip_adapter import IPAdapterTesterMixin
 from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin
 from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin
-from .parallelism import ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin
+from .parallelism import (
+    ContextParallelAttentionBackendsTesterMixin,
+    ContextParallelTesterMixin,
+    TensorParallelTesterMixin,
+)
 from .quantization import (
     AutoRoundCompileTesterMixin,
     AutoRoundConfigMixin,
@@ -63,6 +67,7 @@ __all__ = [
     "CacheTesterMixin",
     "ContextParallelTesterMixin",
     "ContextParallelAttentionBackendsTesterMixin",
+    "TensorParallelTesterMixin",
     "CPUOffloadTesterMixin",
     "FasterCacheConfigMixin",
     "FasterCacheTesterMixin",

tests/models/testing_utils/parallelism.py

@@ -21,13 +21,14 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from diffusers.models._modeling_parallel import ContextParallelConfig
+from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig
 from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
 
 from ...testing_utils import (
     is_attention,
     is_context_parallel,
     is_kernels_available,
+    is_tensor_parallel,
     require_torch_multi_accelerator,
     torch_device,
 )
@@ -241,6 +242,109 @@ def _custom_mesh_worker(
             dist.destroy_process_group()
 
 
+def _tensor_parallel_worker(
+    rank, world_size, master_port, model_class, init_dict, inputs_dict, return_dict, state_dict
+):
+    """Worker function for tensor parallel inference testing.
+
+    Each rank builds the (identical, `state_dict`-loaded) model, sets up the accelerator device, shards the model
+    with `enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))` and runs a forward pass. Rank 0
+    reports its output so the caller can compare it against a single-device reference (TP is mathematically equivalent
+    to the unsharded model up to floating-point reduction order).
+    """
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
+        backend = device_config["backend"]
+        device_module = device_config["module"]
+
+        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
+
+        device_module.set_device(rank)
+        device = torch.device(f"{torch_device}:{rank}")
+
+        model = model_class(**init_dict)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+
+        # Shard the model across all ranks; the device mesh is built from `tp_degree` on the active accelerator.
+        model.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
+
+        with torch.no_grad():
+            output = model(**inputs_on_device, return_dict=False)[0]
+
+        if rank == 0:
+            return_dict["status"] = "success"
+            return_dict["output_shape"] = list(output.shape)
+            # Serialise via nested list so the manager dict can transport it across processes.
+            return_dict["output"] = output.float().cpu().tolist()
+
+    except Exception as e:
+        if rank == 0:
+            return_dict["status"] = "error"
+            return_dict["error"] = str(e)
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+@is_tensor_parallel
+@require_torch_multi_accelerator
+class TensorParallelTesterMixin:
+    def test_tensor_parallel_inference(self, batch_size: int = 1):
+        if not torch.distributed.is_available():
+            pytest.skip("torch.distributed is not available.")
+
+        if getattr(self.model_class, "_tp_plan", None) is None:
+            pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
+
+        world_size = 2
+        init_dict = self.get_init_dict()
+        num_heads = init_dict.get("num_attention_heads")
+        if num_heads is not None and num_heads % world_size != 0:
+            pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
+
+        inputs_dict = self.get_dummy_inputs(batch_size=batch_size)
+
+        # Single-device reference
+        model = self.model_class(**init_dict).eval().to(torch_device)
+        state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
+        with torch.no_grad():
+            ref_output = model(**inputs_dict, return_dict=False)[0].float().cpu()
+
+        # Move all tensors to CPU for multiprocessing
+        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+
+        master_port = _find_free_port()
+        manager = mp.Manager()
+        return_dict = manager.dict()
+
+        mp.spawn(
+            _tensor_parallel_worker,
+            args=(world_size, master_port, self.model_class, init_dict, inputs_dict, return_dict, state_dict),
+            nprocs=world_size,
+            join=True,
+        )
+
+        assert return_dict.get("status") == "success", (
+            f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}"
+        )
+
+        tp_output = torch.tensor(return_dict["output"])
+        # Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference.
+        torch.testing.assert_close(ref_output, tp_output, atol=1e-3, rtol=1e-3)
+
+    def test_tensor_parallel_batch_inputs(self):
+        self.test_tensor_parallel_inference(batch_size=2)
+
+
 @is_context_parallel
 @require_torch_multi_accelerator
 class ContextParallelTesterMixin:

tests/models/transformers/_neuron_tp_worker.py

@@ -0,0 +1,111 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generic torchrun worker: assert a model's Neuron tensor-parallel output matches its single-device reference.
+
+Model-agnostic. The model under test is supplied as a `module:function` spec reference on the command line; the
+referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors, so all model-specific test data lives
+with the launching test rather than here.
+
+Launched as a subprocess by a `@require_torch_neuron` test (and runnable directly for debugging)::
+
+    torchrun --nproc_per_node=2 _neuron_tp_worker.py \\
+        tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec
+
+Each rank builds an identical (seeded) model on CPU, computes a single-device reference, then shards it with
+`enable_parallelism(TensorParallelConfig(mesh=neuron_mesh))` — which auto-selects the Neuron pre-shard backend — runs
+a forward pass on the Neuron device, and asserts the gathered output matches the reference. Exit code 0 means the TP
+path is numerically equivalent to the unsharded model; non-zero means failure.
+"""
+
+import argparse
+import importlib
+import os
+import sys
+import traceback
+
+
+# Make the in-repo `diffusers` and `tests` packages importable when run via torchrun from an arbitrary CWD.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+
+import torch
+import torch.distributed as dist
+import torch_neuronx  # noqa: F401 — registers torch.neuron
+from torch.distributed.device_mesh import DeviceMesh
+
+from diffusers import TensorParallelConfig
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Neuron tensor-parallel correctness worker.")
+    parser.add_argument(
+        "spec",
+        help="`module:function` reference returning (model_class, init_dict, cpu_inputs) for the model under test.",
+    )
+    args = parser.parse_args()
+    module_name, _, fn_name = args.spec.partition(":")
+    model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)()
+
+    dist.init_process_group(backend="neuron")
+    rank = dist.get_rank()
+    tp_size = dist.get_world_size()
+    device = torch.neuron.current_device()
+    tp_mesh = DeviceMesh("neuron", list(range(tp_size)))
+
+    # Identical weights on every rank (same seed), kept on CPU as the Neuron pre-shard backend requires.
+    torch.manual_seed(0)
+    model = model_class(**init_dict).eval()
+
+    # Single-device (unsharded) reference on CPU, computed before TP mutates the weights in place.
+    with torch.no_grad():
+        ref_output = model(**inputs, return_dict=False)[0].float().cpu()
+
+    # Shard across all ranks; the Neuron backend is auto-selected from the mesh device type.
+    model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
+    model = model.to(device)
+    torch.neuron.synchronize()
+
+    inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+    with torch.no_grad():
+        tp_output = model(**inputs_on_device, return_dict=False)[0]
+    torch.neuron.synchronize()
+    tp_output = tp_output.float().cpu()
+
+    if rank == 0:
+        assert tp_output.shape == ref_output.shape, f"shape mismatch: {tp_output.shape} vs {ref_output.shape}"
+        assert torch.isfinite(tp_output).all(), "TP output contains non-finite values"
+        max_abs = (tp_output - ref_output).abs().max().item()
+        denom = ref_output.abs().max().item() + 1e-6
+        print(
+            f"[rank0] tp_size={tp_size} output_shape={tuple(tp_output.shape)} "
+            f"max_abs_diff={max_abs:.4e} max_rel_diff={max_abs / denom:.4e}"
+        )
+        # Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance. A wrong shard
+        # plan produces grossly different output and is caught comfortably within this bound.
+        torch.testing.assert_close(tp_output, ref_output, atol=2e-2, rtol=2e-2)
+        print("[rank0] PASS: Neuron tensor-parallel output matches single-device reference.")
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+        # Ensure a non-zero exit so the launching pytest sees the failure.
+        os._exit(1)

tests/models/transformers/test_models_transformer_flux.py

@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
+import sys
 import tempfile
 from typing import Any
 
@@ -24,7 +27,7 @@ from diffusers.models.embeddings import ImageProjection
 from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
 from ..testing_utils import (
     AttentionBackendTesterMixin,
     AttentionTesterMixin,
@@ -51,6 +54,7 @@ from ..testing_utils import (
     SDNQTesterMixin,
     SingleFileTesterMixin,
     TaylorSeerCacheTesterMixin,
+    TensorParallelTesterMixin,
     TorchAoCompileTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
@@ -157,7 +161,7 @@ class FluxTransformerTesterConfig(BaseModelTesterConfig):
             "axes_dims_rope": [4, 4, 8],
         }
 
-    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
+    def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]:
         height = width = 4
         num_latent_channels = 4
         num_image_channels = 3
@@ -168,34 +172,34 @@ class FluxTransformerTesterConfig(BaseModelTesterConfig):
             "hidden_states": randn_tensor(
                 (batch_size, height * width, num_latent_channels),
                 generator=self.generator,
-                device=torch_device,
+                device=device,
                 dtype=self.torch_dtype,
             ),
             "encoder_hidden_states": randn_tensor(
                 (batch_size, sequence_length, embedding_dim),
                 generator=self.generator,
-                device=torch_device,
+                device=device,
                 dtype=self.torch_dtype,
             ),
             "pooled_projections": randn_tensor(
                 (batch_size, embedding_dim),
                 generator=self.generator,
-                device=torch_device,
+                device=device,
                 dtype=self.torch_dtype,
             ),
             "img_ids": randn_tensor(
                 (height * width, num_image_channels),
                 generator=self.generator,
-                device=torch_device,
+                device=device,
                 dtype=self.torch_dtype,
             ),
             "txt_ids": randn_tensor(
                 (sequence_length, num_image_channels),
                 generator=self.generator,
-                device=torch_device,
+                device=device,
                 dtype=self.torch_dtype,
             ),
-            "timestep": torch.tensor([1.0]).to(torch_device, self.torch_dtype).expand(batch_size),
+            "timestep": torch.tensor([1.0]).to(device, self.torch_dtype).expand(batch_size),
         }
 
 
@@ -262,6 +266,43 @@ class TestFluxTransformerContextParallelAttnBackends(
     """Context Parallel inference x attention backends tests for Flux Transformer"""
 
 
+class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParallelTesterMixin):
+    """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator)."""
+
+
+def make_neuron_tp_spec():
+    """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`).
+
+    Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux-specific test data lives in this file
+    while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of
+    the Flux tests.
+    """
+    config = FluxTransformerTesterConfig()
+    return FluxTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
+
+
+@is_tensor_parallel
+@require_torch_neuron
+class TestFluxTransformerTensorParallelNeuron:
+    """Tensor Parallel inference test for Flux Transformer on AWS Neuron.
+
+    Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the
+    `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker
+    with the Flux model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device
+    reference, and the test checks its exit code.
+    """
+
+    def test_tensor_parallel_neuron_inference(self):
+        worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py")
+        spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_tp_spec"
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        assert result.returncode == 0, (
+            f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n"
+            f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+        )
+
+
 class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin):
     """IP Adapter tests for Flux Transformer."""
 

tests/models/transformers/test_models_transformer_flux2.py

@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
+import sys
+
 import torch
 
 from diffusers import Flux2Transformer2DModel
@@ -24,7 +28,7 @@ from diffusers.models.transformers.transformer_flux2 import (
 )
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
@@ -37,6 +41,7 @@ from ..testing_utils import (
     MemoryTesterMixin,
     ModelTesterMixin,
     SingleFileTesterMixin,
+    TensorParallelTesterMixin,
     TorchAoCompileTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
@@ -91,16 +96,18 @@ class Flux2TransformerTesterConfig(BaseModelTesterConfig):
             "axes_dims_rope": [4, 4, 4, 4],
         }
 
-    def get_dummy_inputs(self, height: int = 4, width: int = 4, batch_size: int = 1) -> dict[str, torch.Tensor]:
+    def get_dummy_inputs(
+        self, height: int = 4, width: int = 4, batch_size: int = 1, device: str = torch_device
+    ) -> dict[str, torch.Tensor]:
         num_latent_channels = 4
         sequence_length = 48
         embedding_dim = 32
 
         hidden_states = randn_tensor(
-            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=device
         )
         encoder_hidden_states = randn_tensor(
-            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=device
         )
 
         t_coords = torch.arange(1)
@@ -108,17 +115,17 @@ class Flux2TransformerTesterConfig(BaseModelTesterConfig):
         w_coords = torch.arange(width)
         l_coords = torch.arange(1)
         image_ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords)  # [height * width, 4]
-        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+        image_ids = image_ids.unsqueeze(0).expand(batch_size, -1, -1).to(device)
 
         text_t_coords = torch.arange(1)
         text_h_coords = torch.arange(1)
         text_w_coords = torch.arange(1)
         text_l_coords = torch.arange(sequence_length)
         text_ids = torch.cartesian_prod(text_t_coords, text_h_coords, text_w_coords, text_l_coords)
-        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(torch_device)
+        text_ids = text_ids.unsqueeze(0).expand(batch_size, -1, -1).to(device)
 
-        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
-        guidance = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+        timestep = torch.tensor([1.0]).to(device).expand(batch_size)
+        guidance = torch.tensor([1.0]).to(device).expand(batch_size)
 
         return {
             "hidden_states": hidden_states,
@@ -154,6 +161,43 @@ class TestFlux2TransformerContextParallel(Flux2TransformerTesterConfig, ContextP
     """Context Parallel inference tests for Flux2 Transformer."""
 
 
+class TestFlux2TransformerTensorParallel(Flux2TransformerTesterConfig, TensorParallelTesterMixin):
+    """Tensor Parallel inference tests for Flux2 Transformer (CUDA/XPU multi-accelerator)."""
+
+
+def make_neuron_tp_spec():
+    """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`).
+
+    Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux2-specific test data lives in this file
+    while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of
+    the Flux2 tests.
+    """
+    config = Flux2TransformerTesterConfig()
+    return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
+
+
+@is_tensor_parallel
+@require_torch_neuron
+class TestFlux2TransformerTensorParallelNeuron:
+    """Tensor Parallel inference test for Flux2 Transformer on AWS Neuron.
+
+    Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the
+    `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker
+    with the Flux2 model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device
+    reference, and the test checks its exit code.
+    """
+
+    def test_tensor_parallel_neuron_inference(self):
+        worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py")
+        spec = "tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec"
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        assert result.returncode == 0, (
+            f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n"
+            f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+        )
+
+
 class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin):
     """LoRA adapter tests for Flux2 Transformer."""
 

tests/models/transformers/test_models_transformer_qwenimage.py

@@ -13,6 +13,10 @@
 # limitations under the License.
 
 
+import os
+import subprocess
+import sys
+
 import pytest
 import torch
 
@@ -20,7 +24,7 @@ from diffusers import QwenImageTransformer2DModel
 from diffusers.models.transformers.transformer_qwenimage import compute_text_seq_len_from_mask
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
 from ..testing_utils import (
     AttentionBackendTesterMixin,
     AttentionTesterMixin,
@@ -32,6 +36,7 @@ from ..testing_utils import (
     LoraTesterMixin,
     MemoryTesterMixin,
     ModelTesterMixin,
+    TensorParallelTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
     TrainingTesterMixin,
@@ -79,20 +84,20 @@ class QwenImageTransformerTesterConfig(BaseModelTesterConfig):
             "axes_dims_rope": (8, 4, 4),
         }
 
-    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
+    def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]:
         num_latent_channels = embedding_dim = 16
         height = width = 4
         sequence_length = 8
         vae_scale_factor = 4
 
         hidden_states = randn_tensor(
-            (batch_size, height * width, num_latent_channels), generator=self.generator, device=torch_device
+            (batch_size, height * width, num_latent_channels), generator=self.generator, device=device
         )
         encoder_hidden_states = randn_tensor(
-            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+            (batch_size, sequence_length, embedding_dim), generator=self.generator, device=device
         )
-        encoder_hidden_states_mask = torch.ones((batch_size, sequence_length)).to(torch_device, torch.long)
-        timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
+        encoder_hidden_states_mask = torch.ones((batch_size, sequence_length)).to(device, torch.long)
+        timestep = torch.tensor([1.0]).to(device).expand(batch_size)
         orig_height = height * 2 * vae_scale_factor
         orig_width = width * 2 * vae_scale_factor
         img_shapes = [(1, orig_height // vae_scale_factor // 2, orig_width // vae_scale_factor // 2)] * batch_size
@@ -298,6 +303,43 @@ class TestQwenImageTransformerContextParallelAttnBackends(
         return inputs
 
 
+class TestQwenImageTransformerTensorParallel(QwenImageTransformerTesterConfig, TensorParallelTesterMixin):
+    """Tensor Parallel inference tests for QwenImage Transformer (CUDA/XPU multi-accelerator)."""
+
+
+def make_neuron_tp_spec():
+    """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``).
+
+    Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all QwenImage-specific test data lives in this
+    file while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest
+    of the QwenImage tests.
+    """
+    config = QwenImageTransformerTesterConfig()
+    return QwenImageTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
+
+
+@is_tensor_parallel
+@require_torch_neuron
+class TestQwenImageTransformerTensorParallelNeuron:
+    """Tensor Parallel inference test for QwenImage Transformer on AWS Neuron.
+
+    Neuron TP runs through ``torchrun`` with the ``"neuron"`` distributed backend, so it cannot use the
+    ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker
+    with the QwenImage model spec (``make_neuron_tp_spec``); the worker asserts the sharded output matches a
+    single-device reference, and the test checks its exit code.
+    """
+
+    def test_tensor_parallel_neuron_inference(self):
+        worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py")
+        spec = "tests.models.transformers.test_models_transformer_qwenimage:make_neuron_tp_spec"
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        assert result.returncode == 0, (
+            f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n"
+            f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+        )
+
+
 class TestQwenImageTransformerLoRA(QwenImageTransformerTesterConfig, LoraTesterMixin):
     """LoRA adapter tests for QwenImage Transformer."""
 

tests/testing_utils.py

@@ -486,6 +486,14 @@ def is_context_parallel(test_case):
     return pytest.mark.context_parallel(test_case)
 
 
+def is_tensor_parallel(test_case):
+    """
+    Decorator marking a test as a tensor parallel inference test. These tests can be filtered using:
+        pytest -m "not tensor_parallel" to skip pytest -m tensor_parallel to run only these tests
+    """
+    return pytest.mark.tensor_parallel(test_case)
+
+
 def is_cache(test_case):
     """
     Decorator marking a test as a cache test. These tests can be filtered using: