Qwen-Image-2.1 inference on AMD ROCm GPUs (R9700 FP8 / W7900 INT8)

configs/platforms/amd_rocm/qwen_image_21_bf16.json

@@ -0,0 +1,16 @@
+{
+    "infer_steps": 40,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "torch_sdpa",
+    "rope_type": "torch_real_rope",
+    "rms_norm_type": "fp32_variance",
+    "layer_norm_type": "torch",
+    "modulate_type": "torch",
+    "fused_qk_rms_norm": false,
+    "fused_block_ops": false,
+    "enable_cfg": false,
+    "sample_guide_scale": 1.0,
+    "warmup": false,
+    "text_encoder_cpu_offload": true
+}

configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json

@@ -0,0 +1,19 @@
+{
+    "infer_steps": 40,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2_amd_rocm",
+    "rope_type": "torch_real_rope",
+    "rms_norm_type": "fp32_variance",
+    "layer_norm_type": "torch",
+    "modulate_type": "torch",
+    "fused_qk_rms_norm": false,
+    "fused_block_ops": false,
+    "enable_cfg": false,
+    "sample_guide_scale": 1.0,
+    "warmup": true,
+    "text_encoder_cpu_offload": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-rocm",
+    "use_compile": true
+}

configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json

@@ -0,0 +1,19 @@
+{
+    "infer_steps": 40,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2_amd_rocm",
+    "rope_type": "torch_real_rope",
+    "rms_norm_type": "fp32_variance",
+    "layer_norm_type": "torch",
+    "modulate_type": "torch",
+    "fused_qk_rms_norm": false,
+    "fused_block_ops": false,
+    "enable_cfg": false,
+    "sample_guide_scale": 1.0,
+    "warmup": true,
+    "text_encoder_cpu_offload": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "int8-rocm",
+    "use_compile": true
+}

lightx2v/models/networks/base_model.py

@@ -113,6 +113,8 @@ def _check_dit_quantized(self):
             bool: True if the model is quantized
         """
         assert self.config.get("dit_quant_scheme", "Default") in [
+            "fp8-rocm",
+            "int8-rocm",
             "fp8-pertensor",
             "fp8-musa",
             "fp8-triton",

lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py

@@ -16,6 +16,7 @@ def __init__(self, config):
         else:
             tp_size = 1
         self.heads = config["num_attention_heads"] // tp_size
+        self.head_dim = config["attention_head_dim"]
         self.use_fused_qk_rms_norm = config.get("fused_qk_rms_norm", True)
         self.use_fused_block_ops = config.get("fused_block_ops", True) and config.get("modulate_type", "triton") == "triton"
         self.fused_block_min_tokens = config.get("fused_block_min_tokens", 1024)
@@ -46,7 +47,7 @@ def _modulation(self, state):
 
     def _qkv(self, block, x, scale, rotary, rotary_positions):
         h = block.norm1.apply(x) * scale
-        q = block.q.apply(h).reshape(-1, self.heads, self.config["attention_head_dim"])
+        q = block.q.apply(h).reshape(-1, self.heads, self.head_dim)
         k = block.k.apply(h).reshape_as(q)
         v = block.v.apply(h).reshape_as(q)
         q, k = apply_qk_rms_norm(q, k, block.norm_q, block.norm_k, use_triton=self.use_fused_qk_rms_norm)
@@ -70,7 +71,7 @@ def prefill_block(self, index, block, x, modulation, state, cache):
         """Run one condition-prefix block and store its K/V."""
         q, k, v = self._qkv(block, x, modulation[0], state.rotary, state.rotary_positions)
         cache.store_kv(k, v, index)
-        attention = x.new_empty((x.shape[0], self.heads * self.config["attention_head_dim"]))
+        attention = x.new_empty((x.shape[0], self.heads * self.head_dim))
         for begin, end, is_text in state.layout.segments:
             mask = None
             if is_text:
@@ -134,5 +135,6 @@ def infer(self, weights, state, cache):
         """Denoise only target tokens, attending to the prefilled condition K/V."""
         if not cache.is_ready():
             raise RuntimeError("Condition KV must be prefilled before denoising")
+        # Shared across every block; kept outside the compiled block graph.
         modulation = self._modulation(state)
         return self._infer_blocks(weights.blocks, state.hidden_states, modulation, state, cache)

lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py

@@ -81,9 +81,10 @@ def __init__(self, config):
             if parallel.get("seq_p_quant_scheme") not in (None, "fp8", "fp4"):
                 raise ValueError("qwen_image_21 seq_p_quant_scheme supports only fp8 and fp4")
         if config.get("dit_quantized"):
-            if config.get("dit_quant_scheme") not in ("fp8-sgl", "fp8-f16-accum"):
-                raise ValueError("qwen_image_21 DiT quantization supports only fp8-sgl and fp8-f16-accum")
-            if not config.get("dit_quantized_ckpt"):
+            if config.get("dit_quant_scheme") not in ("fp8-sgl", "fp8-f16-accum", "fp8-rocm", "int8-rocm"):
+                raise ValueError("qwen_image_21 DiT quantization supports only fp8-sgl, fp8-f16-accum, fp8-rocm and int8-rocm")
+            # fp8-rocm / int8-rocm quantize the released BF16 weights on load, no prequantized checkpoint.
+            if config["dit_quant_scheme"] not in ("fp8-rocm", "int8-rocm") and not config.get("dit_quantized_ckpt"):
                 raise ValueError("qwen_image_21 FP8 requires dit_quantized_ckpt")
             if config["dit_quant_scheme"] == "fp8-f16-accum":
                 validate_fp8_f16_accum_qmax(config.get("dit_fp8_activation_qmax", ACTIVATION_QMAX))

lightx2v_platform/base/amd_rocm.py

@@ -4,7 +4,7 @@
 AMD ROCm provides CUDA-compatible APIs through HIP (Heterogeneous-computing Interface for Portability).
 This module handles AMD-specific optimizations including:
 - Disabling cudnn for faster VAE convolution
-- sgl_kernel compatibility layer using aiter library (required on AMD)
+- sgl_kernel compatibility layer using aiter library (optional; only for the aiter-backed GEMM/RMSNorm path)
 """
 
 import sys
