From 2864e99e5141186277dcaf96a45ee5c98a546653 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 01:19:19 +0000
Subject: [PATCH 01/21] fix(qwen-image-2.1): MIOpen-free VAE conv via im2col
 GEMM on RDNA4

MIOpen convolution kernels return NaN nondeterministically for the
Qwen-Image-2.1 VAE on gfx1201 (RDNA4). Route every VAE convolution
(ImageConv, the Resample upsampler, and AttentionBlock) through
im2col + GEMM (hipBLASLt), which is numerically reliable there.

Add opt-in vae_conv_im2col plus vae_device / vae_dtype so the VAE can
also run on CPU or in a chosen dtype as a fallback.
---
 .../video_encoders/hf/qwen_image_21/vae.py    | 62 ++++++++++++++++---
 1 file changed, 53 insertions(+), 9 deletions(-)

diff --git a/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py b/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
index 81f3ba8d7..48a727a59 100644
--- a/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
+++ b/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
@@ -18,14 +18,48 @@
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 
+# MIOpen convolutions on gfx1201 (RDNA4) return NaN nondeterministically for this
+# VAE's shapes. Express every conv as im2col + GEMM (hipBLASLt), which is reliable
+# on this hardware. The input is padded by the caller.
+def _im2col_conv2d(x, weight, bias, stride):
+    kh, kw = weight.shape[-2:]
+    sh, sw = stride
+    b, _, h, w = x.shape
+    cols = F.unfold(x, (kh, kw), stride=(sh, sw))  # (B, C*kh*kw, L)
+    out = torch.matmul(weight.reshape(weight.shape[0], -1), cols)  # (B, O, L)
+    out_h = (h - kh) // sh + 1
+    out_w = (w - kw) // sw + 1
+    out = out.reshape(b, weight.shape[0], out_h, out_w)
+    if bias is not None:
+        out = out + bias.reshape(1, -1, 1, 1)
+    return out
+
+
+class GemmConv2d(nn.Conv2d):
+    """nn.Conv2d that runs through im2col+GEMM (MIOpen-free) when enabled."""
+
+    im2col_conv = False
+
+    def forward(self, x):
+        if self.im2col_conv:
+            ph, pw = self.padding
+            padded = F.pad(x, (pw, pw, ph, ph)) if (ph or pw) else x
+            return _im2col_conv2d(padded, self.weight, self.bias, self.stride)
+        return super().forward(x)
+
+
 class ImageConv(nn.Conv2d):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.image_padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0])
         self.padding = (0, 0)
+        self.im2col_conv = False
 
     def forward(self, x):
-        return super().forward(F.pad(x.squeeze(2), self.image_padding)).unsqueeze(2)
+        padded = F.pad(x.squeeze(2), self.image_padding)
+        if self.im2col_conv:
+            return _im2col_conv2d(padded, self.weight, self.bias, self.stride).unsqueeze(2)
+        return super().forward(padded).unsqueeze(2)
 
 
 class ImageRMSNorm(nn.Module):
@@ -56,8 +90,8 @@ class AttentionBlock(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.norm = ImageRMSNorm(dim, images=True)
-        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
-        self.proj = nn.Conv2d(dim, dim, 1)
+        self.to_qkv = GemmConv2d(dim, dim * 3, 1)
+        self.proj = GemmConv2d(dim, dim, 1)
 
     def forward(self, x):
         b, c, _, h, w = x.shape
@@ -87,9 +121,9 @@ class Resample(nn.Module):
     def __init__(self, dim, up, temporal):
         super().__init__()
         if up:
-            self.resample = nn.Sequential(Upsample(), nn.Conv2d(dim, dim, 3, padding=1))
+            self.resample = nn.Sequential(Upsample(), GemmConv2d(dim, dim, 3, padding=1))
         else:
-            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=2))
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), GemmConv2d(dim, dim, 3, stride=2))
         if temporal:
             self.time_conv = ImageConv(dim, dim * 2 if up else dim, 1)
 
@@ -181,6 +215,11 @@ def __init__(self, config):
         if not cfg["is_residual"] or cfg.get("patch_size") is not None or cfg["attn_scales"]:
             raise ValueError("Expected the released residual, unpatched Qwen-Image-2.1 image VAE")
         self.scale_factor = config["vae_scale_factor"]
