FLUX.2: wire e2e pipeline into benchmark CI + 4-chip nightly

Mirrors the image-gen e2e pattern (#5044 Playground, #5244 SDXL-Lightning,
#5291 Janus-Pro), extended for the first multichip (tensor-parallel) image-gen
model. FLUX.2 runs on 4 Blackhole chips: ~24B Mistral3 text encoder + ~32B
Flux2 transformer SPMD-sharded (model-parallel degree 4), VAE replicated.

Benchmark CI:
- imagegen_benchmark.py: make per-component perf reporting component-set
  agnostic (a model reports only the components it runs; FLUX.2 has a single
  text encoder, so te2 is omitted), with an optional step label and mesh_shape
  reporting. SDXL/Playground output is unchanged (verified on hardware).
- benchmarks/flux2_pipeline.py: new self-contained multichip pipeline. Sets
  CONVERT_SHLO_TO_SHARDY=1 before use_spmd() (required so tt-mlir gets shardy
  annotations), mesh from MESH_SHAPES[num_devices].
- test_imagegen.py: add test_flux2.
- perf-bench-matrix.json: flux2 entry on qb2-blackhole (4-chip).

Nightly test (tests/torch/models/flux2/test_flux2_pipeline.py):
- Add @pytest.mark.qb2_blackhole and make the mesh adaptive via MESH_SHAPES
  (no-op at 4 chips).
- Add a qb2-blackhole standalone-models job (scoped to the flux2 dir) to
  model-test-lb-blackhole-nightly.json so the test is collected on 4 chips.

Validated on a 4-chip qb2 Blackhole: FLUX.2 benchmark passes with a coherent
1024x1024 image; SDXL single-chip benchmark passes unchanged.

Author: ashokkumarkannan1

.github/workflows/perf-bench-matrix.json

@@ -70,6 +70,11 @@
         "name": "sdxl_lightning",
         "pytest": "tests/benchmark/test_imagegen.py::test_sdxl_lightning"
       },
+      {
+        "name": "flux2",
+        "pytest": "tests/benchmark/test_imagegen.py::test_flux2",
+        "runs-on": "qb2-blackhole"
+      },
       {
         "name": "bge_m3_encode",
         "pyreq": "transformers==4.57.1 FlagEmbedding",

.github/workflows/test-matrix-presets/model-test-lb-blackhole-nightly.json

@@ -2,5 +2,6 @@
   {"runs-on": "lb-blackhole",   "name": "run_forge_models_lb_blackhole_multichip",    "dir": "./tests/runner/test_models.py::test_all_models_torch", "test-mark": "lb-blackhole and tensor_parallel and expected_passing", "forge-models": true},
   {"runs-on": "lb-blackhole",   "name": "run_forge_models_lb_blackhole_multichip",    "dir": "./tests/runner/test_models.py::test_all_models_jax", "test-mark": "lb-blackhole and tensor_parallel and expected_passing", "forge-models": true},
   {"runs-on": "lb-blackhole",   "name": "run_forge_models_lb_blackhole_llm_multichip", "dir": "./tests/runner/test_models.py::test_llms_torch", "test-mark": "lb-blackhole and tensor_parallel and expected_passing", "forge-models": true},
-  { "runs-on": "lb-blackhole",  "name": "run_standalone_models_lb_blackhole_multichip", "dir": "./tests/torch/models", "test-mark": "lb_blackhole" }
+  { "runs-on": "lb-blackhole",  "name": "run_standalone_models_lb_blackhole_multichip", "dir": "./tests/torch/models", "test-mark": "lb_blackhole" },
+  { "runs-on": "qb2-blackhole", "name": "run_standalone_models_qb2_blackhole_flux2", "dir": "./tests/torch/models/flux2", "test-mark": "qb2_blackhole" }
 ]

