diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml
index 525734a2e4cb..9156ca5e505c 100644
--- a/docs/source/en/_toctree.yml
+++ b/docs/source/en/_toctree.yml
@@ -124,6 +124,8 @@
       title: Intel Gaudi
     - local: optimization/neuron
       title: AWS Neuron
+    - local: optimization/tpu
+      title: TPU
     title: Hardware-specific acceleration
   - isExpanded: false
     sections:
diff --git a/docs/source/en/api/parallel.md b/docs/source/en/api/parallel.md
index 5f300d5dd566..d1e9b7d4177c 100644
--- a/docs/source/en/api/parallel.md
+++ b/docs/source/en/api/parallel.md
@@ -27,4 +27,4 @@ Parallelism strategies help speed up diffusion transformers by distributing comp
 
 [[autodoc]] TensorParallelConfig
 
-[[autodoc]] hooks.apply_tensor_parallel
+[[autodoc]] hooks.apply_tensor_parallel
\ No newline at end of file
diff --git a/docs/source/en/optimization/tpu.md b/docs/source/en/optimization/tpu.md
new file mode 100644
index 000000000000..1a24bd60b930
--- /dev/null
+++ b/docs/source/en/optimization/tpu.md
@@ -0,0 +1,145 @@
+<!--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.
+-->
+
+# TorchTPU
+
+[TorchTPU](https://github.com/google-pytorch/torch_tpu/) is a PyTorch backend for Google's Tensor Processing Units (TPUs), which lets you run Diffusers pipelines on Cloud TPUs (v6e, v5p, etc.) with minimal code changes.
+
+Two execution modes are available:
+
+| Mode | Constant | How to activate | Notes |
+|---|---|---|---|
+| Strict eager (default) | `EagerMode.DEFER_NEVER` | `import torch_tpu` | Operations dispatched one at a time, asynchronous |
+| Compile | — | `torch.compile(module, backend="tpu")` | AOT compilation with `TpuBackend` |
+
+Follow the [TorchTPU installation guide](https://github.com/google-pytorch/torch_tpu/). After installation,
+`import torch_tpu` registers the `"tpu"` device automatically.
+
+## Eager mode
+
+```python
+import gc
+import torch
+import torch_tpu  # noqa: F401
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+
+# 1. Encode on TPU.
+pipe.text_encoder.to("tpu")
+pipe.text_encoder_2.to("tpu")
+with torch.no_grad():
+    prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(
+        prompt="a golden retriever surfing a wave, photorealistic",
+        prompt_2="a golden retriever surfing a wave, photorealistic",
+        device=torch.device("tpu"),
+        max_sequence_length=512,
+    )
+
+# 2. Free the text encoders — nothing below needs them.
+pipe.text_encoder = None
+pipe.text_encoder_2 = None
+gc.collect()
+
+# 3. Move the transformer and VAE in, then denoise with the precomputed embeddings.
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+image = pipe(
+    prompt_embeds=prompt_embeds,
+    pooled_prompt_embeds=pooled_prompt_embeds,
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+image.save("output.png")
+```
+
+If the text encoder alone is too large for a single chip(eg. FLUX.2-dev's Mistral-3-Small is ~45GB),
+shard it across multiple chips with [`~diffusers.hooks.tensor_parallel.apply_tensor_parallel`], the
+same mechanism [`~ModelMixin.enable_parallelism`] uses for the transformer (see [Tensor
+parallelism](../training/distributed_inference#tensor-parallelism)). It only requires `model:
+torch.nn.Module`, so it works directly on a `transformers.PreTrainedModel` text encoder too, not
+just a diffusers `ModelMixin`. The text encoder doesn't define a `_tp_plan`, so supply one: pair
+each attention/MLP projection that expands the hidden dimension (`"colwise"`) with the one that
+contracts it back (`"rowwise"`), matching the `transformers` model's actual module names.
+
+## Compiled mode
+
+`import torch_tpu` registers `"tpu"` as a `torch.compile` backend name (`TpuBackend` under the hood), so
+components compile like any other `torch.compile` target — no diffusers-specific method needed. The first
+call (warmup) is slow because it compiles; later calls with the same shapes reuse the compiled graph.
+
+> [!IMPORTANT]
+> TorchTPU requires **static shapes** — pass `dynamic=False`. Every time `height`, `width`, or
+> `num_inference_steps` changes, the graph is recompiled from scratch. Keep these values constant
+> across all calls after warmup, or run another warmup pass before changing them.
+
+```python
+import torch
+import torch_tpu  # noqa: F401 — registers the "tpu" torch.compile backend
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+)
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+
+pipe.transformer = torch.compile(pipe.transformer, backend="tpu", fullgraph=True, dynamic=False)
+pipe.vae = torch.compile(pipe.vae, backend="tpu", fullgraph=True, dynamic=False)
+
+# Warmup — triggers static graph compilation.
+with torch.no_grad():
+    pipe(
+        prompt="warmup",
+        height=1024,
+        width=1024,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+    )
+
+# Timed inference reuses the compiled graph.
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+image.save("output.png")
+```
+
+## Tensor parallelism
+
+Shard a transformer too large for one chip across several with [`~ModelMixin.enable_parallelism`]. Pass a `TensorParallelConfig` with a TPU `DeviceMesh`. For general TP details, (`_tp_plan`, colwise/rowwise), see the [Tensor parallelism](../training/distributed_inference#tensor-parallelism) guide. Set `backend="tpu_dist"` and `DeviceMesh("tpu", ...)` here to enable tensor parallelism.
+
+```python
+import torch
+import torch.distributed as dist
+import torch_tpu  # noqa: F401
+from torch.distributed.device_mesh import DeviceMesh
+
+from diffusers import DiffusionPipeline, TensorParallelConfig
+
+dist.init_process_group(backend="tpu_dist")
+tp_mesh = DeviceMesh("tpu", list(range(dist.get_world_size())))
+
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16)
+pipe.transformer.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
+pipe.transformer.to("tpu")
+```
diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py
index b90a5761d043..3397d53f2e37 100644
--- a/src/diffusers/hooks/tensor_parallel.py
+++ b/src/diffusers/hooks/tensor_parallel.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import NamedTuple
+
 import torch
 
 from ..models._modeling_parallel import TensorParallelConfig
@@ -20,7 +22,7 @@
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name
 
-_SUPPORTED_TP_DEVICES = ("cuda", "neuron")
+_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu")
 
 
 class PackedColwiseParallel:
@@ -65,6 +67,89 @@ def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]":
     return [b * unit for b in blocks]
 
 
+class TPShardSpec(NamedTuple):
+    """How one parameter is laid out across the tensor-parallel ranks.
+
+    `dim` is the dimension sharded across ranks, or `None` when the parameter is replicated on every rank (a rowwise
+    bias, which is added after the all-reduce). `block_sizes` partitions `dim` into independently sharded blocks; a
+    plain `"colwise"` / `"rowwise"` style has a single block covering the whole dimension, and packed styles have one
+    per fused projection.
+    """
+
+    dim: "int | None"
+    block_sizes: "list[int] | None"
+
+
+def _local_shard(tensor: torch.Tensor, dim: int, block_sizes: "list[int]", tp_mesh) -> torch.Tensor:
+    """Extract this rank's slice of `tensor` along `dim`.
+
+    Each block is sliced independently and the pieces concatenated, so every fused projection of a packed weight
+    contributes its own contiguous chunk to the rank's shard. A `dim`-1 slice comes back strided, hence the final
+    `contiguous()` — `DTensor.from_local` needs a contiguous local tensor.
+    """
+    rank = tp_mesh.get_local_rank()
+    tp_size = tp_mesh.size()
+
+    parts, offset = [], 0
+    for block_size in block_sizes:
+        # An uneven split is rejected rather than handed to `Shard`, which pads the tail and would break both the
+        # paired colwise/rowwise matmul and the attention head split.
+        if block_size % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard a block of size {block_size} across {tp_size} tensor-parallel ranks: "
+                f"{block_size} is not divisible by {tp_size}."
+            )
+        chunk = block_size // tp_size
+        index = [slice(None)] * tensor.dim()
+        index[dim] = slice(offset + rank * chunk, offset + (rank + 1) * chunk)
+        parts.append(tensor[tuple(index)])
+        offset += block_size
+
+    local = parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
+    return local.contiguous()
+
+
+def _block_shard_specs(
+    block: torch.nn.Module, relative_plan: dict
+) -> "list[tuple[torch.nn.Module, str, TPShardSpec]]":
+    """Resolve one block's plan to `(module, param_name, spec)` triples, covering both `weight` and `bias`.
+
+    Modules are returned directly rather than by name so the caller can place the shards without a second lookup.
+    """
+    resolved = []
+    for relative_path, style in relative_plan.items():
+        submodule = block
+        for atom in relative_path.split("."):
+            submodule = getattr(submodule, atom)
+
+        # `_tp_packed_*_blocks` hold absolute sizes rather than proportions; that works because they sum to the
+        # full dimension, so `_blocks_to_block_sizes` computes `unit == 1`.
+        if style == "colwise":
+            weight_spec = TPShardSpec(0, [submodule.weight.shape[0]])
+            bias_spec = weight_spec
+        elif style == "rowwise":
+            weight_spec = TPShardSpec(1, [submodule.weight.shape[1]])
+            bias_spec = TPShardSpec(None, None)
+        elif isinstance(style, PackedColwiseParallel):
+            blocks = style.blocks if style.blocks is not None else submodule._tp_packed_col_blocks
+            weight_spec = TPShardSpec(0, _blocks_to_block_sizes(submodule.weight.shape[0], blocks))
+            bias_spec = weight_spec
+        elif isinstance(style, PackedRowwiseParallel):
+            blocks = style.blocks if style.blocks is not None else submodule._tp_packed_row_blocks
+            weight_spec = TPShardSpec(1, _blocks_to_block_sizes(submodule.weight.shape[1], blocks))
+            bias_spec = TPShardSpec(None, None)
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{relative_path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+
+        resolved.append((submodule, "weight", weight_spec))
+        if submodule.bias is not None:
+            resolved.append((submodule, "bias", bias_spec))
+    return resolved
+
+
 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.
 
@@ -193,7 +278,7 @@ def _partition_linear_fn(self, name, module, device_mesh):
     # `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.
+    # matching what `_local_shard` already does for the packed styles.
     def _make_checked_col(path: str) -> ColwiseParallel:
         class _CheckedColwiseImpl(ColwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
@@ -240,6 +325,78 @@ def _partition_linear_fn(self, name, module, device_mesh):
     return resolved
 
 
+def _hooks_only_styles(relative_plan: dict) -> dict:
+    """Map a `{relative_path: style}` plan to styles that partition nothing.
+
+    Used when the caller has already placed every planned parameter as a `DTensor`. `parallelize_module` then runs only
+    to register the forward input/output hooks; `_partition_linear_fn` must not re-partition. Packed and plain styles
+    share hook behaviour, so both collapse onto the two styles here.
+
+    Note this is not purely additive: `distribute_module` still replicates any *remaining* plain parameter of the
+    targeted module into a `Replicate()` DTensor via a broadcast, so callers must place every planned parameter
+    themselves.
+    """
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    class _NoPartitionColwise(ColwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(0)
+
+    class _NoPartitionRowwise(RowwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(1)
+
+    resolved = {}
+    for path, style in relative_plan.items():
+        if style == "colwise" or isinstance(style, PackedColwiseParallel):
+            resolved[path] = _NoPartitionColwise()
+        elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
+            resolved[path] = _NoPartitionRowwise()
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+    return resolved
+
+
+def _pre_shard_and_parallelize(tp_mesh, groups: list, device: torch.device) -> None:
+    """Slice every planned parameter on CPU and move only this rank's shard to `device`.
+
+    The default path lets `parallelize_module` distribute the weights, which materializes each full weight on every
+    rank before scattering it, so peak memory per rank is the size of the whole weight even though only a shard
+    survives. Slicing first and handing `DTensor.from_local` just this rank's piece keeps the full tensor off the
+    accelerator, which is what allows sharding a model larger than one device's memory. `parallelize_module` then runs
+    over `_hooks_only_styles`, distributing nothing and registering only the hooks the forward pass needs.
+
+    `device` is the only backend-specific input; the slicing itself is the same on any accelerator. Unlike the default
+    path this does not broadcast from a single rank, so every rank must already hold the same weights — true after
+    loading a checkpoint, not after a random init.
+
+    Model weights must be on CPU when this is called.
+    """
+    import torch.nn as nn
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    for block, relative_plan in groups:
+        for submodule, param_name, spec in _block_shard_specs(block, relative_plan):
+            param = getattr(submodule, param_name)
+            if spec.dim is None:
+                # A rowwise bias is added after the all-reduce, so every rank needs the whole vector.
+                local, placement = param.data, Replicate()
+            else:
+                local, placement = _local_shard(param.data, spec.dim, spec.block_sizes, tp_mesh), Shard(spec.dim)
+            submodule.register_parameter(
+                param_name,
+                nn.Parameter(
+                    DTensor.from_local(local.to(device), tp_mesh, [placement]),
+                    requires_grad=param.requires_grad,
+                ),
+            )
+        parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan))
+
+
 def apply_tensor_parallel(
     model: torch.nn.Module,
     config: TensorParallelConfig,
@@ -257,7 +414,12 @@ def apply_tensor_parallel(
             f"or from the active accelerator when the mesh is built from `tp_degree`."
         )
 
-    backend = "neuron" if tp_mesh.device_type == "neuron" else "default"
+    if tp_mesh.device_type == "neuron":
+        backend = "neuron"
+    elif tp_mesh.device_type == "tpu":
+        backend = "tpu"
+    else:
+        backend = "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}.")
 
@@ -267,6 +429,14 @@ def apply_tensor_parallel(
         _apply_tp_neuron(model, tp_mesh, groups)
         return
 
+    if backend == "tpu":
+        # Pre-shard rather than let `parallelize_module` distribute: it materializes each full weight on every chip
+        # before scattering it, which exhausts HBM for a large diffusion transformer. Address the chip as "tpu" with
+        # no index — "tpu:rank" would mean chip `rank` from this process's view, but each torchrun worker only has
+        # access to its own assigned chip.
+        _pre_shard_and_parallelize(tp_mesh, groups, torch.device("tpu"))
+        return
+
     from torch.distributed.tensor.parallel import parallelize_module
 
     for submodule, relative_plan in groups:
diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py
index 86627284e078..9d00dc888b6d 100644
--- a/src/diffusers/models/_modeling_parallel.py
+++ b/src/diffusers/models/_modeling_parallel.py
@@ -161,7 +161,7 @@ class TensorParallelConfig:
     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"`.
+    device types are `"cuda"`, `"neuron"` and `"tpu"`.
 
     Args:
         tp_degree (`int`, defaults to `1`):
diff --git a/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py b/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py
index 11fce6a204bf..df5b27ace653 100644
--- a/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py
+++ b/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py
@@ -114,7 +114,7 @@ def _enhance_prompt_with_pe(
             tokenize=False,
             add_generation_prompt=False,  # "Output:" is already in the user block
         )
-        inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(device)
+        inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(self.pe.device)
         output_ids = self.pe.generate(
             **inputs,
             max_new_tokens=self.pe_tokenizer.model_max_length,
@@ -155,7 +155,7 @@ def encode_prompt(
                 else:
                     ids = [0]
 
-            input_ids = torch.tensor([ids], device=device)
+            input_ids = torch.tensor([ids], device=self.text_encoder.device)
             with torch.no_grad():
                 outputs = self.text_encoder(
                     input_ids=input_ids,
diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py
index 5c63a4bc7661..0b1362f769d4 100644
--- a/src/diffusers/utils/__init__.py
+++ b/src/diffusers/utils/__init__.py
@@ -115,6 +115,7 @@
     is_torch_mlu_available,
     is_torch_neuronx_available,
     is_torch_npu_available,
+    is_torch_tpu_available,
     is_torch_version,
     is_torch_xla_available,
     is_torch_xla_version,
diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py
index d2cf394cd9a7..e9c8a9e1858d 100644
--- a/src/diffusers/utils/import_utils.py
+++ b/src/diffusers/utils/import_utils.py
@@ -178,6 +178,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b
 _torch_xla_available, _torch_xla_version = _is_package_available("torch_xla")
 _torch_npu_available, _torch_npu_version = _is_package_available("torch_npu")
 _torch_mlu_available, _torch_mlu_version = _is_package_available("torch_mlu")
+_torch_tpu_available, _torch_tpu_version = _is_package_available("torch_tpu")
 _torch_neuronx_available, _torch_neuronx_version = _is_package_available("torch_neuronx")
 _transformers_available, _transformers_version = _is_package_available("transformers")
 _hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub")
@@ -238,6 +239,10 @@ def is_torch_mlu_available():
     return _torch_mlu_available
 
 
+def is_torch_tpu_available():
+    return _torch_tpu_available
+
+
 def is_torch_neuronx_available():
     return _torch_neuronx_available
 
@@ -553,6 +558,11 @@ def is_av_available():
 torchao`
 """
 
+TORCH_TPU_IMPORT_ERROR = """
+{0} requires the torch_tpu library but it was not found in your environment. Please follow the installation
+instructions at https://github.com/pytorch/tpu
+"""
+
 QUANTO_IMPORT_ERROR = """
 {0} requires the optimum-quanto library but it was not found in your environment. You can install it with pip: `pip
 install optimum-quanto`
@@ -613,6 +623,7 @@ def is_av_available():
         ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)),
         ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)),
         ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)),
+        ("torch_tpu", (is_torch_tpu_available, TORCH_TPU_IMPORT_ERROR)),
         ("torch_neuronx", (is_torch_neuronx_available, TORCH_NEURONX_IMPORT_ERROR)),
     ]
 )
diff --git a/src/diffusers/utils/torch_utils.py b/src/diffusers/utils/torch_utils.py
index b3292d5cf0d2..cddc7d667bf0 100644
--- a/src/diffusers/utils/torch_utils.py
+++ b/src/diffusers/utils/torch_utils.py
@@ -46,6 +46,7 @@
         "cpu": True,
         "mps": False,
         "neuron": False,
+        "tpu": False,
         "default": True,
     }
     BACKEND_EMPTY_CACHE = {
@@ -53,6 +54,7 @@
         "xpu": torch.xpu.empty_cache,
         "cpu": None,
         "mps": torch.mps.empty_cache,
+        "tpu": getattr(getattr(torch, "tpu", None), "empty_cache", None),
         "neuron": None,
         "default": None,
     }
@@ -61,6 +63,7 @@
         "xpu": torch.xpu.device_count,
         "cpu": lambda: 0,
         "mps": lambda: 0,
+        "tpu": lambda: getattr(getattr(torch, "tpu", None), "device_count", lambda: 0)(),
         "neuron": lambda: getattr(getattr(torch, "neuron", None), "device_count", lambda: 0)(),
         "default": 0,
     }
@@ -69,6 +72,9 @@
         "xpu": torch.xpu.manual_seed,
         "cpu": torch.manual_seed,
         "mps": torch.mps.manual_seed,
+        # TPU latents are always generated on CPU (TPU RNG has unaligned DUS bug),
+        # so CPU seeding is the correct behaviour here.
+        "tpu": torch.manual_seed,
         "neuron": torch.manual_seed,
         "default": torch.manual_seed,
     }
@@ -77,6 +83,7 @@
         "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
         "cpu": None,
         "mps": None,
+        "tpu": None,
         "neuron": None,
         "default": None,
     }
@@ -85,6 +92,7 @@
         "xpu": getattr(torch.xpu, "reset_peak_memory_stats", None),
         "cpu": None,
         "mps": None,
+        "tpu": None,
         "neuron": None,
         "default": None,
     }
@@ -93,6 +101,7 @@
         "xpu": getattr(torch.xpu, "max_memory_allocated", None),
         "cpu": 0,
         "mps": 0,
+        "tpu": 0,
         "neuron": 0,
         "default": 0,
     }
@@ -101,6 +110,7 @@
         "xpu": getattr(torch.xpu, "synchronize", None),
         "cpu": None,
         "mps": None,
+        "tpu": getattr(getattr(torch, "tpu", None), "synchronize", None),
         "neuron": getattr(getattr(torch, "neuron", None), "synchronize", None),
         "default": None,
     }
@@ -198,6 +208,11 @@ def randn_tensor(
     rand_device = device
     batch_size = shape[0]
 
+    # TPU RNG has an unaligned DUS (dynamic-update-slice) bug — generate on CPU
+    # and move to TPU via the existing .to(device) call at the end.
+    if device is not None and device.type == "tpu":
+        rand_device = torch.device("cpu")
+
     layout = layout or torch.strided
     device = device or torch.device("cpu")
 
diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py
index 2d7d5ae23257..0932e46f9b93 100644
--- a/tests/models/testing_utils/__init__.py
+++ b/tests/models/testing_utils/__init__.py
@@ -23,6 +23,7 @@
     ContextParallelAttentionBackendsTesterMixin,
     ContextParallelTesterMixin,
     TensorParallelTesterMixin,
+    TensorParallelTPUTesterMixin,
 )
 from .quantization import (
     AutoRoundCompileTesterMixin,
@@ -67,6 +68,7 @@
     "ContextParallelTesterMixin",
     "ContextParallelAttentionBackendsTesterMixin",
     "TensorParallelTesterMixin",
+    "TensorParallelTPUTesterMixin",
     "CPUOffloadTesterMixin",
     "FasterCacheConfigMixin",
     "FasterCacheTesterMixin",
diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py
index 63575abf6b7b..11b343bd6fcd 100644
--- a/tests/models/testing_utils/parallelism.py
+++ b/tests/models/testing_utils/parallelism.py
@@ -15,6 +15,8 @@
 
 import os
 import socket
+import subprocess
+import sys
 
 import pytest
 import torch
@@ -30,6 +32,7 @@
     is_kernels_available,
     is_tensor_parallel,
     require_torch_multi_accelerator,
+    require_torch_tpu,
     torch_device,
 )
 from .utils import _maybe_cast_to_bf16
@@ -345,6 +348,99 @@ def test_tensor_parallel_batch_inputs(self):
         self.test_tensor_parallel_inference(batch_size=2)
 
 
+def _run_tp_worker_subprocess(worker_filename: str, spec: str, world_size: int, timeout_s: int = 900) -> None:
+    """Launch a `torchrun` TP-correctness worker subprocess and assert it exits cleanly.
+
+    Args:
+        worker_filename: Name of the worker script, resolved relative to `tests/models/transformers/` (e.g.
+            `"_tpu_tp_worker.py"`).
+        spec: `module:function` reference forwarded to the worker, see `_tp_worker_common.run_tp_correctness_worker`.
+        world_size: Number of ranks to launch (`torchrun --nproc_per_node`).
+        timeout_s: Seconds to wait for the subprocess before failing the test. The worker itself only needs a couple
+            of minutes even from a cold compile; this generously bounds it so a real hang (e.g. a distributed-runtime
+            barrier timeout) fails the test loudly instead of stalling the run.
+    """
+    worker = os.path.join(os.path.dirname(__file__), "..", "transformers", worker_filename)
+    cmd = [sys.executable, "-m", "torch.distributed.run", f"--nproc_per_node={world_size}", worker, spec]
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_s)
+    except subprocess.TimeoutExpired as e:
+        raise AssertionError(
+            f"TP worker did not finish within {timeout_s}s (likely stuck on a distributed-runtime barrier).\n"
+            f"--- stdout ---\n{e.stdout}\n--- stderr ---\n{e.stderr}"
+        ) from e
+    assert result.returncode == 0, (
+        f"TP worker failed (exit {result.returncode}).\n--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+    )
+
+
+@is_tensor_parallel
+@require_torch_tpu
+class TensorParallelTPUTesterMixin:
+    """Mixin for a tensor-parallel correctness test on TPU, run via `_tpu_tp_worker.py`.
+
+    TPU TP runs through `torchrun` with the `"tpu_dist"` distributed backend, so — like `TestFlux2TransformerTensorParallelNeuron`
+    for Neuron — it cannot use `TensorParallelTesterMixin`'s `torch.multiprocessing.spawn`/NCCL path above and instead
+    launches a subprocess worker script and checks its exit code.
+
+    Subclasses set `TP_SPEC` to a `module:function` reference (see
+    `_tp_worker_common.run_tp_correctness_worker`'s `spec` argument) and, only if the model spec's head count
+    doesn't divide 4, override `WORLD_SIZE`.
+
+    `WORLD_SIZE` defaults to 4 rather than an arbitrary rank count: `torch_tpu`'s per-generation topology table
+    (`torch_tpu._internal.utils.hardware`) only enumerates whole-pod-slice chip counts (1/4/8 for v6e, for example),
+    not arbitrary sub-slices of a larger single host. A rank count with no matching whole-slice topology has
+    nothing to advertise and the PJRT client never completes its start-session barrier — the test would hang for
+    the barrier's full multi-minute timeout instead of failing. 4 is the smallest whole-slice count every current
+    TPU generation defines (see `_V4_TOPOLOGY` / `_V5E_TOPOLOGY` / `_V6E_TOPOLOGY` / `_V7_TOPOLOGY` in
+    `torch_tpu._internal.utils.hardware`). `skip_if_unsupported` below still checks the actual host up front and
+    skips fast instead of hanging when it doesn't have exactly that many chips.
+
+    Requires `TORCH_TPU_TOPOLOGY` and `TORCH_TPU_SLICEBUILDER_ADDRESSES` to be set. Source them via::
+
+        eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
+    """
+
+    WORLD_SIZE = 4
+    # The worker itself only needs a couple of minutes even from a cold XLA compile; this generously bounds the
+    # subprocess so a real hang (e.g. a barrier timeout this skip failed to catch) fails the test loudly instead of
+    # stalling the run.
+    TIMEOUT_S = 900
+    TP_SPEC: str = ""
+
+    def skip_if_unsupported(self):
+        """Skip unless the host has exactly `WORLD_SIZE` TPU chips.
+
+        A topology *string* existing for a chip count (`hardware.get_tpu_topology`) isn't enough to guarantee the
+        PJRT client can actually form that session: a sub-slice of a larger single host (e.g. claiming 2 of a
+        4-chip v6e-4's chips via `TORCH_TPU_TOPOLOGY`/`TORCH_TPU_SLICEBUILDER_ADDRESSES`) can still fail with a
+        low-level `START_SESSION` GRPC error, since the runtime's session setup is tied to the host's actual
+        provisioned slice, not just a topology label. The only combination verified to work is running with exactly
+        as many ranks as the host has chips.
+        """
+        from torch_tpu._internal.utils import hardware
+
+        try:
+            device_count = hardware.get_tpu_device_count()
+        except Exception as e:  # pragma: no cover - defensive, hardware detection is best-effort
+            pytest.skip(f"Could not determine local TPU chip count: {e}")
+            return
+
+        if device_count != self.WORLD_SIZE:
+            pytest.skip(
+                f"This host exposes {device_count} TPU chip(s), but this test requires exactly "
+                f"{self.WORLD_SIZE} (a TPU single-host tensor-parallel job must use all chips on the host; "
+                f"sub-slicing a larger host is not reliably supported by the runtime). Run this test on a host "
+                f"with exactly {self.WORLD_SIZE} TPU chips."
+            )
+
+    def test_tensor_parallel_tpu_inference(self):
+        self.skip_if_unsupported()
+        _run_tp_worker_subprocess(
+            "_tpu_tp_worker.py", self.TP_SPEC, world_size=self.WORLD_SIZE, timeout_s=self.TIMEOUT_S
+        )
+
+
 @is_context_parallel
 @require_torch_multi_accelerator
 class ContextParallelTesterMixin:
diff --git a/tests/models/transformers/_neuron_tp_worker.py b/tests/models/transformers/_neuron_tp_worker.py
index 681f5686bee9..1306335e2928 100644
--- a/tests/models/transformers/_neuron_tp_worker.py
+++ b/tests/models/transformers/_neuron_tp_worker.py
@@ -99,7 +99,6 @@ def main():
         print("[rank0] PASS: Neuron tensor-parallel output matches single-device reference.")
 
     dist.barrier()
-    dist.destroy_process_group()
 
 
 if __name__ == "__main__":
diff --git a/tests/models/transformers/_tp_worker_common.py b/tests/models/transformers/_tp_worker_common.py
new file mode 100644
index 000000000000..750da2a94e7c
--- /dev/null
+++ b/tests/models/transformers/_tp_worker_common.py
@@ -0,0 +1,109 @@
+# 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.
+
+"""Shared logic for the per-backend TP-correctness `torchrun` workers (`_tpu_tp_worker.py`, `_neuron_tp_worker.py`).
+
+Both workers follow the same recipe — build an identical model on every rank, compute a single-device reference,
+shard with `enable_parallelism`, run on-device, and compare — differing only in backend-specific details (device
+string, sync call, whether the reference itself needs to run on-device, and numerical tolerance). This module holds
+that shared recipe; each `_<backend>_tp_worker.py` is a thin wrapper supplying those details.
+"""
+
+import copy
+import importlib
+from typing import Callable
+
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+
+from diffusers import TensorParallelConfig
+
+
+def run_tp_correctness_worker(
+    spec: str,
+    *,
+    mesh_device_type: str,
+    to_device,
+    backend_label: str,
+    synchronize: Callable[[], None],
+    reference_on_device: bool,
+    atol: float,
+    rtol: float,
+) -> None:
+    """Assert a model's tensor-parallel output matches its single-device reference, on the given backend.
+
+    Args:
+        spec: `module:function` reference returning `(model_class, init_dict, cpu_inputs)` for the model under test.
+        mesh_device_type: The `DeviceMesh` device type (e.g. `"tpu"`, `"neuron"`).
+        to_device: The value passed to `.to(...)` to move the model/inputs onto the accelerator. Usually the same as
+            `mesh_device_type`, but some backends (e.g. Neuron) need a more specific device handle here.
+        backend_label: Human-readable backend name for log messages (e.g. `"TPU"`, `"Neuron"`).
+        synchronize: Callable that blocks until pending device work completes.
+        reference_on_device: If `True`, the unsharded reference forward pass also runs on the accelerator (before TP
+            mutates the weights in place), so it uses the same kernels as the TP forward pass and only sharding
+            differs. If `False`, the reference runs on CPU.
+        atol: Absolute tolerance for the final `torch.testing.assert_close` comparison.
+        rtol: Relative tolerance for the final `torch.testing.assert_close` comparison.
+    """
+    module_name, _, fn_name = spec.partition(":")
+    model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)()
+
+    rank = dist.get_rank()
+    tp_size = dist.get_world_size()
+    tp_mesh = DeviceMesh(mesh_device_type, list(range(tp_size)))
+
+    # Identical weights on every rank (same seed), kept on CPU as the pre-shard backends require.
+    torch.manual_seed(0)
+    model = model_class(**init_dict).eval()
+
+    if reference_on_device:
+        # Single-device (unsharded) reference on the accelerator, computed before TP mutates the weights in place.
+        ref_model = copy.deepcopy(model).to(to_device)
+        synchronize()
+        inputs_on_device = {k: v.to(to_device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+        with torch.no_grad():
+            ref_output = ref_model(**inputs_on_device, return_dict=False)[0]
+        synchronize()
+        ref_output = ref_output.float().cpu()
+        del ref_model
+    else:
+        with torch.no_grad():
+            ref_output = model(**inputs, return_dict=False)[0].float().cpu()
+        inputs_on_device = {k: v.to(to_device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+
+    # Shard across all ranks; the backend is auto-selected from the mesh device type.
+    model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
+    model = model.to(to_device)
+    synchronize()
+
+    with torch.no_grad():
+        tp_output = model(**inputs_on_device, return_dict=False)[0]
+    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}"
+        )
+        torch.testing.assert_close(tp_output, ref_output, atol=atol, rtol=rtol)
+        print(f"[rank0] PASS: {backend_label} tensor-parallel output matches single-device reference.")
+
+    dist.barrier()
diff --git a/tests/models/transformers/_tpu_tp_worker.py b/tests/models/transformers/_tpu_tp_worker.py
new file mode 100644
index 000000000000..e219fb1a6f45
--- /dev/null
+++ b/tests/models/transformers/_tpu_tp_worker.py
@@ -0,0 +1,80 @@
+# 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.
+
+"""TPU entry point for the generic TP-correctness worker (see `_tp_worker_common.py`).
+
+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_tpu`` test (and runnable directly for debugging)::
+
+    eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
+    torchrun --nproc_per_node=4 _tpu_tp_worker.py \\
+        tests.models.transformers.test_models_transformer_flux2:make_tpu_tp_spec
+
+Exit code 0 means the TP path is numerically equivalent to the unsharded model; non-zero means failure.
+"""
+
+import argparse
+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.distributed as dist
+import torch_tpu  # noqa: F401 — registers "tpu" device and "tpu_dist" backend
+from torch_tpu._internal import sync as tpu_sync
+
+from tests.models.transformers._tp_worker_common import run_tp_correctness_worker
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TPU tensor-parallel correctness worker.")
+    parser.add_argument(
+        "spec",
+        help="`module:function` reference returning (model_class, init_dict, tpu_inputs) for the model under test.",
+    )
+    args = parser.parse_args()
+
+    dist.init_process_group(backend="tpu_dist")
+    # The reference runs on the TPU (not CPU) so both the reference and the TP pass use the same Flash Attention
+    # kernel; the only difference between them is sharding, not numerical implementation. TPU Flash Attention has
+    # bf16-level numerics, so the tolerance is wider than fp32 — but a wrong shard plan produces grossly different
+    # output and is caught comfortably within this bound.
+    run_tp_correctness_worker(
+        args.spec,
+        mesh_device_type="tpu",
+        to_device="tpu",
+        backend_label="TPU",
+        synchronize=lambda: tpu_sync.synchronize(None, wait=True),
+        reference_on_device=True,
+        atol=0.1,
+        rtol=0.1,
+    )
+    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)
diff --git a/tests/models/transformers/run_flux2_tp_tpu.py b/tests/models/transformers/run_flux2_tp_tpu.py
new file mode 100644
index 000000000000..08a0420a65e5
--- /dev/null
+++ b/tests/models/transformers/run_flux2_tp_tpu.py
@@ -0,0 +1,333 @@
+#!/usr/bin/env python3
+"""Verify Flux2 transformer forward pass with tensor parallelism on TPU.
+
+Without --model-id (default):
+    Builds a small Flux2 model with random weights, runs a TP forward pass, and compares
+    it against a single-device TPU reference (same model, same device, no sharding). Both
+    PASS/FAIL and a max_abs_diff are printed.
+
+With --model-id (e.g. black-forest-labs/FLUX.2-dev):
+    Loads the real model from the Hub (CPU), applies TP, runs one forward pass on TPU, and
+    checks the output is finite and has the expected shape. No reference comparison (too slow).
+
+The script self-relaunches under torchrun when it is not already a distributed worker, so a
+single ``python run_flux2_tp_tpu.py`` invocation is enough. The TPU topology env-vars must be
+set before the torchrun relaunch; pass them via --topology / --addresses or export them first:
+
+    eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
+    python run_flux2_tp_tpu.py --tp-degree 4
+
+    # or, to test against real weights:
+    python run_flux2_tp_tpu.py --tp-degree 4 --model-id black-forest-labs/FLUX.2-dev
+
+To run end-to-end image generation, use run_flux2_tp_tpu_pipeline.py instead.
+
+The default sequence lengths (latent_h=16, latent_w=16, txt_len=256) give a joint sequence of
+512, which satisfies the TPU Flash Attention requirement of seq_len divisible by 512.
+"""
+
+import argparse
+import copy
+import os
+import sys
+import time
+import traceback
+
+
+# Make in-repo packages importable when run from an arbitrary CWD via torchrun.
+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_tpu  # noqa: F401 — registers "tpu" device and "tpu_dist" backend
+from torch.distributed.device_mesh import DeviceMesh
+from torch_tpu._internal import sync as tpu_sync
+
+from diffusers import Flux2Transformer2DModel, TensorParallelConfig
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+
+def log(rank: int, msg: str) -> None:
+    if rank == 0:
+        print(f"[flux2-tp-tpu] {msg}", flush=True)
+
+
+def relaunch_via_torchrun(tp_degree: int, topology: str | None, addresses: str | None) -> None:
+    """Re-invoke this script under torch.distributed.run if not already a worker."""
+    if os.environ.get("LOCAL_RANK") is not None:
+        return  # already running inside torchrun — nothing to do
+
+    if tp_degree == 1:
+        return  # single-process, torchrun not needed
+
+    if topology:
+        os.environ["TORCH_TPU_TOPOLOGY"] = topology
+    if addresses:
+        os.environ["TORCH_TPU_SLICEBUILDER_ADDRESSES"] = addresses
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        f"--nproc-per-node={tp_degree}",
+        os.path.abspath(__file__),
+    ] + sys.argv[1:]  # forward all original flags to workers
+    raise SystemExit(os.execvp(sys.executable, cmd))
+
+
+def _make_inputs(
+    in_channels: int,
+    joint_attention_dim: int,
+    latent_height: int,
+    latent_width: int,
+    txt_len: int,
+    batch_size: int,
+    device: str | torch.device,
+    dtype: torch.dtype,
+) -> dict:
+    """Build a forward-pass input dict from model config dimensions."""
+    seq_len = latent_height * latent_width
+
+    hidden_states = torch.randn(batch_size, seq_len, in_channels, dtype=dtype, device=device)
+    encoder_hidden_states = torch.randn(batch_size, txt_len, joint_attention_dim, dtype=dtype, device=device)
+
+    t_c = torch.arange(1)
+    h_c = torch.arange(latent_height)
+    w_c = torch.arange(latent_width)
+    l_c = torch.arange(1)
+    img_ids = torch.cartesian_prod(t_c, h_c, w_c, l_c).unsqueeze(0).expand(batch_size, -1, -1).to(device)
+
+    txt_ids = (
+        torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(txt_len))
+        .unsqueeze(0)
+        .expand(batch_size, -1, -1)
+        .to(device)
+    )
+
+    timestep = torch.tensor([500.0], dtype=dtype, device=device).expand(batch_size)
+    guidance = torch.tensor([3.5], dtype=dtype, device=device).expand(batch_size)
+
+    return {
+        "hidden_states": hidden_states,
+        "encoder_hidden_states": encoder_hidden_states,
+        "img_ids": img_ids,
+        "txt_ids": txt_ids,
+        "timestep": timestep,
+        "guidance": guidance,
+    }
+
+
+# ── main worker ───────────────────────────────────────────────────────────────
+
+
+def run(args: argparse.Namespace) -> int:
+    dist.init_process_group(backend="tpu_dist")
+    rank = dist.get_rank()
+    tp_size = dist.get_world_size()
+    tp_mesh = DeviceMesh("tpu", list(range(tp_size)))
+
+    log(rank, f"tp_size={tp_size}  dtype=bfloat16")
+
+    # ── load model ────────────────────────────────────────────────────────────
+    t0 = time.perf_counter()
+    try:
+        if args.model_id:
+            log(rank, f"loading from Hub: {args.model_id}")
+            model = Flux2Transformer2DModel.from_pretrained(
+                args.model_id,
+                subfolder="transformer",
+                torch_dtype=torch.bfloat16,
+            )
+        else:
+            log(rank, "building small dummy model (random weights)")
+            model = Flux2Transformer2DModel(
+                patch_size=1,
+                in_channels=4,
+                num_layers=1,
+                num_single_layers=1,
+                attention_head_dim=16,
+                num_attention_heads=4,  # must be divisible by tp_degree
+                joint_attention_dim=32,
+                timestep_guidance_channels=256,
+                axes_dims_rope=[4, 4, 4, 4],
+            ).to(torch.bfloat16)
+    except Exception:
+        log(rank, "LOAD FAILED")
+        if rank == 0:
+            traceback.print_exc()
+        return 1
+    log(rank, f"load OK ({time.perf_counter() - t0:.1f}s)")
+
+    cfg = model.config
+    latent_h = args.latent_height
+    latent_w = args.latent_width
+    txt_len = args.txt_len
+
+    # ── single-device reference (dummy mode only, before TP mutates weights) ──
+    # The reference runs on the TPU device (not CPU) so both the reference and the TP forward
+    # use the same kernels (e.g. Flash Attention). The only difference between them is sharding.
+    ref_output = None
+    if not args.model_id:
+        ref_model = copy.deepcopy(model).to("tpu")
+        tpu_sync.synchronize(None, wait=True)
+        torch.manual_seed(0)
+        ref_inputs = _make_inputs(
+            cfg.in_channels,
+            cfg.joint_attention_dim,
+            latent_h,
+            latent_w,
+            txt_len,
+            batch_size=1,
+            device="tpu",
+            dtype=torch.bfloat16,
+        )
+        ref_model.eval()
+        with torch.no_grad():
+            ref_output_tpu = ref_model(**ref_inputs, return_dict=False)[0]
+        tpu_sync.synchronize(None, wait=True)
+        ref_output = ref_output_tpu.float().cpu()
+        del ref_model, ref_inputs, ref_output_tpu
+        tpu_sync.synchronize(None, wait=True)
+        log(rank, f"TPU ref computed  max_abs={ref_output.abs().max():.4f}")
+
+    # ── apply tensor parallelism ──────────────────────────────────────────────
+    try:
+        model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
+    except Exception:
+        log(rank, "enable_parallelism FAILED")
+        if rank == 0:
+            traceback.print_exc()
+        return 1
+    log(rank, "TP applied")
+
+    model = model.to("tpu")
+    tpu_sync.synchronize(None, wait=True)
+    log(rank, "model on TPU, triggering compilation ...")
+
+    # ── build inputs ──────────────────────────────────────────────────────────
+    torch.manual_seed(0)
+    tpu_inputs = _make_inputs(
+        cfg.in_channels,
+        cfg.joint_attention_dim,
+        latent_h,
+        latent_w,
+        txt_len,
+        batch_size=1,
+        device="tpu",
+        dtype=torch.bfloat16,
+    )
+
+    # ── warm-up forward (triggers XLA compilation) ────────────────────────────
+    model.eval()
+    try:
+        with torch.no_grad():
+            _ = model(**tpu_inputs, return_dict=False)[0]
+        tpu_sync.synchronize(None, wait=True)
+        log(rank, "warm-up OK (graph compiled)")
+    except Exception:
+        log(rank, "warm-up FAILED")
+        if rank == 0:
+            traceback.print_exc()
+        return 1
+
+    # ── timed forward ─────────────────────────────────────────────────────────
+    try:
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            tp_out = model(**tpu_inputs, return_dict=False)[0]
+        tpu_sync.synchronize(None, wait=True)
+        elapsed = time.perf_counter() - t0
+    except Exception:
+        log(rank, "timed forward FAILED")
+        if rank == 0:
+            traceback.print_exc()
+        return 1
+
+    tp_out_cpu = tp_out.float().cpu()
+
+    # ── verify ────────────────────────────────────────────────────────────────
+    expected_shape = (1, latent_h * latent_w, cfg.in_channels)
+    if tp_out_cpu.shape != torch.Size(expected_shape):
+        log(rank, f"FAIL: shape {tuple(tp_out_cpu.shape)} != expected {expected_shape}")
+        return 1
+
+    if not torch.isfinite(tp_out_cpu).all():
+        log(rank, "FAIL: output contains non-finite values")
+        return 1
+
+    if rank == 0:
+        stats = f"shape={tuple(tp_out_cpu.shape)}  max_abs={tp_out_cpu.abs().max():.4f}  time={elapsed * 1000:.1f}ms"
+
+        if ref_output is not None:
+            # dummy-model mode: compare TP output against single-device TPU reference
+            max_abs_diff = (tp_out_cpu - ref_output).abs().max().item()
+            denom = ref_output.abs().max().item() + 1e-6
+            max_rel_diff = max_abs_diff / denom
+            stats += f"  max_abs_diff={max_abs_diff:.4e}  max_rel_diff={max_rel_diff:.4e}"
+
+            # TPU Flash Attention introduces bf16-level rounding vs standard SDPA.
+            # A wrong shard plan produces grossly different output (off by ~10x) and is caught
+            # easily within this tolerance; the bound is wider than the MLP case because
+            # Flash Attention rounds differently than sequential matmul+softmax.
+            if max_abs_diff > 0.1:
+                log(rank, f"FAIL: max_abs_diff={max_abs_diff:.4e} exceeds tolerance 0.1")
+                return 1
+
+        log(rank, f"PASS  {stats}")
+
+    dist.barrier()
+    dist.destroy_process_group()
+    return 0
+
+
+# ── entry point ───────────────────────────────────────────────────────────────
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Flux2 TP-on-TPU verification script.")
+    p.add_argument("--tp-degree", type=int, default=4, help="number of TPU chips to shard across")
+    p.add_argument("--model-id", type=str, default="", help="HuggingFace model ID (empty = random weights)")
+    p.add_argument(
+        "--latent-height",
+        type=int,
+        default=16,
+        help="latent grid height (default 16; combined with txt-len=256 gives joint-seq=512)",
+    )
+    p.add_argument(
+        "--latent-width",
+        type=int,
+        default=16,
+        help="latent grid width  (default 16; combined with txt-len=256 gives joint-seq=512)",
+    )
+    p.add_argument(
+        "--txt-len",
+        type=int,
+        default=256,
+        help="text sequence length (default 256; combined with 16x16 image gives joint-seq=512)",
+    )
+    p.add_argument("--topology", type=str, default="", help="TORCH_TPU_TOPOLOGY (e.g. '2,2,1')")
+    p.add_argument("--addresses", type=str, default="", help="TORCH_TPU_SLICEBUILDER_ADDRESSES")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    relaunch_via_torchrun(
+        args.tp_degree,
+        args.topology or None,
+        args.addresses or None,
+    )
+    raise SystemExit(run(args))
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except SystemExit:
+        raise
+    except Exception:
+        traceback.print_exc()
+        sys.exit(1)
diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py
index 53af9eedc50c..db6ecaced54f 100644
--- a/tests/models/transformers/test_models_transformer_flux.py
+++ b/tests/models/transformers/test_models_transformer_flux.py
@@ -27,7 +27,12 @@
 from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
+from ...testing_utils import (
+    enable_full_determinism,
+    is_tensor_parallel,
+    require_torch_neuron,
+    torch_device,
+)
 from ..testing_utils import (
     AttentionBackendTesterMixin,
     AttentionTesterMixin,
@@ -53,6 +58,7 @@
     SingleFileTesterMixin,
     TaylorSeerCacheTesterMixin,
     TensorParallelTesterMixin,
+    TensorParallelTPUTesterMixin,
     TorchAoCompileTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
@@ -268,6 +274,35 @@ class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParal
     """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator)."""
 
 
+def make_tpu_tp_spec():
+    """Model spec consumed by the generic TPU TP worker (``_tpu_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.
+
+    Overrides ``num_attention_heads`` to 4 (instead of reusing the shared tester config's 2) so
+    ``TensorParallelTPUTesterMixin``'s default 4-rank ``WORLD_SIZE`` divides the head count — see
+    ``make_tpu_tp_spec`` in ``test_models_transformer_flux2.py`` for the full rationale (TPU can't shard across an
+    arbitrary rank count the way CUDA/XPU can). Every other field still comes from the shared config so the rest of
+    the spec doesn't drift from the other Flux tests.
+    """
+    config = FluxTransformerTesterConfig()
+    init_dict = {**config.get_init_dict(), "num_attention_heads": 4}
+    return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu")
+
+
+class TestFluxTransformerTensorParallelTPU(TensorParallelTPUTesterMixin):
+    """Tensor Parallel inference test for Flux Transformer on TPU.
+
+    TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` 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_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts the
+    sharded output matches a single-device reference, and the test checks its exit code.
+    """
+
+    TP_SPEC = "tests.models.transformers.test_models_transformer_flux:make_tpu_tp_spec"
+
+
 def make_neuron_tp_spec():
     """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`).
 
diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py
index 3263ce68202c..dae72a863f24 100644
--- a/tests/models/transformers/test_models_transformer_flux2.py
+++ b/tests/models/transformers/test_models_transformer_flux2.py
@@ -28,7 +28,12 @@
 )
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
+from ...testing_utils import (
+    enable_full_determinism,
+    is_tensor_parallel,
+    require_torch_neuron,
+    torch_device,
+)
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
@@ -42,6 +47,7 @@
     ModelTesterMixin,
     SingleFileTesterMixin,
     TensorParallelTesterMixin,
+    TensorParallelTPUTesterMixin,
     TorchAoCompileTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
@@ -176,6 +182,40 @@ def make_neuron_tp_spec():
     return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
 
 
+def make_tpu_tp_spec():
+    """Model spec consumed by the generic TPU TP worker (``_tpu_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.
+
+    Overrides ``num_attention_heads`` to 4 (instead of reusing the shared tester config's 2) so the TP degree divides
+    the head count on a whole-pod-slice TPU host: `torch_tpu`'s per-generation topology table
+    (``torch_tpu._internal.utils.hardware``) only enumerates whole-slice chip counts (1/4/8 for v6e, for example), not
+    arbitrary sub-slices of a larger single host, and ``TestFlux2TransformerTensorParallelTPU`` shards across
+    ``WORLD_SIZE`` ranks to match. Every other field still comes from the shared config so the rest of the spec
+    doesn't drift from the other Flux2 tests.
+    """
+    config = Flux2TransformerTesterConfig()
+    init_dict = {**config.get_init_dict(), "num_attention_heads": 4}
+    return Flux2Transformer2DModel, init_dict, config.get_dummy_inputs(device="cpu")
+
+
+class TestFlux2TransformerTensorParallelTPU(TensorParallelTPUTesterMixin):
+    """Tensor Parallel inference test for Flux2 Transformer on TPU.
+
+    TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` 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_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts the sharded
+    output matches a single-device reference, and the test checks its exit code.
+
+    ``make_tpu_tp_spec`` overrides ``num_attention_heads`` to 4 so that ``TensorParallelTPUTesterMixin``'s default
+    4-rank ``WORLD_SIZE`` divides the head count — unlike the CUDA/XPU ``TensorParallelTesterMixin``, which hardcodes
+    ``world_size = 2`` to match `Flux2TransformerTesterConfig`'s 2 heads.
+    """
+
+    TP_SPEC = "tests.models.transformers.test_models_transformer_flux2:make_tpu_tp_spec"
+
+
 @is_tensor_parallel
 @require_torch_neuron
 class TestFlux2TransformerTensorParallelNeuron:
diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py
index 7a03a8fe2353..0b9e0dc50cbb 100644
--- a/tests/models/transformers/test_models_transformer_qwenimage.py
+++ b/tests/models/transformers/test_models_transformer_qwenimage.py
@@ -24,7 +24,12 @@
 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, is_tensor_parallel, require_torch_neuron, torch_device
+from ...testing_utils import (
+    enable_full_determinism,
+    is_tensor_parallel,
+    require_torch_neuron,
+    torch_device,
+)
 from ..testing_utils import (
     AttentionBackendTesterMixin,
     AttentionTesterMixin,
@@ -37,6 +42,7 @@
     MemoryTesterMixin,
     ModelTesterMixin,
     TensorParallelTesterMixin,
+    TensorParallelTPUTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
     TrainingTesterMixin,
@@ -307,6 +313,30 @@ class TestQwenImageTransformerTensorParallel(QwenImageTransformerTesterConfig, T
     """Tensor Parallel inference tests for QwenImage Transformer (CUDA/XPU multi-accelerator)."""
 
 
+def make_tpu_tp_spec():
+    """Model spec consumed by the generic TPU TP worker (``_tpu_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. ``QwenImageTransformerTesterConfig``'s default ``num_attention_heads``
+    (4) already divides ``TensorParallelTPUTesterMixin``'s default 4-rank ``WORLD_SIZE``, so no override is needed
+    here (contrast Flux/Flux2, whose shared config defaults to 2 heads and does need one).
+    """
+    config = QwenImageTransformerTesterConfig()
+    return QwenImageTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
+
+
+class TestQwenImageTransformerTensorParallelTPU(TensorParallelTPUTesterMixin):
+    """Tensor Parallel inference test for QwenImage Transformer on TPU.
+
+    TPU TP runs through ``torchrun`` with the ``"tpu_dist"`` 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_tpu_tp_spec``) via ``TensorParallelTPUTesterMixin``; the worker asserts
+    the sharded output matches a single-device reference, and the test checks its exit code.
+    """
+
+    TP_SPEC = "tests.models.transformers.test_models_transformer_qwenimage:make_tpu_tp_spec"
+
+
 def make_neuron_tp_spec():
     """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``).
 
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 9da89e626198..d9671828314c 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -46,6 +46,7 @@
     is_timm_available,
     is_torch_available,
     is_torch_neuronx_available,
+    is_torch_tpu_available,
     is_torch_version,
     is_torchao_available,
     is_torchsde_available,
@@ -565,6 +566,14 @@ def require_torch_neuron(test_case):
     )(test_case)
 
 
+def require_torch_tpu(test_case):
+    """Decorator marking a test that requires a TPU device (torch_tpu)."""
+    return pytest.mark.skipif(
+        not is_torch_tpu_available(),
+        reason="test requires TPU device (torch_tpu)",
+    )(test_case)
+
+
 def require_torch_multi_gpu(test_case):
     """
     Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without