+        # MIOpen convolutions on gfx1201 (RDNA4) return NaN for this decoder's
+        # large convs, so allow running the VAE on CPU (always fp32 there).
+        self.vae_device = torch.device("cpu") if str(config.get("vae_device", "")).lower() == "cpu" else torch.device(AI_DEVICE)
+        _vd = str(config.get("vae_dtype", "")).lower()
+        self.vae_dtype = torch.float32 if self.vae_device.type == "cpu" else {"fp32": torch.float32, "float32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(_vd, GET_DTYPE())
         with torch.device("meta"):
             self.encoder = Encoder(cfg)
             self.decoder = Decoder(cfg)
@@ -192,16 +231,21 @@ def __init__(self, config):
         self.load_state_dict(state, strict=True, assign=True)
         self.register_buffer("latents_mean", torch.tensor(cfg["latents_mean"]).reshape(1, -1, 1, 1, 1), persistent=False)
         self.register_buffer("latents_std", torch.tensor(cfg["latents_std"]).reshape(1, -1, 1, 1, 1), persistent=False)
-        self.to(dtype=GET_DTYPE(), device=AI_DEVICE)
+        self.to(dtype=self.vae_dtype, device=self.vae_device)
+        if self.vae_device.type != "cpu" and config.get("vae_conv_im2col", False):
+            for module in self.modules():
+                if isinstance(module, (ImageConv, GemmConv2d)):
+                    module.im2col_conv = True
         self.eval().requires_grad_(False)
 
     @torch.inference_mode()
     def encode(self, image):
-        mean = self.quant_conv(self.encoder(image.to(device=AI_DEVICE, dtype=GET_DTYPE()))).chunk(2, dim=1)[0]
+        mean = self.quant_conv(self.encoder(image.to(device=self.vae_device, dtype=self.vae_dtype))).chunk(2, dim=1)[0]
         return ((mean - self.latents_mean) / self.latents_std).flatten(2).transpose(1, 2)
 
     @torch.inference_mode()
     def decode(self, latents, size):
         h, w = size
-        z = latents.transpose(1, 2).reshape(1, self.config["z_dim"], 1, h // self.scale_factor, w // self.scale_factor).to(GET_DTYPE())
-        return self.decoder(self.post_quant_conv(z * self.latents_std + self.latents_mean)).clamp(-1, 1)[:, :, 0]
+        z = latents.transpose(1, 2).reshape(1, self.config["z_dim"], 1, h // self.scale_factor, w // self.scale_factor).to(device=self.vae_device, dtype=self.vae_dtype)
+        out = self.decoder(self.post_quant_conv(z * self.latents_std + self.latents_mean)).clamp(-1, 1)[:, :, 0]
+        return out.to(AI_DEVICE)

From 2a39bebc9630d4c245e7747e34b49a76671b07c0 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 01:19:58 +0000
Subject: [PATCH 02/21] fix(sage-attn): work around Triton ROCm num_stages
 pipeliner crash on gfx1201

SageAttention's INT8 WMMA Triton kernels fail to compile on ROCm/gfx1201
at num_stages>=3 ("operation destroyed but still has uses" in the AMD
pipeliner pass). Clamp num_stages<=2 on both SageAttention's own kernels
and every torch.compile / Inductor Triton compilation, and run Inductor
single-threaded so the clamp reaches its worker (only on the SageAttention
path). With this, dense SageAttention2 runs and is faster than torch SDPA
flash on this hardware.
---
 lightx2v/common/ops/attn/sage_attn.py | 105 ++++++++++++++++++++++++++
 1 file changed, 105 insertions(+)

diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index 07945b9f5..6e999e241 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -39,6 +39,110 @@
         sageattn = None
 
 
+def _clamp_sageattn_triton_num_stages_on_rocm(max_stages=2):
+    """Work around a Triton ROCm software-pipeliner bug.
+
+    On gfx1201 (ROCm/Triton), SageAttention's INT8 WMMA kernels fail to compile
+    at num_stages>=3 ("operation destroyed but still has uses" during the AMD
+    pipeliner pass). Clamp the launch-time num_stages on every SageAttention
+    Triton kernel so they compile; num_stages=2 runs correctly and is faster
+    than torch SDPA flash on this hardware.
+    """
+    if sageattn is None or getattr(torch.version, "hip", None) is None:
+        return
+    import importlib
+    import pkgutil
+
+    try:
+        import sageattention
+        from triton.runtime.jit import JITFunction
+    except ImportError:
+        return
+
+    patched = 0
+    for module_info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
+        try:
+            module = importlib.import_module(module_info.name)
+        except Exception:
+            continue
+        for value in vars(module).values():
+            if not isinstance(value, JITFunction) or getattr(value, "_rocm_num_stages_clamped", False):
+                continue
+            original_run = value.run
+
+            def run(*args, _original_run=original_run, **kwargs):
+                stages = kwargs.get("num_stages")
+                if isinstance(stages, int) and stages > max_stages:
+                    kwargs["num_stages"] = max_stages
+                return _original_run(*args, **kwargs)
+
+            try:
+                value.run = run
+                value._rocm_num_stages_clamped = True
+                patched += 1
+            except Exception:
+                continue
+    if patched:
+        logger.info(f"[SageAttention] ROCm: clamped Triton num_stages<={max_stages} on {patched} kernels")
+
+
+def _clamp_triton_compile_num_stages_on_rocm(max_stages=2):
+    """Clamp num_stages for every Triton compilation (including torch.compile /
+    Inductor kernels) on ROCm, where num_stages>=3 crashes the AMD pipeliner.
+
+    Inductor compiles in worker subprocesses; requires compile_threads=1 so this
+    main-process patch is in effect at compile time.
+    """
+    if getattr(torch.version, "hip", None) is None:
+        return
+    try:
+        import triton
+    except ImportError:
+        return
+    if getattr(triton.compile, "_rocm_num_stages_clamped", False):
+        return
+    _original_compile = triton.compile
+
+    def compile(*args, **kwargs):
+        options = kwargs.get("options")
+        if isinstance(options, dict) and isinstance(options.get("num_stages"), int) and options["num_stages"] > max_stages:
+            options["num_stages"] = max_stages
+        elif options is not None and isinstance(getattr(options, "num_stages", None), int) and options.num_stages > max_stages:
+            try:
+                options.num_stages = max_stages
+            except Exception:
+                pass
+        return _original_compile(*args, **kwargs)
+
+    compile._rocm_num_stages_clamped = True
+    triton.compile = compile
+
+
+def _force_rocm_inductor_single_thread():
+    """Make torch.compile / Inductor compile in-process on ROCm.
+
+    The num_stages clamp on ``triton.compile`` only takes effect in the process
+    that installed it. Inductor's default multi-threaded pool compiles kernels
+    in worker subprocesses that never see the clamp, so an Inductor-generated
+    kernel can still hit the num_stages>=3 pipeliner crash. Compiling in-process
+    keeps the clamp effective. Only pays the (one-time, warmup) single-threaded
+    compile cost on the SageAttention path that needs it.
+    """
+    if getattr(torch.version, "hip", None) is None:
+        return
+    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: set inductor compile_threads=1 so the Triton num_stages clamp reaches torch.compile kernels")
+
+
+_clamp_sageattn_triton_num_stages_on_rocm()
+_clamp_triton_compile_num_stages_on_rocm()
+
+
 if magi_register_custom_op is not None and sageattn is not None:
 
     @magi_register_custom_op(
@@ -75,6 +179,7 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
+        _force_rocm_inductor_single_thread()
 
     def apply(
         self,

From 89a144b98848013c0594227b686134972a474cb4 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 01:20:13 +0000
Subject: [PATCH 03/21] =?UTF-8?q?optimize(qwen-image-2.1):=20R9700=20ROCm?=
 =?UTF-8?q?=20DiT=20=E2=80=94=20fp8=20=5Fscaled=5Fmm,=20torch.compile,=20c?=
 =?UTF-8?q?ross-GPU=20text=20encoder,=20QKV/MLP=20fusion?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- fp8-rocm quant scheme: per-channel weight + per-token activation FP8 via
  torch._scaled_mm, auto-quantized from the released BF16 weights (no
  prequantized checkpoint; releases the BF16 source on load to halve the
  load-time peak). CUTLASS/sgl FP8 kernels are CUDA-only.
- use_compile for qwen_image_21 via BaseTransformerInfer (infer_block /
  run_block); the block graph passes only tensors, keeping modulation and
  KV-cache lookups outside it.
- Run the text encoder on a separate GPU (text_encoder_device_index): ROCm
  empty_cache does not return the onloaded ~15G to the driver, which would
  otherwise starve the VAE decode on the main GPU.
- Optional dit_fuse_qkv_mlp: fuse q/k/v and gate/up into single FP8 GEMMs
  (numerically identical; one shared activation quantization).
---
 lightx2v/common/ops/mm/mm_weight.py           | 96 +++++++++++++++++++
 lightx2v/models/networks/base_model.py        |  1 +
 .../qwen_image_21/infer/transformer_infer.py  | 46 ++++++---
 .../weights/transformer_weights.py            | 29 ++++--
 .../qwen_image_21/qwen_image_21_runner.py     | 37 ++++---
 5 files changed, 175 insertions(+), 34 deletions(-)

diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py
index e6c1ca2e6..39f293009 100755
--- a/lightx2v/common/ops/mm/mm_weight.py
+++ b/lightx2v/common/ops/mm/mm_weight.py
@@ -771,6 +771,102 @@ def apply(self, input_tensor):
         return output_tensor
 
 
+@MM_WEIGHT_REGISTER("fp8-rocm")
+class MMWeightWfp8channelAfp8channeldynamicRocm(MMWeightQuantTemplate):
+    """W-fp8-e4m3-channel-sym / A-fp8-e4m3-token-dynamic via torch._scaled_mm.
+
+    MIOpen/CUTLASS/sgl FP8 kernels are unavailable on ROCm gfx1201 (RDNA4), but
+    torch._scaled_mm drives a working FP8 GEMM there. Weights are quantized from
+    the released BF16 checkpoint on load (no pre-quantized checkpoint needed).
+    """
+
+    FP8_MAX = 448.0
+
+    def __init__(
+        self,
+        weight_name,
+        bias_name,
+        create_cuda_buffer=False,
+        create_cpu_buffer=False,
+        lazy_load=False,
+        lazy_load_file=None,
+        is_post_adapter=False,
+        lora_prefix="diffusion_model.blocks",
+        lora_path="",
+    ):
+        super().__init__(
+            weight_name,
+            bias_name,
+            create_cuda_buffer,
+            create_cpu_buffer,
+            lazy_load,
+            lazy_load_file,
+            is_post_adapter,
+            lora_prefix,
+            lora_path,
+        )
+        # Weight stored (N, K) fp8; a column-major (K, N) view is taken in apply().
+        self.weight_need_transpose = False
+        self.scale_force_fp32 = True
+
+    def load(self, weight_dict):
+        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (N, 1)
+        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        # Release the BF16 source immediately so it does not coexist with every
+        # quantized weight (halves the load-time peak on GPU).
+        weight_dict.pop(self.weight_name, None)
+        if self.bias_name is not None:
+            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
+        else:
+            self.bias = None
+
+    def apply(self, input_tensor):
+        x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX).float()  # (M, 1)
+        x_q = (input_tensor / x_scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)  # (M, K)
+        output_tensor = torch._scaled_mm(
+            x_q,
+            self.weight.t(),  # (K, N) column-major
+            scale_a=x_scale,
+            scale_b=self.weight_scale.t(),  # (1, N)
+            bias=self.bias,
+            out_dtype=input_tensor.dtype,
+        )
+        if self.has_lora_branch:
+            return output_tensor + self.apply_lora(input_tensor)
+        return output_tensor
+
+
+class MMWeightFusedFp8Rocm(MMWeightWfp8channelAfp8channeldynamicRocm):
+    """Fuse several row-concatenatable linears (q/k/v, gate/up) into one fp8 GEMM.
+
+    Per-output-channel weight scales make this numerically identical to the
+    separate GEMMs, while cutting GEMM count and quantizing the (shared) input
+    activation once instead of once per linear. ``apply()`` returns the
+    concatenated output; the caller splits it with ``split_sizes``.
+    """
+
+    def __init__(self, weight_names, bias_name=None, **kwargs):
+        super().__init__(weight_names[0], bias_name, **kwargs)
+        self.fused_weight_names = list(weight_names)
+        self.split_sizes = None
+
+    def load(self, weight_dict):
+        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]  # each (N_i, K)
+        self.split_sizes = [w.shape[0] for w in weights]
+        weight = torch.cat(weights, dim=0)  # (sum_N, K)
+        del weights
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (sum_N, 1)
+        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        for name in self.fused_weight_names:
+            weight_dict.pop(name, None)
+        self.bias = None
+
+
 @MM_WEIGHT_REGISTER("int8-vllm")
 class MMWeightWint8channelAint8channeldynamicVllm(MMWeightQuantTemplate):
     """
diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py
index a572bc214..15f3885e2 100644
--- a/lightx2v/models/networks/base_model.py
+++ b/lightx2v/models/networks/base_model.py
@@ -113,6 +113,7 @@ def _check_dit_quantized(self):
             bool: True if the model is quantized
         """
         assert self.config.get("dit_quant_scheme", "Default") in [
+            "fp8-rocm",
             "fp8-pertensor",
             "fp8-musa",
             "fp8-triton",
diff --git a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
index cf5521b08..16c815150 100644
--- a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
+++ b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
@@ -2,13 +2,15 @@
 import torch.nn.functional as F
 
 from lightx2v.common.ops.norm.rms_norm_weight import apply_qk_rms_norm
+from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer
 from lightx2v.utils.registry_factory import ROPE_REGISTER
 
 
-class QwenImage21TransformerInfer:
+class QwenImage21TransformerInfer(BaseTransformerInfer):
     def __init__(self, config):
         self.config = config
         self.heads = config["num_attention_heads"]
+        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)
@@ -18,6 +20,7 @@ def __init__(self, config):
             self.block_ops = triton_ops
         self.rope = ROPE_REGISTER[config.get("rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float32)
         self.rope.set_config(config)
+        self.init_compile(config)
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
@@ -27,15 +30,28 @@ def _modulation(self, state):
         scale1, gate1, scale2, gate2 = state.modulation.chunk(4, -1)
         return 1 + scale1, gate1.tanh(), 1 + scale2, gate2.tanh()
 
-    def _qkv(self, block, x, scale, 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"])
-        k = block.k.apply(h).reshape_as(q)
-        v = block.v.apply(h).reshape_as(q)
+        if hasattr(block, "qkv"):
+            q, k, v = block.qkv.apply(h).split(block.qkv.split_sizes, dim=-1)
+            q = q.reshape(-1, self.heads, self.head_dim)
+            k = k.reshape_as(q)
+            v = v.reshape_as(q)
+        else:
+            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)
-        q, k = self.rope.apply(q, k, state.rotary, positions=state.rotary_positions)
+        q, k = self.rope.apply(q, k, rotary, positions=rotary_positions)
         return q, k, v
 
+    def infer_block(self, block, x, scale1, gate1, scale2, gate2, rotary, rotary_positions, k_cache, v_cache):
+        q, k, v = self._qkv(block, x, scale1, rotary, rotary_positions)
+        k = torch.cat((k_cache, k))
+        v = torch.cat((v_cache, v))
+        attention = block.attention.apply(q, k, v)
+        return self._finish_block(block, x, attention, (scale1, gate1, scale2, gate2))
+
     def _finish_block(self, block, x, attention, modulation):
         _, gate1, scale2, gate2 = modulation
         # Short sequences are launch-bound; eager Torch is faster there.
@@ -46,7 +62,12 @@ def _finish_block(self, block, x, attention, modulation):
 
         x = x + gate1 * block.out.apply(attention)
         h = block.norm2.apply(x) * scale2
-        x = x + gate2 * block.down.apply(F.silu(block.gate.apply(h)) * block.up.apply(h))
+        if hasattr(block, "gate_up"):
+            gate_out, up_out = block.gate_up.apply(h).split(block.gate_up.split_sizes, dim=-1)
+            mlp = F.silu(gate_out) * up_out
+        else:
+            mlp = F.silu(block.gate.apply(h)) * block.up.apply(h)
+        x = x + gate2 * block.down.apply(mlp)
         return x.clamp(-65504, 65504) if x.dtype == torch.float16 else x
 
     def prefill(self, weights, state, cache):
@@ -54,7 +75,7 @@ def prefill(self, weights, state, cache):
         x = state.hidden_states
         modulation = self._modulation(state)
         for index, block in enumerate(weights.blocks):
-            q, k, v = self._qkv(block, x, modulation[0], state)
+            q, k, v = self._qkv(block, x, modulation[0], state.rotary, state.rotary_positions)
             cache.store_kv(k, v, index)
             attention = torch.empty_like(x)
             for begin, end, is_text in state.layout.segments:
@@ -70,11 +91,8 @@ def infer(self, weights, state, cache):
         if not cache.is_ready():
             raise RuntimeError("Condition KV must be prefilled before denoising")
         x = state.hidden_states
-        modulation = self._modulation(state)
+        # Shared across every block; kept outside the compiled block graph.
+        scale1, gate1, scale2, gate2 = self._modulation(state)
         for index, block in enumerate(weights.blocks):
-            q, k, v = self._qkv(block, x, modulation[0], state)
-            k = torch.cat((cache.k_cache(index), k))
-            v = torch.cat((cache.v_cache(index), v))
-            attention = block.attention.apply(q, k, v)
-            x = self._finish_block(block, x, attention, modulation)
+            x = self.run_block(index, block, x, scale1, gate1, scale2, gate2, state.rotary, state.rotary_positions, cache.k_cache(index), cache.v_cache(index))
         return x
diff --git a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
index 26c67d14c..2cf3a9211 100644
--- a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
+++ b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
@@ -8,15 +8,26 @@ def __init__(self, index, config):
         super().__init__()
         prefix = f"transformer_blocks.{index}"
         mm_type = config["dit_quant_scheme"] if config.get("dit_quantized", False) else "Default"
-        for name, key in {
-            "q": "attn.to_q",
-            "k": "attn.to_k",
-            "v": "attn.to_v",
-            "out": "attn.to_out.0",
-            "up": "img_mlp.proj",
-            "gate": "img_mlp.gate_layer",
-            "down": "img_mlp.out",
-        }.items():
+        # Fuse q/k/v and gate/up into single fp8 GEMMs (fewer, larger GEMMs and
+        # one shared activation quantization instead of one per linear).
+        self.fused_qkv_mlp = bool(config.get("dit_fuse_qkv_mlp", False)) and mm_type == "fp8-rocm"
+        if self.fused_qkv_mlp:
+            from lightx2v.common.ops.mm.mm_weight import MMWeightFusedFp8Rocm
+
+            self.add_module("qkv", MMWeightFusedFp8Rocm([f"{prefix}.attn.to_q.weight", f"{prefix}.attn.to_k.weight", f"{prefix}.attn.to_v.weight"], bias_name=None))
+            self.add_module("gate_up", MMWeightFusedFp8Rocm([f"{prefix}.img_mlp.gate_layer.weight", f"{prefix}.img_mlp.proj.weight"], bias_name=None))
+            linear_keys = {"out": "attn.to_out.0", "down": "img_mlp.out"}
+        else:
+            linear_keys = {
+                "q": "attn.to_q",
+                "k": "attn.to_k",
+                "v": "attn.to_v",
+                "out": "attn.to_out.0",
+                "up": "img_mlp.proj",
+                "gate": "img_mlp.gate_layer",
+                "down": "img_mlp.out",
+            }
+        for name, key in linear_keys.items():
             linear = MM_WEIGHT_REGISTER[mm_type](f"{prefix}.{key}.weight", bias_name=None)
             if mm_type == "fp8-f16-accum":
                 linear.enable_fp8_f16_accum(config.get("dit_fp8_activation_qmax", ACTIVATION_QMAX))
diff --git a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
index f4f49cdfc..6b70def69 100644
--- a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
+++ b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
@@ -1,3 +1,4 @@
+import contextlib
 import math
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -51,15 +52,15 @@ def __init__(self, config):
             "disagg_mode",
             "lora_configs",
             "vae_tiling",
-            "use_compile",
         )
         for key in unsupported:
             if config.get(key):
                 raise ValueError(f"qwen_image_21 does not yet support {key}")
         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"):
+                raise ValueError("qwen_image_21 DiT quantization supports only fp8-sgl, fp8-f16-accum and fp8-rocm")
+            # fp8-rocm quantizes the released BF16 weights on load (torch._scaled_mm), no prequantized checkpoint.
+            if config["dit_quant_scheme"] != "fp8-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))
@@ -184,18 +185,32 @@ def run_text_encoder(self, images, shapes):
             branches.append(("uncond", info.negative_prompt or ""))
         encoder = self.text_encoders[0]
         offload = self.config.get("text_encoder_cpu_offload", False)
+        # ROCm/R9700: empty_cache() does not physically return freed VRAM to the
+        # driver, so the ~15G text-encoder onload stays resident on the main GPU
+        # and starves the VAE decoder. Run it on a separate GPU when configured.
+        te_index = self.config.get("text_encoder_device_index", None)
+        dev_ctx = torch_device_module.device(te_index) if te_index is not None else contextlib.nullcontext()
+        # Resolve an explicit indexed device: a bare "cuda" would re-resolve to the
+        # current device (the text-encoder GPU) inside the context above.
+        main_device = torch.device(AI_DEVICE, torch_device_module.current_device())
         try:
-            if offload:
-                with ProfilingContext4DebugL1("Onload Text Encoder"):
-                    encoder.to_cuda()
-            for name, prompt in branches:
-                branch = encoder.infer(prompt, images)
-                branch["layout"] = build_token_layout(branch["image_mask"], shapes, self.config["axes_dims_rope"], rope=self.model.transformer_infer.rope)
-                outputs[name] = branch
+            with dev_ctx:
+                if offload:
+                    with ProfilingContext4DebugL1("Onload Text Encoder"):
+                        encoder.to_cuda()
+                for name, prompt in branches:
+                    branch = encoder.infer(prompt, images)
+                    # Pull encoder outputs onto the main GPU before leaving the
+                    # text-encoder device context.
+                    outputs[name] = {k: (v.to(main_device) if torch.is_tensor(v) else v) for k, v in branch.items()}
         finally:
             if offload:
                 with ProfilingContext4DebugL1("Offload Text Encoder"):
                     encoder.to_cpu()
+        # Build the DiT rope layout on the main GPU (outside the TE device context,
+        # otherwise freshly created rope tensors would land on the TE device).
+        for branch in outputs.values():
+            branch["layout"] = build_token_layout(branch["image_mask"], shapes, self.config["axes_dims_rope"], rope=self.model.transformer_infer.rope)
         return outputs
 
     @ProfilingContext4DebugL1("Run VAE Encoder")

From ab4a95b508b03f17e4e1e1410bec1afd39683304 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 01:20:25 +0000
Subject: [PATCH 04/21] chore(qwen-image-2.1): add R9700 (gfx1201) ROCm
 inference configs

Torch-op inference path for RDNA4 (torch_sdpa, torch_real_rope, im2col VAE,
text encoder on cuda:1):
- qwen_image_21_r9700.json          : bf16 baseline
- ..._fp8_compile.json              : fp8 + torch.compile
- ..._fp8_compile_sage.json         : fp8 + compile + SageAttention2 (fastest, 40 steps)
- ..._fp8_compile_fused.json        : + QKV/MLP fusion (opt-in)
- ..._fp8_sage_25steps.json         : 25-step practical config (~1.6x fewer steps, no quality loss)
---
 .../qwen_image_21/qwen_image_21_r9700.json    | 18 +++++++++++++++
 .../qwen_image_21_r9700_fp8_compile.json      | 21 ++++++++++++++++++
 ...qwen_image_21_r9700_fp8_compile_fused.json | 22 +++++++++++++++++++
 .../qwen_image_21_r9700_fp8_compile_sage.json | 21 ++++++++++++++++++
 .../qwen_image_21_r9700_fp8_sage_25steps.json | 21 ++++++++++++++++++
 5 files changed, 103 insertions(+)
 create mode 100644 configs/qwen_image_21/qwen_image_21_r9700.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json

diff --git a/configs/qwen_image_21/qwen_image_21_r9700.json b/configs/qwen_image_21/qwen_image_21_r9700.json
new file mode 100644
index 000000000..aaac144fe
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_r9700.json
@@ -0,0 +1,18 @@
+{
+    "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,
+    "text_encoder_device_index": 1,
+    "vae_conv_im2col": true
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
new file mode 100644
index 000000000..be4ff748f
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
@@ -0,0 +1,21 @@
+{
+    "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": true,
+    "text_encoder_cpu_offload": true,
+    "text_encoder_device_index": 1,
+    "vae_conv_im2col": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-rocm",
+    "use_compile": true
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
new file mode 100644
index 000000000..295f456c5
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
@@ -0,0 +1,22 @@
+{
+    "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": true,
+    "text_encoder_cpu_offload": true,
+    "text_encoder_device_index": 1,
+    "vae_conv_im2col": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-rocm",
+    "use_compile": true,
+    "dit_fuse_qkv_mlp": true
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
new file mode 100644
index 000000000..2ac288bc5
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
@@ -0,0 +1,21 @@
+{
+    "infer_steps": 40,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2",
+    "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,
+    "text_encoder_device_index": 1,
+    "vae_conv_im2col": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-rocm",
+    "use_compile": true
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
new file mode 100644
index 000000000..d84331957
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
@@ -0,0 +1,21 @@
+{
+    "infer_steps": 25,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2",
+    "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,
+    "text_encoder_device_index": 1,
+    "vae_conv_im2col": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-rocm",
+    "use_compile": true
+}
\ No newline at end of file

From b27987d15314ab3e44cb8427fb281976e58a1a1f Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 04:41:38 +0000
Subject: [PATCH 05/21] optimize(qwen-image-2.1): int8-rocm DiT quantization
 for RDNA3 (gfx1100 / W7900)

RDNA3 has no FP8 tensor path, so add an INT8 scheme mirroring fp8-rocm:
- int8-rocm: per-channel weight + per-token activation INT8 via torch._int_mm,
  auto-quantized from the released BF16 weights (no prequantized checkpoint).
  Falls back to a BF16 GEMM for the small prefill segments where torch._int_mm
  requires M > 16.
- dit_fuse_qkv_mlp now also fuses the int8 q/k/v and gate/up GEMMs.
- Allow int8-rocm in the qwen_image_21 runner and base-model quant whitelists.

On gfx1100 native MIOpen VAE conv is fine (no im2col needed) and SageAttention2
INT8 is ~1.75x over torch SDPA flash, so the practical W7900 stack is
int8-rocm + torch.compile + SageAttention2.
---
 lightx2v/common/ops/mm/mm_weight.py           | 80 +++++++++++++++++++
 lightx2v/models/networks/base_model.py        |  1 +
 .../weights/transformer_weights.py            |  9 ++-
 .../qwen_image_21/qwen_image_21_runner.py     |  8 +-
 4 files changed, 90 insertions(+), 8 deletions(-)

diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py
index 39f293009..c67a042f9 100755
--- a/lightx2v/common/ops/mm/mm_weight.py
+++ b/lightx2v/common/ops/mm/mm_weight.py
@@ -867,6 +867,86 @@ def load(self, weight_dict):
         self.bias = None
 
 
+@MM_WEIGHT_REGISTER("int8-rocm")
+class MMWeightWint8channelAint8tokenRocm(MMWeightQuantTemplate):
+    """W-int8-channel-sym / A-int8-token-dynamic via torch._int_mm.
+
+    RDNA3 (gfx1100, e.g. W7900) has no FP8 tensor path, but torch._int_mm drives
+    a working INT8 WMMA GEMM there. Weights are quantized from the released BF16
+    checkpoint on load (no prequantized checkpoint needed).
+    """
+
+    INT8_MAX = 127.0
+
+    def __init__(
+        self,
+        weight_name,
+        bias_name=None,
+        create_cuda_buffer=False,
+        create_cpu_buffer=False,
+        lazy_load=False,
+        lazy_load_file=None,
+        is_post_adapter=False,
+        lora_prefix="diffusion_model.blocks",
+        lora_path="",
+    ):
+        super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, lora_prefix, lora_path)
+        self.weight_need_transpose = False
+        self.scale_force_fp32 = True
+
+    def load(self, weight_dict):
+        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX  # (N, 1)
+        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        weight_dict.pop(self.weight_name, None)
+        if self.bias_name is not None:
+            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
+        else:
+            self.bias = None
+
+    def apply(self, input_tensor):
+        if input_tensor.shape[0] <= 16:
+            # torch._int_mm requires M > 16; dequantize to BF16 for the small
+            # prefill segments (rare, negligible cost).
+            weight = (self.weight.to(torch.float32) * self.weight_scale).to(input_tensor.dtype)  # (N, K)
+            out = torch.mm(input_tensor, weight.t())
+        else:
+            x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX).float()  # (M, 1)
+            x_q = (input_tensor / x_scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)  # (M, K)
+            out = torch._int_mm(x_q, self.weight.t())  # (M, N) int32
+            out = (out.to(torch.float32) * x_scale * self.weight_scale.t()).to(input_tensor.dtype)  # dequant
+        if self.bias is not None:
+            out = out + self.bias
+        if self.has_lora_branch:
+            out = out + self.apply_lora(input_tensor)
+        return out
+
+
+class MMWeightFusedInt8Rocm(MMWeightWint8channelAint8tokenRocm):
+    """Fuse q/k/v, gate/up into one INT8 GEMM; apply() returns the concatenated
+    output, split by ``split_sizes``. See MMWeightFusedFp8Rocm."""
+
+    def __init__(self, weight_names, bias_name=None, **kwargs):
+        super().__init__(weight_names[0], bias_name, **kwargs)
+        self.fused_weight_names = list(weight_names)
+        self.split_sizes = None
+
+    def load(self, weight_dict):
+        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]
+        self.split_sizes = [w.shape[0] for w in weights]
+        weight = torch.cat(weights, dim=0)
+        del weights
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX
+        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        for name in self.fused_weight_names:
+            weight_dict.pop(name, None)
+        self.bias = None
+
+
 @MM_WEIGHT_REGISTER("int8-vllm")
 class MMWeightWint8channelAint8channeldynamicVllm(MMWeightQuantTemplate):
     """
diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py
index 15f3885e2..bbbbf538e 100644
--- a/lightx2v/models/networks/base_model.py
+++ b/lightx2v/models/networks/base_model.py
@@ -114,6 +114,7 @@ def _check_dit_quantized(self):
         """
         assert self.config.get("dit_quant_scheme", "Default") in [
             "fp8-rocm",
+            "int8-rocm",
             "fp8-pertensor",
             "fp8-musa",
             "fp8-triton",
diff --git a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
index 2cf3a9211..5535ecfe1 100644
--- a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
+++ b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
@@ -10,12 +10,13 @@ def __init__(self, index, config):
         mm_type = config["dit_quant_scheme"] if config.get("dit_quantized", False) else "Default"
         # Fuse q/k/v and gate/up into single fp8 GEMMs (fewer, larger GEMMs and
         # one shared activation quantization instead of one per linear).
-        self.fused_qkv_mlp = bool(config.get("dit_fuse_qkv_mlp", False)) and mm_type == "fp8-rocm"
+        self.fused_qkv_mlp = bool(config.get("dit_fuse_qkv_mlp", False)) and mm_type in ("fp8-rocm", "int8-rocm")
         if self.fused_qkv_mlp:
-            from lightx2v.common.ops.mm.mm_weight import MMWeightFusedFp8Rocm
+            from lightx2v.common.ops.mm.mm_weight import MMWeightFusedFp8Rocm, MMWeightFusedInt8Rocm
 
-            self.add_module("qkv", MMWeightFusedFp8Rocm([f"{prefix}.attn.to_q.weight", f"{prefix}.attn.to_k.weight", f"{prefix}.attn.to_v.weight"], bias_name=None))
-            self.add_module("gate_up", MMWeightFusedFp8Rocm([f"{prefix}.img_mlp.gate_layer.weight", f"{prefix}.img_mlp.proj.weight"], bias_name=None))
+            fused_cls = MMWeightFusedInt8Rocm if mm_type == "int8-rocm" else MMWeightFusedFp8Rocm
+            self.add_module("qkv", fused_cls([f"{prefix}.attn.to_q.weight", f"{prefix}.attn.to_k.weight", f"{prefix}.attn.to_v.weight"], bias_name=None))
+            self.add_module("gate_up", fused_cls([f"{prefix}.img_mlp.gate_layer.weight", f"{prefix}.img_mlp.proj.weight"], bias_name=None))
             linear_keys = {"out": "attn.to_out.0", "down": "img_mlp.out"}
         else:
             linear_keys = {
diff --git a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
index 6b70def69..f1d12c9c4 100644
--- a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
+++ b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
@@ -57,10 +57,10 @@ def __init__(self, config):
             if config.get(key):
                 raise ValueError(f"qwen_image_21 does not yet support {key}")
         if config.get("dit_quantized"):
-            if config.get("dit_quant_scheme") not in ("fp8-sgl", "fp8-f16-accum", "fp8-rocm"):
-                raise ValueError("qwen_image_21 DiT quantization supports only fp8-sgl, fp8-f16-accum and fp8-rocm")
-            # fp8-rocm quantizes the released BF16 weights on load (torch._scaled_mm), no prequantized checkpoint.
-            if config["dit_quant_scheme"] != "fp8-rocm" and 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))

From 4f1d0a4a21bb673f30cf7db432da2939d5366759 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 04:41:51 +0000
Subject: [PATCH 06/21] chore(qwen-image-2.1): add W7900 (gfx1100) int8
 inference configs

Native-conv VAE + CPU-offload text encoder (no gfx1201 workarounds):
- qwen_image_21_w7900_int8.json               : int8 baseline
- ..._int8_compile.json                       : int8 + torch.compile
- ..._int8_compile_sage.json                  : int8 + compile + SageAttention2 (fastest, 40 steps)
- ..._int8_sage_25steps.json                  : 25-step practical config (~29s end-to-end)
---
 .../qwen_image_21_w7900_int8.json             | 18 ++++++++++++++++++
 .../qwen_image_21_w7900_int8_compile.json     | 19 +++++++++++++++++++
 ...qwen_image_21_w7900_int8_compile_sage.json | 19 +++++++++++++++++++
 ...qwen_image_21_w7900_int8_sage_25steps.json | 19 +++++++++++++++++++
 4 files changed, 75 insertions(+)
 create mode 100644 configs/qwen_image_21/qwen_image_21_w7900_int8.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
 create mode 100644 configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json

diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8.json b/configs/qwen_image_21/qwen_image_21_w7900_int8.json
new file mode 100644
index 000000000..52998196a
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8.json
@@ -0,0 +1,18 @@
+{
+    "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,
+    "dit_quantized": true,
+    "dit_quant_scheme": "int8-rocm"
+}
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
new file mode 100644
index 000000000..e683666e4
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
@@ -0,0 +1,19 @@
+{
+    "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": true,
+    "text_encoder_cpu_offload": true,
+    "dit_quantized": true,
+    "dit_quant_scheme": "int8-rocm",
+    "use_compile": true
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
new file mode 100644
index 000000000..42ffb7a8a
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
@@ -0,0 +1,19 @@
+{
+    "infer_steps": 40,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2",
+    "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
+}
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
new file mode 100644
index 000000000..6079e7f78
--- /dev/null
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
@@ -0,0 +1,19 @@
+{
+    "infer_steps": 25,
+    "resolution": 1024,
+    "vae_scale_factor": 16,
+    "attn_type": "sage_attn2",
+    "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
+}
\ No newline at end of file

From 31a578c7436f7e112d10a1ab993db2c0eb11e99f Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 04:54:53 +0000
Subject: [PATCH 07/21] chore(qwen-image-2.1): R9700 configs run on a single
 GPU

Drop text_encoder_device_index and use text_encoder_cpu_offload instead. The
im2col GEMM VAE decode reuses PyTorch's allocator pool (where the offloaded
text-encoder memory is cached), so the VAE no longer starves for MIOpen
workspace on the main GPU. A single 32G R9700 fits the whole pipeline
(verified: bf16 and fp8, same speed as the previous two-GPU setup).
---
 configs/qwen_image_21/qwen_image_21_r9700.json                   | 1 -
 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json       | 1 -
 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json | 1 -
 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json  | 1 -
 configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json  | 1 -
 5 files changed, 5 deletions(-)

diff --git a/configs/qwen_image_21/qwen_image_21_r9700.json b/configs/qwen_image_21/qwen_image_21_r9700.json
index aaac144fe..44f45884d 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700.json
@@ -13,6 +13,5 @@
     "sample_guide_scale": 1.0,
     "warmup": false,
     "text_encoder_cpu_offload": true,
-    "text_encoder_device_index": 1,
     "vae_conv_im2col": true
 }
\ No newline at end of file
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
index be4ff748f..566b5b8e1 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "text_encoder_device_index": 1,
     "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
index 295f456c5..81d4cf2be 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "text_encoder_device_index": 1,
     "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
index 2ac288bc5..50b34bdbf 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "text_encoder_device_index": 1,
     "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
index d84331957..984a18d4b 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "text_encoder_device_index": 1,
     "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",

From 09c01e8a8395235ba5a3e5ffdf8b7d6312bf27b2 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 05:42:17 +0000
Subject: [PATCH 08/21] docs(qwen-image-2.1): document R9700 / W7900 ROCm
 inference (EN + ZH)

Add a section (in both README.md and README_zh.md) covering the AMD RDNA runs:
deps, single-GPU run command, per-GPU configs (fp8-rocm on R9700 / int8-rocm on
W7900), ROCm-specific notes (BF16-on-load quantization, im2col VAE on gfx1201,
SageAttention num_stages workaround), and measured end-to-end latency at 40/25 steps.
---
 scripts/qwen_image_21/README.md    | 43 ++++++++++++++++++++++++++++++
 scripts/qwen_image_21/README_zh.md | 43 ++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index d888f112f..96ed7ca3f 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -84,6 +84,49 @@ The RTX 5090 config retains the applicable general optimizations and uses:
 
 Both tasks were measured on the current code with 40 steps, seed 42, CFG disabled, and the median latency of three consecutive requests. I2I uses one 1024×1024 reference image and produces a 1024×1024 image. End-to-end latency covers input encoding, condition-KV prefill, denoising, VAE decoding, post-processing, and PNG saving; it excludes model loading and one-time runner initialization.
 
+### 3.3 AMD ROCm GPUs (Radeon AI PRO R9700 / Radeon PRO W7900)
+
+The 5090 stack above (FlashAttention3, FlashInfer RoPE, CUTLASS FP8, SageAttention2 CUDA kernels) is CUDA-only. The configs below run Qwen-Image-2.1 on AMD RDNA GPUs with a torch-op inference path (`torch_sdpa`, `torch_real_rope`, torch LayerNorm/modulation) plus ROCm-friendly quantization, `torch.compile`, and SageAttention2. They quantize the **released BF16 weights on load** — no `tools/convert` step and no `dit_quantized_ckpt`.
+
+Validated in the `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` image (PyTorch 2.10, Triton 3.6). Install the runtime deps not baked into that image:
+
+```bash
+pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
+    qtorch langdetect tqdm imageio imageio-ffmpeg opencv-python-headless comfy-kitchen \
+    peft gguf prometheus_client fastapi uvicorn pydantic aiohttp pyzmq python-multipart \
+    PyJWT jsonschema sageattention
+```
+
+Run directly with `python -m lightx2v.infer` (a single GPU is enough — the text encoder is CPU-offloaded):
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
+    --model_cls qwen_image_21 --task t2i \
+    --model_path /path/to/Qwen-Image-2.1 \
+    --config_json configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json \
+    --prompt "A capybara wearing a wizard hat, oil painting" \
+    --size 1024 1024 --seed 42 \
+    --save_result_path ./save_results/qwen_image_21_t2i.png
+```
+
+Per-GPU config (pick one for `--config_json`):
+
+| GPU (arch) | Quantization | Configs |
+| --- | --- | --- |
+| R9700 (gfx1201 / RDNA4) | FP8 via `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json` (fastest), `..._fp8_compile.json`, `..._fp8_sage_25steps.json`, `qwen_image_21_r9700.json` (BF16 baseline) |
+| W7900 (gfx1100 / RDNA3) | INT8 via `torch._int_mm` (no FP8 on RDNA3) | `qwen_image_21_w7900_int8_compile_sage.json` (fastest), `..._int8_compile.json`, `..._int8_sage_25steps.json`, `qwen_image_21_w7900_int8.json` (baseline) |
+
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900): per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. On gfx1201 the VAE runs through an im2col GEMM (`vae_conv_im2col`) to avoid a nondeterministic MIOpen convolution defect; gfx1100 uses the native conv. The required SageAttention/Inductor Triton `num_stages` workaround for ROCm is applied automatically.
+
+| GPU | Quant | Steps | End-to-end |
+| --- | --- | ---: | ---: |
+| R9700 | FP8 | 40 | ~24 s |
+| R9700 | FP8 | 25 | ~17 s |
+| W7900 | INT8 | 40 | ~46 s |
+| W7900 | INT8 | 25 | ~29 s |
+
+Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup. 25 steps (the ComfyUI default for this model) has no visible quality loss versus 40. Compilation is a one-time warmup cost (~1.5–3 min); SageAttention needs `TORCHINDUCTOR_COMPILE_THREADS=1`, which the code sets automatically on the SageAttention path.
+
 ## 4. Service Deployment and API Usage
 
 Set the repository path, model path, and GPU ID in `server/start_server.sh`, then start the server:
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index c4fb17bf1..2b52631d3 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -84,6 +84,49 @@ RTX 5090 配置保留适用的通用优化，并使用：
 
 两类任务均使用当前代码，测试条件为 40 steps、seed 42、关闭 CFG，连续执行三次请求并取中位数。I2I 使用单张 1024×1024 参考图，输出尺寸为 1024×1024。端到端耗时包含输入编码、condition-KV prefill、去噪、VAE 解码、后处理和 PNG 保存，不包含模型加载及 Runner 的一次性初始化。
 
+### 3.3 AMD ROCm 显卡（Radeon AI PRO R9700 / Radeon PRO W7900）
+
+上面 5090 那套（FlashAttention3、FlashInfer RoPE、CUTLASS FP8、SageAttention2 的 CUDA kernel）都是 CUDA 专用。下面的配置让 Qwen-Image-2.1 跑在 AMD RDNA 显卡上，走 torch 算子推理路径（`torch_sdpa`、`torch_real_rope`、torch LayerNorm/modulation），配合 ROCm 可用的量化、`torch.compile` 和 SageAttention2。它们**在加载时直接量化官方 BF16 权重**——无需 `tools/convert` 转换、无需 `dit_quantized_ckpt`。
+
+验证环境为 `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` 镜像（PyTorch 2.10、Triton 3.6）。安装该镜像未自带的运行时依赖：
+
+```bash
+pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
+    qtorch langdetect tqdm imageio imageio-ffmpeg opencv-python-headless comfy-kitchen \
+    peft gguf prometheus_client fastapi uvicorn pydantic aiohttp pyzmq python-multipart \
+    PyJWT jsonschema sageattention
+```
+
+直接用 `python -m lightx2v.infer` 运行（单卡即可——条件编码器走 CPU offload）：
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
+    --model_cls qwen_image_21 --task t2i \
+    --model_path /path/to/Qwen-Image-2.1 \
+    --config_json configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json \
+    --prompt "A capybara wearing a wizard hat, oil painting" \
+    --size 1024 1024 --seed 42 \
+    --save_result_path ./save_results/qwen_image_21_t2i.png
+```
+
+各显卡对应配置（`--config_json` 选其一）：
+
+| 显卡（架构） | 量化 | 配置 |
+| --- | --- | --- |
+| R9700（gfx1201 / RDNA4） | FP8，走 `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json`（最快）、`..._fp8_compile.json`、`..._fp8_sage_25steps.json`、`qwen_image_21_r9700.json`（BF16 基线） |
+| W7900（gfx1100 / RDNA3） | INT8，走 `torch._int_mm`（RDNA3 无 FP8） | `qwen_image_21_w7900_int8_compile_sage.json`（最快）、`..._int8_compile.json`、`..._int8_sage_25steps.json`、`qwen_image_21_w7900_int8.json`（基线） |
+
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900）：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。gfx1201 上 VAE 走 im2col GEMM（`vae_conv_im2col`）以规避 MIOpen 卷积的非确定性缺陷；gfx1100 用原生卷积。SageAttention/Inductor 在 ROCm 上所需的 Triton `num_stages` 规避会自动生效。
+
+| 显卡 | 量化 | 步数 | 端到端 |
+| --- | --- | ---: | ---: |
+| R9700 | FP8 | 40 | ~24 s |
+| R9700 | FP8 | 25 | ~17 s |
+| W7900 | INT8 | 40 | ~46 s |
+| W7900 | INT8 | 25 | ~29 s |
+
+测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态。25 步（该模型在 ComfyUI 的默认步数）相比 40 步无可见画质损失。编译为一次性 warmup 开销（约 1.5–3 分钟）；SageAttention 需要 `TORCHINDUCTOR_COMPILE_THREADS=1`，代码在 SageAttention 路径上会自动设置。
+
 ## 4. 服务化部署与 API 调用
 
 先修改 `server/start_server.sh` 中的仓库路径、模型路径和 GPU 编号，然后启动：

From b08085fcbf9751ee048d4fdf42cfe4119e40b75f Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 08:09:16 +0000
Subject: [PATCH 09/21] fix(qwen-image-2.1): add trailing newline to ROCm
 configs for end-of-file-fixer

---
 configs/qwen_image_21/qwen_image_21_r9700.json                  | 2 +-
 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json      | 2 +-
 .../qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json    | 2 +-
 configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json | 2 +-
 configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json | 2 +-
 configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json     | 2 +-
 .../qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json    | 2 +-
 .../qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json    | 2 +-
 8 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/configs/qwen_image_21/qwen_image_21_r9700.json b/configs/qwen_image_21/qwen_image_21_r9700.json
index 44f45884d..fe53c9a0c 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700.json
@@ -14,4 +14,4 @@
     "warmup": false,
     "text_encoder_cpu_offload": true,
     "vae_conv_im2col": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
index 566b5b8e1..bc2dbfbcd 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
@@ -17,4 +17,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
index 81d4cf2be..a4a10023b 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
@@ -18,4 +18,4 @@
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true,
     "dit_fuse_qkv_mlp": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
index 50b34bdbf..8acb536fb 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
@@ -17,4 +17,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
index 984a18d4b..672e32052 100644
--- a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
+++ b/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
@@ -17,4 +17,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
index e683666e4..e3fefa9dc 100644
--- a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
@@ -16,4 +16,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "int8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
index 42ffb7a8a..459eb2eca 100644
--- a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
@@ -16,4 +16,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "int8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json b/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
index 6079e7f78..1775e1f70 100644
--- a/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
+++ b/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
@@ -16,4 +16,4 @@
     "dit_quantized": true,
     "dit_quant_scheme": "int8-rocm",
     "use_compile": true
-}
\ No newline at end of file
+}

From 7f60da42085a6705e1aba5aeee3063acc2b00ff4 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 09:54:47 +0000
Subject: [PATCH 10/21] refactor(qwen-image-2.1): move AMD ROCm ops and configs
 into lightx2v_platform

Address review: keep platform-specific code under the platform layer.

- MM: fp8-rocm/int8-rocm GEMM (torch._scaled_mm / torch._int_mm) moved to
  lightx2v_platform/ops/mm/amd_rocm and registered via PLATFORM_MM_WEIGHT_REGISTER;
  removed from the shared lightx2v/common/ops/mm/mm_weight.py.
- Attn: SageAttention/Triton ROCm num_stages workaround moved to
  lightx2v_platform/ops/attn/amd_rocm/sage_patch.py; removed from the shared
  sage_attn.py.
- Configs moved to configs/platforms/amd_rocm; run under PLATFORM=amd_rocm.
- Make aiter optional in AmdRocmDevice so the torch-native FP8/INT8 path needs
  no aiter build (aiter is not required on R9700/W7900).

Validated end-to-end (1024x1024, 40 steps, seed 42): R9700 fp8 ~24.4s,
W7900 int8 ~44.9s; both produce correct images.
---
 .../amd_rocm}/qwen_image_21_r9700.json        |   0
 .../qwen_image_21_r9700_fp8_compile.json      |   0
 ...qwen_image_21_r9700_fp8_compile_fused.json |   0
 .../qwen_image_21_r9700_fp8_compile_sage.json |   0
 .../qwen_image_21_r9700_fp8_sage_25steps.json |   0
 .../amd_rocm}/qwen_image_21_w7900_int8.json   |   0
 .../qwen_image_21_w7900_int8_compile.json     |   0
 ...qwen_image_21_w7900_int8_compile_sage.json |   0
 ...qwen_image_21_w7900_int8_sage_25steps.json |   0
 lightx2v/common/ops/attn/sage_attn.py         | 105 ----------
 lightx2v/common/ops/mm/mm_weight.py           | 176 -----------------
 .../weights/transformer_weights.py            |   2 +-
 lightx2v_platform/base/amd_rocm.py            |  42 ++--
 lightx2v_platform/ops/__init__.py             |   5 +
 .../ops/attn/amd_rocm/__init__.py             |   1 +
 .../ops/attn/amd_rocm/sage_patch.py           | 111 +++++++++++
 lightx2v_platform/ops/mm/amd_rocm/__init__.py |   1 +
 .../ops/mm/amd_rocm/mm_weight.py              | 183 ++++++++++++++++++
 scripts/qwen_image_21/README.md               |  10 +-
 scripts/qwen_image_21/README_zh.md            |  10 +-
 20 files changed, 336 insertions(+), 310 deletions(-)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_r9700.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_r9700_fp8_compile.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_r9700_fp8_compile_fused.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_r9700_fp8_compile_sage.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_r9700_fp8_sage_25steps.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_w7900_int8.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_w7900_int8_compile.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_w7900_int8_compile_sage.json (100%)
 rename configs/{qwen_image_21 => platforms/amd_rocm}/qwen_image_21_w7900_int8_sage_25steps.json (100%)
 create mode 100644 lightx2v_platform/ops/attn/amd_rocm/sage_patch.py
 create mode 100644 lightx2v_platform/ops/mm/amd_rocm/__init__.py
 create mode 100644 lightx2v_platform/ops/mm/amd_rocm/mm_weight.py