tests/benchmark/benchmarks/flux2_pipeline.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""FLUX.2-dev — multichip benchmark-side pipeline for the imagegen harness.
+
+Mirrors the nightly e2e test (tests/torch/models/flux2/test_flux2_pipeline.py):
+the diffusers Flux2Pipeline orchestrates, the Mistral3 text encoder and Flux2
+transformer are SPMD-sharded (model axis), the VAE is replicated. The text
+encoder is evicted before the transformer is placed and the VAE is placed
+lazily, so peak DRAM ≈ max(component). Per-component times go into ``self._perf``.
+"""
+
+import os
+import time
+from typing import Optional
+
+import numpy as np
+import torch
+import torch_xla
+import torch_xla.distributed.spmd as xs
+import torch_xla.runtime as xr
+from diffusers import Flux2Pipeline
+from loguru import logger
+from torch_xla.distributed.spmd import Mesh
+
+from third_party.tt_forge_models.flux2.pytorch.src.model_utils import (
+    DTYPE,
+    GUIDANCE_SCALE,
+    HEIGHT,
+    MESH_NAMES,
+    MESH_SHAPES,
+    PROMPT,
+    REPO_ID,
+    SEED,
+    WIDTH,
+    Mistral3TextEncoderWrapper,
+    shard_text_encoder_specs,
+    shard_transformer_specs,
+    tokenize_prompt,
+)
+
+
+class _DeviceDenoiser:
+    """Routes the transformer to the TP-sharded model on TT; each call is one
+    denoise step, timed into ``perf["unet_steps"]``."""
+
+    def __init__(self, transformer, mesh, perf):
+        self._dev = torch_xla.device()
+        self._perf = perf
+        self.config = transformer.config
+        self.dtype = next(transformer.parameters()).dtype
+
+        transformer = transformer.to(self._dev)
+        if hasattr(transformer, "tie_weights"):
+            transformer.tie_weights()
+        for tensor, spec in shard_transformer_specs(transformer).items():
+            xs.mark_sharding(tensor, mesh, spec)
+        self._compiled = torch.compile(transformer, backend="tt")
+
+    def __call__(self, **kwargs):
+        moved = {
+            k: (v.to(self._dev) if torch.is_tensor(v) else v) for k, v in kwargs.items()
+        }
+        t0 = time.perf_counter()
+        out = self._compiled(**moved)
+        torch_xla.sync()
+        # cpu cast forces a sync — the timer ends after the result lands on host.
+        if isinstance(out, (tuple, list)):
+            result = type(out)(o.cpu() if torch.is_tensor(o) else o for o in out)
+        else:
+            result = out.cpu()
+        self._perf["unet_steps"].append(time.perf_counter() - t0)
+        return result
+
+
+class _DeviceVAEDecoder:
+    """Routes vae.decode() to TT (replicated), placed lazily. Decode time goes
+    into ``perf["vae"]``; raw pixels ([-1, 1]) are stashed on ``last_pixels`` so
+    the harness can save them without the pipeline's PIL postprocess."""
+
+    def __init__(self, vae, mesh, perf):
+        self._dev = torch_xla.device()
+        self._perf = perf
+        self.config = vae.config
+        self.dtype = next(vae.parameters()).dtype
+        self.bn = vae.bn  # stays on CPU; pipeline reads it host-side for denorm
+        self._vae = vae
+        self._compiled = None
+        self.last_pixels = None
+
+    def decode(self, latents, return_dict=False):
+        # Lazy device placement: keep the VAE off-device during the denoise loop
+        # so it does not inflate the denoiser's peak DRAM; place it only now.
+        if self._compiled is None:
+            vae = self._vae.to(self._dev)
+            self._compiled = torch.compile(
+                lambda z: vae.decode(z, return_dict=False)[0], backend="tt"
+            )
+        t0 = time.perf_counter()
+        out = self._compiled(latents.to(self._dev))
+        torch_xla.sync()
+        image = out.cpu() if torch.is_tensor(out) else out
+        self._perf["vae"] = time.perf_counter() - t0
+        self.last_pixels = image
+        return (image,)
+
+
+class Flux2Config:
+    def __init__(
+        self,
+        height: int = HEIGHT,
+        width: int = WIDTH,
+        compile_options: Optional[dict] = None,
+    ):
+        self.height = height
+        self.width = width
+        # Forwarded for parity with the other imagegen pipelines; unused inline.
+        self.compile_options = compile_options or {}
+
+
+class Flux2Pipeline_TT:
+    """Flux2Pipeline with every module on TT, tensor-parallel sharded. Built
+    once; generate() runs twice (warmup + steady). Raw transformer/VAE are kept
+    so the TT wrappers can be rebuilt fresh each call."""
+
+    def __init__(self, config: Flux2Config):
+        self.config = config
+        self._perf = {}
+
+    def setup(self):
+        # CONVERT_SHLO_TO_SHARDY=1 + use_spmd() (mirrors infra.enable_spmd):
+        # required so tt-mlir gets shardy annotations, else presharded args lose
+        # their @Sharding custom call and compilation fails.
+        os.environ["CONVERT_SHLO_TO_SHARDY"] = "1"
+        xr.use_spmd()
+        self.num_devices = xr.global_runtime_device_count()
+        self.mesh_shape = MESH_SHAPES[self.num_devices]
+        device_ids = np.array(range(self.num_devices))
+        self.mesh = Mesh(device_ids, self.mesh_shape, MESH_NAMES)
+        logger.info(
+            f"FLUX.2 mesh {self.mesh_shape} (names={MESH_NAMES}) "
+            f"over {self.num_devices} devices"
+        )
+        self.pipe = Flux2Pipeline.from_pretrained(REPO_ID, torch_dtype=DTYPE)
+        # Keep the raw modules so wrappers can be rebuilt on each generate().
+        self._raw_transformer = self.pipe.transformer
+        self._raw_vae = self.pipe.vae
+
+    def generate(
+        self,
+        prompt: str,
+        num_inference_steps: int,
+        seed: Optional[int] = SEED,
+    ):
+        dev = torch_xla.device()
+        self._perf = {
+            "te1": None,
+            "unet_steps": [],
+            "vae": None,
+            "total": None,
+            "step_label": "Transformer",
+        }
+        t_total_start = time.perf_counter()
+
+        # ── Stage 1: text encoder (sharded, compiled) → prompt embeds, evict ──
+        logger.info("[STAGE] Text encoder: start")
+        text_encoder = self.pipe.text_encoder
+        encoder_wrapper = Mistral3TextEncoderWrapper(text_encoder).eval()
+        input_ids, attention_mask = tokenize_prompt(prompt)
+
+        text_encoder = text_encoder.to(dev)
+        if hasattr(text_encoder, "tie_weights"):
+            text_encoder.tie_weights()
+        te_specs = shard_text_encoder_specs(text_encoder)
+        assert te_specs, "text-encoder shard spec is empty — descent failed (would OOM)"
+        for tensor, spec in te_specs.items():
+            xs.mark_sharding(tensor, self.mesh, spec)
+        te_compiled = torch.compile(encoder_wrapper, backend="tt")
+
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            prompt_embeds = te_compiled(input_ids.to(dev), attention_mask.to(dev))
+        torch_xla.sync()
+        prompt_embeds = prompt_embeds.cpu()
+        self._perf["te1"] = time.perf_counter() - t0
+
+        # Free the 24B encoder from device before placing the 32B denoiser.
+        self.pipe.text_encoder = text_encoder.to("cpu")
+        del te_compiled, encoder_wrapper
+        import gc
+
+        gc.collect()
+        torch_xla.sync()
+        logger.info("[STAGE] Text encoder: done")
+
+        # ── Stage 2: denoiser (sharded) + VAE (replicated, lazy) → image ─────
+        logger.info("[STAGE] Transformer + VAE: start")
+        self.pipe.transformer = _DeviceDenoiser(
+            self._raw_transformer, self.mesh, self._perf
+        )
+        vae_wrapper = _DeviceVAEDecoder(self._raw_vae, self.mesh, self._perf)
+        self.pipe.vae = vae_wrapper
+
+        generator = torch.Generator().manual_seed(seed) if seed is not None else None
+        self.pipe(
+            prompt=None,
+            prompt_embeds=prompt_embeds,
+            height=self.config.height,
+            width=self.config.width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=GUIDANCE_SCALE,
+            generator=generator,
+        )
+        logger.info("[STAGE] Transformer + VAE: done")
+
+        self._perf["total"] = time.perf_counter() - t_total_start
+        # Raw VAE pixels in [-1, 1], shape (1, 3, H, W) — the harness's
+        # save_image() expects this range, so return it instead of the PIL.
+        return vae_wrapper.last_pixels