@@ -85,10 +85,12 @@ def _get_aiter_sgl_kernel():
 
         return AiterSglKernelCompat(aiter)
     except ImportError:
-        logger.error(
-            f"\n{'=' * 60}\nERROR: AMD ROCm detected but aiter is not installed.\naiter is REQUIRED for LightX2V to work on AMD GPUs.\n\nPlease install aiter:\n{AITER_INSTALL_CMD}\n{'=' * 60}\n"
-        )
-        raise ImportError(f"aiter is required for AMD ROCm support. Please install: pip install git+{AITER_REPO}@{AITER_COMMIT}")
+        # aiter is optional: it only backs the aiter_attn / sgl_kernel GEMM +
+        # RMSNorm path. Torch-native paths (e.g. torch._scaled_mm / torch._int_mm
+        # fp8/int8 GEMM, torch_sdpa / SageAttention) run without it. Warn and let
+        # the caller skip the sgl_kernel injection rather than aborting setup.
+        logger.warning(f"aiter not installed; skipping the aiter sgl_kernel/RMSNorm compatibility layer. This is fine for torch-native paths. To enable the aiter-backed path:\n{AITER_INSTALL_CMD}")
+        return None
 
 
 @PLATFORM_DEVICE_REGISTER("amd_rocm")
@@ -109,27 +111,31 @@ def init_device_env():
 
         This is called from lightx2v_platform.set_ai_device when platform is amd_rocm.
         1. Disable cudnn for faster VAE convolution