diff --git a/configs/qwen_image_21/qwen_image_21_r9700.json b/configs/platforms/amd_rocm/qwen_image_21_r9700.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_r9700.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700.json
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_r9700_fp8_compile.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_fused.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
diff --git a/configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_sage_25steps.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_r9700_fp8_sage_25steps.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_sage_25steps.json
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_w7900_int8.json
rename to configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_w7900_int8_compile.json
rename to configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_w7900_int8_compile_sage.json
rename to configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json
diff --git a/configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_sage_25steps.json
similarity index 100%
rename from configs/qwen_image_21/qwen_image_21_w7900_int8_sage_25steps.json
rename to configs/platforms/amd_rocm/qwen_image_21_w7900_int8_sage_25steps.json
diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index 6e999e241..07945b9f5 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -39,110 +39,6 @@
         sageattn = None
 
 
-def _clamp_sageattn_triton_num_stages_on_rocm(max_stages=2):
-    """Work around a Triton ROCm software-pipeliner bug.
-
-    On gfx1201 (ROCm/Triton), SageAttention's INT8 WMMA kernels fail to compile
-    at num_stages>=3 ("operation destroyed but still has uses" during the AMD
-    pipeliner pass). Clamp the launch-time num_stages on every SageAttention
-    Triton kernel so they compile; num_stages=2 runs correctly and is faster
-    than torch SDPA flash on this hardware.
-    """
-    if sageattn is None or getattr(torch.version, "hip", None) is None:
-        return
-    import importlib
-    import pkgutil
-
-    try:
-        import sageattention
-        from triton.runtime.jit import JITFunction
-    except ImportError:
-        return
-
-    patched = 0
-    for module_info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
-        try:
-            module = importlib.import_module(module_info.name)
-        except Exception:
-            continue
-        for value in vars(module).values():
-            if not isinstance(value, JITFunction) or getattr(value, "_rocm_num_stages_clamped", False):
-                continue
-            original_run = value.run
-
-            def run(*args, _original_run=original_run, **kwargs):
-                stages = kwargs.get("num_stages")
-                if isinstance(stages, int) and stages > max_stages:
-                    kwargs["num_stages"] = max_stages
-                return _original_run(*args, **kwargs)
-
-            try:
-                value.run = run
-                value._rocm_num_stages_clamped = True
-                patched += 1
-            except Exception:
-                continue
-    if patched:
-        logger.info(f"[SageAttention] ROCm: clamped Triton num_stages<={max_stages} on {patched} kernels")
-
-
-def _clamp_triton_compile_num_stages_on_rocm(max_stages=2):
-    """Clamp num_stages for every Triton compilation (including torch.compile /
-    Inductor kernels) on ROCm, where num_stages>=3 crashes the AMD pipeliner.
-
-    Inductor compiles in worker subprocesses; requires compile_threads=1 so this
-    main-process patch is in effect at compile time.
-    """
-    if getattr(torch.version, "hip", None) is None:
-        return
-    try:
-        import triton
-    except ImportError:
-        return
-    if getattr(triton.compile, "_rocm_num_stages_clamped", False):
-        return
-    _original_compile = triton.compile
-
-    def compile(*args, **kwargs):
-        options = kwargs.get("options")
-        if isinstance(options, dict) and isinstance(options.get("num_stages"), int) and options["num_stages"] > max_stages:
-            options["num_stages"] = max_stages
-        elif options is not None and isinstance(getattr(options, "num_stages", None), int) and options.num_stages > max_stages:
-            try:
-                options.num_stages = max_stages
-            except Exception:
-                pass
-        return _original_compile(*args, **kwargs)
-
-    compile._rocm_num_stages_clamped = True
-    triton.compile = compile
-
-
-def _force_rocm_inductor_single_thread():
-    """Make torch.compile / Inductor compile in-process on ROCm.
-
-    The num_stages clamp on ``triton.compile`` only takes effect in the process
-    that installed it. Inductor's default multi-threaded pool compiles kernels
-    in worker subprocesses that never see the clamp, so an Inductor-generated
-    kernel can still hit the num_stages>=3 pipeliner crash. Compiling in-process
-    keeps the clamp effective. Only pays the (one-time, warmup) single-threaded
-    compile cost on the SageAttention path that needs it.
-    """
-    if getattr(torch.version, "hip", None) is None:
-        return
-    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: set inductor compile_threads=1 so the Triton num_stages clamp reaches torch.compile kernels")
-
-
-_clamp_sageattn_triton_num_stages_on_rocm()
-_clamp_triton_compile_num_stages_on_rocm()
-
-
 if magi_register_custom_op is not None and sageattn is not None:
 
     @magi_register_custom_op(
@@ -179,7 +75,6 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
-        _force_rocm_inductor_single_thread()
 
     def apply(
         self,
diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py
index c67a042f9..e6c1ca2e6 100755
--- a/lightx2v/common/ops/mm/mm_weight.py
+++ b/lightx2v/common/ops/mm/mm_weight.py
@@ -771,182 +771,6 @@ def apply(self, input_tensor):
         return output_tensor
 
 
-@MM_WEIGHT_REGISTER("fp8-rocm")
-class MMWeightWfp8channelAfp8channeldynamicRocm(MMWeightQuantTemplate):
-    """W-fp8-e4m3-channel-sym / A-fp8-e4m3-token-dynamic via torch._scaled_mm.
-
-    MIOpen/CUTLASS/sgl FP8 kernels are unavailable on ROCm gfx1201 (RDNA4), but
-    torch._scaled_mm drives a working FP8 GEMM there. Weights are quantized from
-    the released BF16 checkpoint on load (no pre-quantized checkpoint needed).
-    """
-
-    FP8_MAX = 448.0
-
-    def __init__(
-        self,
-        weight_name,
-        bias_name,
-        create_cuda_buffer=False,
-        create_cpu_buffer=False,
-        lazy_load=False,
-        lazy_load_file=None,
-        is_post_adapter=False,
-        lora_prefix="diffusion_model.blocks",
-        lora_path="",
-    ):
-        super().__init__(
-            weight_name,
-            bias_name,
-            create_cuda_buffer,
-            create_cpu_buffer,
-            lazy_load,
-            lazy_load_file,
-            is_post_adapter,
-            lora_prefix,
-            lora_path,
-        )
-        # Weight stored (N, K) fp8; a column-major (K, N) view is taken in apply().
-        self.weight_need_transpose = False
-        self.scale_force_fp32 = True
-
-    def load(self, weight_dict):
-        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (N, 1)
-        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        # Release the BF16 source immediately so it does not coexist with every
-        # quantized weight (halves the load-time peak on GPU).
-        weight_dict.pop(self.weight_name, None)
-        if self.bias_name is not None:
-            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
-        else:
-            self.bias = None
-
-    def apply(self, input_tensor):
-        x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX).float()  # (M, 1)
-        x_q = (input_tensor / x_scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)  # (M, K)
-        output_tensor = torch._scaled_mm(
-            x_q,
-            self.weight.t(),  # (K, N) column-major
-            scale_a=x_scale,
-            scale_b=self.weight_scale.t(),  # (1, N)
-            bias=self.bias,
-            out_dtype=input_tensor.dtype,
-        )
-        if self.has_lora_branch:
-            return output_tensor + self.apply_lora(input_tensor)
-        return output_tensor
-
-
-class MMWeightFusedFp8Rocm(MMWeightWfp8channelAfp8channeldynamicRocm):
-    """Fuse several row-concatenatable linears (q/k/v, gate/up) into one fp8 GEMM.
-
-    Per-output-channel weight scales make this numerically identical to the
-    separate GEMMs, while cutting GEMM count and quantizing the (shared) input
-    activation once instead of once per linear. ``apply()`` returns the
-    concatenated output; the caller splits it with ``split_sizes``.
-    """
-
-    def __init__(self, weight_names, bias_name=None, **kwargs):
-        super().__init__(weight_names[0], bias_name, **kwargs)
-        self.fused_weight_names = list(weight_names)
-        self.split_sizes = None
-
-    def load(self, weight_dict):
-        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]  # each (N_i, K)
-        self.split_sizes = [w.shape[0] for w in weights]
-        weight = torch.cat(weights, dim=0)  # (sum_N, K)
-        del weights
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (sum_N, 1)
-        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        for name in self.fused_weight_names:
-            weight_dict.pop(name, None)
-        self.bias = None
-
-
-@MM_WEIGHT_REGISTER("int8-rocm")
-class MMWeightWint8channelAint8tokenRocm(MMWeightQuantTemplate):
-    """W-int8-channel-sym / A-int8-token-dynamic via torch._int_mm.
-
-    RDNA3 (gfx1100, e.g. W7900) has no FP8 tensor path, but torch._int_mm drives
-    a working INT8 WMMA GEMM there. Weights are quantized from the released BF16
-    checkpoint on load (no prequantized checkpoint needed).
-    """
-
-    INT8_MAX = 127.0
-
-    def __init__(
-        self,
-        weight_name,
-        bias_name=None,
-        create_cuda_buffer=False,
-        create_cpu_buffer=False,
-        lazy_load=False,
-        lazy_load_file=None,
-        is_post_adapter=False,
-        lora_prefix="diffusion_model.blocks",
-        lora_path="",
-    ):
-        super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, lora_prefix, lora_path)
-        self.weight_need_transpose = False
-        self.scale_force_fp32 = True
-
-    def load(self, weight_dict):
-        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX  # (N, 1)
-        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        weight_dict.pop(self.weight_name, None)
-        if self.bias_name is not None:
-            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
-        else:
-            self.bias = None
-
-    def apply(self, input_tensor):
-        if input_tensor.shape[0] <= 16:
-            # torch._int_mm requires M > 16; dequantize to BF16 for the small
-            # prefill segments (rare, negligible cost).
-            weight = (self.weight.to(torch.float32) * self.weight_scale).to(input_tensor.dtype)  # (N, K)
-            out = torch.mm(input_tensor, weight.t())
-        else:
-            x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX).float()  # (M, 1)
-            x_q = (input_tensor / x_scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)  # (M, K)
-            out = torch._int_mm(x_q, self.weight.t())  # (M, N) int32
-            out = (out.to(torch.float32) * x_scale * self.weight_scale.t()).to(input_tensor.dtype)  # dequant
-        if self.bias is not None:
-            out = out + self.bias
-        if self.has_lora_branch:
-            out = out + self.apply_lora(input_tensor)
-        return out
-
-
-class MMWeightFusedInt8Rocm(MMWeightWint8channelAint8tokenRocm):
-    """Fuse q/k/v, gate/up into one INT8 GEMM; apply() returns the concatenated
-    output, split by ``split_sizes``. See MMWeightFusedFp8Rocm."""
-
-    def __init__(self, weight_names, bias_name=None, **kwargs):
-        super().__init__(weight_names[0], bias_name, **kwargs)
-        self.fused_weight_names = list(weight_names)
-        self.split_sizes = None
-
-    def load(self, weight_dict):
-        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]
-        self.split_sizes = [w.shape[0] for w in weights]
-        weight = torch.cat(weights, dim=0)
-        del weights
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX
-        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        for name in self.fused_weight_names:
-            weight_dict.pop(name, None)
-        self.bias = None
-
-
 @MM_WEIGHT_REGISTER("int8-vllm")
 class MMWeightWint8channelAint8channeldynamicVllm(MMWeightQuantTemplate):
     """
diff --git a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
index 5535ecfe1..5cb8c7ab8 100644
--- a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
+++ b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
@@ -12,7 +12,7 @@ def __init__(self, index, config):
         # one shared activation quantization instead of one per linear).
         self.fused_qkv_mlp = bool(config.get("dit_fuse_qkv_mlp", False)) and mm_type in ("fp8-rocm", "int8-rocm")
         if self.fused_qkv_mlp:
-            from lightx2v.common.ops.mm.mm_weight import MMWeightFusedFp8Rocm, MMWeightFusedInt8Rocm
+            from lightx2v_platform.ops.mm.amd_rocm.mm_weight import MMWeightFusedFp8Rocm, MMWeightFusedInt8Rocm
 
             fused_cls = MMWeightFusedInt8Rocm if mm_type == "int8-rocm" else MMWeightFusedFp8Rocm
             self.add_module("qkv", fused_cls([f"{prefix}.attn.to_q.weight", f"{prefix}.attn.to_k.weight", f"{prefix}.attn.to_v.weight"], bias_name=None))
diff --git a/lightx2v_platform/base/amd_rocm.py b/lightx2v_platform/base/amd_rocm.py
index 21aeb93e4..94f154522 100644
--- a/lightx2v_platform/base/amd_rocm.py
+++ b/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,7 +111,8 @@ 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...")
 
@@ -117,19 +120,22 @@ def init_device_env():
         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:
diff --git a/lightx2v_platform/ops/__init__.py b/lightx2v_platform/ops/__init__.py
index 175d4deb1..4ac8949de 100755
--- a/lightx2v_platform/ops/__init__.py
+++ b/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
diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index 9ecd99131..e9fbba530 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1 +1,2 @@
+from . import sage_patch  # noqa: F401  applies ROCm Triton num_stages workarounds on import
 from .flash_attn import *
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_patch.py b/lightx2v_platform/ops/attn/amd_rocm/sage_patch.py
new file mode 100644
index 000000000..1f20ff026
--- /dev/null
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_patch.py
@@ -0,0 +1,111 @@
+"""ROCm-only SageAttention / Triton num_stages workarounds.
+
+On gfx1201 (ROCm/Triton), Triton's AMD software-pipeliner pass miscompiles
+kernels at num_stages>=3 ("operation destroyed but still has uses"). These
+patches clamp num_stages<=2 on both SageAttention's own Triton kernels and any
+torch.compile / Inductor Triton kernels, and force Inductor to compile
+in-process so the clamp actually reaches those kernels.
+
+Imported (and applied) only on the amd_rocm platform via
+``lightx2v_platform/ops/attn/amd_rocm/__init__.py``. Each function is a no-op
+off ROCm, so importing it elsewhere is harmless.
+"""
+
+import torch
+from loguru import logger
+
+
+def clamp_sageattn_triton_num_stages(max_stages=2):
+    """Clamp launch-time num_stages on every SageAttention Triton kernel."""
+    if getattr(torch.version, "hip", None) is None:
+        return
+    import importlib
+    import pkgutil
+
+    try:
+        import sageattention
+        from triton.runtime.jit import JITFunction
+    except ImportError:
+        return
+
+    patched = 0
+    for module_info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
+        try:
+            module = importlib.import_module(module_info.name)
+        except Exception:
+            continue
+        for value in vars(module).values():
+            if not isinstance(value, JITFunction) or getattr(value, "_rocm_num_stages_clamped", False):
+                continue
+            original_run = value.run
+
+            def run(*args, _original_run=original_run, **kwargs):
+                stages = kwargs.get("num_stages")
+                if isinstance(stages, int) and stages > max_stages:
+                    kwargs["num_stages"] = max_stages
+                return _original_run(*args, **kwargs)
+
+            try:
+                value.run = run
+                value._rocm_num_stages_clamped = True
+                patched += 1
+            except Exception:
+                continue
+    if patched:
+        logger.info(f"[SageAttention] ROCm: clamped Triton num_stages<={max_stages} on {patched} kernels")
+
+
+def clamp_triton_compile_num_stages(max_stages=2):
+    """Clamp num_stages for every Triton compilation (incl. torch.compile /
+    Inductor kernels) on ROCm, where num_stages>=3 crashes the AMD pipeliner."""
+    if getattr(torch.version, "hip", None) is None:
+        return
+    try:
+        import triton
+    except ImportError:
+        return
+    if getattr(triton.compile, "_rocm_num_stages_clamped", False):
+        return
+    _original_compile = triton.compile
+
+    def compile(*args, **kwargs):
+        options = kwargs.get("options")
+        if isinstance(options, dict) and isinstance(options.get("num_stages"), int) and options["num_stages"] > max_stages:
+            options["num_stages"] = max_stages
+        elif options is not None and isinstance(getattr(options, "num_stages", None), int) and options.num_stages > max_stages:
+            try:
+                options.num_stages = max_stages
+            except Exception:
+                pass
+        return _original_compile(*args, **kwargs)
+
+    compile._rocm_num_stages_clamped = True
+    triton.compile = compile
+
+
+def force_rocm_inductor_single_thread():
+    """Make torch.compile / Inductor compile in-process on ROCm.
+
+    The num_stages clamp on ``triton.compile`` only takes effect in the process
+    that installed it. Inductor's default multi-threaded pool compiles kernels
+    in worker subprocesses that never see the clamp, so an Inductor-generated
+    kernel can still hit the num_stages>=3 pipeliner crash. Compiling in-process
+    keeps the clamp effective (one-time, warmup-only single-threaded cost)."""
+    if getattr(torch.version, "hip", None) is None:
+        return
+    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: set inductor compile_threads=1 so the Triton num_stages clamp reaches torch.compile kernels")
+
+
+def apply_rocm_sage_patches():
+    clamp_sageattn_triton_num_stages()
+    clamp_triton_compile_num_stages()
+    force_rocm_inductor_single_thread()
+
+
+apply_rocm_sage_patches()
diff --git a/lightx2v_platform/ops/mm/amd_rocm/__init__.py b/lightx2v_platform/ops/mm/amd_rocm/__init__.py
new file mode 100644
index 000000000..3f9898101
--- /dev/null
+++ b/lightx2v_platform/ops/mm/amd_rocm/__init__.py
@@ -0,0 +1 @@
+from .mm_weight import *
diff --git a/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
new file mode 100644
index 000000000..a91e128c3
--- /dev/null
+++ b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
@@ -0,0 +1,183 @@
+"""AMD ROCm quantized GEMM weights (aiter-free).
+
+FP8 (RDNA4 / gfx1201) drives torch._scaled_mm; INT8 (RDNA3 / gfx1100) drives
+torch._int_mm. Both quantize the released BF16 checkpoint on load, so no
+pre-quantized checkpoint and no aiter build are required. CUTLASS / sgl / vLLM
+FP8 kernels are CUDA-only, which is why these platform-native paths exist.
+"""
+
+import torch
+
+from lightx2v_platform.base.global_var import AI_DEVICE
+from lightx2v_platform.ops.mm.template import MMWeightQuantTemplate
+from lightx2v_platform.registry_factory import PLATFORM_MM_WEIGHT_REGISTER
+
+
+@PLATFORM_MM_WEIGHT_REGISTER("fp8-rocm")
+class MMWeightWfp8channelAfp8channeldynamicRocm(MMWeightQuantTemplate):
+    """W-fp8-e4m3-channel-sym / A-fp8-e4m3-token-dynamic via torch._scaled_mm.
+
+    MIOpen/CUTLASS/sgl FP8 kernels are unavailable on ROCm gfx1201 (RDNA4), but
+    torch._scaled_mm drives a working FP8 GEMM there. Weights are quantized from
+    the released BF16 checkpoint on load (no pre-quantized checkpoint needed).
+    """
+
+    FP8_MAX = 448.0
+
+    def __init__(
+        self,
+        weight_name,
+        bias_name,
+        create_cuda_buffer=False,
+        create_cpu_buffer=False,
+        lazy_load=False,
+        lazy_load_file=None,
+        is_post_adapter=False,
+        lora_prefix="diffusion_model.blocks",
+        lora_path="",
+    ):
+        super().__init__(
+            weight_name,
+            bias_name,
+            create_cuda_buffer,
+            create_cpu_buffer,
+            lazy_load,
+            lazy_load_file,
+            is_post_adapter,
+            lora_prefix,
+            lora_path,
+        )
+        # Weight stored (N, K) fp8; a column-major (K, N) view is taken in apply().
+        self.weight_need_transpose = False
+
+    def load(self, weight_dict):
+        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (N, 1)
+        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        # Release the BF16 source immediately so it does not coexist with every
+        # quantized weight (halves the load-time peak on GPU).
+        weight_dict.pop(self.weight_name, None)
+        if self.bias_name is not None:
+            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
+        else:
+            self.bias = None
+
+    def apply(self, input_tensor):
+        x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX).float()  # (M, 1)
+        x_q = (input_tensor / x_scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)  # (M, K)
+        output_tensor = torch._scaled_mm(
+            x_q,
+            self.weight.t(),  # (K, N) column-major
+            scale_a=x_scale,
+            scale_b=self.weight_scale.t(),  # (1, N)
+            bias=self.bias,
+            out_dtype=input_tensor.dtype,
+        )
+        return output_tensor
+
+
+class MMWeightFusedFp8Rocm(MMWeightWfp8channelAfp8channeldynamicRocm):
+    """Fuse several row-concatenatable linears (q/k/v, gate/up) into one fp8 GEMM.
+
+    Per-output-channel weight scales make this numerically identical to the
+    separate GEMMs, while cutting GEMM count and quantizing the (shared) input
+    activation once instead of once per linear. ``apply()`` returns the
+    concatenated output; the caller splits it with ``split_sizes``.
+    """
+
+    def __init__(self, weight_names, bias_name=None, **kwargs):
+        super().__init__(weight_names[0], bias_name, **kwargs)
+        self.fused_weight_names = list(weight_names)
+        self.split_sizes = None
+
+    def load(self, weight_dict):
+        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]  # each (N_i, K)
+        self.split_sizes = [w.shape[0] for w in weights]
+        weight = torch.cat(weights, dim=0)  # (sum_N, K)
+        del weights
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (sum_N, 1)
+        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        for name in self.fused_weight_names:
+            weight_dict.pop(name, None)
+        self.bias = None
+
+
+@PLATFORM_MM_WEIGHT_REGISTER("int8-rocm")
+class MMWeightWint8channelAint8tokenRocm(MMWeightQuantTemplate):
+    """W-int8-channel-sym / A-int8-token-dynamic via torch._int_mm.
+
+    RDNA3 (gfx1100, e.g. W7900) has no FP8 tensor path, but torch._int_mm drives
+    a working INT8 WMMA GEMM there. Weights are quantized from the released BF16
+    checkpoint on load (no prequantized checkpoint needed).
+    """
+
+    INT8_MAX = 127.0
+
+    def __init__(
+        self,
+        weight_name,
+        bias_name=None,
+        create_cuda_buffer=False,
+        create_cpu_buffer=False,
+        lazy_load=False,
+        lazy_load_file=None,
+        is_post_adapter=False,
+        lora_prefix="diffusion_model.blocks",
+        lora_path="",
+    ):
+        super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, lora_prefix, lora_path)
+        self.weight_need_transpose = False
+
+    def load(self, weight_dict):
+        weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX  # (N, 1)
+        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        weight_dict.pop(self.weight_name, None)
+        if self.bias_name is not None:
+            self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
+        else:
+            self.bias = None
+
+    def apply(self, input_tensor):
+        if input_tensor.shape[0] <= 16:
+            # torch._int_mm requires M > 16; dequantize to BF16 for the small
+            # prefill segments (rare, negligible cost).
+            weight = (self.weight.to(torch.float32) * self.weight_scale).to(input_tensor.dtype)  # (N, K)
+            out = torch.mm(input_tensor, weight.t())
+        else:
+            x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX).float()  # (M, 1)
+            x_q = (input_tensor / x_scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)  # (M, K)
+            out = torch._int_mm(x_q, self.weight.t())  # (M, N) int32
+            out = (out.to(torch.float32) * x_scale * self.weight_scale.t()).to(input_tensor.dtype)  # dequant
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+class MMWeightFusedInt8Rocm(MMWeightWint8channelAint8tokenRocm):
+    """Fuse q/k/v, gate/up into one INT8 GEMM; apply() returns the concatenated
+    output, split by ``split_sizes``. See MMWeightFusedFp8Rocm."""
+
+    def __init__(self, weight_names, bias_name=None, **kwargs):
+        super().__init__(weight_names[0], bias_name, **kwargs)
+        self.fused_weight_names = list(weight_names)
+        self.split_sizes = None
+
+    def load(self, weight_dict):
+        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]
+        self.split_sizes = [w.shape[0] for w in weights]
+        weight = torch.cat(weights, dim=0)
+        del weights
+        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX
+        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
+        self.weight_scale = scale.to(torch.float32)
+        del weight
+        for name in self.fused_weight_names:
+            weight_dict.pop(name, None)
+        self.bias = None
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index 96ed7ca3f..36a662846 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -97,26 +97,26 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-Run directly with `python -m lightx2v.infer` (a single GPU is enough — the text encoder is CPU-offloaded):
+Run with `PLATFORM=amd_rocm` (this activates the AMD ROCm platform backend under `lightx2v_platform/`, which registers the ROCm FP8/INT8 GEMM and applies the SageAttention Triton workaround). A single GPU is enough — the text encoder is CPU-offloaded. No `aiter` build is required for this torch-native path:
 
 ```bash
-CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
+PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
     --model_cls qwen_image_21 --task t2i \
     --model_path /path/to/Qwen-Image-2.1 \
-    --config_json configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json \
+    --config_json configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json \
     --prompt "A capybara wearing a wizard hat, oil painting" \
     --size 1024 1024 --seed 42 \
     --save_result_path ./save_results/qwen_image_21_t2i.png
 ```
 
-Per-GPU config (pick one for `--config_json`):
+Per-GPU config (pick one for `--config_json`, all under `configs/platforms/amd_rocm/`):
 
 | GPU (arch) | Quantization | Configs |
 | --- | --- | --- |
 | R9700 (gfx1201 / RDNA4) | FP8 via `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json` (fastest), `..._fp8_compile.json`, `..._fp8_sage_25steps.json`, `qwen_image_21_r9700.json` (BF16 baseline) |
 | W7900 (gfx1100 / RDNA3) | INT8 via `torch._int_mm` (no FP8 on RDNA3) | `qwen_image_21_w7900_int8_compile_sage.json` (fastest), `..._int8_compile.json`, `..._int8_sage_25steps.json`, `qwen_image_21_w7900_int8.json` (baseline) |
 
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900): per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. On gfx1201 the VAE runs through an im2col GEMM (`vae_conv_im2col`) to avoid a nondeterministic MIOpen convolution defect; gfx1100 uses the native conv. The required SageAttention/Inductor Triton `num_stages` workaround for ROCm is applied automatically.
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. On gfx1201 the VAE runs through an im2col GEMM (`vae_conv_im2col`) to avoid a nondeterministic MIOpen convolution defect; gfx1100 uses the native conv. The required SageAttention/Inductor Triton `num_stages` workaround for ROCm is applied automatically.
 
 | GPU | Quant | Steps | End-to-end |
 | --- | --- | ---: | ---: |
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 2b52631d3..47ebef962 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -97,26 +97,26 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-直接用 `python -m lightx2v.infer` 运行（单卡即可——条件编码器走 CPU offload）：
+运行时设置 `PLATFORM=amd_rocm`（这会激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端：注册 ROCm 的 FP8/INT8 GEMM，并自动应用 SageAttention 的 Triton 规避）。单卡即可——条件编码器走 CPU offload。该 torch 原生路径**无需编译 aiter**：
 
 ```bash
-CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
+PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
     --model_cls qwen_image_21 --task t2i \
     --model_path /path/to/Qwen-Image-2.1 \
-    --config_json configs/qwen_image_21/qwen_image_21_r9700_fp8_compile_sage.json \
+    --config_json configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json \
     --prompt "A capybara wearing a wizard hat, oil painting" \
     --size 1024 1024 --seed 42 \
     --save_result_path ./save_results/qwen_image_21_t2i.png
 ```
 
-各显卡对应配置（`--config_json` 选其一）：
+各显卡对应配置（`--config_json` 选其一，均在 `configs/platforms/amd_rocm/` 下）：
 
 | 显卡（架构） | 量化 | 配置 |
 | --- | --- | --- |
 | R9700（gfx1201 / RDNA4） | FP8，走 `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json`（最快）、`..._fp8_compile.json`、`..._fp8_sage_25steps.json`、`qwen_image_21_r9700.json`（BF16 基线） |
 | W7900（gfx1100 / RDNA3） | INT8，走 `torch._int_mm`（RDNA3 无 FP8） | `qwen_image_21_w7900_int8_compile_sage.json`（最快）、`..._int8_compile.json`、`..._int8_sage_25steps.json`、`qwen_image_21_w7900_int8.json`（基线） |
 
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900）：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。gfx1201 上 VAE 走 im2col GEMM（`vae_conv_im2col`）以规避 MIOpen 卷积的非确定性缺陷；gfx1100 用原生卷积。SageAttention/Inductor 在 ROCm 上所需的 Triton `num_stages` 规避会自动生效。
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。gfx1201 上 VAE 走 im2col GEMM（`vae_conv_im2col`）以规避 MIOpen 卷积的非确定性缺陷；gfx1100 用原生卷积。SageAttention/Inductor 在 ROCm 上所需的 Triton `num_stages` 规避会自动生效。
 
 | 显卡 | 量化 | 步数 | 端到端 |
 | --- | --- | ---: | ---: |

From 22655cd825036b3beee100ef2d0082446cf1740f Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 10:05:25 +0000
Subject: [PATCH 11/21] refactor(qwen-image-2.1): rename amd_rocm sage_patch.py
 to sage_attn.py

---
 lightx2v_platform/ops/attn/amd_rocm/__init__.py                 | 2 +-
 .../ops/attn/amd_rocm/{sage_patch.py => sage_attn.py}           | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename lightx2v_platform/ops/attn/amd_rocm/{sage_patch.py => sage_attn.py} (100%)

diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index e9fbba530..5bb7ca05d 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1,2 +1,2 @@
-from . import sage_patch  # noqa: F401  applies ROCm Triton num_stages workarounds on import
+from . import sage_attn  # noqa: F401  applies ROCm Triton num_stages workarounds on import
 from .flash_attn import *
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_patch.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
similarity index 100%
rename from lightx2v_platform/ops/attn/amd_rocm/sage_patch.py
rename to lightx2v_platform/ops/attn/amd_rocm/sage_attn.py

From 8d6e078c8a564573ad2ccf42c86bf449b83c3328 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 12:02:19 +0000
Subject: [PATCH 12/21] =?UTF-8?q?refactor(qwen-image-2.1):=20address=20ROC?=
 =?UTF-8?q?m=20PR=20review=20=E2=80=94=20decouple,=20shrink=20scope,=20add?=
 =?UTF-8?q?=20tests?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- VAE: drop the im2col/GemmConv2d workaround and vae_device/vae_dtype; the AMD
  platform already disables cuDNN/MIOpen, so gfx1201 runs native conv without the
  nondeterministic NaN (verified on R9700). vae.py returns to plain nn.Conv2d.
- Remove QKV / gate-up fusion (measured ~0%, compute-bound) and its AMD-specific
  fused MM classes + config; the model no longer imports AMD backend classes and
  the infer path no longer branches on hasattr(block, "qkv").
- SageAttention Triton num_stages workaround now installs only when sage_attn2 is
  selected (from SageAttn2Weight.__init__) and only on validated archs
  (gfx1201/gfx1100), instead of unconditionally for every ROCm workload.
- fp8/int8 MM: fail-fast if torch._scaled_mm/_int_mm is missing or arch is wrong;
  int8 GEMM inputs made contiguous.
- Configs/README: rename *_sage_25steps -> *_compile_sage_25steps, relabel INT8
  eager baseline, drop the experimental fused config, pin tested transformers
  version, soften the 25-step quality claim, add benchmark methodology.
- Add lightx2v_platform/test/test_amd_rocm_quant.py (fp8/int8 vs bf16 across
  M=1/16/17/31/64, incl. the int8 fallback boundary; registration without aiter).

Validated on R9700 (fp8) and W7900 (int8): quant test ALL PASS, t2i produces
correct images (native conv, no fusion), 1024x1024.
---
 .../amd_rocm/qwen_image_21_r9700.json         |  3 +-
 .../qwen_image_21_r9700_fp8_compile.json      |  1 -
 ...qwen_image_21_r9700_fp8_compile_fused.json | 21 ------
 .../qwen_image_21_r9700_fp8_compile_sage.json |  1 -
 ...ge_21_r9700_fp8_compile_sage_25steps.json} |  1 -
 ...e_21_w7900_int8_compile_sage_25steps.json} |  0
 lightx2v/common/ops/attn/sage_attn.py         |  8 +++
 .../qwen_image_21/infer/transformer_infer.py  | 18 ++---
 .../weights/transformer_weights.py            | 29 +++-----
 .../video_encoders/hf/qwen_image_21/vae.py    | 62 +++-------------
 .../ops/attn/amd_rocm/__init__.py             |  1 -
 .../ops/attn/amd_rocm/sage_attn.py            | 27 ++++++-
 .../ops/mm/amd_rocm/mm_weight.py              | 72 +++++--------------
 lightx2v_platform/test/test_amd_rocm_quant.py | 71 ++++++++++++++++++
 scripts/qwen_image_21/README.md               | 10 +--
 scripts/qwen_image_21/README_zh.md            | 10 +--
 16 files changed, 153 insertions(+), 182 deletions(-)
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json
 rename configs/platforms/amd_rocm/{qwen_image_21_r9700_fp8_sage_25steps.json => qwen_image_21_r9700_fp8_compile_sage_25steps.json} (94%)
 rename configs/platforms/amd_rocm/{qwen_image_21_w7900_int8_sage_25steps.json => qwen_image_21_w7900_int8_compile_sage_25steps.json} (100%)
 create mode 100644 lightx2v_platform/test/test_amd_rocm_quant.py

diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700.json b/configs/platforms/amd_rocm/qwen_image_21_r9700.json
index fe53c9a0c..b4ab63018 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_r9700.json
@@ -12,6 +12,5 @@
     "enable_cfg": false,
     "sample_guide_scale": 1.0,
     "warmup": false,
-    "text_encoder_cpu_offload": true,
-    "vae_conv_im2col": true
+    "text_encoder_cpu_offload": true
 }
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
index bc2dbfbcd..1b1be254e 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json
deleted file mode 100644
index a4a10023b..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_fused.json
+++ /dev/null
@@ -1,21 +0,0 @@
-{
-    "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": true,
-    "text_encoder_cpu_offload": true,
-    "vae_conv_im2col": true,
-    "dit_quantized": true,
-    "dit_quant_scheme": "fp8-rocm",
-    "use_compile": true,
-    "dit_fuse_qkv_mlp": true
-}
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
index 8acb536fb..6cc41a24c 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
similarity index 94%
rename from configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_sage_25steps.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
index 672e32052..142a6e8e6 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_sage_25steps.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
@@ -13,7 +13,6 @@
     "sample_guide_scale": 1.0,
     "warmup": true,
     "text_encoder_cpu_offload": true,
-    "vae_conv_im2col": true,
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-rocm",
     "use_compile": true
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json
similarity index 100%
rename from configs/platforms/amd_rocm/qwen_image_21_w7900_int8_sage_25steps.json
rename to configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json
diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index 07945b9f5..ea94b191a 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -75,6 +75,14 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
+        if getattr(torch.version, "hip", None) is not None:
+            # SageAttention's Triton kernels miscompile at num_stages>=3 on some
+            # RDNA archs; install the clamp before any compile. No-op off the
+            # validated archs. Scoped here so it only runs when SageAttention is
+            # actually selected, not for every ROCm workload.
+            from lightx2v_platform.ops.attn.amd_rocm.sage_attn import apply_rocm_sage_patches
+
+            apply_rocm_sage_patches()
 
     def apply(
         self,
diff --git a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
index 16c815150..940d4c083 100644
--- a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
+++ b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
@@ -32,15 +32,9 @@ def _modulation(self, state):
 
     def _qkv(self, block, x, scale, rotary, rotary_positions):
         h = block.norm1.apply(x) * scale
-        if hasattr(block, "qkv"):
-            q, k, v = block.qkv.apply(h).split(block.qkv.split_sizes, dim=-1)
-            q = q.reshape(-1, self.heads, self.head_dim)
-            k = k.reshape_as(q)
-            v = v.reshape_as(q)
-        else:
-            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 = 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)
         q, k = self.rope.apply(q, k, rotary, positions=rotary_positions)
         return q, k, v
@@ -62,11 +56,7 @@ def _finish_block(self, block, x, attention, modulation):
 
         x = x + gate1 * block.out.apply(attention)
         h = block.norm2.apply(x) * scale2
-        if hasattr(block, "gate_up"):
-            gate_out, up_out = block.gate_up.apply(h).split(block.gate_up.split_sizes, dim=-1)
-            mlp = F.silu(gate_out) * up_out
-        else:
-            mlp = F.silu(block.gate.apply(h)) * block.up.apply(h)
+        mlp = F.silu(block.gate.apply(h)) * block.up.apply(h)
         x = x + gate2 * block.down.apply(mlp)
         return x.clamp(-65504, 65504) if x.dtype == torch.float16 else x
 
diff --git a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
index 5cb8c7ab8..37e629251 100644
--- a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
+++ b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
@@ -8,26 +8,15 @@ def __init__(self, index, config):
         super().__init__()
         prefix = f"transformer_blocks.{index}"
         mm_type = config["dit_quant_scheme"] if config.get("dit_quantized", False) else "Default"
-        # Fuse q/k/v and gate/up into single fp8 GEMMs (fewer, larger GEMMs and
-        # one shared activation quantization instead of one per linear).
-        self.fused_qkv_mlp = bool(config.get("dit_fuse_qkv_mlp", False)) and mm_type in ("fp8-rocm", "int8-rocm")
-        if self.fused_qkv_mlp:
-            from lightx2v_platform.ops.mm.amd_rocm.mm_weight import MMWeightFusedFp8Rocm, MMWeightFusedInt8Rocm
-
-            fused_cls = MMWeightFusedInt8Rocm if mm_type == "int8-rocm" else MMWeightFusedFp8Rocm
-            self.add_module("qkv", fused_cls([f"{prefix}.attn.to_q.weight", f"{prefix}.attn.to_k.weight", f"{prefix}.attn.to_v.weight"], bias_name=None))
-            self.add_module("gate_up", fused_cls([f"{prefix}.img_mlp.gate_layer.weight", f"{prefix}.img_mlp.proj.weight"], bias_name=None))
-            linear_keys = {"out": "attn.to_out.0", "down": "img_mlp.out"}
-        else:
-            linear_keys = {
-                "q": "attn.to_q",
-                "k": "attn.to_k",
-                "v": "attn.to_v",
-                "out": "attn.to_out.0",
-                "up": "img_mlp.proj",
-                "gate": "img_mlp.gate_layer",
-                "down": "img_mlp.out",
-            }
+        linear_keys = {
+            "q": "attn.to_q",
+            "k": "attn.to_k",
+            "v": "attn.to_v",
+            "out": "attn.to_out.0",
+            "up": "img_mlp.proj",
+            "gate": "img_mlp.gate_layer",
+            "down": "img_mlp.out",
+        }
         for name, key in linear_keys.items():
             linear = MM_WEIGHT_REGISTER[mm_type](f"{prefix}.{key}.weight", bias_name=None)
             if mm_type == "fp8-f16-accum":
diff --git a/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py b/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
index 48a727a59..81f3ba8d7 100644
--- a/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
+++ b/lightx2v/models/video_encoders/hf/qwen_image_21/vae.py
@@ -18,48 +18,14 @@
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 
-# MIOpen convolutions on gfx1201 (RDNA4) return NaN nondeterministically for this
-# VAE's shapes. Express every conv as im2col + GEMM (hipBLASLt), which is reliable
-# on this hardware. The input is padded by the caller.
-def _im2col_conv2d(x, weight, bias, stride):
-    kh, kw = weight.shape[-2:]
-    sh, sw = stride
-    b, _, h, w = x.shape
-    cols = F.unfold(x, (kh, kw), stride=(sh, sw))  # (B, C*kh*kw, L)
-    out = torch.matmul(weight.reshape(weight.shape[0], -1), cols)  # (B, O, L)
-    out_h = (h - kh) // sh + 1
-    out_w = (w - kw) // sw + 1
-    out = out.reshape(b, weight.shape[0], out_h, out_w)
-    if bias is not None:
-        out = out + bias.reshape(1, -1, 1, 1)
-    return out
-
-
-class GemmConv2d(nn.Conv2d):
-    """nn.Conv2d that runs through im2col+GEMM (MIOpen-free) when enabled."""
-
-    im2col_conv = False
-
-    def forward(self, x):
-        if self.im2col_conv:
-            ph, pw = self.padding
-            padded = F.pad(x, (pw, pw, ph, ph)) if (ph or pw) else x
-            return _im2col_conv2d(padded, self.weight, self.bias, self.stride)
-        return super().forward(x)
-
-
 class ImageConv(nn.Conv2d):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.image_padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0])
         self.padding = (0, 0)
-        self.im2col_conv = False
 
     def forward(self, x):
-        padded = F.pad(x.squeeze(2), self.image_padding)
-        if self.im2col_conv:
-            return _im2col_conv2d(padded, self.weight, self.bias, self.stride).unsqueeze(2)
-        return super().forward(padded).unsqueeze(2)
+        return super().forward(F.pad(x.squeeze(2), self.image_padding)).unsqueeze(2)
 
 
 class ImageRMSNorm(nn.Module):
@@ -90,8 +56,8 @@ class AttentionBlock(nn.Module):
     def __init__(self, dim):
         super().__init__()
         self.norm = ImageRMSNorm(dim, images=True)
-        self.to_qkv = GemmConv2d(dim, dim * 3, 1)
-        self.proj = GemmConv2d(dim, dim, 1)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
 
     def forward(self, x):
         b, c, _, h, w = x.shape
@@ -121,9 +87,9 @@ class Resample(nn.Module):
     def __init__(self, dim, up, temporal):
         super().__init__()
         if up:
-            self.resample = nn.Sequential(Upsample(), GemmConv2d(dim, dim, 3, padding=1))
+            self.resample = nn.Sequential(Upsample(), nn.Conv2d(dim, dim, 3, padding=1))
         else:
-            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), GemmConv2d(dim, dim, 3, stride=2))
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=2))
         if temporal:
             self.time_conv = ImageConv(dim, dim * 2 if up else dim, 1)
 
@@ -215,11 +181,6 @@ def __init__(self, config):
         if not cfg["is_residual"] or cfg.get("patch_size") is not None or cfg["attn_scales"]:
             raise ValueError("Expected the released residual, unpatched Qwen-Image-2.1 image VAE")
         self.scale_factor = config["vae_scale_factor"]
-        # MIOpen convolutions on gfx1201 (RDNA4) return NaN for this decoder's
-        # large convs, so allow running the VAE on CPU (always fp32 there).
-        self.vae_device = torch.device("cpu") if str(config.get("vae_device", "")).lower() == "cpu" else torch.device(AI_DEVICE)
-        _vd = str(config.get("vae_dtype", "")).lower()
-        self.vae_dtype = torch.float32 if self.vae_device.type == "cpu" else {"fp32": torch.float32, "float32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(_vd, GET_DTYPE())
         with torch.device("meta"):
             self.encoder = Encoder(cfg)
             self.decoder = Decoder(cfg)
@@ -231,21 +192,16 @@ def __init__(self, config):
         self.load_state_dict(state, strict=True, assign=True)
         self.register_buffer("latents_mean", torch.tensor(cfg["latents_mean"]).reshape(1, -1, 1, 1, 1), persistent=False)
         self.register_buffer("latents_std", torch.tensor(cfg["latents_std"]).reshape(1, -1, 1, 1, 1), persistent=False)
-        self.to(dtype=self.vae_dtype, device=self.vae_device)
-        if self.vae_device.type != "cpu" and config.get("vae_conv_im2col", False):
-            for module in self.modules():
-                if isinstance(module, (ImageConv, GemmConv2d)):
-                    module.im2col_conv = True
+        self.to(dtype=GET_DTYPE(), device=AI_DEVICE)
         self.eval().requires_grad_(False)
 
     @torch.inference_mode()
     def encode(self, image):
-        mean = self.quant_conv(self.encoder(image.to(device=self.vae_device, dtype=self.vae_dtype))).chunk(2, dim=1)[0]
+        mean = self.quant_conv(self.encoder(image.to(device=AI_DEVICE, dtype=GET_DTYPE()))).chunk(2, dim=1)[0]
         return ((mean - self.latents_mean) / self.latents_std).flatten(2).transpose(1, 2)
 
     @torch.inference_mode()
     def decode(self, latents, size):
         h, w = size
-        z = latents.transpose(1, 2).reshape(1, self.config["z_dim"], 1, h // self.scale_factor, w // self.scale_factor).to(device=self.vae_device, dtype=self.vae_dtype)
-        out = self.decoder(self.post_quant_conv(z * self.latents_std + self.latents_mean)).clamp(-1, 1)[:, :, 0]
-        return out.to(AI_DEVICE)
+        z = latents.transpose(1, 2).reshape(1, self.config["z_dim"], 1, h // self.scale_factor, w // self.scale_factor).to(GET_DTYPE())
+        return self.decoder(self.post_quant_conv(z * self.latents_std + self.latents_mean)).clamp(-1, 1)[:, :, 0]
diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index 5bb7ca05d..9ecd99131 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1,2 +1 @@
-from . import sage_attn  # noqa: F401  applies ROCm Triton num_stages workarounds on import
 from .flash_attn import *
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index 1f20ff026..c65165a01 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -102,10 +102,31 @@ def force_rocm_inductor_single_thread():
         logger.info("[SageAttention] ROCm: set inductor compile_threads=1 so the Triton num_stages clamp reaches torch.compile kernels")
 
 
+_VALIDATED_ARCHS = ("gfx1201", "gfx1100")
+
+
+def _is_validated_rocm_arch():
+    """True only on the RDNA arch(es) where this workaround is validated.
+
+    The num_stages>=3 pipeliner crash was reproduced and the clamp validated on
+    gfx1201 (R9700) and gfx1100 (W7900). Other archs (e.g. MI300) are left
+    untouched so their kernels keep their own num_stages tuning.
+    """
+    if getattr(torch.version, "hip", None) is None or not torch.cuda.is_available():
+        return False
+    try:
+        arch = torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:
+        return False
+    return any(arch.startswith(a) for a in _VALIDATED_ARCHS)
+
+
 def apply_rocm_sage_patches():
+    """Install the Triton num_stages workaround. No-op unless running on a
+    validated RDNA arch; call this only when SageAttention is actually selected
+    (from SageAttn2Weight.__init__), not for every ROCm workload."""
+    if not _is_validated_rocm_arch():
+        return
     clamp_sageattn_triton_num_stages()
     clamp_triton_compile_num_stages()
     force_rocm_inductor_single_thread()
-
-
-apply_rocm_sage_patches()
diff --git a/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
index a91e128c3..99961b9bc 100644
--- a/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
+++ b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
@@ -13,6 +13,13 @@
 from lightx2v_platform.registry_factory import PLATFORM_MM_WEIGHT_REGISTER
 
 
+def _gcn_arch():
+    try:
+        return torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:
+        return ""
+
+
 @PLATFORM_MM_WEIGHT_REGISTER("fp8-rocm")
 class MMWeightWfp8channelAfp8channeldynamicRocm(MMWeightQuantTemplate):
     """W-fp8-e4m3-channel-sym / A-fp8-e4m3-token-dynamic via torch._scaled_mm.