tests/benchmark/benchmarks/imagegen_benchmark.py

@@ -120,12 +120,24 @@ def benchmark_imagegen_torch_xla(
     total_samples = 1
     samples_per_sec = total_samples / steady_state_time
 
-    # Per-component forward+sync times from the pipeline's own instrumentation
-    # (steady-state pass).
+    # Per-component steady-state times. A model reports only the components it
+    # runs (e.g. FLUX.2 has a single text encoder, so te2 is absent); missing
+    # ones are treated as 0 s and dropped from the breakdown + measurements.
     perf = pipeline._perf
-    unet_steps = perf["unet_steps"]
-    unet_step_mean_s = sum(unet_steps) / len(unet_steps)
-    tt_components_total = perf["te1"] + perf["te2"] + sum(unet_steps) + perf["vae"]
+    unet_steps = perf.get("unet_steps") or []
+    unet_step_mean_s = sum(unet_steps) / len(unet_steps) if unet_steps else 0.0
+
+    def _component_s(key):
+        value = perf.get(key)
+        return value if value is not None else 0.0
+
+    te1_s = _component_s("te1")
+    te2_s = _component_s("te2")
+    vae_s = _component_s("vae")
+    # Per-step label (UNet / Transformer); measurement name stays unet_step_mean_s.
+    step_label = perf.get("step_label", "UNet")
+
+    tt_components_total = te1_s + te2_s + sum(unet_steps) + vae_s
     cpu_overhead = max(0.0, perf["total"] - tt_components_total)
 
     metadata = get_benchmark_metadata()
@@ -149,25 +161,40 @@ def benchmark_imagegen_torch_xla(
         data_format="bfloat16",
         input_size=input_size,
     )
-    print(
-        f"| Num inference steps: {num_inference_steps}\n"
-        f"| Steady-state:\n"
-        f"|   Text encoder 1 (s):  {perf['te1']:.3f}\n"
-        f"|   Text encoder 2 (s):  {perf['te2']:.3f}\n"
-        f"|   UNet step mean (s):  {unet_step_mean_s:.3f}\n"
-        f"|   VAE (s):             {perf['vae']:.3f}\n"
-        f"|   CPU overhead (s):    {cpu_overhead:.3f}"
-    )
+    breakdown_lines = [
+        f"| Num inference steps: {num_inference_steps}",
+        "| Steady-state:",
+    ]
+    if perf.get("te1") is not None:
+        breakdown_lines.append(f"|   Text encoder 1 (s):  {te1_s:.3f}")
+    if perf.get("te2") is not None:
+        breakdown_lines.append(f"|   Text encoder 2 (s):  {te2_s:.3f}")
+    breakdown_lines.append(f"|   {step_label} step mean (s):  {unet_step_mean_s:.3f}")
+    if perf.get("vae") is not None:
+        breakdown_lines.append(f"|   VAE (s):             {vae_s:.3f}")
+    breakdown_lines.append(f"|   CPU overhead (s):    {cpu_overhead:.3f}")
+    print("\n".join(breakdown_lines))
 
     custom_measurements = [
         {"measurement_name": "images_per_second", "value": samples_per_sec},
         {"measurement_name": "e2e_latency", "value": steady_state_time},
-        {"measurement_name": "text_encoder_1_s", "value": perf["te1"]},
-        {"measurement_name": "text_encoder_2_s", "value": perf["te2"]},
-        {"measurement_name": "unet_step_mean_s", "value": unet_step_mean_s},
-        {"measurement_name": "vae_s", "value": perf["vae"]},
-        {"measurement_name": "cpu_overhead_s", "value": cpu_overhead},
     ]
+    if perf.get("te1") is not None:
+        custom_measurements.append(
+            {"measurement_name": "text_encoder_1_s", "value": te1_s}
+        )
+    if perf.get("te2") is not None:
+        custom_measurements.append(
+            {"measurement_name": "text_encoder_2_s", "value": te2_s}
+        )
+    custom_measurements.append(
+        {"measurement_name": "unet_step_mean_s", "value": unet_step_mean_s}
+    )
+    if perf.get("vae") is not None:
+        custom_measurements.append({"measurement_name": "vae_s", "value": vae_s})
+    custom_measurements.append(
+        {"measurement_name": "cpu_overhead_s", "value": cpu_overhead}
+    )
 
     result = create_benchmark_result(
         full_model_name=full_model_name,
@@ -192,6 +219,7 @@ def benchmark_imagegen_torch_xla(
         device_name=socket.gethostname(),
         arch=get_xla_device_arch(),
         device_count=xr.global_runtime_device_count(),
+        mesh_shape=getattr(pipeline, "mesh_shape", None),
         input_is_image=True,
     )
 

tests/benchmark/test_imagegen.py

@@ -191,3 +191,49 @@ def generate_fn(prompt, steps):
         optimization_level=0,
         output_image_path="test_sdxl_lightning_output.png",
     )
+
+
+def test_flux2(output_file, request):
+    from benchmarks.flux2_pipeline import Flux2Config, Flux2Pipeline_TT
+    from third_party.tt_forge_models.flux2.pytorch.src.model_utils import (
+        HEIGHT,
+        NUM_INFERENCE_STEPS,
+        PROMPT,
+        WIDTH,
+    )
+
+    # FLUX.2-dev: ~24B Mistral3 text encoder + ~32B Flux2 transformer (both
+    # tensor-parallel sharded across the mesh's model axis) + replicated VAE.
+    # Multichip — wired to the 4-chip blackhole (qb2) in perf-bench-matrix.json.
+    prompt = PROMPT
+    num_inference_steps = NUM_INFERENCE_STEPS
+    height = HEIGHT
+    width = WIDTH
+
+    def build_pipeline_fn(compile_options):
+        pipeline = Flux2Pipeline_TT(
+            config=Flux2Config(compile_options=compile_options)
+        )
+        pipeline.setup()
+
+        def generate_fn(prompt, steps):
+            return pipeline.generate(
+                prompt=prompt,
+                num_inference_steps=steps,
+                seed=DEFAULT_SEED,
+            )
+
+        return pipeline, generate_fn
+
+    test_imagegen(
+        build_pipeline_fn=build_pipeline_fn,
+        model_info_name="flux2",
+        output_file=output_file,
+        request=request,
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        height=height,
+        width=width,
+        optimization_level=0,
+        output_image_path="test_flux2_output.png",
+    )