Z-Image on AMD GPUs: the 1344x768 bucket captured, batch sizes 1-8

configs/z_image_turbo.yaml

@@ -18,26 +18,29 @@ model_kwargs:
   # [norms] for reference-faithful output.
   compile_exact_ops: false
   accelerator_graph: true
-  capture_sizes: [[1024, 1024]]
+  # [height, width]: 1024^2 and 1344x768, the two sizes of the image benchmark set; an uncaptured size compiles per
+  # batch size inside its first requests and decodes on the eager VAE.
+  capture_sizes: [[1024, 1024], [768, 1344]]
   # Captions pad to multiples of 32 up to max_seq_len, so cap_len has 16 possible
   # values -- and it is an INPUT, not a shape: the DiT substitutes a learned
   # cap_pad_token at every padded position and attends over it, so a prompt cannot
   # be padded up into a wider bucket without changing the image. A cap_len with no
   # captured graph falls back to the eager batched path.
   #
-  # The capture budget is len(capture_caption_lengths) x len(capture_batch_sizes) graphs,
-  # 16 here. A step of 3 or 5-7 rows replays the next captured size with dummy rows.
-  # Capturing every batch size up to max_batch_size (the commented line) skips that padding,
-  # which pays under concurrency (requests at different steps, one or two outside the dit).
-  # H200, 1024^2, 8 steps, 2 rounds (before the fused q/k norm + rotary):
+  # The capture budget per size is len(capture_caption_lengths) x len(capture_batch_sizes) graphs.
+  # A step of 3 or 5-7 rows would replay the next captured size with dummy rows; at 16 concurrent
+  # requests steps average 5.7 rows, so with only 1/2/4/8 captured about 30% of a step's rows would
+  # be dummies. Capturing every batch size pays under concurrency (requests at different steps, one
+  # or two outside the dit). H200, 1024^2, 8 steps, 2 rounds (before the fused q/k norm + rotary):
   # c=4 1.05-1.08 -> 1.22 images/s, c=8 1.12-1.14 -> 1.21-1.22,
   # c=16 1.17-1.19 -> 1.23, c=1 unchanged, same peak VRAM. But each graph is a compile of its own:
-  # 32 graphs start ~100 s slower than 16 on a warm Inductor cache, ~8 min slower on a cold one.
+  # 32 graphs start ~100 s slower than 16 on a warm Inductor cache, ~8 min slower on a cold one,
+  # and the 64 graphs of the two sizes cost 5.3 GiB of VRAM and take a warm start from about 96 s to about 220 s on MI325X.
   # The text encoder captures its own geometric_batch_sizes: powers of two below max_batch_size,
   # then max_batch_size (4 graphs here, any caption length).
   capture_caption_lengths: [32, 64, 96, 128]
-  capture_batch_sizes: [1, 2, 4, 8]
-  # capture_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8]  # for deployments that batch 3, 5, 6 or 7
+  # every batch size up to max_batch_size; [1, 2, 4, 8] halves the graphs for deployments that start cold often
+  capture_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8]
   max_batch_size: 8
   # compiled VAE decode (max-autotune), not bit-exact with the eager decode: before the folded upsample and split
   # GroupNorm it was 3x faster at 1024^2 and ~56 dB from eager, and the autotuner may pick other conv kernels in another

docs/architecture.rst

@@ -209,6 +209,14 @@ step.
   ``os.environ.setdefault`` in the conductor before it spawns workers, so values you
   export win. ROCm turns off MIOpen's split-K conv solver there (see
   :doc:`environment_variables`); the other platforms set nothing.
+- **Compile pool.** ``Platform.compile_pool_overrides()`` gives
+  ``torch._inductor.config`` values in force while a worker builds Inductor's compile
+  pool, before it touches the model, since a submodule can compile at load. The pool reads
+  them once, so they stay out of every graph's cache key, which the same setting in the
+  environment would not (that recompiles every graph a golden compile cache holds). ROCm
+  turns off the pool's idle quiesce: the sidecar
+  retires a pool without waiting for it, and a worker of the next pool forked meanwhile can
+  hang forever on the old pool's shutdown lock, along with the compile that gave it a job.
 - **Capture refusal.** A resource whose kernel can't run some capture bucket inside a
   graph returns a reason from ``Resource.capture_refusal(spec)``. Both graph runners ask
   before capturing, and a refused bucket runs eagerly. Answer from the spec and the