@@ -49,6 +56,11 @@ def __init__(
         )
         # Weight stored (N, K) fp8; a column-major (K, N) view is taken in apply().
         self.weight_need_transpose = False
+        if not hasattr(torch, "_scaled_mm"):
+            raise RuntimeError("fp8-rocm requires torch._scaled_mm, which is absent from this PyTorch/ROCm build.")
+        arch = _gcn_arch()
+        if not arch.startswith("gfx1201"):
+            raise RuntimeError(f"fp8-rocm is validated only on gfx1201 (RDNA4); detected '{arch}'. Use int8-rocm instead.")
 
     def load(self, weight_dict):
         weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
@@ -56,8 +68,7 @@ def load(self, weight_dict):
         self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
         self.weight_scale = scale.to(torch.float32)
         del weight
-        # Release the BF16 source immediately so it does not coexist with every
-        # quantized weight (halves the load-time peak on GPU).
+        # Drop the source weight so it does not coexist with every quantized one.
         weight_dict.pop(self.weight_name, None)
         if self.bias_name is not None:
             self.bias = weight_dict[self.bias_name].to(AI_DEVICE).to(self.infer_dtype)
@@ -78,34 +89,6 @@ def apply(self, input_tensor):
         return output_tensor
 
 
-class MMWeightFusedFp8Rocm(MMWeightWfp8channelAfp8channeldynamicRocm):
-    """Fuse several row-concatenatable linears (q/k/v, gate/up) into one fp8 GEMM.
-
-    Per-output-channel weight scales make this numerically identical to the
-    separate GEMMs, while cutting GEMM count and quantizing the (shared) input
-    activation once instead of once per linear. ``apply()`` returns the
-    concatenated output; the caller splits it with ``split_sizes``.
-    """
-
-    def __init__(self, weight_names, bias_name=None, **kwargs):
-        super().__init__(weight_names[0], bias_name, **kwargs)
-        self.fused_weight_names = list(weight_names)
-        self.split_sizes = None
-
-    def load(self, weight_dict):
-        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]  # each (N_i, K)
-        self.split_sizes = [w.shape[0] for w in weights]
-        weight = torch.cat(weights, dim=0)  # (sum_N, K)
-        del weights
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.FP8_MAX  # (sum_N, 1)
-        self.weight = (weight / scale).clamp(-self.FP8_MAX, self.FP8_MAX).to(torch.float8_e4m3fn)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        for name in self.fused_weight_names:
-            weight_dict.pop(name, None)
-        self.bias = None
-
-
 @PLATFORM_MM_WEIGHT_REGISTER("int8-rocm")
 class MMWeightWint8channelAint8tokenRocm(MMWeightQuantTemplate):
     """W-int8-channel-sym / A-int8-token-dynamic via torch._int_mm.
@@ -131,6 +114,8 @@ def __init__(
     ):
         super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, lora_prefix, lora_path)
         self.weight_need_transpose = False
+        if not hasattr(torch, "_int_mm"):
+            raise RuntimeError("int8-rocm requires torch._int_mm, which is absent from this PyTorch/ROCm build.")
 
     def load(self, weight_dict):
         weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
@@ -152,32 +137,9 @@ def apply(self, input_tensor):
             out = torch.mm(input_tensor, weight.t())
         else:
             x_scale = (input_tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX).float()  # (M, 1)
-            x_q = (input_tensor / x_scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)  # (M, K)
-            out = torch._int_mm(x_q, self.weight.t())  # (M, N) int32
+            x_q = (input_tensor / x_scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8).contiguous()  # (M, K)
+            out = torch._int_mm(x_q, self.weight.t().contiguous())  # (M, N) int32
             out = (out.to(torch.float32) * x_scale * self.weight_scale.t()).to(input_tensor.dtype)  # dequant
         if self.bias is not None:
             out = out + self.bias
         return out
-
-
-class MMWeightFusedInt8Rocm(MMWeightWint8channelAint8tokenRocm):
-    """Fuse q/k/v, gate/up into one INT8 GEMM; apply() returns the concatenated
-    output, split by ``split_sizes``. See MMWeightFusedFp8Rocm."""
-
-    def __init__(self, weight_names, bias_name=None, **kwargs):
-        super().__init__(weight_names[0], bias_name, **kwargs)
-        self.fused_weight_names = list(weight_names)
-        self.split_sizes = None
-
-    def load(self, weight_dict):
-        weights = [weight_dict[name].to(AI_DEVICE).to(torch.float32) for name in self.fused_weight_names]
-        self.split_sizes = [w.shape[0] for w in weights]
-        weight = torch.cat(weights, dim=0)
-        del weights
-        scale = weight.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / self.INT8_MAX
-        self.weight = (weight / scale).round().clamp(-self.INT8_MAX, self.INT8_MAX).to(torch.int8)
-        self.weight_scale = scale.to(torch.float32)
-        del weight
-        for name in self.fused_weight_names:
-            weight_dict.pop(name, None)
-        self.bias = None
diff --git a/lightx2v_platform/test/test_amd_rocm_quant.py b/lightx2v_platform/test/test_amd_rocm_quant.py
new file mode 100644
index 000000000..f8b119cba
--- /dev/null
+++ b/lightx2v_platform/test/test_amd_rocm_quant.py
@@ -0,0 +1,71 @@
+"""AMD ROCm quantized GEMM correctness checks.
+
+Run inside the ROCm container:
+    PYTHONPATH=/path-to-LightX2V PLATFORM=amd_rocm python lightx2v_platform/test/test_amd_rocm_quant.py
+
+What it checks:
+- The fp8-rocm / int8-rocm scheme keys register (importing set_ai_device with
+  aiter and sageattention absent must still succeed).
+- Quantized GEMM output matches a BF16 matmul reference across several M,
+  including the int8 M<=16 BF16-fallback boundary (M = 1, 16, 17, 31, 64).
+
+fp8-rocm requires gfx1201 (RDNA4) and is skipped on other archs; int8-rocm runs
+on gfx1100 / gfx1201.
+"""
+
+# isort: off
+import lightx2v_platform.set_ai_device  # noqa: F401  must run first: initializes AI device + registers platform ops
+import torch
+
+from lightx2v_platform.registry_factory import PLATFORM_MM_WEIGHT_REGISTER
+
+# isort: on
+
+
+def _arch():
+    try:
+        return torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:
+        return ""
+
+
+def _rel_err(out, ref):
+    return ((out.float() - ref.float()).abs().mean() / ref.float().abs().mean().clamp_min(1e-6)).item()
+
+
+def _check_scheme(scheme, tol):
+    n, k = 512, 1024
+    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
+    linear = PLATFORM_MM_WEIGHT_REGISTER[scheme]("w.weight", None)
+    linear.load({"w.weight": weight.clone()})
+    ref_w = weight.float()
+    for m in (1, 16, 17, 31, 64):
+        x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+        ref = (x.float() @ ref_w.t()).to(torch.bfloat16)
+        out = linear.apply(x)
+        err = _rel_err(out, ref)
+        assert err < tol, f"{scheme} M={m}: rel_err {err:.4f} >= tol {tol}"
+        print(f"  {scheme:<10} M={m:>3}  rel_err={err:.4f}  OK")
+
+
+def main():
+    assert "fp8-rocm" in PLATFORM_MM_WEIGHT_REGISTER, "fp8-rocm not registered"
+    assert "int8-rocm" in PLATFORM_MM_WEIGHT_REGISTER, "int8-rocm not registered"
+    print("registration OK (no aiter / sageattention required)")
+
+    if not torch.cuda.is_available():
+        print("no GPU available; skipping numerical checks")
+        return
+
+    arch = _arch()
+    print(f"arch: {arch}")
+    if arch.startswith("gfx1201"):
+        _check_scheme("fp8-rocm", tol=0.08)
+    else:
+        print("  fp8-rocm skipped (requires gfx1201)")
+    _check_scheme("int8-rocm", tol=0.04)
+    print("ALL PASS")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index 36a662846..aa94c84b6 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -88,7 +88,7 @@ Both tasks were measured on the current code with 40 steps, seed 42, CFG disable
 
 The 5090 stack above (FlashAttention3, FlashInfer RoPE, CUTLASS FP8, SageAttention2 CUDA kernels) is CUDA-only. The configs below run Qwen-Image-2.1 on AMD RDNA GPUs with a torch-op inference path (`torch_sdpa`, `torch_real_rope`, torch LayerNorm/modulation) plus ROCm-friendly quantization, `torch.compile`, and SageAttention2. They quantize the **released BF16 weights on load** — no `tools/convert` step and no `dit_quantized_ckpt`.
 
-Validated in the `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` image (PyTorch 2.10, Triton 3.6). Install the runtime deps not baked into that image:
+Validated in the `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` image (PyTorch 2.10, Triton 3.6) with the `transformers` 4.57.x line (`Qwen3VLProcessor` needs a recent release; pin to a version you have tested rather than only a lower bound). Install the runtime deps not baked into that image:
 
 ```bash
 pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
@@ -113,10 +113,10 @@ Per-GPU config (pick one for `--config_json`, all under `configs/platforms/amd_r
 
 | GPU (arch) | Quantization | Configs |
 | --- | --- | --- |
-| R9700 (gfx1201 / RDNA4) | FP8 via `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json` (fastest), `..._fp8_compile.json`, `..._fp8_sage_25steps.json`, `qwen_image_21_r9700.json` (BF16 baseline) |
-| W7900 (gfx1100 / RDNA3) | INT8 via `torch._int_mm` (no FP8 on RDNA3) | `qwen_image_21_w7900_int8_compile_sage.json` (fastest), `..._int8_compile.json`, `..._int8_sage_25steps.json`, `qwen_image_21_w7900_int8.json` (baseline) |
+| R9700 (gfx1201 / RDNA4) | FP8 via `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json` (fastest), `..._fp8_compile.json`, `..._fp8_compile_sage_25steps.json`, `qwen_image_21_r9700.json` (BF16 eager baseline) |
+| W7900 (gfx1100 / RDNA3) | INT8 via `torch._int_mm` (no FP8 on RDNA3) | `qwen_image_21_w7900_int8_compile_sage.json` (fastest), `..._int8_compile.json`, `..._int8_compile_sage_25steps.json`, `qwen_image_21_w7900_int8.json` (INT8 eager baseline) |
 
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. On gfx1201 the VAE runs through an im2col GEMM (`vae_conv_im2col`) to avoid a nondeterministic MIOpen convolution defect; gfx1100 uses the native conv. The required SageAttention/Inductor Triton `num_stages` workaround for ROCm is applied automatically.
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when `attn_type: sage_attn2` is selected.
 
 | GPU | Quant | Steps | End-to-end |
 | --- | --- | ---: | ---: |
@@ -125,7 +125,7 @@ The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered
 | W7900 | INT8 | 40 | ~46 s |
 | W7900 | INT8 | 25 | ~29 s |
 
-Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup. 25 steps (the ComfyUI default for this model) has no visible quality loss versus 40. Compilation is a one-time warmup cost (~1.5–3 min); SageAttention needs `TORCHINDUCTOR_COMPILE_THREADS=1`, which the code sets automatically on the SageAttention path.
+Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup (median of a few runs). End-to-end covers text encoding, condition-KV prefill, denoising, VAE decode and PNG save; it excludes model load and one-time initialization. In the tested prompts and seeds, 25 steps (a common ComfyUI setting for this model) showed no obvious visual degradation versus 40 — this is not a systematic quality evaluation. Compilation is a one-time warmup cost (~1.5–3 min); `TORCHINDUCTOR_COMPILE_THREADS=1` is set automatically on the SageAttention path.
 
 ## 4. Service Deployment and API Usage
 
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 47ebef962..3296b3cca 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -88,7 +88,7 @@ RTX 5090 配置保留适用的通用优化，并使用：
 
 上面 5090 那套（FlashAttention3、FlashInfer RoPE、CUTLASS FP8、SageAttention2 的 CUDA kernel）都是 CUDA 专用。下面的配置让 Qwen-Image-2.1 跑在 AMD RDNA 显卡上，走 torch 算子推理路径（`torch_sdpa`、`torch_real_rope`、torch LayerNorm/modulation），配合 ROCm 可用的量化、`torch.compile` 和 SageAttention2。它们**在加载时直接量化官方 BF16 权重**——无需 `tools/convert` 转换、无需 `dit_quantized_ckpt`。
 
-验证环境为 `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` 镜像（PyTorch 2.10、Triton 3.6）。安装该镜像未自带的运行时依赖：
+验证环境为 `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` 镜像（PyTorch 2.10、Triton 3.6），transformers 使用 4.57.x（`Qwen3VLProcessor` 需要较新版本；请固定到你实测过的版本，而非只写下限）。安装该镜像未自带的运行时依赖：
 
 ```bash
 pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
@@ -113,10 +113,10 @@ PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
 
 | 显卡（架构） | 量化 | 配置 |
 | --- | --- | --- |
-| R9700（gfx1201 / RDNA4） | FP8，走 `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json`（最快）、`..._fp8_compile.json`、`..._fp8_sage_25steps.json`、`qwen_image_21_r9700.json`（BF16 基线） |
-| W7900（gfx1100 / RDNA3） | INT8，走 `torch._int_mm`（RDNA3 无 FP8） | `qwen_image_21_w7900_int8_compile_sage.json`（最快）、`..._int8_compile.json`、`..._int8_sage_25steps.json`、`qwen_image_21_w7900_int8.json`（基线） |
+| R9700（gfx1201 / RDNA4） | FP8，走 `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json`（最快）、`..._fp8_compile.json`、`..._fp8_compile_sage_25steps.json`、`qwen_image_21_r9700.json`（BF16 eager 基线） |
+| W7900（gfx1100 / RDNA3） | INT8，走 `torch._int_mm`（RDNA3 无 FP8） | `qwen_image_21_w7900_int8_compile_sage.json`（最快）、`..._int8_compile.json`、`..._int8_compile_sage_25steps.json`、`qwen_image_21_w7900_int8.json`（INT8 eager 基线） |
 
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。gfx1201 上 VAE 走 im2col GEMM（`vae_conv_im2col`）以规避 MIOpen 卷积的非确定性缺陷；gfx1100 用原生卷积。SageAttention/Inductor 在 ROCm 上所需的 Triton `num_stages` 规避会自动生效。
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在选择 `attn_type: sage_attn2` 时生效。
 
 | 显卡 | 量化 | 步数 | 端到端 |
 | --- | --- | ---: | ---: |
@@ -125,7 +125,7 @@ PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
 | W7900 | INT8 | 40 | ~46 s |
 | W7900 | INT8 | 25 | ~29 s |
 
-测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态。25 步（该模型在 ComfyUI 的默认步数）相比 40 步无可见画质损失。编译为一次性 warmup 开销（约 1.5–3 分钟）；SageAttention 需要 `TORCHINDUCTOR_COMPILE_THREADS=1`，代码在 SageAttention 路径上会自动设置。
+测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态（多次取中位数）。端到端包含文本编码、condition-KV prefill、去噪、VAE 解码与 PNG 保存，不含模型加载与一次性初始化。在所测试的 prompt/seed 上，25 步（该模型在 ComfyUI 的常用步数）相比 40 步未观察到明显画质退化——这不是系统性画质评估。编译为一次性 warmup 开销（约 1.5–3 分钟）；`TORCHINDUCTOR_COMPILE_THREADS=1` 在 SageAttention 路径上会自动设置。
 
 ## 4. 服务化部署与 API 调用
 

From 05ed342e9903392ce16a8fd5bd06210e3c3fe9e2 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 12:54:09 +0000
Subject: [PATCH 13/21] refactor(qwen-image-2.1): keep the ROCm sage workaround
 in the platform layer + add run scripts

- Move the SageAttention/Triton num_stages workaround entirely into the amd_rocm
  platform ops (applied when amd_rocm attention loads, arch-gated to gfx1201 /
  gfx1100). Removes the backend (torch.version.hip) branch from the shared
  SageAttn2Weight.__init__, keeping shared inference code backend-agnostic.
- Add scripts/platforms/amd_rocm/qwen_image_21_{r9700,w7900}_t2i.sh (export
  PLATFORM=amd_rocm + source scripts/base/base.sh), matching the other platforms;
  point the README at them instead of an inline PLATFORM= command.

Re-validated on R9700: patch applies at platform load, fp8+compile+sage 25 steps
~15.4s, output unchanged.
---
 lightx2v/common/ops/attn/sage_attn.py         |  8 -------
 .../ops/attn/amd_rocm/__init__.py             |  1 +
 .../ops/attn/amd_rocm/sage_attn.py            | 11 ++++++---
 .../amd_rocm/qwen_image_21_r9700_t2i.sh       | 23 +++++++++++++++++++
 .../amd_rocm/qwen_image_21_w7900_t2i.sh       | 23 +++++++++++++++++++
 scripts/qwen_image_21/README.md               | 17 +++++++-------
 scripts/qwen_image_21/README_zh.md            | 17 +++++++-------
 7 files changed, 71 insertions(+), 29 deletions(-)
 create mode 100755 scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
 create mode 100755 scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh

diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index ea94b191a..07945b9f5 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -75,14 +75,6 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
-        if getattr(torch.version, "hip", None) is not None:
-            # SageAttention's Triton kernels miscompile at num_stages>=3 on some
-            # RDNA archs; install the clamp before any compile. No-op off the
-            # validated archs. Scoped here so it only runs when SageAttention is
-            # actually selected, not for every ROCm workload.
-            from lightx2v_platform.ops.attn.amd_rocm.sage_attn import apply_rocm_sage_patches
-
-            apply_rocm_sage_patches()
 
     def apply(
         self,
diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index 9ecd99131..481a9cd1e 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1 +1,2 @@
+from . import sage_attn  # noqa: F401  installs the ROCm Triton num_stages workaround (arch-gated, no-op elsewhere)
 from .flash_attn import *
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index c65165a01..61fabc62f 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -122,11 +122,16 @@ def _is_validated_rocm_arch():
 
 
 def apply_rocm_sage_patches():
-    """Install the Triton num_stages workaround. No-op unless running on a
-    validated RDNA arch; call this only when SageAttention is actually selected
-    (from SageAttn2Weight.__init__), not for every ROCm workload."""
+    """Install the Triton num_stages workaround. Applied automatically when the
+    amd_rocm attention ops load (i.e. under PLATFORM=amd_rocm). No-op unless
+    running on a validated RDNA arch, where Triton's pipeliner miscompiles at
+    num_stages>=3 — so this stays entirely inside the platform layer rather than
+    branching on the backend in shared inference code."""
     if not _is_validated_rocm_arch():
         return
     clamp_sageattn_triton_num_stages()
     clamp_triton_compile_num_stages()
     force_rocm_inductor_single_thread()
+
+
+apply_rocm_sage_patches()
diff --git a/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh b/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
new file mode 100755
index 000000000..7135c4446
--- /dev/null
+++ b/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# System management interface: amd-smi
+
+# set path firstly
+lightx2v_path=
+model_path=
+
+export PLATFORM=amd_rocm
+export CUDA_VISIBLE_DEVICES=0
+
+# set environment variables
+source "${lightx2v_path}/scripts/base/base.sh"
+
+python -m lightx2v.infer \
+    --model_cls qwen_image_21 \
+    --task t2i \
+    --model_path "${model_path}" \
+    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json" \
+    --prompt "A capybara wearing a wizard hat, oil painting" \
+    --size 1024 1024 \
+    --seed 42 \
+    --save_result_path "${lightx2v_path}/save_results/qwen_image_21_r9700_t2i.png"
diff --git a/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh b/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
new file mode 100755
index 000000000..b4207436a
--- /dev/null
+++ b/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# System management interface: amd-smi
+
+# set path firstly
+lightx2v_path=
+model_path=
+
+export PLATFORM=amd_rocm
+export CUDA_VISIBLE_DEVICES=0
+
+# set environment variables
+source "${lightx2v_path}/scripts/base/base.sh"
+
+python -m lightx2v.infer \
+    --model_cls qwen_image_21 \
+    --task t2i \
+    --model_path "${model_path}" \
+    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json" \
+    --prompt "A capybara wearing a wizard hat, oil painting" \
+    --size 1024 1024 \
+    --seed 42 \
+    --save_result_path "${lightx2v_path}/save_results/qwen_image_21_w7900_t2i.png"
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index aa94c84b6..6c6ec7196 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -97,19 +97,18 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-Run with `PLATFORM=amd_rocm` (this activates the AMD ROCm platform backend under `lightx2v_platform/`, which registers the ROCm FP8/INT8 GEMM and applies the SageAttention Triton workaround). A single GPU is enough — the text encoder is CPU-offloaded. No `aiter` build is required for this torch-native path:
+Run through the platform scripts under `scripts/platforms/amd_rocm/`, which `export PLATFORM=amd_rocm` and source the common env. This activates the AMD ROCm platform backend under `lightx2v_platform/` (registers the ROCm FP8/INT8 GEMM and applies the SageAttention Triton workaround). A single GPU is enough — the text encoder is CPU-offloaded, and no `aiter` build is required.
+
+Set `lightx2v_path` and `model_path` at the top of the script, then run:
 
 ```bash
-PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
-    --model_cls qwen_image_21 --task t2i \
-    --model_path /path/to/Qwen-Image-2.1 \
-    --config_json configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json \
-    --prompt "A capybara wearing a wizard hat, oil painting" \
-    --size 1024 1024 --seed 42 \
-    --save_result_path ./save_results/qwen_image_21_t2i.png
+# R9700 (FP8)
+bash scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
+# W7900 (INT8)
+bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 ```
 
-Per-GPU config (pick one for `--config_json`, all under `configs/platforms/amd_rocm/`):
+Each script selects the fastest config for that GPU; edit `--config_json` in the script to pick another. Per-GPU config (all under `configs/platforms/amd_rocm/`):
 
 | GPU (arch) | Quantization | Configs |
 | --- | --- | --- |
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 3296b3cca..9adf7b815 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -97,19 +97,18 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-运行时设置 `PLATFORM=amd_rocm`（这会激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端：注册 ROCm 的 FP8/INT8 GEMM，并自动应用 SageAttention 的 Triton 规避）。单卡即可——条件编码器走 CPU offload。该 torch 原生路径**无需编译 aiter**：
+通过 `scripts/platforms/amd_rocm/` 下的平台脚本运行——脚本里会 `export PLATFORM=amd_rocm` 并 source 公共环境，从而激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端（注册 ROCm 的 FP8/INT8 GEMM，并自动应用 SageAttention 的 Triton 规避）。单卡即可——条件编码器走 CPU offload，且该路径**无需编译 aiter**。
+
+先在脚本顶部设置 `lightx2v_path` 和 `model_path`，然后运行：
 
 ```bash
-PLATFORM=amd_rocm CUDA_VISIBLE_DEVICES=0 python -m lightx2v.infer \
-    --model_cls qwen_image_21 --task t2i \
-    --model_path /path/to/Qwen-Image-2.1 \
-    --config_json configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json \
-    --prompt "A capybara wearing a wizard hat, oil painting" \
-    --size 1024 1024 --seed 42 \
-    --save_result_path ./save_results/qwen_image_21_t2i.png
+# R9700（FP8）
+bash scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
+# W7900（INT8）
+bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 ```
 
-各显卡对应配置（`--config_json` 选其一，均在 `configs/platforms/amd_rocm/` 下）：
+每个脚本默认选该卡最快的配置；改脚本里的 `--config_json` 可切换到其它配置。各配置（均在 `configs/platforms/amd_rocm/` 下）：
 
 | 显卡（架构） | 量化 | 配置 |
 | --- | --- | --- |

From 9d1c8bcd42645a85e8d53a692f29d12aa448aff2 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 14:14:41 +0000
Subject: [PATCH 14/21] =?UTF-8?q?refactor(qwen-image-2.1):=20address=202nd?=
 =?UTF-8?q?=20ROCm=20review=20=E2=80=94=20decouple=20sage/TE,=20slim=20con?=
 =?UTF-8?q?figs,=20real=20pytest?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- SageAttention Triton workaround now fires only when a sage_attn2 backend is
  actually constructed: shared SageAttn2Weight calls a uniform platform hook
  (get_platform_device().on_sage_attn2_init), which the amd_rocm device
  implements (arch-gated). No import-time side effect on non-sage ROCm runs, and
  no backend branch in shared inference code.
- Revert the multi-GPU text_encoder_device_index change in the Qwen runner
  (unused, backend-specific); drop the now-dead `import contextlib`.
- Revert the incidental linear_keys / mlp refactors so the Qwen transformer files
  differ from base only by the (necessary) torch.compile support.
- Reduce configs to 3 user-facing ones: qwen_image_21_r9700_fp8 / _w7900_int8 /
  _bf16; point the two run scripts + README at them.
- mm: _gcn_arch() uses current_device(); fp8 error no longer prescribes INT8;
  sage patch logs skipped modules at debug.
- Rewrite the quant test as real pytest (collectable, skips off-ROCm), exercising
  the production MM_WEIGHT_REGISTER with a fixed seed.

Validated: pytest passes on R9700 (3) and W7900 (2, fp8 skipped); t2i correct on
both (R9700 ~24s, W7900 ~47s) via the run scripts.
---
 ..._21_r9700.json => qwen_image_21_bf16.json} |  0
 ...sage.json => qwen_image_21_r9700_fp8.json} |  0
 .../qwen_image_21_r9700_fp8_compile.json      | 19 -----
 ...age_21_r9700_fp8_compile_sage_25steps.json | 19 -----
 .../amd_rocm/qwen_image_21_w7900_int8.json    |  7 +-
 .../qwen_image_21_w7900_int8_compile.json     | 19 -----
 ...qwen_image_21_w7900_int8_compile_sage.json | 19 -----
 ...ge_21_w7900_int8_compile_sage_25steps.json | 19 -----
 lightx2v/common/ops/attn/sage_attn.py         |  8 ++
 .../qwen_image_21/infer/transformer_infer.py  |  3 +-
 .../weights/transformer_weights.py            |  5 +-
 .../qwen_image_21/qwen_image_21_runner.py     | 29 ++-----
 lightx2v_platform/base/amd_rocm.py            |  8 ++
 lightx2v_platform/base/base.py                |  9 +++
 .../ops/attn/amd_rocm/__init__.py             |  1 -
 .../ops/attn/amd_rocm/sage_attn.py            | 16 ++--
 .../ops/mm/amd_rocm/mm_weight.py              |  4 +-
 lightx2v_platform/test/test_amd_rocm_quant.py | 78 ++++++++++---------
 .../amd_rocm/qwen_image_21_r9700_t2i.sh       |  2 +-
 .../amd_rocm/qwen_image_21_w7900_t2i.sh       |  2 +-
 scripts/qwen_image_21/README.md               | 11 +--
 scripts/qwen_image_21/README_zh.md            | 11 +--
 22 files changed, 103 insertions(+), 186 deletions(-)
 rename configs/platforms/amd_rocm/{qwen_image_21_r9700.json => qwen_image_21_bf16.json} (100%)
 rename configs/platforms/amd_rocm/{qwen_image_21_r9700_fp8_compile_sage.json => qwen_image_21_r9700_fp8.json} (100%)
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json
 delete mode 100644 configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json

diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700.json b/configs/platforms/amd_rocm/qwen_image_21_bf16.json
similarity index 100%
rename from configs/platforms/amd_rocm/qwen_image_21_r9700.json
rename to configs/platforms/amd_rocm/qwen_image_21_bf16.json
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json
similarity index 100%
rename from configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json
rename to configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
deleted file mode 100644
index 1b1be254e..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "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": true,
-    "text_encoder_cpu_offload": true,
-    "dit_quantized": true,
-    "dit_quant_scheme": "fp8-rocm",
-    "use_compile": true
-}
diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
deleted file mode 100644
index 142a6e8e6..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage_25steps.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "infer_steps": 25,
-    "resolution": 1024,
-    "vae_scale_factor": 16,
-    "attn_type": "sage_attn2",
-    "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
-}
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
index 52998196a..459eb2eca 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
@@ -2,7 +2,7 @@
     "infer_steps": 40,
     "resolution": 1024,
     "vae_scale_factor": 16,
-    "attn_type": "torch_sdpa",
+    "attn_type": "sage_attn2",
     "rope_type": "torch_real_rope",
     "rms_norm_type": "fp32_variance",
     "layer_norm_type": "torch",
@@ -11,8 +11,9 @@
     "fused_block_ops": false,
     "enable_cfg": false,
     "sample_guide_scale": 1.0,
-    "warmup": false,
+    "warmup": true,
     "text_encoder_cpu_offload": true,
     "dit_quantized": true,
-    "dit_quant_scheme": "int8-rocm"
+    "dit_quant_scheme": "int8-rocm",
+    "use_compile": true
 }
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json
deleted file mode 100644
index e3fefa9dc..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "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": true,
-    "text_encoder_cpu_offload": true,
-    "dit_quantized": true,
-    "dit_quant_scheme": "int8-rocm",
-    "use_compile": true
-}
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json
deleted file mode 100644
index 459eb2eca..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "infer_steps": 40,
-    "resolution": 1024,
-    "vae_scale_factor": 16,
-    "attn_type": "sage_attn2",
-    "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
-}
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json
deleted file mode 100644
index 1775e1f70..000000000
--- a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage_25steps.json
+++ /dev/null
@@ -1,19 +0,0 @@
-{
-    "infer_steps": 25,
-    "resolution": 1024,
-    "vae_scale_factor": 16,
-    "attn_type": "sage_attn2",
-    "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
-}
diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index 07945b9f5..64108beea 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -75,6 +75,14 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
+        # Let the active platform run any SageAttention-specific setup (e.g. the
+        # ROCm Triton num_stages workaround). Uniform hook — no backend branch.
+        from lightx2v_platform.base.base import get_platform_device
+
+        device = get_platform_device()
+        hook = getattr(device, "on_sage_attn2_init", None)
+        if hook is not None:
+            hook()
 
     def apply(
         self,
diff --git a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
index 940d4c083..84babfc37 100644
--- a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
+++ b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
@@ -56,8 +56,7 @@ def _finish_block(self, block, x, attention, modulation):
 
         x = x + gate1 * block.out.apply(attention)
         h = block.norm2.apply(x) * scale2
-        mlp = F.silu(block.gate.apply(h)) * block.up.apply(h)
-        x = x + gate2 * block.down.apply(mlp)
+        x = x + gate2 * block.down.apply(F.silu(block.gate.apply(h)) * block.up.apply(h))
         return x.clamp(-65504, 65504) if x.dtype == torch.float16 else x
 
     def prefill(self, weights, state, cache):
diff --git a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
index 37e629251..26c67d14c 100644
--- a/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
+++ b/lightx2v/models/networks/qwen_image_21/weights/transformer_weights.py
@@ -8,7 +8,7 @@ def __init__(self, index, config):
         super().__init__()
         prefix = f"transformer_blocks.{index}"
         mm_type = config["dit_quant_scheme"] if config.get("dit_quantized", False) else "Default"
-        linear_keys = {
+        for name, key in {
             "q": "attn.to_q",
             "k": "attn.to_k",
             "v": "attn.to_v",
@@ -16,8 +16,7 @@ def __init__(self, index, config):
             "up": "img_mlp.proj",
             "gate": "img_mlp.gate_layer",
             "down": "img_mlp.out",
-        }
-        for name, key in linear_keys.items():
+        }.items():
             linear = MM_WEIGHT_REGISTER[mm_type](f"{prefix}.{key}.weight", bias_name=None)
             if mm_type == "fp8-f16-accum":
                 linear.enable_fp8_f16_accum(config.get("dit_fp8_activation_qmax", ACTIVATION_QMAX))
diff --git a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
index f1d12c9c4..3d81b7e82 100644
--- a/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
+++ b/lightx2v/models/runners/qwen_image_21/qwen_image_21_runner.py
@@ -1,4 +1,3 @@
-import contextlib
 import math
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -185,32 +184,18 @@ def run_text_encoder(self, images, shapes):
             branches.append(("uncond", info.negative_prompt or ""))
         encoder = self.text_encoders[0]
         offload = self.config.get("text_encoder_cpu_offload", False)
-        # ROCm/R9700: empty_cache() does not physically return freed VRAM to the
-        # driver, so the ~15G text-encoder onload stays resident on the main GPU
-        # and starves the VAE decoder. Run it on a separate GPU when configured.
-        te_index = self.config.get("text_encoder_device_index", None)
-        dev_ctx = torch_device_module.device(te_index) if te_index is not None else contextlib.nullcontext()
-        # Resolve an explicit indexed device: a bare "cuda" would re-resolve to the
-        # current device (the text-encoder GPU) inside the context above.
-        main_device = torch.device(AI_DEVICE, torch_device_module.current_device())
         try:
-            with dev_ctx:
-                if offload:
-                    with ProfilingContext4DebugL1("Onload Text Encoder"):
-                        encoder.to_cuda()
-                for name, prompt in branches:
-                    branch = encoder.infer(prompt, images)
-                    # Pull encoder outputs onto the main GPU before leaving the
-                    # text-encoder device context.
-                    outputs[name] = {k: (v.to(main_device) if torch.is_tensor(v) else v) for k, v in branch.items()}
+            if offload:
+                with ProfilingContext4DebugL1("Onload Text Encoder"):
+                    encoder.to_cuda()
+            for name, prompt in branches:
+                branch = encoder.infer(prompt, images)
+                branch["layout"] = build_token_layout(branch["image_mask"], shapes, self.config["axes_dims_rope"], rope=self.model.transformer_infer.rope)
+                outputs[name] = branch
         finally:
             if offload:
                 with ProfilingContext4DebugL1("Offload Text Encoder"):
                     encoder.to_cpu()
-        # Build the DiT rope layout on the main GPU (outside the TE device context,
-        # otherwise freshly created rope tensors would land on the TE device).
-        for branch in outputs.values():
-            branch["layout"] = build_token_layout(branch["image_mask"], shapes, self.config["axes_dims_rope"], rope=self.model.transformer_infer.rope)
         return outputs
 
     @ProfilingContext4DebugL1("Run VAE Encoder")
diff --git a/lightx2v_platform/base/amd_rocm.py b/lightx2v_platform/base/amd_rocm.py
index 94f154522..089259a86 100644
--- a/lightx2v_platform/base/amd_rocm.py
+++ b/lightx2v_platform/base/amd_rocm.py
@@ -137,6 +137,14 @@ def init_device_env():
                     setattr(mod, "sgl_kernel", sgl_kernel)
             logger.info("  - aiter sgl_kernel compatibility layer enabled (RMSNorm, GEMM)")
 
+    @staticmethod
+    def on_sage_attn2_init():
+        """Install the ROCm Triton num_stages workaround when a SageAttention2
+        backend is constructed. Arch-gated (no-op off gfx1201 / gfx1100)."""
+        from lightx2v_platform.ops.attn.amd_rocm.sage_attn import apply_rocm_sage_patches
+
+        apply_rocm_sage_patches()
+
     @staticmethod
     def is_available() -> bool:
         """Check if AMD ROCm is available."""
diff --git a/lightx2v_platform/base/base.py b/lightx2v_platform/base/base.py
index 72c133a34..e7c949eac 100644
--- a/lightx2v_platform/base/base.py
+++ b/lightx2v_platform/base/base.py
@@ -6,6 +6,15 @@
 from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER
 
 
+def get_platform_device():
+    """Return the active platform device class (or None before init).
+
+    Lets shared code query the current platform for optional hooks without
+    branching on a specific backend.
+    """
+    return PLATFORM_DEVICE_REGISTER.get(global_var.PLATFORM)
+
+
 def init_ai_device(platform="cuda"):
     platform_device = PLATFORM_DEVICE_REGISTER.get(platform, None)
     if platform_device is None:
diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index 481a9cd1e..9ecd99131 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1,2 +1 @@
-from . import sage_attn  # noqa: F401  installs the ROCm Triton num_stages workaround (arch-gated, no-op elsewhere)
 from .flash_attn import *
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index 61fabc62f..8c09feec9 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -32,7 +32,8 @@ def clamp_sageattn_triton_num_stages(max_stages=2):
     for module_info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
         try:
             module = importlib.import_module(module_info.name)
-        except Exception:
+        except Exception as e:
+            logger.debug(f"[SageAttention] ROCm: skipped {module_info.name} while installing num_stages clamp: {e}")
             continue
         for value in vars(module).values():
             if not isinstance(value, JITFunction) or getattr(value, "_rocm_num_stages_clamped", False):
@@ -122,16 +123,13 @@ def _is_validated_rocm_arch():
 
 
 def apply_rocm_sage_patches():
-    """Install the Triton num_stages workaround. Applied automatically when the
-    amd_rocm attention ops load (i.e. under PLATFORM=amd_rocm). No-op unless
-    running on a validated RDNA arch, where Triton's pipeliner miscompiles at
-    num_stages>=3 — so this stays entirely inside the platform layer rather than
-    branching on the backend in shared inference code."""
+    """Install the Triton num_stages workaround. Called only when a SageAttention2
+    backend is actually constructed (via the amd_rocm platform's
+    on_sage_attn2_init hook), so non-SageAttention ROCm workloads are untouched.
+    No-op unless running on a validated RDNA arch, where Triton's pipeliner
+    miscompiles at num_stages>=3."""
     if not _is_validated_rocm_arch():
         return
     clamp_sageattn_triton_num_stages()
     clamp_triton_compile_num_stages()
     force_rocm_inductor_single_thread()
-
-
-apply_rocm_sage_patches()
diff --git a/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
index 99961b9bc..99afeab7c 100644
--- a/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
+++ b/lightx2v_platform/ops/mm/amd_rocm/mm_weight.py
@@ -15,7 +15,7 @@
 
 def _gcn_arch():
     try:
-        return torch.cuda.get_device_properties(0).gcnArchName
+        return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName
     except Exception:
         return ""
 
@@ -60,7 +60,7 @@ def __init__(
             raise RuntimeError("fp8-rocm requires torch._scaled_mm, which is absent from this PyTorch/ROCm build.")
         arch = _gcn_arch()
         if not arch.startswith("gfx1201"):
-            raise RuntimeError(f"fp8-rocm is validated only on gfx1201 (RDNA4); detected '{arch}'. Use int8-rocm instead.")
+            raise RuntimeError(f"fp8-rocm (torch._scaled_mm) is validated only on gfx1201 (RDNA4); detected '{arch}'.")
 
     def load(self, weight_dict):
         weight = weight_dict[self.weight_name].to(AI_DEVICE).to(torch.float32)  # (N, K)
diff --git a/lightx2v_platform/test/test_amd_rocm_quant.py b/lightx2v_platform/test/test_amd_rocm_quant.py
index f8b119cba..f53321842 100644
--- a/lightx2v_platform/test/test_amd_rocm_quant.py
+++ b/lightx2v_platform/test/test_amd_rocm_quant.py
@@ -1,71 +1,75 @@
 """AMD ROCm quantized GEMM correctness checks.
 