-        2. Inject aiter as sgl_kernel compatibility layer (REQUIRED on AMD)
+        2. Inject aiter as sgl_kernel compatibility layer (optional; only the
+           aiter-backed GEMM/RMSNorm path needs it)
         """
         logger.info("AMD ROCm platform detected, initializing optimizations...")
 
         # Disable cudnn for faster VAE conv computation
         torch.backends.cudnn.enabled = False
         logger.info("  - cudnn disabled for faster VAE convolution")
 
-        # Inject aiter as sgl_kernel compatibility layer (REQUIRED)
+        # Inject aiter as sgl_kernel compatibility layer when available. Optional:
+        # torch-native paths (torch._scaled_mm / torch._int_mm fp8/int8 GEMM,
+        # torch_sdpa / SageAttention) run without aiter.
         sgl_kernel = _get_aiter_sgl_kernel()
-        sys.modules["sgl_kernel"] = sgl_kernel
-        # Update any module that already imported sgl_kernel. torch.ops and
-        # torch.classes are ModuleType subclasses living in sys.modules whose
-        # attributes are operator namespaces rather than imported modules, so
-        # overwriting them would replace torch.ops.sgl_kernel itself.
-        for mod_name, mod in list(sys.modules.items()):
-            if mod_name in ("torch.ops", "torch.classes"):
-                continue
-            if mod is not None and hasattr(mod, "sgl_kernel"):
-                setattr(mod, "sgl_kernel", sgl_kernel)
-        logger.info("  - aiter sgl_kernel compatibility layer enabled (RMSNorm, GEMM)")
+        if sgl_kernel is not None:
+            sys.modules["sgl_kernel"] = sgl_kernel
+            # Update any module that already imported sgl_kernel. torch.ops and
+            # torch.classes are ModuleType subclasses living in sys.modules whose
+            # attributes are operator namespaces rather than imported modules, so
+            # overwriting them would replace torch.ops.sgl_kernel itself.
+            for mod_name, mod in list(sys.modules.items()):
+                if mod_name in ("torch.ops", "torch.classes"):
+                    continue
+                if mod is not None and hasattr(mod, "sgl_kernel"):
+                    setattr(mod, "sgl_kernel", sgl_kernel)
+            logger.info("  - aiter sgl_kernel compatibility layer enabled (RMSNorm, GEMM)")
 
     @staticmethod
     def is_available() -> bool:

lightx2v_platform/ops/__init__.py

@@ -12,7 +12,12 @@
     from .attn.hygon_dcu import *
     from .mm.hygon_dcu import *
 elif PLATFORM == "amd_rocm":
+    # Register MM before attention so the fp8/int8 GEMM keys are present before
+    # any framework registry snapshot the attention import may trigger.
+    # isort: off
+    from .mm.amd_rocm import *
     from .attn.amd_rocm import *
+    # isort: on
 elif PLATFORM == "ascend_npu":
     # Register A2A first because Ascend attention may import the common
     # Ulysses backend factory. Keep MoE before MM: MM imports common utilities

lightx2v_platform/ops/attn/amd_rocm/__init__.py

@@ -1 +1,2 @@
 from .flash_attn import *
+from .sage_attn import AmdSageAttn2Weight

lightx2v_platform/ops/attn/amd_rocm/sage_attn.py

@@ -0,0 +1,172 @@
+"""AMD ROCm SageAttention2 backend and Triton num_stages workaround.
+
+On gfx1201 / gfx1100, Triton's AMD software-pipeliner miscompiles kernels at
+num_stages>=3 ("operation destroyed but still has uses"). We clamp num_stages<=2
+on SageAttention's Triton kernels and on any torch.compile / Inductor Triton
+kernels, and force Inductor to compile in-process so the clamp reaches them.
+
+Select with ``attn_type=sage_attn2_amd_rocm``. The workaround is installed only
+when this backend is constructed, and only on the validated archs.
+
+Toggle with the ``LIGHTX2V_ROCM_TRITON_MAX_STAGES`` env var (default 2): it is
+the clamp ceiling, and setting it to 0 (or a negative value) disables the whole
+workaround — e.g. once Triton fixes the pipeliner bug, so kernels can use their
+own num_stages again without a code change.
+"""
+
+import os
+
+import torch
+from loguru import logger
+
+from lightx2v_platform.ops.attn.template import AttnWeightTemplate
+from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER
+
+_MAX_STAGES = int(os.getenv("LIGHTX2V_ROCM_TRITON_MAX_STAGES", "2"))
+_VALIDATED_ARCHS = ("gfx1201", "gfx1100")
+
+
+@PLATFORM_ATTN_WEIGHT_REGISTER("sage_attn2_amd_rocm")
+class AmdSageAttn2Weight(AttnWeightTemplate):
+    """Dense SageAttention2 with opt-in ROCm compiler workarounds."""
+
+    def __init__(self):
+        self.config = {}
+        if getattr(torch.version, "hip", None) is None:
+            raise RuntimeError("sage_attn2_amd_rocm requires AMD ROCm (torch.version.hip is not set).")
+
+        # Keep SageAttention optional during platform operator registration.
+        from sageattention import sageattn
+
+        self.sageattn = sageattn
+        apply_rocm_sage_patches()
+
+    def apply(
+        self,
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_kv=None,
+        max_seqlen_q=None,
+        max_seqlen_kv=None,
+        **kwargs,
+    ):
+        q, k, v = self._prepare_qkv(q, k, v)
+        output = self.sageattn(q, k, v, tensor_layout="NHD")
+        return output.reshape(q.shape[0] * q.shape[1], -1)
+
+    def apply_with_lse(self, q, k, v, softmax_scale=None):
+        """Return dense attention and LSE as [tokens, heads]."""
+        q, k, v = self._prepare_qkv(q, k, v)
+        output, lse = self.sageattn(
+            q,
+            k,
+            v,
+            tensor_layout="NHD",
+            is_causal=False,
+            sm_scale=softmax_scale,
+            return_lse=True,
+        )
+        output = output.reshape(q.shape[0] * q.shape[1], -1)
+        lse = lse.transpose(1, 2).reshape(q.shape[0] * q.shape[1], q.shape[2])
+        return output, lse
+
+    @staticmethod
+    def _prepare_qkv(q, k, v):
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+        if q.ndim == 3:
+            return q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)
+        if q.ndim != 4:
+            raise ValueError(f"AMD SageAttention2 expects 3D or 4D Q/K/V, got q.ndim={q.ndim}.")
+        return q, k, v
+
+
+def _is_validated_rocm_arch():
+    if getattr(torch.version, "hip", None) is None or not torch.cuda.is_available():
+        return False
+    try:
+        return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.startswith(_VALIDATED_ARCHS)
+    except Exception:
+        return False
+
+
+def _clamp_sageattn_kernels():
+    """Clamp num_stages on SageAttention's own Triton kernels."""
+    import importlib
+    import pkgutil
+
+    try:
+        import sageattention
+        from triton.runtime.jit import JITFunction
+    except ImportError:
+        return
+    patched = 0
+    for info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
+        try:
+            module = importlib.import_module(info.name)
+        except Exception as e:
+            logger.debug(f"[SageAttention] ROCm: skipped {info.name}: {e}")
+            continue
+        for fn in vars(module).values():
+            if not isinstance(fn, JITFunction) or getattr(fn, "_rocm_clamped", False):
+                continue
+
+            def run(*args, _run=fn.run, **kwargs):
+                if isinstance(kwargs.get("num_stages"), int):
+                    kwargs["num_stages"] = min(kwargs["num_stages"], _MAX_STAGES)
+                return _run(*args, **kwargs)
+
+            fn.run, fn._rocm_clamped = run, True
+            patched += 1
+    if patched:
+        logger.info(f"[SageAttention] ROCm: clamped num_stages<={_MAX_STAGES} on {patched} kernels")
+
+
+def _clamp_triton_compile():
+    """Clamp num_stages for every Triton compilation (incl. Inductor kernels)."""
+    try:
+        import triton
+    except ImportError:
+        return
+    if getattr(triton.compile, "_rocm_clamped", False):
+        return
+    _original = triton.compile
+
+    def compile(*args, **kwargs):
+        opts = kwargs.get("options")
+        if isinstance(opts, dict) and isinstance(opts.get("num_stages"), int):
+            opts["num_stages"] = min(opts["num_stages"], _MAX_STAGES)
+        elif isinstance(getattr(opts, "num_stages", None), int):
+            try:
+                opts.num_stages = min(opts.num_stages, _MAX_STAGES)
+            except Exception:
+                pass
+        return _original(*args, **kwargs)
+
+    compile._rocm_clamped = True
+    triton.compile = compile
+
+
+def _force_inductor_single_thread():
+    """Compile Inductor kernels in-process so the triton.compile clamp reaches them."""
+    try:
+        from torch._inductor import config as inductor_config
+    except ImportError:
+        return
+    if getattr(inductor_config, "compile_threads", 1) != 1:
+        inductor_config.compile_threads = 1
+        logger.info("[SageAttention] ROCm: inductor compile_threads=1 so the num_stages clamp reaches compiled kernels")
+
+
+def apply_rocm_sage_patches():
+    """Install the workaround. No-op when disabled (LIGHTX2V_ROCM_TRITON_MAX_STAGES<=0)
+    or off the validated RDNA archs."""
+    if _MAX_STAGES <= 0:
+        logger.info("[SageAttention] ROCm: num_stages workaround disabled via LIGHTX2V_ROCM_TRITON_MAX_STAGES")
+        return
+    if not _is_validated_rocm_arch():
+        return
+    _clamp_sageattn_kernels()
+    _clamp_triton_compile()
+    _force_inductor_single_thread()

lightx2v_platform/ops/mm/amd_rocm/__init__.py

@@ -0,0 +1 @@
+from .mm_weight import *