mstar/model/z_image/submodules.py

@@ -23,11 +23,9 @@
 
 from mstar.communication.tensors import NameToTensorList
 from mstar.conductor.request_info import CurrentForwardPassInfo
+from mstar.engine import recompile_limit
 from mstar.engine.accelerator_graph_config import geometric_batch_sizes
-from mstar.model.components.diffusion.compile_utils import (
-    compile_transformer_forward,
-    step_schedule_recompile_limit,
-)
+from mstar.model.components.diffusion.compile_utils import compile_transformer_forward
 from mstar.model.components.diffusion.decode_utils import (
     VAE_MAX_DECODE_BATCH,
     compile_vae_decode,
@@ -145,9 +143,16 @@ def __init__(
                 if module is not None:
                     compile_transformer_forward(
                         module, eager_rounding=compile_eager_rounding, exact_ops=compile_exact_ops,
-                        recompile_limit=step_schedule_recompile_limit(step_dtypes),
+                        recompile_limit=self._recompile_limit(step_dtypes),
                     )
 
+    def _recompile_limit(self, step_dtypes) -> int:
+        # every block runs one code object, so dynamo's per-code limit counts its shapes over all three spans of every
+        # captured graph, at most two per graph; a step schedule compiles each dtype on its own
+        per_dtype = max(recompile_limit(), 2 * len(self.capture_buckets) * len(self.capture_batch_sizes))
+        dtypes = 1 if step_dtypes is None else len(step_dtypes.dtypes)
+        return min(per_dtype * dtypes, 256 * dtypes)
+
     def bucket_key_for(self, fwd_info: CurrentForwardPassInfo) -> ZShape:
         return shape_from_metadata(self.config, fwd_info.step_metadata)
 

mstar/platforms/compile.py

@@ -1,4 +1,4 @@
-"""The platform's inductor settings, applied to one compiled callable."""
+"""The platform's inductor settings, applied to one compiled callable or to the compile pool."""
 
 from __future__ import annotations
 
@@ -39,6 +39,23 @@ def forward(*args, **kwargs):
     return forward
 
 
+def start_compile_pool() -> None:
+    """Build Inductor's compile pool under the platform's ``compile_pool_overrides``.
+
+    Call before the process's first compile, since the first caller builds the
+    one pool. The pool reads these when built, and keeps them after the patch.
+    A value set in the environment would stay in the config instead, and so in
+    every graph's cache key.
+    """
+    overrides = current_platform.compile_pool_overrides()
+    if not overrides:
+        return
+
+    import torch._inductor.config as inductor_config
+    from torch._inductor.async_compile import AsyncCompile
+
+    with inductor_config.patch(**overrides):
+        AsyncCompile.warm_pool()
 @functools.cache
 def compiled(fn: Callable[..., Any], dynamic: bool | None = None) -> Callable[..., Any]:
     """``torch.compile(fn)`` under the platform's inductor overrides, built on

mstar/platforms/interface.py

@@ -41,6 +41,18 @@ class HardwareQueueColocation(WorkerColocation):
     def start(self) -> None:
         logger.info("mps: true on ROCm: nothing to start, the GPU schedules the workers' queues itself")
 
+BF16_HEADER_FLAG = "-include hip/hip_bf16.h"
+
+
+def _with_bf16_header(flags: str | None) -> str:
+    """``FLASHINFER_EXTRA_CUDAFLAGS`` with the header amd-flashinfer 0.6.18+amd.1's
+    AITER prefill sources need: their generated config names ``__hip_bfloat16``
+    before anything declares it, since FlashInfer's ``-DUSE_ROCM=1`` turns off
+    the include in torch's ``BFloat16.h``. Later releases include it themselves."""
+    if flags is None:
+        return BF16_HEADER_FLAG
+    return flags if BF16_HEADER_FLAG in flags else f"{flags} {BF16_HEADER_FLAG}"
+
 
 class Platform:
     """One accelerator stack; ``current_platform`` is the one this process runs.
@@ -113,6 +125,12 @@ def inductor_overrides(
         """
         return {}
 
+    def compile_pool_overrides(self) -> dict[str, object]:
+        """``torch._inductor.config`` values in force while a worker builds
+        Inductor's compile pool. Unlike ``inductor_overrides`` these never reach
+        a graph's cache key. Empty keeps inductor's defaults."""
+        return {}
+
     def process_env_defaults(self) -> dict[str, str | Callable[[str | None], str]]:
         """Environment variables for settings a runtime or library reads once
         at init. Workers apply them with ``os.environ.setdefault`` before their
@@ -305,8 +323,16 @@ def inductor_overrides(
             overrides["post_grad_custom_post_pass"] = post
         return overrides
 
-    def process_env_defaults(self) -> dict[str, str]:
+    def compile_pool_overrides(self) -> dict[str, object]:
+        # Inductor's sidecar retires a pool that returned nothing for a minute
+        # without waiting for it. A worker of the next pool forked meanwhile
+        # inherits the old pool's shutdown lock held by a thread it doesn't
+        # have, and hangs at its first full GC with the first job it took.
+        return {"quiesce_async_compile_pool": False}
+
+    def process_env_defaults(self) -> dict[str, str | Callable[[str | None], str]]:
         return {
+            "FLASHINFER_EXTRA_CUDAFLAGS": _with_bf16_header,
             # MIOpen's split-K implicit-GEMM conv solvers are not deterministic. Turning
             # them off costs about 1% and makes same-seed outputs repeat. The backward
             # one runs transposed convs (LTX-2.5's vocoder upsamplers).

mstar/worker/engine_manager.py

@@ -16,6 +16,7 @@
 from mstar.graph.runtime.base import GraphRuntime
 from mstar.model.base import Model, node_moe_parallel_configs, parse_autocast_dtype, serving_resources
 from mstar.model.submodule_base import NodeSubmodule
+from mstar.platforms.compile import start_compile_pool
 from mstar.utils.device import DeviceCapability
 from mstar.utils.streams import reset_device_scheduling
 
@@ -269,6 +270,10 @@ def build(
         the latter, so there is no separate cache config. The conductor's
         ``validate_deployment`` has already checked the specs and overrides.
         """
+        # Before the model does anything: the process's first compile builds the
+        # pool, and a submodule can compile at load (klein warms its VAE decode).
+        if (device_capability or DeviceCapability.of(device)).is_available:
+            start_compile_pool()
         specs = serving_resources(model, model_config)
 
         # YAML autocast_dtype over the model's own, but never re-enabling

test/modular/test_platforms.py

@@ -1,16 +1,21 @@
 """``mstar.platforms``: which hardware the process runs on."""
 
 import copy
+import importlib.util
 import multiprocessing
 import os
 import pickle
+import shutil
+import subprocess
+import sys
 from types import SimpleNamespace
 
 import pytest
 import torch
 
 from mstar import platforms
 from mstar.platforms import CpuPlatform, CudaPlatform, RocmPlatform, XpuPlatform, current_platform
+from mstar.platforms.interface import BF16_HEADER_FLAG
 
 
 @pytest.fixture
@@ -163,6 +168,60 @@ def test_a_callable_env_default_adds_to_the_value_already_set(monkeypatch):
     assert os.environ["PLATFORM_ENV_TEST_SET"] == "user-default"
 
 
+@pytest.mark.parametrize(
+    ("exported", "expected"),
+    [
+        (None, BF16_HEADER_FLAG),
+        ("-DUSER_FLAG", f"-DUSER_FLAG {BF16_HEADER_FLAG}"),
+        (f"{BF16_HEADER_FLAG} -DUSER_FLAG", f"{BF16_HEADER_FLAG} -DUSER_FLAG"),
+    ],
+)
+def test_rocm_adds_the_bf16_header_to_flashinfer_flags(monkeypatch, exported, expected):
+    monkeypatch.setattr(torch.version, "hip", "7.0")
+    monkeypatch.setenv("FLASHINFER_EXTRA_CUDAFLAGS", "")
+    if exported is None:
+        monkeypatch.delenv("FLASHINFER_EXTRA_CUDAFLAGS")
+    else:
+        monkeypatch.setenv("FLASHINFER_EXTRA_CUDAFLAGS", exported)
+
+    # the conductor applies the defaults, then each worker again
+    platforms.apply_process_env_defaults()
+    platforms.apply_process_env_defaults()
+    assert os.environ["FLASHINFER_EXTRA_CUDAFLAGS"] == expected
+
+
+# aiter's import asks rocminfo for a live GPU; the module generator reads only its path
+_BF16_PREFILL_JIT = """
+import importlib.util, sys, types
+import torch
+aiter = types.ModuleType("aiter")
+aiter.__file__ = importlib.util.find_spec("aiter").origin
+sys.modules["aiter"] = aiter
+from flashinfer.jit.rocm.modules import gen_batch_prefill_module
+bf16 = torch.bfloat16
+gen_batch_prefill_module("aiter", bf16, bf16, bf16, torch.int32, 64, 64, 0, False, False, False).build()
+"""
+
+
+def test_flashinfer_builds_its_bf16_aiter_prefill_with_the_rocm_flags(tmp_path):
+    if importlib.util.find_spec("aiter") is None:
+        pytest.skip("needs AITER")
+    pytest.importorskip("flashinfer.jit.rocm.modules")
+    if shutil.which("amdclang++") is None:
+        pytest.skip("needs amdclang++")
+    env = {
+        **os.environ,
+        "FLASHINFER_EXTRA_CUDAFLAGS": BF16_HEADER_FLAG,
+        "FLASHINFER_ROCM_ARCH_LIST": "gfx942",
+        "FLASHINFER_WORKSPACE_BASE": str(tmp_path),
+    }
+    built = subprocess.run(
+        [sys.executable, "-c", _BF16_PREFILL_JIT],
+        env=env, capture_output=True, text=True, check=False,
+    )
+    assert built.returncode == 0, built.stderr[-2000:] + built.stdout[-2000:]
+
+
 def test_a_worker_sets_its_env_defaults_before_anything_else(monkeypatch):
     from mstar.conductor import conductor
 
@@ -246,6 +305,74 @@ def test_only_rocm_turns_off_miopens_split_k_conv_solvers(solver):
         assert solver not in platform.process_env_defaults()
 
 
+def test_only_rocm_builds_inductors_compile_pool_without_its_quiesce():
+    assert RocmPlatform().compile_pool_overrides() == {"quiesce_async_compile_pool": False}
+    for platform in (CudaPlatform(), XpuPlatform(), CpuPlatform()):
+        assert platform.compile_pool_overrides() == {}
+
+
+def test_the_compile_pool_keeps_its_settings_out_of_the_cache_key():
+    # a real sidecar, in a process of its own since torch keeps the first pool it builds
+    code = """
+import torch._inductor.config as config
+from torch._inductor.async_compile import AsyncCompile
+from mstar.platforms import RocmPlatform, compile as platform_compile
+
+key = config.save_config_portable()
+platform_compile.current_platform = RocmPlatform()
+platform_compile.start_compile_pool()
+pool = AsyncCompile.process_pool()
+assert pool.timer is None, "the pool would still quiesce"
+assert config.save_config_portable() == key, "a graph's cache key moved"
+pool.shutdown()
+"""
+    subprocess.run([sys.executable, "-c", code], env={**os.environ, "TORCHINDUCTOR_COMPILE_THREADS": "2"}, check=True)
+
+
+@pytest.mark.parametrize(
+    ("platform", "available", "built"),
+    [("RocmPlatform", True, True), ("RocmPlatform", False, False), ("CudaPlatform", True, False)],
+)
+def test_a_worker_builds_the_compile_pool_before_it_touches_the_model(platform, available, built):
+    # a model can compile at load, and the first compile builds the pool under inductor's defaults
+    code = f"""
+from types import SimpleNamespace
+
+import torch
+from torch._inductor.async_compile import AsyncCompile
+
+from mstar import platforms
+from mstar.platforms import compile as platform_compile
+from mstar.worker.engine_manager import EngineManager
+
+
+class ModelTouched(Exception):
+    pass
+
+
+class Model:
+    def __getattr__(self, name):
+        raise ModelTouched
+
+
+platform_compile.current_platform = platforms.{platform}()
+try:
+    EngineManager.build(
+        node_names={{"node"}}, device=torch.device("cpu"), model_config={{}}, parallel_groups=None,
+        transfer_engine_info=None, model=Model(), graph_runtime=None,
+        device_capability=SimpleNamespace(is_available={available}),
+    )
+except ModelTouched:
+    pass
+assert AsyncCompile.process_pool.cache_info().currsize == {int(built)}
+if {built}:
+    pool = AsyncCompile.process_pool()
+    assert pool.timer is None, "the pool would still quiesce"
+    pool.shutdown()
+"""
+    subprocess.run([sys.executable, "-c", code], env={**os.environ, "TORCHINDUCTOR_COMPILE_THREADS": "2"}, check=True)
+
+
 def test_platform_knob_defaults_are_empty_but_for_rocm():
     assert RocmPlatform().knob_defaults() == {"dit_exact_ops": "0"}
     for platform in (CudaPlatform(), XpuPlatform(), CpuPlatform()):

test/modular/test_z_image_model.py

@@ -116,6 +116,31 @@ def test_capture_shapes_cover_sizes_times_caption_buckets():
     assert _make_model(accelerator_graph=False).capture_buckets() == []
 
 
+@pytest.mark.parametrize(("capture_sizes", "batch_sizes", "limit"), [
+    ([[1024, 1024]], [1, 2, 4, 8], None),  # 2 x 4 x 4 graphs fit the process-wide limit
+    ([[1024, 1024], [768, 1344]], list(range(1, 9)), 128),
+])
+def test_the_transformers_recompile_limit_covers_the_capture_plan(monkeypatch, capture_sizes, batch_sizes, limit):
+    from mstar.engine import recompile_limit
+    from mstar.model.components.diffusion import compile_utils
+    from mstar.model.z_image.submodules import ZImageDenoiseSubmodule
+
+    seen = {}
+
+    def fake_compile(fn, **kwargs):
+        seen.update(kwargs)
+        return fn
+
+    monkeypatch.setattr(compile_utils.torch, "compile", fake_compile)
+    model = _make_model(capture_sizes=capture_sizes, capture_caption_lengths=[32, 64, 96, 128],
+                        capture_batch_sizes=batch_sizes)
+    ZImageDenoiseSubmodule(
+        torch.nn.Linear(2, 2), ZImageConfig(), loop_name=DENOISE_LOOP, attn_resource_key=None,
+        capture_buckets=model.capture_buckets(), capture_batch_sizes=model.capture_batch_sizes,
+    )
+    assert seen["recompile_limit"] == (limit or recompile_limit())
+
+
 def test_postprocess_png_and_adapter():
     model = _make_model()
     assert model.postprocess(torch.zeros(1, 3, 16, 16, dtype=torch.uint8), "image")[:4] == b"\x89PNG"

test/rocm/test_varlen_parity_shapes.py

@@ -346,10 +346,11 @@ def test_patch_embed_conv(case, worker_env):
 
     cin, cout, kernel, shape, dtype, tol = _PATCH_CONVS[case]
     env = dict(os.environ)
-    for key in _built_for().process_env_defaults():
+    defaults = _built_for().process_env_defaults()
+    for key in defaults:
         env.pop(key, None)
     if worker_env:
-        env.update(_built_for().process_env_defaults())
+        env.update({k: v(None) if callable(v) else v for k, v in defaults.items()})
     res = subprocess.run(
         [sys.executable, "-c", _CONV_CHECK, json.dumps([cin, cout, kernel, shape, dtype])],
         env=env, capture_output=True, text=True, timeout=300, check=True,