-Run inside the ROCm container:
-    PYTHONPATH=/path-to-LightX2V PLATFORM=amd_rocm python lightx2v_platform/test/test_amd_rocm_quant.py
+Run inside the ROCm container (a ROCm GPU and PLATFORM=amd_rocm are required;
+otherwise every test is skipped, so this stays safe to collect on CPU CI):
 
-What it checks:
-- The fp8-rocm / int8-rocm scheme keys register (importing set_ai_device with
-  aiter and sageattention absent must still succeed).
-- Quantized GEMM output matches a BF16 matmul reference across several M,
-  including the int8 M<=16 BF16-fallback boundary (M = 1, 16, 17, 31, 64).
+    PYTHONPATH=/path-to-LightX2V PLATFORM=amd_rocm python -m pytest \\
+        lightx2v_platform/test/test_amd_rocm_quant.py -v
 
-fp8-rocm requires gfx1201 (RDNA4) and is skipped on other archs; int8-rocm runs
-on gfx1100 / gfx1201.
+It checks that fp8-rocm / int8-rocm register into the production
+``MM_WEIGHT_REGISTER`` (the platform ops must merge into the same registry
+inference looks them up in), and that their GEMM matches a BF16 reference across
+several M — including the int8 M<=16 BF16-fallback boundary. These torch-native
+ops do not import aiter or sageattention, so registration works without them.
 """
 
-# isort: off
-import lightx2v_platform.set_ai_device  # noqa: F401  must run first: initializes AI device + registers platform ops
+import os
+
+import pytest
 import torch
 
-from lightx2v_platform.registry_factory import PLATFORM_MM_WEIGHT_REGISTER
+_ON_ROCM = getattr(torch.version, "hip", None) is not None and os.getenv("PLATFORM") == "amd_rocm" and torch.cuda.is_available()
 
-# isort: on
+pytestmark = pytest.mark.skipif(not _ON_ROCM, reason="requires PLATFORM=amd_rocm on a ROCm GPU")
 
 
-def _arch():
+def _gcn_arch():
     try:
-        return torch.cuda.get_device_properties(0).gcnArchName
+        return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName
     except Exception:
         return ""
 
 
+def _mm_register():
+    # Import via the lightx2v package (lazily, so collecting this module on a
+    # non-ROCm host never fails): lightx2v/__init__ runs set_ai_device before
+    # registry_factory takes its one-time platform snapshot — the same order
+    # production uses, so the platform schemes are merged into MM_WEIGHT_REGISTER.
+    from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER
+
+    return MM_WEIGHT_REGISTER
+
+
 def _rel_err(out, ref):
     return ((out.float() - ref.float()).abs().mean() / ref.float().abs().mean().clamp_min(1e-6)).item()
 
 
-def _check_scheme(scheme, tol):
+def _check(scheme, tol):
+    torch.manual_seed(42)
+    register = _mm_register()
+    assert scheme in register, f"{scheme} not merged into the production MM_WEIGHT_REGISTER"
     n, k = 512, 1024
     weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    linear = PLATFORM_MM_WEIGHT_REGISTER[scheme]("w.weight", None)
+    linear = register[scheme]("w.weight", None)
     linear.load({"w.weight": weight.clone()})
     ref_w = weight.float()
     for m in (1, 16, 17, 31, 64):
         x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
         ref = (x.float() @ ref_w.t()).to(torch.bfloat16)
-        out = linear.apply(x)
-        err = _rel_err(out, ref)
+        err = _rel_err(linear.apply(x), ref)
         assert err < tol, f"{scheme} M={m}: rel_err {err:.4f} >= tol {tol}"
-        print(f"  {scheme:<10} M={m:>3}  rel_err={err:.4f}  OK")
 
 
-def main():
-    assert "fp8-rocm" in PLATFORM_MM_WEIGHT_REGISTER, "fp8-rocm not registered"
-    assert "int8-rocm" in PLATFORM_MM_WEIGHT_REGISTER, "int8-rocm not registered"
-    print("registration OK (no aiter / sageattention required)")
+def test_schemes_registered():
+    register = _mm_register()
+    assert "fp8-rocm" in register
+    assert "int8-rocm" in register
 
-    if not torch.cuda.is_available():
-        print("no GPU available; skipping numerical checks")
-        return
 
-    arch = _arch()
-    print(f"arch: {arch}")
-    if arch.startswith("gfx1201"):
-        _check_scheme("fp8-rocm", tol=0.08)
-    else:
-        print("  fp8-rocm skipped (requires gfx1201)")
-    _check_scheme("int8-rocm", tol=0.04)
-    print("ALL PASS")
+def test_int8_rocm():
+    _check("int8-rocm", tol=0.04)
 
 
-if __name__ == "__main__":
-    main()
+@pytest.mark.skipif(not _gcn_arch().startswith("gfx1201"), reason="fp8-rocm requires gfx1201 (RDNA4)")
+def test_fp8_rocm():
+    _check("fp8-rocm", tol=0.08)
diff --git a/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh b/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
index 7135c4446..8bb92e667 100755
--- a/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
+++ b/scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
@@ -16,7 +16,7 @@ python -m lightx2v.infer \
     --model_cls qwen_image_21 \
     --task t2i \
     --model_path "${model_path}" \
-    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8_compile_sage.json" \
+    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json" \
     --prompt "A capybara wearing a wizard hat, oil painting" \
     --size 1024 1024 \
     --seed 42 \
diff --git a/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh b/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
index b4207436a..6d3662816 100755
--- a/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
+++ b/scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
@@ -16,7 +16,7 @@ python -m lightx2v.infer \
     --model_cls qwen_image_21 \
     --task t2i \
     --model_path "${model_path}" \
-    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_w7900_int8_compile_sage.json" \
+    --config_json "${lightx2v_path}/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json" \
     --prompt "A capybara wearing a wizard hat, oil painting" \
     --size 1024 1024 \
     --seed 42 \
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index 6c6ec7196..ee3269ec2 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -110,12 +110,13 @@ bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 
 Each script selects the fastest config for that GPU; edit `--config_json` in the script to pick another. Per-GPU config (all under `configs/platforms/amd_rocm/`):
 
-| GPU (arch) | Quantization | Configs |
+| GPU (arch) | Config | Contents |
 | --- | --- | --- |
-| R9700 (gfx1201 / RDNA4) | FP8 via `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json` (fastest), `..._fp8_compile.json`, `..._fp8_compile_sage_25steps.json`, `qwen_image_21_r9700.json` (BF16 eager baseline) |
-| W7900 (gfx1100 / RDNA3) | INT8 via `torch._int_mm` (no FP8 on RDNA3) | `qwen_image_21_w7900_int8_compile_sage.json` (fastest), `..._int8_compile.json`, `..._int8_compile_sage_25steps.json`, `qwen_image_21_w7900_int8.json` (INT8 eager baseline) |
+| R9700 (gfx1201 / RDNA4) | `qwen_image_21_r9700_fp8.json` | FP8 (`torch._scaled_mm`) + `torch.compile` + SageAttention2 |
+| W7900 (gfx1100 / RDNA3) | `qwen_image_21_w7900_int8.json` | INT8 (`torch._int_mm`) + `torch.compile` + SageAttention2 |
+| either | `qwen_image_21_bf16.json` | BF16 eager baseline |
 
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. Both use `use_compile` and `attn_type: sage_attn2`. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when `attn_type: sage_attn2` is selected.
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when a SageAttention2 backend is actually constructed. To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
 
 | GPU | Quant | Steps | End-to-end |
 | --- | --- | ---: | ---: |
@@ -124,7 +125,7 @@ The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered
 | W7900 | INT8 | 40 | ~46 s |
 | W7900 | INT8 | 25 | ~29 s |
 
-Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup (median of a few runs). End-to-end covers text encoding, condition-KV prefill, denoising, VAE decode and PNG save; it excludes model load and one-time initialization. In the tested prompts and seeds, 25 steps (a common ComfyUI setting for this model) showed no obvious visual degradation versus 40 — this is not a systematic quality evaluation. Compilation is a one-time warmup cost (~1.5–3 min); `TORCHINDUCTOR_COMPILE_THREADS=1` is set automatically on the SageAttention path.
+Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup (median of 3 runs). End-to-end covers text encoding, condition-KV prefill, denoising, VAE decode and PNG save; it excludes model load and one-time initialization. In the tested prompts and seeds, 25 steps showed no obvious visual degradation versus 40 — this is not a systematic quality evaluation. Compilation is a one-time warmup cost (~1.5–3 min); `TORCHINDUCTOR_COMPILE_THREADS=1` is set automatically on the SageAttention path.
 
 ## 4. Service Deployment and API Usage
 
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 9adf7b815..607ccb19b 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -110,12 +110,13 @@ bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 
 每个脚本默认选该卡最快的配置；改脚本里的 `--config_json` 可切换到其它配置。各配置（均在 `configs/platforms/amd_rocm/` 下）：
 
-| 显卡（架构） | 量化 | 配置 |
+| 显卡（架构） | 配置 | 内容 |
 | --- | --- | --- |
-| R9700（gfx1201 / RDNA4） | FP8，走 `torch._scaled_mm` | `qwen_image_21_r9700_fp8_compile_sage.json`（最快）、`..._fp8_compile.json`、`..._fp8_compile_sage_25steps.json`、`qwen_image_21_r9700.json`（BF16 eager 基线） |
-| W7900（gfx1100 / RDNA3） | INT8，走 `torch._int_mm`（RDNA3 无 FP8） | `qwen_image_21_w7900_int8_compile_sage.json`（最快）、`..._int8_compile.json`、`..._int8_compile_sage_25steps.json`、`qwen_image_21_w7900_int8.json`（INT8 eager 基线） |
+| R9700（gfx1201 / RDNA4） | `qwen_image_21_r9700_fp8.json` | FP8（`torch._scaled_mm`）+ `torch.compile` + SageAttention2 |
+| W7900（gfx1100 / RDNA3） | `qwen_image_21_w7900_int8.json` | INT8（`torch._int_mm`）+ `torch.compile` + SageAttention2 |
+| 通用 | `qwen_image_21_bf16.json` | BF16 eager 基线 |
 
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。两者都开启 `use_compile` 和 `attn_type: sage_attn2`。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在选择 `attn_type: sage_attn2` 时生效。
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在真正构造 SageAttention2 backend 时才生效。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
 
 | 显卡 | 量化 | 步数 | 端到端 |
 | --- | --- | ---: | ---: |
@@ -124,7 +125,7 @@ bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 | W7900 | INT8 | 40 | ~46 s |
 | W7900 | INT8 | 25 | ~29 s |
 
-测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态（多次取中位数）。端到端包含文本编码、condition-KV prefill、去噪、VAE 解码与 PNG 保存，不含模型加载与一次性初始化。在所测试的 prompt/seed 上，25 步（该模型在 ComfyUI 的常用步数）相比 40 步未观察到明显画质退化——这不是系统性画质评估。编译为一次性 warmup 开销（约 1.5–3 分钟）；`TORCHINDUCTOR_COMPILE_THREADS=1` 在 SageAttention 路径上会自动设置。
+测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态（3 次取中位数）。端到端包含文本编码、condition-KV prefill、去噪、VAE 解码与 PNG 保存，不含模型加载与一次性初始化。在所测试的 prompt/seed 上，25 步相比 40 步未观察到明显画质退化——这不是系统性画质评估。编译为一次性 warmup 开销（约 1.5–3 分钟）；`TORCHINDUCTOR_COMPILE_THREADS=1` 在 SageAttention 路径上会自动设置。
 
 ## 4. 服务化部署与 API 调用
 

From 7502bd19c36f7980acfdc44e2088e306c3c7d961 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 14:37:59 +0000
Subject: [PATCH 15/21] refactor(qwen-image-2.1): pass modulation as one tuple
 to infer_block (10 args -> 7)

The four modulation tensors were split into separate params for the (now
removed) fusion path; the base code already passed them as one tuple. Regroup
them so the compiled block takes modulation/rotary/positions/k_cache/v_cache.
Re-validated on R9700 (fp8+compile+sage, ~24s).
---
 .../networks/qwen_image_21/infer/transformer_infer.py  | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
index 84babfc37..ef619afa6 100644
--- a/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
+++ b/lightx2v/models/networks/qwen_image_21/infer/transformer_infer.py
@@ -39,12 +39,12 @@ def _qkv(self, block, x, scale, rotary, rotary_positions):
         q, k = self.rope.apply(q, k, rotary, positions=rotary_positions)
         return q, k, v
 
-    def infer_block(self, block, x, scale1, gate1, scale2, gate2, rotary, rotary_positions, k_cache, v_cache):
-        q, k, v = self._qkv(block, x, scale1, rotary, rotary_positions)
+    def infer_block(self, block, x, modulation, rotary, rotary_positions, k_cache, v_cache):
+        q, k, v = self._qkv(block, x, modulation[0], rotary, rotary_positions)
         k = torch.cat((k_cache, k))
         v = torch.cat((v_cache, v))
         attention = block.attention.apply(q, k, v)
-        return self._finish_block(block, x, attention, (scale1, gate1, scale2, gate2))
+        return self._finish_block(block, x, attention, modulation)
 
     def _finish_block(self, block, x, attention, modulation):
         _, gate1, scale2, gate2 = modulation
@@ -81,7 +81,7 @@ def infer(self, weights, state, cache):
             raise RuntimeError("Condition KV must be prefilled before denoising")
         x = state.hidden_states
         # Shared across every block; kept outside the compiled block graph.
-        scale1, gate1, scale2, gate2 = self._modulation(state)
+        modulation = self._modulation(state)
         for index, block in enumerate(weights.blocks):
-            x = self.run_block(index, block, x, scale1, gate1, scale2, gate2, state.rotary, state.rotary_positions, cache.k_cache(index), cache.v_cache(index))
+            x = self.run_block(index, block, x, modulation, state.rotary, state.rotary_positions, cache.k_cache(index), cache.v_cache(index))
         return x

From 8518d2c23535546db25e739f62c50e3a0165e141 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 14:45:15 +0000
Subject: [PATCH 16/21] refactor: slim down amd_rocm sage_attn.py (drop
 redundant hip guards, tighten helpers)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Functions are private and only run after apply_rocm_sage_patches() has already
checked the arch, so the per-function hip guards were redundant; use min() /
startswith(tuple) / a _MAX_STAGES constant and trim docstrings. Behavior
unchanged — re-validated on R9700 (num_stages clamp on 15 kernels, ~24s).
---
 .../ops/attn/amd_rocm/sage_attn.py            | 135 +++++++-----------
 1 file changed, 51 insertions(+), 84 deletions(-)

diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index 8c09feec9..a0f8b3b61 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -1,24 +1,32 @@
-"""ROCm-only SageAttention / Triton num_stages workarounds.
+"""ROCm-only SageAttention / Triton num_stages workaround.
 
-On gfx1201 (ROCm/Triton), Triton's AMD software-pipeliner pass miscompiles
-kernels at num_stages>=3 ("operation destroyed but still has uses"). These
-patches clamp num_stages<=2 on both SageAttention's own Triton kernels and any
-torch.compile / Inductor Triton kernels, and force Inductor to compile
-in-process so the clamp actually reaches those kernels.
+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.
 
-Imported (and applied) only on the amd_rocm platform via
-``lightx2v_platform/ops/attn/amd_rocm/__init__.py``. Each function is a no-op
-off ROCm, so importing it elsewhere is harmless.
+Applied via the amd_rocm platform's ``on_sage_attn2_init`` hook: only when a
+SageAttention2 backend is built, and only on the validated archs.
 """
 
 import torch
 from loguru import logger
 
+_MAX_STAGES = 2
+_VALIDATED_ARCHS = ("gfx1201", "gfx1100")
 
-def clamp_sageattn_triton_num_stages(max_stages=2):
-    """Clamp launch-time num_stages on every SageAttention Triton kernel."""
-    if getattr(torch.version, "hip", None) is None:
-        return
+
+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
 
@@ -27,109 +35,68 @@ def clamp_sageattn_triton_num_stages(max_stages=2):
         from triton.runtime.jit import JITFunction
     except ImportError:
         return
-
     patched = 0
-    for module_info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
+    for info in pkgutil.iter_modules(sageattention.__path__, sageattention.__name__ + "."):
         try:
-            module = importlib.import_module(module_info.name)
+            module = importlib.import_module(info.name)
         except Exception as e:
-            logger.debug(f"[SageAttention] ROCm: skipped {module_info.name} while installing num_stages clamp: {e}")
+            logger.debug(f"[SageAttention] ROCm: skipped {info.name}: {e}")
             continue
-        for value in vars(module).values():
-            if not isinstance(value, JITFunction) or getattr(value, "_rocm_num_stages_clamped", False):
+        for fn in vars(module).values():
+            if not isinstance(fn, JITFunction) or getattr(fn, "_rocm_clamped", False):
                 continue
-            original_run = value.run
 
-            def run(*args, _original_run=original_run, **kwargs):
-                stages = kwargs.get("num_stages")
-                if isinstance(stages, int) and stages > max_stages:
-                    kwargs["num_stages"] = max_stages
-                return _original_run(*args, **kwargs)
+            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)
 
-            try:
-                value.run = run
-                value._rocm_num_stages_clamped = True
-                patched += 1
-            except Exception:
-                continue
+            fn.run, fn._rocm_clamped = run, True
+            patched += 1
     if patched:
-        logger.info(f"[SageAttention] ROCm: clamped Triton num_stages<={max_stages} on {patched} kernels")
+        logger.info(f"[SageAttention] ROCm: clamped num_stages<={_MAX_STAGES} on {patched} kernels")
 
 
-def clamp_triton_compile_num_stages(max_stages=2):
-    """Clamp num_stages for every Triton compilation (incl. torch.compile /
-    Inductor kernels) on ROCm, where num_stages>=3 crashes the AMD pipeliner."""
-    if getattr(torch.version, "hip", None) is None:
-        return
+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_num_stages_clamped", False):
+    if getattr(triton.compile, "_rocm_clamped", False):
         return
-    _original_compile = triton.compile
+    _original = triton.compile
 
     def compile(*args, **kwargs):
-        options = kwargs.get("options")
-        if isinstance(options, dict) and isinstance(options.get("num_stages"), int) and options["num_stages"] > max_stages:
-            options["num_stages"] = max_stages
-        elif options is not None and isinstance(getattr(options, "num_stages", None), int) and options.num_stages > max_stages:
+        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:
-                options.num_stages = max_stages
+                opts.num_stages = min(opts.num_stages, _MAX_STAGES)
             except Exception:
                 pass
-        return _original_compile(*args, **kwargs)
+        return _original(*args, **kwargs)
 
-    compile._rocm_num_stages_clamped = True
+    compile._rocm_clamped = True
     triton.compile = compile
 
 
-def force_rocm_inductor_single_thread():
-    """Make torch.compile / Inductor compile in-process on ROCm.
-
-    The num_stages clamp on ``triton.compile`` only takes effect in the process
-    that installed it. Inductor's default multi-threaded pool compiles kernels
-    in worker subprocesses that never see the clamp, so an Inductor-generated
-    kernel can still hit the num_stages>=3 pipeliner crash. Compiling in-process
-    keeps the clamp effective (one-time, warmup-only single-threaded cost)."""
-    if getattr(torch.version, "hip", None) is None:
-        return
+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: set inductor compile_threads=1 so the Triton num_stages clamp reaches torch.compile kernels")
-
-
-_VALIDATED_ARCHS = ("gfx1201", "gfx1100")
-
-
-def _is_validated_rocm_arch():
-    """True only on the RDNA arch(es) where this workaround is validated.
-
-    The num_stages>=3 pipeliner crash was reproduced and the clamp validated on
-    gfx1201 (R9700) and gfx1100 (W7900). Other archs (e.g. MI300) are left
-    untouched so their kernels keep their own num_stages tuning.
-    """
-    if getattr(torch.version, "hip", None) is None or not torch.cuda.is_available():
-        return False
-    try:
-        arch = torch.cuda.get_device_properties(0).gcnArchName
-    except Exception:
-        return False
-    return any(arch.startswith(a) for a in _VALIDATED_ARCHS)
+        logger.info("[SageAttention] ROCm: inductor compile_threads=1 so the num_stages clamp reaches compiled kernels")
 
 
 def apply_rocm_sage_patches():
-    """Install the Triton num_stages workaround. Called only when a SageAttention2
-    backend is actually constructed (via the amd_rocm platform's
-    on_sage_attn2_init hook), so non-SageAttention ROCm workloads are untouched.
-    No-op unless running on a validated RDNA arch, where Triton's pipeliner
-    miscompiles at num_stages>=3."""
+    """Install the workaround. No-op off the validated RDNA archs."""
     if not _is_validated_rocm_arch():
         return
-    clamp_sageattn_triton_num_stages()
-    clamp_triton_compile_num_stages()
-    force_rocm_inductor_single_thread()
+    _clamp_sageattn_kernels()
+    _clamp_triton_compile()
+    _force_inductor_single_thread()

From b3277a40195481e829b4cbc13437eb76f787a385 Mon Sep 17 00:00:00 2001
From: zhangnju <zhangnju@126.com>
Date: Tue, 22 Sep 2026 15:22:48 +0000
Subject: [PATCH 17/21] feat(amd_rocm): make the sage num_stages workaround
 toggleable via env var

LIGHTX2V_ROCM_TRITON_MAX_STAGES (default 2) is the clamp ceiling; <=0 disables
the whole workaround with no code change (e.g. once Triton fixes the pipeliner
bug). Documented in the README. Env var (not JSON config) because the patch is a
process-level, one-time monkeypatch applied before any per-request config.
Re-validated on R9700 (default path still clamps 15 kernels, ~24s).
---
 lightx2v_platform/ops/attn/amd_rocm/sage_attn.py | 15 +++++++++++++--
 scripts/qwen_image_21/README.md                  |  2 +-
 scripts/qwen_image_21/README_zh.md               |  2 +-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index a0f8b3b61..0750661d7 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -7,12 +7,19 @@
 
 Applied via the amd_rocm platform's ``on_sage_attn2_init`` hook: only when a
 SageAttention2 backend is built, 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
 
-_MAX_STAGES = 2
+_MAX_STAGES = int(os.getenv("LIGHTX2V_ROCM_TRITON_MAX_STAGES", "2"))
 _VALIDATED_ARCHS = ("gfx1201", "gfx1100")
 
 
@@ -94,7 +101,11 @@ def _force_inductor_single_thread():
 
 
 def apply_rocm_sage_patches():
-    """Install the workaround. No-op off the validated RDNA archs."""
+    """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()
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index ee3269ec2..968979fe7 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -116,7 +116,7 @@ Each script selects the fastest config for that GPU; edit `--config_json` in the
 | W7900 (gfx1100 / RDNA3) | `qwen_image_21_w7900_int8.json` | INT8 (`torch._int_mm`) + `torch.compile` + SageAttention2 |
 | either | `qwen_image_21_bf16.json` | BF16 eager baseline |
 
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when a SageAttention2 backend is actually constructed. To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when a SageAttention2 backend is actually constructed; set `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` to disable it (e.g. once Triton fixes the pipeliner bug). To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
 
 | GPU | Quant | Steps | End-to-end |
 | --- | --- | ---: | ---: |
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 607ccb19b..d1276469c 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -116,7 +116,7 @@ bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 | W7900（gfx1100 / RDNA3） | `qwen_image_21_w7900_int8.json` | INT8（`torch._int_mm`）+ `torch.compile` + SageAttention2 |
 | 通用 | `qwen_image_21_bf16.json` | BF16 eager 基线 |
 
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在真正构造 SageAttention2 backend 时才生效。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在真正构造 SageAttention2 backend 时才生效；设 `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` 可禁用（例如 Triton 修复该 pipeliner bug 后）。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
 
 | 显卡 | 量化 | 步数 | 端到端 |
 | --- | --- | ---: | ---: |

From 0f59f73a21da0839cb79a69e2f53d8b66d68f0bc Mon Sep 17 00:00:00 2001
From: helloyongyang <yongyang1030@163.com>
Date: Tue, 29 Sep 2026 02:53:06 +0800
Subject: [PATCH 18/21] refactor(amd_rocm): use a dedicated SageAttention2
 backend

---
 .../amd_rocm/qwen_image_21_r9700_fp8.json     |   2 +-
 .../amd_rocm/qwen_image_21_w7900_int8.json    |   2 +-
 lightx2v/common/ops/attn/sage_attn.py         |   8 --
 lightx2v_platform/base/amd_rocm.py            |   8 --
 lightx2v_platform/base/base.py                |   9 --
 .../ops/attn/amd_rocm/__init__.py             |   1 +
 .../ops/attn/amd_rocm/sage_attn.py            |  65 ++++++++++-
 lightx2v_platform/test/test_amd_rocm_quant.py |   3 +
 .../test/test_amd_rocm_sage_attn.py           | 107 ++++++++++++++++++
 scripts/qwen_image_21/README.md               |   4 +-
 scripts/qwen_image_21/README_zh.md            |   4 +-
 11 files changed, 179 insertions(+), 34 deletions(-)
 create mode 100644 lightx2v_platform/test/test_amd_rocm_sage_attn.py

diff --git a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json
index 6cc41a24c..9a2480c59 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_r9700_fp8.json
@@ -2,7 +2,7 @@
     "infer_steps": 40,
     "resolution": 1024,
     "vae_scale_factor": 16,
-    "attn_type": "sage_attn2",
+    "attn_type": "sage_attn2_amd_rocm",
     "rope_type": "torch_real_rope",
     "rms_norm_type": "fp32_variance",
     "layer_norm_type": "torch",
diff --git a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
index 459eb2eca..72a342e94 100644
--- a/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
+++ b/configs/platforms/amd_rocm/qwen_image_21_w7900_int8.json
@@ -2,7 +2,7 @@
     "infer_steps": 40,
     "resolution": 1024,
     "vae_scale_factor": 16,
-    "attn_type": "sage_attn2",
+    "attn_type": "sage_attn2_amd_rocm",
     "rope_type": "torch_real_rope",
     "rms_norm_type": "fp32_variance",
     "layer_norm_type": "torch",
diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py
index c7d07cbab..f34bb94c6 100755
--- a/lightx2v/common/ops/attn/sage_attn.py
+++ b/lightx2v/common/ops/attn/sage_attn.py
@@ -68,14 +68,6 @@ def _sage_attn2(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tens
 class SageAttn2Weight(AttnWeightTemplate):
     def __init__(self):
         self.config = {}
-        # Let the active platform run any SageAttention-specific setup (e.g. the
-        # ROCm Triton num_stages workaround). Uniform hook — no backend branch.
-        from lightx2v_platform.base.base import get_platform_device
-
-        device = get_platform_device()
-        hook = getattr(device, "on_sage_attn2_init", None)
-        if hook is not None:
-            hook()
 
     def apply(
         self,
diff --git a/lightx2v_platform/base/amd_rocm.py b/lightx2v_platform/base/amd_rocm.py
index 089259a86..94f154522 100644
--- a/lightx2v_platform/base/amd_rocm.py
+++ b/lightx2v_platform/base/amd_rocm.py
@@ -137,14 +137,6 @@ def init_device_env():
                     setattr(mod, "sgl_kernel", sgl_kernel)
             logger.info("  - aiter sgl_kernel compatibility layer enabled (RMSNorm, GEMM)")
 
-    @staticmethod
-    def on_sage_attn2_init():
-        """Install the ROCm Triton num_stages workaround when a SageAttention2
-        backend is constructed. Arch-gated (no-op off gfx1201 / gfx1100)."""
-        from lightx2v_platform.ops.attn.amd_rocm.sage_attn import apply_rocm_sage_patches
-
-        apply_rocm_sage_patches()
-
     @staticmethod
     def is_available() -> bool:
         """Check if AMD ROCm is available."""
diff --git a/lightx2v_platform/base/base.py b/lightx2v_platform/base/base.py
index e7c949eac..72c133a34 100644
--- a/lightx2v_platform/base/base.py
+++ b/lightx2v_platform/base/base.py
@@ -6,15 +6,6 @@
 from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER
 
 
-def get_platform_device():
-    """Return the active platform device class (or None before init).
-
-    Lets shared code query the current platform for optional hooks without
-    branching on a specific backend.
-    """
-    return PLATFORM_DEVICE_REGISTER.get(global_var.PLATFORM)
-
-
 def init_ai_device(platform="cuda"):
     platform_device = PLATFORM_DEVICE_REGISTER.get(platform, None)
     if platform_device is None:
diff --git a/lightx2v_platform/ops/attn/amd_rocm/__init__.py b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
index 9ecd99131..3ed111c93 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/__init__.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/__init__.py
@@ -1 +1,2 @@
 from .flash_attn import *
+from .sage_attn import AmdSageAttn2Weight
diff --git a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
index 0750661d7..e90e0c090 100644
--- a/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
+++ b/lightx2v_platform/ops/attn/amd_rocm/sage_attn.py
@@ -1,12 +1,12 @@
-"""ROCm-only SageAttention / Triton num_stages workaround.
+"""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.
 
-Applied via the amd_rocm platform's ``on_sage_attn2_init`` hook: only when a
-SageAttention2 backend is built, and only on the validated archs.
+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
@@ -19,10 +19,69 @@
 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
diff --git a/lightx2v_platform/test/test_amd_rocm_quant.py b/lightx2v_platform/test/test_amd_rocm_quant.py
index f53321842..632026721 100644
--- a/lightx2v_platform/test/test_amd_rocm_quant.py
+++ b/lightx2v_platform/test/test_amd_rocm_quant.py
@@ -61,9 +61,12 @@ def _check(scheme, tol):
 
 
 def test_schemes_registered():
+    from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
+
     register = _mm_register()
     assert "fp8-rocm" in register
     assert "int8-rocm" in register
+    assert "sage_attn2_amd_rocm" in ATTN_WEIGHT_REGISTER
 
 
 def test_int8_rocm():
diff --git a/lightx2v_platform/test/test_amd_rocm_sage_attn.py b/lightx2v_platform/test/test_amd_rocm_sage_attn.py
new file mode 100644
index 000000000..0de87d029
--- /dev/null
+++ b/lightx2v_platform/test/test_amd_rocm_sage_attn.py
@@ -0,0 +1,107 @@
+"""CPU checks for the AMD SageAttention2 wrapper; the GPU kernel is mocked."""
+
+import importlib
+import os
+import sys
+from types import ModuleType
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+@pytest.fixture
+def sage_module():
+    if os.getenv("PLATFORM") == "amd_rocm":
+        # Follow production initialization before platform registries are read.
+        import lightx2v  # noqa: F401
+
+    return importlib.import_module("lightx2v_platform.ops.attn.amd_rocm.sage_attn")
+
+
+@pytest.fixture
+def backend(sage_module, monkeypatch):
+    def sageattn(q, k, v, tensor_layout, is_causal=False, sm_scale=None, return_lse=False):
+        assert tensor_layout == "NHD"
+        assert q.ndim == k.ndim == v.ndim == 4
+        assert all(x.is_contiguous() for x in (q, k, v))
+        q, k, v = (x.transpose(1, 2) for x in (q, k, v))
+        output = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=sm_scale).transpose(1, 2)
+        if return_lse:
+            scale = q.shape[-1] ** -0.5 if sm_scale is None else sm_scale
+            lse = ((q @ k.transpose(-2, -1)) * scale).logsumexp(-1)
+            return output, lse
+        return output
+
+    module = ModuleType("sageattention")
+    module.sageattn = sageattn
+    monkeypatch.setitem(sys.modules, "sageattention", module)
+    monkeypatch.setattr(torch.version, "hip", "test")
+    patches = Mock()
+    monkeypatch.setattr(sage_module, "apply_rocm_sage_patches", patches)
+    instance = sage_module.AmdSageAttn2Weight()
+    patches.assert_called_once_with()
+    return instance
+
+
+def test_dedicated_registration(sage_module):
+    from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER
+
+    assert PLATFORM_ATTN_WEIGHT_REGISTER["sage_attn2_amd_rocm"] is sage_module.AmdSageAttn2Weight
+    assert "sage_attn2" not in PLATFORM_ATTN_WEIGHT_REGISTER
+
+
+def test_rejects_non_rocm(sage_module, monkeypatch):
+    monkeypatch.setattr(torch.version, "hip", None)
+    with pytest.raises(RuntimeError, match="requires AMD ROCm"):
+        sage_module.AmdSageAttn2Weight()
+
+
+def test_sageattention_required_only_at_construction(sage_module, monkeypatch):
+    monkeypatch.setattr(torch.version, "hip", "test")
+    monkeypatch.setitem(sys.modules, "sageattention", None)
+    patches = Mock()
+    monkeypatch.setattr(sage_module, "apply_rocm_sage_patches", patches)
+    with pytest.raises(ImportError):
+        sage_module.AmdSageAttn2Weight()
+    patches.assert_not_called()
+
+
+@pytest.mark.parametrize("batched", [False, True])
+@pytest.mark.parametrize("with_lse", [False, True])
+def test_attention_layout_and_lse(backend, sage_module, batched, with_lse):
+    torch.manual_seed(42)
+    batch, heads, width = 2 if batched else 1, 3, 8
+    # Unequal query/KV lengths cover prefix-cache attention; transpose also
+    # verifies that the wrapper makes noncontiguous inputs contiguous.
+    q, k, v = [torch.randn(batch, heads, length, width).transpose(1, 2) for length in (5, 9, 9)]
+    scale = 0.25 if with_lse else None
+    reference = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale=scale)
+    reference = reference.transpose(1, 2).reshape(batch * 5, heads * width)
+    inputs = (q, k, v) if batched else (q[0], k[0], v[0])
+    if with_lse:
+        output, lse = backend.apply_with_lse(*inputs, softmax_scale=scale)
+        expected_lse = ((q.transpose(1, 2) @ k.transpose(1, 2).transpose(-2, -1)) * scale).logsumexp(-1)
+        torch.testing.assert_close(lse, expected_lse.transpose(1, 2).reshape(batch * 5, heads))
+    else:
+        output = backend.apply(*inputs)
+    torch.testing.assert_close(output, reference)
+    sage_module.apply_rocm_sage_patches.assert_called_once_with()
+
+
+def test_invalid_input_rank(backend):
+    x = torch.zeros(5, 8)
+    with pytest.raises(ValueError, match="expects 3D or 4D"):
+        backend.apply(x, x, x)
+
+
+@pytest.mark.parametrize("max_stages,validated,expected_calls", [(0, True, 0), (2, False, 0), (2, True, 1)])
+def test_workaround_gating(sage_module, monkeypatch, max_stages, validated, expected_calls):
+    monkeypatch.setattr(sage_module, "_MAX_STAGES", max_stages)
+    monkeypatch.setattr(sage_module, "_is_validated_rocm_arch", lambda: validated)
+    patches = [Mock(), Mock(), Mock()]
+    for name, patch in zip(("_clamp_sageattn_kernels", "_clamp_triton_compile", "_force_inductor_single_thread"), patches):
+        monkeypatch.setattr(sage_module, name, patch)
+    sage_module.apply_rocm_sage_patches()
+    assert all(patch.call_count == expected_calls for patch in patches)
diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index e3ffe58db..a6580a4f4 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -124,7 +124,7 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-Run through the platform scripts under `scripts/platforms/amd_rocm/`, which `export PLATFORM=amd_rocm` and source the common env. This activates the AMD ROCm platform backend under `lightx2v_platform/` (registers the ROCm FP8/INT8 GEMM and applies the SageAttention Triton workaround). A single GPU is enough — the text encoder is CPU-offloaded, and no `aiter` build is required.
+Run through the platform scripts under `scripts/platforms/amd_rocm/`, which `export PLATFORM=amd_rocm` and source the common env. This activates the AMD ROCm platform backend under `lightx2v_platform/` (registers the ROCm FP8/INT8 GEMM and the dedicated `sage_attn2_amd_rocm` attention backend). A single GPU is enough — the text encoder is CPU-offloaded, and no `aiter` build is required.
 
 Set `lightx2v_path` and `model_path` at the top of the script, then run:
 
@@ -143,7 +143,7 @@ Each script selects the fastest config for that GPU; edit `--config_json` in the
 | W7900 (gfx1100 / RDNA3) | `qwen_image_21_w7900_int8.json` | INT8 (`torch._int_mm`) + `torch.compile` + SageAttention2 |
 | either | `qwen_image_21_bf16.json` | BF16 eager baseline |
 
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. The SageAttention/Inductor Triton `num_stages` workaround (needed on gfx1201 / gfx1100) is installed automatically, and only when a SageAttention2 backend is actually constructed; set `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` to disable it (e.g. once Triton fixes the pipeliner bug). To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
+The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. Set `attn_type` to `sage_attn2_amd_rocm` to select the AMD-specific SageAttention2 class. Its constructor installs the SageAttention/Inductor Triton `num_stages` workaround on gfx1201 / gfx1100; the shared `sage_attn2` backend does not install ROCm patches. The workaround affects subsequent Triton/Inductor compilation in the process; set `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` to disable it (e.g. once Triton fixes the pipeliner bug). To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
 
 | GPU | Quant | Steps | End-to-end |
 | --- | --- | ---: | ---: |
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index a333e6c75..856e03925 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -124,7 +124,7 @@ pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf eino
     PyJWT jsonschema sageattention
 ```
 
-通过 `scripts/platforms/amd_rocm/` 下的平台脚本运行——脚本里会 `export PLATFORM=amd_rocm` 并 source 公共环境，从而激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端（注册 ROCm 的 FP8/INT8 GEMM，并自动应用 SageAttention 的 Triton 规避）。单卡即可——条件编码器走 CPU offload，且该路径**无需编译 aiter**。
+通过 `scripts/platforms/amd_rocm/` 下的平台脚本运行——脚本里会 `export PLATFORM=amd_rocm` 并 source 公共环境，从而激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端（注册 ROCm 的 FP8/INT8 GEMM 和专用的 `sage_attn2_amd_rocm` attention 后端）。单卡即可——条件编码器走 CPU offload，且该路径**无需编译 aiter**。
 
 先在脚本顶部设置 `lightx2v_path` 和 `model_path`，然后运行：
 
@@ -143,7 +143,7 @@ bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
 | W7900（gfx1100 / RDNA3） | `qwen_image_21_w7900_int8.json` | INT8（`torch._int_mm`）+ `torch.compile` + SageAttention2 |
 | 通用 | `qwen_image_21_bf16.json` | BF16 eager 基线 |
 
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。SageAttention/Inductor 在 gfx1201 / gfx1100 上所需的 Triton `num_stages` 规避会自动装载，且仅在真正构造 SageAttention2 backend 时才生效；设 `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` 可禁用（例如 Triton 修复该 pipeliner bug 后）。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
+`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。将 `attn_type` 设为 `sage_attn2_amd_rocm`，即可选择 AMD 专用 SageAttention2 类。该类在构造时为 gfx1201 / gfx1100 安装 SageAttention/Inductor 的 Triton `num_stages` 规避；通用 `sage_attn2` 后端不再安装 ROCm 补丁。规避生效后会影响该进程后续的 Triton/Inductor 编译；设 `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` 可禁用（例如 Triton 修复该 pipeliner bug 后）。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
 
 | 显卡 | 量化 | 步数 | 端到端 |
 | --- | --- | ---: | ---: |

From f18d059e9d91c4f75c31efa775b684beaf3b8069 Mon Sep 17 00:00:00 2001
From: helloyongyang <yongyang1030@163.com>
Date: Tue, 29 Sep 2026 02:57:02 +0800
Subject: [PATCH 19/21] update

---
 lightx2v_platform/test/test_amd_rocm_quant.py |  78 -------------
 .../test/test_amd_rocm_sage_attn.py           | 107 ------------------
 2 files changed, 185 deletions(-)
 delete mode 100644 lightx2v_platform/test/test_amd_rocm_quant.py
 delete mode 100644 lightx2v_platform/test/test_amd_rocm_sage_attn.py

diff --git a/lightx2v_platform/test/test_amd_rocm_quant.py b/lightx2v_platform/test/test_amd_rocm_quant.py
deleted file mode 100644
index 632026721..000000000
--- a/lightx2v_platform/test/test_amd_rocm_quant.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""AMD ROCm quantized GEMM correctness checks.
-
-Run inside the ROCm container (a ROCm GPU and PLATFORM=amd_rocm are required;
-otherwise every test is skipped, so this stays safe to collect on CPU CI):
-
-    PYTHONPATH=/path-to-LightX2V PLATFORM=amd_rocm python -m pytest \\
-        lightx2v_platform/test/test_amd_rocm_quant.py -v
-
-It checks that fp8-rocm / int8-rocm register into the production
-``MM_WEIGHT_REGISTER`` (the platform ops must merge into the same registry
-inference looks them up in), and that their GEMM matches a BF16 reference across
-several M — including the int8 M<=16 BF16-fallback boundary. These torch-native
-ops do not import aiter or sageattention, so registration works without them.
-"""
-
-import os
-
-import pytest
-import torch
-
-_ON_ROCM = getattr(torch.version, "hip", None) is not None and os.getenv("PLATFORM") == "amd_rocm" and torch.cuda.is_available()
-
-pytestmark = pytest.mark.skipif(not _ON_ROCM, reason="requires PLATFORM=amd_rocm on a ROCm GPU")
-
-
-def _gcn_arch():
-    try:
-        return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName
-    except Exception:
-        return ""
-
-
-def _mm_register():
-    # Import via the lightx2v package (lazily, so collecting this module on a
-    # non-ROCm host never fails): lightx2v/__init__ runs set_ai_device before
-    # registry_factory takes its one-time platform snapshot — the same order
-    # production uses, so the platform schemes are merged into MM_WEIGHT_REGISTER.
-    from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER
-
-    return MM_WEIGHT_REGISTER
-
-
-def _rel_err(out, ref):
-    return ((out.float() - ref.float()).abs().mean() / ref.float().abs().mean().clamp_min(1e-6)).item()
-
-
-def _check(scheme, tol):
-    torch.manual_seed(42)
-    register = _mm_register()
-    assert scheme in register, f"{scheme} not merged into the production MM_WEIGHT_REGISTER"
-    n, k = 512, 1024
-    weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
-    linear = register[scheme]("w.weight", None)
-    linear.load({"w.weight": weight.clone()})
-    ref_w = weight.float()
-    for m in (1, 16, 17, 31, 64):
-        x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
-        ref = (x.float() @ ref_w.t()).to(torch.bfloat16)
-        err = _rel_err(linear.apply(x), ref)
-        assert err < tol, f"{scheme} M={m}: rel_err {err:.4f} >= tol {tol}"
-
-
-def test_schemes_registered():
-    from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
-
-    register = _mm_register()
-    assert "fp8-rocm" in register
-    assert "int8-rocm" in register
-    assert "sage_attn2_amd_rocm" in ATTN_WEIGHT_REGISTER
-
-
-def test_int8_rocm():
-    _check("int8-rocm", tol=0.04)
-
-
-@pytest.mark.skipif(not _gcn_arch().startswith("gfx1201"), reason="fp8-rocm requires gfx1201 (RDNA4)")
-def test_fp8_rocm():
-    _check("fp8-rocm", tol=0.08)
diff --git a/lightx2v_platform/test/test_amd_rocm_sage_attn.py b/lightx2v_platform/test/test_amd_rocm_sage_attn.py
deleted file mode 100644
index 0de87d029..000000000
--- a/lightx2v_platform/test/test_amd_rocm_sage_attn.py
+++ /dev/null
@@ -1,107 +0,0 @@
-"""CPU checks for the AMD SageAttention2 wrapper; the GPU kernel is mocked."""
-
-import importlib
-import os
-import sys
-from types import ModuleType
-from unittest.mock import Mock
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-
-@pytest.fixture
-def sage_module():
-    if os.getenv("PLATFORM") == "amd_rocm":
-        # Follow production initialization before platform registries are read.
-        import lightx2v  # noqa: F401
-
-    return importlib.import_module("lightx2v_platform.ops.attn.amd_rocm.sage_attn")
-
-
-@pytest.fixture
-def backend(sage_module, monkeypatch):
-    def sageattn(q, k, v, tensor_layout, is_causal=False, sm_scale=None, return_lse=False):
-        assert tensor_layout == "NHD"
-        assert q.ndim == k.ndim == v.ndim == 4
-        assert all(x.is_contiguous() for x in (q, k, v))
-        q, k, v = (x.transpose(1, 2) for x in (q, k, v))
-        output = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=sm_scale).transpose(1, 2)
-        if return_lse:
-            scale = q.shape[-1] ** -0.5 if sm_scale is None else sm_scale
-            lse = ((q @ k.transpose(-2, -1)) * scale).logsumexp(-1)
-            return output, lse
-        return output
-
-    module = ModuleType("sageattention")
-    module.sageattn = sageattn
-    monkeypatch.setitem(sys.modules, "sageattention", module)
-    monkeypatch.setattr(torch.version, "hip", "test")
-    patches = Mock()
-    monkeypatch.setattr(sage_module, "apply_rocm_sage_patches", patches)
-    instance = sage_module.AmdSageAttn2Weight()
-    patches.assert_called_once_with()
-    return instance
-
-
-def test_dedicated_registration(sage_module):
-    from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER
-
-    assert PLATFORM_ATTN_WEIGHT_REGISTER["sage_attn2_amd_rocm"] is sage_module.AmdSageAttn2Weight
-    assert "sage_attn2" not in PLATFORM_ATTN_WEIGHT_REGISTER
-
-
-def test_rejects_non_rocm(sage_module, monkeypatch):
-    monkeypatch.setattr(torch.version, "hip", None)
-    with pytest.raises(RuntimeError, match="requires AMD ROCm"):
-        sage_module.AmdSageAttn2Weight()
-
-
-def test_sageattention_required_only_at_construction(sage_module, monkeypatch):
-    monkeypatch.setattr(torch.version, "hip", "test")
-    monkeypatch.setitem(sys.modules, "sageattention", None)
-    patches = Mock()
-    monkeypatch.setattr(sage_module, "apply_rocm_sage_patches", patches)
-    with pytest.raises(ImportError):
-        sage_module.AmdSageAttn2Weight()
-    patches.assert_not_called()
-
-
-@pytest.mark.parametrize("batched", [False, True])
-@pytest.mark.parametrize("with_lse", [False, True])
-def test_attention_layout_and_lse(backend, sage_module, batched, with_lse):
-    torch.manual_seed(42)
-    batch, heads, width = 2 if batched else 1, 3, 8
-    # Unequal query/KV lengths cover prefix-cache attention; transpose also
-    # verifies that the wrapper makes noncontiguous inputs contiguous.
-    q, k, v = [torch.randn(batch, heads, length, width).transpose(1, 2) for length in (5, 9, 9)]
-    scale = 0.25 if with_lse else None
-    reference = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), scale=scale)
-    reference = reference.transpose(1, 2).reshape(batch * 5, heads * width)
-    inputs = (q, k, v) if batched else (q[0], k[0], v[0])
-    if with_lse:
-        output, lse = backend.apply_with_lse(*inputs, softmax_scale=scale)
-        expected_lse = ((q.transpose(1, 2) @ k.transpose(1, 2).transpose(-2, -1)) * scale).logsumexp(-1)
-        torch.testing.assert_close(lse, expected_lse.transpose(1, 2).reshape(batch * 5, heads))
-    else:
-        output = backend.apply(*inputs)
-    torch.testing.assert_close(output, reference)
-    sage_module.apply_rocm_sage_patches.assert_called_once_with()
-
-
-def test_invalid_input_rank(backend):
-    x = torch.zeros(5, 8)
-    with pytest.raises(ValueError, match="expects 3D or 4D"):
-        backend.apply(x, x, x)
-
-
-@pytest.mark.parametrize("max_stages,validated,expected_calls", [(0, True, 0), (2, False, 0), (2, True, 1)])
-def test_workaround_gating(sage_module, monkeypatch, max_stages, validated, expected_calls):
-    monkeypatch.setattr(sage_module, "_MAX_STAGES", max_stages)
-    monkeypatch.setattr(sage_module, "_is_validated_rocm_arch", lambda: validated)
-    patches = [Mock(), Mock(), Mock()]
-    for name, patch in zip(("_clamp_sageattn_kernels", "_clamp_triton_compile", "_force_inductor_single_thread"), patches):
-        monkeypatch.setattr(sage_module, name, patch)
-    sage_module.apply_rocm_sage_patches()
-    assert all(patch.call_count == expected_calls for patch in patches)

From 94f782d35f59b5da0c7ebc9d991ebf64d435a196 Mon Sep 17 00:00:00 2001
From: helloyongyang <yongyang1030@163.com>
Date: Tue, 29 Sep 2026 03:01:17 +0800
Subject: [PATCH 20/21] update

---
 scripts/qwen_image_21/README.md    | 67 ------------------------------
 scripts/qwen_image_21/README_zh.md | 67 ------------------------------
 2 files changed, 134 deletions(-)

diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index a6580a4f4..24eaddfaf 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -110,70 +110,3 @@ Tests use PyTorch 2.11.0+cu130, 40 steps, seed 42, CFG disabled, and the example
 Latency is the median of three consecutive requests after the first request in the same instance; dual-GPU samples use the slower rank. End-to-end timing includes image saving and excludes model loading and one-time initialization.
 
 Search the runtime log for **`RUN pipeline cost`** to find the pipeline latency in seconds.
-
-### 3.3 AMD ROCm GPUs (Radeon AI PRO R9700 / Radeon PRO W7900)
-
-The 5090 stack above (FlashAttention3, FlashInfer RoPE, CUTLASS FP8, SageAttention2 CUDA kernels) is CUDA-only. The configs below run Qwen-Image-2.1 on AMD RDNA GPUs with a torch-op inference path (`torch_sdpa`, `torch_real_rope`, torch LayerNorm/modulation) plus ROCm-friendly quantization, `torch.compile`, and SageAttention2. They quantize the **released BF16 weights on load** — no `tools/convert` step and no `dit_quantized_ckpt`.
-
-Validated in the `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` image (PyTorch 2.10, Triton 3.6) with the `transformers` 4.57.x line (`Qwen3VLProcessor` needs a recent release; pin to a version you have tested rather than only a lower bound). Install the runtime deps not baked into that image:
-
-```bash
-pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
-    qtorch langdetect tqdm imageio imageio-ffmpeg opencv-python-headless comfy-kitchen \
-    peft gguf prometheus_client fastapi uvicorn pydantic aiohttp pyzmq python-multipart \
-    PyJWT jsonschema sageattention
-```
-
-Run through the platform scripts under `scripts/platforms/amd_rocm/`, which `export PLATFORM=amd_rocm` and source the common env. This activates the AMD ROCm platform backend under `lightx2v_platform/` (registers the ROCm FP8/INT8 GEMM and the dedicated `sage_attn2_amd_rocm` attention backend). A single GPU is enough — the text encoder is CPU-offloaded, and no `aiter` build is required.
-
-Set `lightx2v_path` and `model_path` at the top of the script, then run:
-
-```bash
-# R9700 (FP8)
-bash scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
-# W7900 (INT8)
-bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
-```
-
-Each script selects the fastest config for that GPU; edit `--config_json` in the script to pick another. Per-GPU config (all under `configs/platforms/amd_rocm/`):
-
-| GPU (arch) | Config | Contents |
-| --- | --- | --- |
-| R9700 (gfx1201 / RDNA4) | `qwen_image_21_r9700_fp8.json` | FP8 (`torch._scaled_mm`) + `torch.compile` + SageAttention2 |
-| W7900 (gfx1100 / RDNA3) | `qwen_image_21_w7900_int8.json` | INT8 (`torch._int_mm`) + `torch.compile` + SageAttention2 |
-| either | `qwen_image_21_bf16.json` | BF16 eager baseline |
-
-The `dit_quant_scheme` is `fp8-rocm` (R9700) or `int8-rocm` (W7900), registered by the AMD ROCm platform ops in `lightx2v_platform/ops/mm/amd_rocm/`: per-channel symmetric weight + per-token dynamic activation, quantized from BF16 on load. The VAE runs native convolution — the AMD ROCm platform disables cuDNN/MIOpen, which sidesteps the nondeterministic-NaN convolution path observed on gfx1201. Set `attn_type` to `sage_attn2_amd_rocm` to select the AMD-specific SageAttention2 class. Its constructor installs the SageAttention/Inductor Triton `num_stages` workaround on gfx1201 / gfx1100; the shared `sage_attn2` backend does not install ROCm patches. The workaround affects subsequent Triton/Inductor compilation in the process; set `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` to disable it (e.g. once Triton fixes the pipeliner bug). To trade a little quality for speed, lower `infer_steps` (e.g. 25) in the config or pass `--infer_steps 25`.
-
-| GPU | Quant | Steps | End-to-end |
-| --- | --- | ---: | ---: |
-| R9700 | FP8 | 40 | ~24 s |
-| R9700 | FP8 | 25 | ~17 s |
-| W7900 | INT8 | 40 | ~46 s |
-| W7900 | INT8 | 25 | ~29 s |
-
-Measured at 1024×1024, seed 42, CFG disabled, single GPU, steady state after warmup (median of 3 runs). End-to-end covers text encoding, condition-KV prefill, denoising, VAE decode and PNG save; it excludes model load and one-time initialization. In the tested prompts and seeds, 25 steps showed no obvious visual degradation versus 40 — this is not a systematic quality evaluation. Compilation is a one-time warmup cost (~1.5–3 min); `TORCHINDUCTOR_COMPILE_THREADS=1` is set automatically on the SageAttention path.
-
-## 4. Service Deployment and API Usage
-
-Set the repository path, model path, and GPU ID in `server/start_server.sh`, then start the server:
-
-```bash
-bash scripts/qwen_image_21/server/start_server.sh
-```
-
-The default port is `8000`. The server supports both text-to-image and image-to-image requests.
-
-Once the server is ready, open a new terminal in the same container and send requests:
-
-```bash
-# Text-to-image
-python scripts/qwen_image_21/server/post_t2i.py
-
-# Image-to-image
-python scripts/qwen_image_21/server/post_i2i.py
-```
-
-Edit `url`, `message`, and `output_path` directly; for image-to-image, also set `image_path`. The scripts already specify `task: "t2i"` and `task: "i2i"`, respectively.
-
-Service requests follow the same size rules. For automatic image-to-image sizing, remove `size` from `message`.
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 856e03925..14f3c14b0 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -110,70 +110,3 @@ bash scripts/qwen_image_21/qwen_image_21_i2i_5090_2k_sp2.sh
 同一实例中首次请求后，连续三次请求取中位数；双卡每次取较慢 rank 的耗时。端到端耗时包含图片保存，不包含模型加载和一次性初始化。
 
 运行时在日志中搜索 **`RUN pipeline cost`**，查看单次 pipeline 耗时，单位为秒。
-
-### 3.3 AMD ROCm 显卡（Radeon AI PRO R9700 / Radeon PRO W7900）
-
-上面 5090 那套（FlashAttention3、FlashInfer RoPE、CUTLASS FP8、SageAttention2 的 CUDA kernel）都是 CUDA 专用。下面的配置让 Qwen-Image-2.1 跑在 AMD RDNA 显卡上，走 torch 算子推理路径（`torch_sdpa`、`torch_real_rope`、torch LayerNorm/modulation），配合 ROCm 可用的量化、`torch.compile` 和 SageAttention2。它们**在加载时直接量化官方 BF16 权重**——无需 `tools/convert` 转换、无需 `dit_quantized_ckpt`。
-
-验证环境为 `rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0` 镜像（PyTorch 2.10、Triton 3.6），transformers 使用 4.57.x（`Qwen3VLProcessor` 需要较新版本；请固定到你实测过的版本，而非只写下限）。安装该镜像未自带的运行时依赖：
-
-```bash
-pip install "transformers>=4.57" diffusers ftfy accelerate loguru omegaconf einops \
-    qtorch langdetect tqdm imageio imageio-ffmpeg opencv-python-headless comfy-kitchen \
-    peft gguf prometheus_client fastapi uvicorn pydantic aiohttp pyzmq python-multipart \
-    PyJWT jsonschema sageattention
-```
-
-通过 `scripts/platforms/amd_rocm/` 下的平台脚本运行——脚本里会 `export PLATFORM=amd_rocm` 并 source 公共环境，从而激活 `lightx2v_platform/` 下的 AMD ROCm 平台后端（注册 ROCm 的 FP8/INT8 GEMM 和专用的 `sage_attn2_amd_rocm` attention 后端）。单卡即可——条件编码器走 CPU offload，且该路径**无需编译 aiter**。
-
-先在脚本顶部设置 `lightx2v_path` 和 `model_path`，然后运行：
-
-```bash
-# R9700（FP8）
-bash scripts/platforms/amd_rocm/qwen_image_21_r9700_t2i.sh
-# W7900（INT8）
-bash scripts/platforms/amd_rocm/qwen_image_21_w7900_t2i.sh
-```
-
-每个脚本默认选该卡最快的配置；改脚本里的 `--config_json` 可切换到其它配置。各配置（均在 `configs/platforms/amd_rocm/` 下）：
-
-| 显卡（架构） | 配置 | 内容 |
-| --- | --- | --- |
-| R9700（gfx1201 / RDNA4） | `qwen_image_21_r9700_fp8.json` | FP8（`torch._scaled_mm`）+ `torch.compile` + SageAttention2 |
-| W7900（gfx1100 / RDNA3） | `qwen_image_21_w7900_int8.json` | INT8（`torch._int_mm`）+ `torch.compile` + SageAttention2 |
-| 通用 | `qwen_image_21_bf16.json` | BF16 eager 基线 |
-
-`dit_quant_scheme` 为 `fp8-rocm`（R9700）或 `int8-rocm`（W7900），由 `lightx2v_platform/ops/mm/amd_rocm/` 下的 AMD ROCm 平台算子注册：权重 per-channel 对称量化 + 激活 per-token 动态量化，均在加载时从 BF16 量化得到。VAE 走原生卷积——AMD ROCm 平台会关闭 cuDNN/MIOpen，从而绕开 gfx1201 上观察到的非确定性 NaN 卷积路径。将 `attn_type` 设为 `sage_attn2_amd_rocm`，即可选择 AMD 专用 SageAttention2 类。该类在构造时为 gfx1201 / gfx1100 安装 SageAttention/Inductor 的 Triton `num_stages` 规避；通用 `sage_attn2` 后端不再安装 ROCm 补丁。规避生效后会影响该进程后续的 Triton/Inductor 编译；设 `LIGHTX2V_ROCM_TRITON_MAX_STAGES=0` 可禁用（例如 Triton 修复该 pipeliner bug 后）。想拿画质换速度，可把配置里的 `infer_steps` 调小（如 25）或传 `--infer_steps 25`。
-
-| 显卡 | 量化 | 步数 | 端到端 |
-| --- | --- | ---: | ---: |
-| R9700 | FP8 | 40 | ~24 s |
-| R9700 | FP8 | 25 | ~17 s |
-| W7900 | INT8 | 40 | ~46 s |
-| W7900 | INT8 | 25 | ~29 s |
-
-测试条件：1024×1024、seed 42、关闭 CFG、单卡、warmup 后稳态（3 次取中位数）。端到端包含文本编码、condition-KV prefill、去噪、VAE 解码与 PNG 保存，不含模型加载与一次性初始化。在所测试的 prompt/seed 上，25 步相比 40 步未观察到明显画质退化——这不是系统性画质评估。编译为一次性 warmup 开销（约 1.5–3 分钟）；`TORCHINDUCTOR_COMPILE_THREADS=1` 在 SageAttention 路径上会自动设置。
-
-## 4. 服务化部署与 API 调用
-
-先修改 `server/start_server.sh` 中的仓库路径、模型路径和 GPU 编号，然后启动：
-
-```bash
-bash scripts/qwen_image_21/server/start_server.sh
-```
-
-默认端口为 `8000`，同一个服务支持文生图和图生图。
-
-服务启动完成后，在同一容器新开终端并发送请求：
-
-```bash
-# 文生图
-python scripts/qwen_image_21/server/post_t2i.py
-
-# 图生图
-python scripts/qwen_image_21/server/post_i2i.py
-```
-
-直接修改代码中的 `url`、`message`、`output_path`；图生图还需设置 `image_path`。两个脚本已分别设置 `task: "t2i"` 和 `task: "i2i"`。
-
-服务请求使用相同的尺寸规则。图生图需自动确定尺寸时，删除 `message` 中的 `size` 字段。

From 377a8a3df54ece80de6828bf88d7f8120f02d20b Mon Sep 17 00:00:00 2001
From: helloyongyang <yongyang1030@163.com>
Date: Tue, 29 Sep 2026 03:03:28 +0800
Subject: [PATCH 21/21] docs: remove Qwen-Image README changes from ROCm PR

---
 scripts/qwen_image_21/README.md    | 24 ++++++++++++++++++++++++
 scripts/qwen_image_21/README_zh.md | 24 ++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/scripts/qwen_image_21/README.md b/scripts/qwen_image_21/README.md
index 24eaddfaf..a970f1c84 100644
--- a/scripts/qwen_image_21/README.md
+++ b/scripts/qwen_image_21/README.md
@@ -110,3 +110,27 @@ Tests use PyTorch 2.11.0+cu130, 40 steps, seed 42, CFG disabled, and the example
 Latency is the median of three consecutive requests after the first request in the same instance; dual-GPU samples use the slower rank. End-to-end timing includes image saving and excludes model loading and one-time initialization.
 
 Search the runtime log for **`RUN pipeline cost`** to find the pipeline latency in seconds.
+
+## 4. Service Deployment and API Usage
+
+Set the repository path, model path, and GPU ID in `server/start_server.sh`, then start the server:
+
+```bash
+bash scripts/qwen_image_21/server/start_server.sh
+```
+
+The default port is `8000`. The server supports both text-to-image and image-to-image requests.
+
+Once the server is ready, open a new terminal in the same container and send requests:
+
+```bash
+# Text-to-image
+python scripts/qwen_image_21/server/post_t2i.py
+
+# Image-to-image
+python scripts/qwen_image_21/server/post_i2i.py
+```
+
+Edit `url`, `message`, and `output_path` directly; for image-to-image, also set `image_path`. The scripts already specify `task: "t2i"` and `task: "i2i"`, respectively.
+
+Service requests follow the same size rules. For automatic image-to-image sizing, remove `size` from `message`.
diff --git a/scripts/qwen_image_21/README_zh.md b/scripts/qwen_image_21/README_zh.md
index 14f3c14b0..245dec371 100644
--- a/scripts/qwen_image_21/README_zh.md
+++ b/scripts/qwen_image_21/README_zh.md
@@ -110,3 +110,27 @@ bash scripts/qwen_image_21/qwen_image_21_i2i_5090_2k_sp2.sh
 同一实例中首次请求后，连续三次请求取中位数；双卡每次取较慢 rank 的耗时。端到端耗时包含图片保存，不包含模型加载和一次性初始化。
 
 运行时在日志中搜索 **`RUN pipeline cost`**，查看单次 pipeline 耗时，单位为秒。
+
+## 4. 服务化部署与 API 调用
+
+先修改 `server/start_server.sh` 中的仓库路径、模型路径和 GPU 编号，然后启动：
+
+```bash
+bash scripts/qwen_image_21/server/start_server.sh
+```
+
+默认端口为 `8000`，同一个服务支持文生图和图生图。
+
+服务启动完成后，在同一容器新开终端并发送请求：
+
+```bash
+# 文生图
+python scripts/qwen_image_21/server/post_t2i.py
+
+# 图生图
+python scripts/qwen_image_21/server/post_i2i.py
+```
+
+直接修改代码中的 `url`、`message`、`output_path`；图生图还需设置 `image_path`。两个脚本已分别设置 `task: "t2i"` 和 `task: "i2i"`。
+
+服务请求使用相同的尺寸规则。图生图需自动确定尺寸时，删除 `message` 中的 `size` 字段。