From 90fff812d2db5afd04bbd0419bd67e1576008981 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 15:47:44 +0000 Subject: [PATCH 1/6] feat: support INT8 Ulysses communication --- .../common/ops/attn/kernels/ulysses_layout.py | 202 ++++++++++++++---- lightx2v/common/ops/attn/ulysses_prepost.py | 27 ++- lightx2v/common/ops/attn/utils/seq_p.py | 39 +++- .../networks/wan/infer/transformer_infer.py | 4 +- 4 files changed, 223 insertions(+), 49 deletions(-) diff --git a/lightx2v/common/ops/attn/kernels/ulysses_layout.py b/lightx2v/common/ops/attn/kernels/ulysses_layout.py index 5d0fbb5d0..36c707df2 100644 --- a/lightx2v/common/ops/attn/kernels/ulysses_layout.py +++ b/lightx2v/common/ops/attn/kernels/ulysses_layout.py @@ -3,8 +3,8 @@ The names intentionally mirror the two Ulysses communications: - ``qkv_pre`` / ``qkv_post`` wrap the all-to-all before attention. - ``attn_pre`` / ``attn_post`` wrap the all-to-all after attention. -- ``*_fp8`` variants keep the same layout contract, but return separate - FP8 payload and FP32 scale tensors. Keeping these as separate collectives +- ``*_fp8`` and ``*_int8`` variants keep the same layout contract, but return separate + quantized payload and FP32 scale tensors. Keeping these as separate collectives matches the faster NCCL message shape used by the legacy implementation. The public pre/post wrappers each launch one Triton kernel on the hot path. @@ -13,6 +13,7 @@ import torch import triton import triton.language as tl +from triton.language.extra import libdevice # Validation and launch-shape helpers @@ -71,8 +72,8 @@ def _numel_from_shape(shape): return numel -def _fp8_split_layout(payload_shape, scale_shape, name="fp8 split"): - """Return per-rank split metadata for separate FP8 payload/scale tensors.""" +def _quantized_split_layout(payload_shape, scale_shape, name="quantized split"): + """Return per-rank split metadata for separate quantized payload/scale tensors.""" if int(payload_shape[0]) != int(scale_shape[0]): raise ValueError(f"{name}: payload and scale must share the all-to-all split dimension.") payload_elems_per_rank = _numel_from_shape(payload_shape[1:]) @@ -80,18 +81,18 @@ def _fp8_split_layout(payload_shape, scale_shape, name="fp8 split"): return payload_elems_per_rank, scale_elems_per_rank -def _check_fp8_split(payload, scale, payload_shape, scale_shape, scale_dtype=torch.float32, name="fp8 split"): - """Validate separate FP8 payload and FP32 scale tensors used by post kernels.""" +def _check_quantized_split(payload, scale, payload_shape, scale_shape, payload_dtype, scale_dtype=torch.float32, name="quantized split"): + """Validate separate quantized payload and FP32 scale tensors.""" _check_cuda_contiguous(payload, scale) - if payload.dtype != torch.float8_e4m3fn: - raise ValueError(f"{name}: payload dtype must be torch.float8_e4m3fn, got {payload.dtype}.") + if payload.dtype != payload_dtype: + raise ValueError(f"{name}: payload dtype must be {payload_dtype}, got {payload.dtype}.") if scale.dtype != scale_dtype: raise ValueError(f"{name}: scale dtype must be {scale_dtype}, got {scale.dtype}.") if tuple(payload.shape) != tuple(payload_shape): raise ValueError(f"{name}: payload shape {tuple(payload.shape)} does not match expected {tuple(payload_shape)}.") if tuple(scale.shape) != tuple(scale_shape): raise ValueError(f"{name}: scale shape {tuple(scale.shape)} does not match expected {tuple(scale_shape)}.") - return _fp8_split_layout(payload_shape, scale_shape, name) + return _quantized_split_layout(payload_shape, scale_shape, name) # BF16/FP16 layout kernels @@ -341,15 +342,15 @@ def _attn_post_kernel( tl.store(out + out_offsets, tl.load(attn + input_offsets, mask=mask), mask=mask) -# FP8 quantized layout kernels +# Quantized layout kernels @triton.jit -def _qkv_pre_fp8_kernel( +def _qkv_pre_quant_kernel( q, k, v, - payload_fp8, + payload, scale_ptr, local_len: tl.constexpr, shard_heads: tl.constexpr, @@ -360,6 +361,7 @@ def _qkv_pre_fp8_kernel( payload_elems_per_rank: tl.constexpr, scale_base_offset: tl.constexpr, scale_elems_per_rank: tl.constexpr, + int8_quant: tl.constexpr, block_d: tl.constexpr, ): # One program handles the q/k/v triplet for one [dst_rank, local_s, shard_head]. @@ -381,12 +383,29 @@ def _qkv_pre_fp8_kernel( k_vals = tl.load(k + src_offsets, mask=mask, other=0.0).to(tl.float32) v_vals = tl.load(v + src_offsets, mask=mask, other=0.0).to(tl.float32) - q_amax = tl.maximum(tl.max(tl.abs(q_vals), axis=0), 0.001953125) - k_amax = tl.maximum(tl.max(tl.abs(k_vals), axis=0), 0.001953125) - v_amax = tl.maximum(tl.max(tl.abs(v_vals), axis=0), 0.001953125) - q_scale = q_amax / 448.0 - k_scale = k_amax / 448.0 - v_scale = v_amax / 448.0 + q_amax = tl.max(tl.abs(q_vals), axis=0) + k_amax = tl.max(tl.abs(k_vals), axis=0) + v_amax = tl.max(tl.abs(v_vals), axis=0) + if int8_quant: + q_scale = tl.where(q_amax > 0.0, q_amax / 127.0, 1.0) + k_scale = tl.where(k_amax > 0.0, k_amax / 127.0, 1.0) + v_scale = tl.where(v_amax > 0.0, v_amax / 127.0, 1.0) + q_quant = libdevice.rint(q_vals / q_scale) + k_quant = libdevice.rint(k_vals / k_scale) + v_quant = libdevice.rint(v_vals / v_scale) + q_quant = tl.minimum(tl.maximum(q_quant, -127.0), 127.0).to(tl.int8) + k_quant = tl.minimum(tl.maximum(k_quant, -127.0), 127.0).to(tl.int8) + v_quant = tl.minimum(tl.maximum(v_quant, -127.0), 127.0).to(tl.int8) + else: + q_amax = tl.maximum(q_amax, 0.001953125) + k_amax = tl.maximum(k_amax, 0.001953125) + v_amax = tl.maximum(v_amax, 0.001953125) + q_scale = q_amax / 448.0 + k_scale = k_amax / 448.0 + v_scale = v_amax / 448.0 + q_quant = (q_vals / q_scale).to(tl.float8e4nv) + k_quant = (k_vals / k_scale).to(tl.float8e4nv) + v_quant = (v_vals / v_scale).to(tl.float8e4nv) q_row = s * 3 * shard_heads + h k_row = q_row + shard_heads @@ -394,9 +413,9 @@ def _qkv_pre_fp8_kernel( q_payload = rank * payload_elems_per_rank + q_row * hidden_dims + offs k_payload = rank * payload_elems_per_rank + k_row * hidden_dims + offs v_payload = rank * payload_elems_per_rank + v_row * hidden_dims + offs - tl.store(payload_fp8 + q_payload, (q_vals / q_scale).to(tl.float8e4nv), mask=mask) - tl.store(payload_fp8 + k_payload, (k_vals / k_scale).to(tl.float8e4nv), mask=mask) - tl.store(payload_fp8 + v_payload, (v_vals / v_scale).to(tl.float8e4nv), mask=mask) + tl.store(payload + q_payload, q_quant, mask=mask) + tl.store(payload + k_payload, k_quant, mask=mask) + tl.store(payload + v_payload, v_quant, mask=mask) q_scale_offset = rank * scale_elems_per_rank + scale_base_offset + q_row k_scale_offset = rank * scale_elems_per_rank + scale_base_offset + k_row @@ -407,7 +426,7 @@ def _qkv_pre_fp8_kernel( @triton.jit -def _qkv_post_fp8_kernel( +def _qkv_post_quant_kernel( payload_fp8, scale_ptr, q_source, @@ -491,7 +510,7 @@ def _qkv_post_fp8_kernel( @triton.jit -def _qonly_qkv_post_fp8_kernel( +def _qonly_qkv_post_quant_kernel( payload_fp8, scale_ptr, k_source, @@ -578,9 +597,9 @@ def _qonly_qkv_post_fp8_kernel( @triton.jit -def _attn_pre_fp8_kernel( +def _attn_pre_quant_kernel( attn, - payload_fp8, + payload, scale_ptr, local_len: tl.constexpr, shard_heads: tl.constexpr, @@ -588,6 +607,7 @@ def _attn_pre_fp8_kernel( payload_elems_per_rank: tl.constexpr, scale_base_offset: tl.constexpr, scale_elems_per_rank: tl.constexpr, + int8_quant: tl.constexpr, block_d: tl.constexpr, ): # One program handles one [shard_head, local_s] row and writes payload plus scale. @@ -605,18 +625,24 @@ def _attn_pre_fp8_kernel( src_offsets = src_s * (shard_heads * hidden_dims) + h * hidden_dims + offs vals = tl.load(attn + src_offsets, mask=mask, other=0.0).to(tl.float32) - amax = tl.maximum(tl.max(tl.abs(vals), axis=0), 0.001953125) - scale = amax / 448.0 - quant = (vals / scale).to(tl.float8e4nv) + amax = tl.max(tl.abs(vals), axis=0) + if int8_quant: + scale = tl.where(amax > 0.0, amax / 127.0, 1.0) + quant = libdevice.rint(vals / scale) + quant = tl.minimum(tl.maximum(quant, -127.0), 127.0).to(tl.int8) + else: + amax = tl.maximum(amax, 0.001953125) + scale = amax / 448.0 + quant = (vals / scale).to(tl.float8e4nv) payload_offsets = rank * payload_elems_per_rank + row_in_rank * hidden_dims + offs - tl.store(payload_fp8 + payload_offsets, quant, mask=mask) + tl.store(payload + payload_offsets, quant, mask=mask) scale_offset = rank * scale_elems_per_rank + scale_base_offset + row_in_rank tl.store(scale_ptr + scale_offset, scale) @triton.jit -def _attn_post_fp8_kernel( +def _attn_post_quant_kernel( payload_fp8, scale_ptr, out, @@ -859,13 +885,13 @@ def qkv_pre_fp8(q, k, v, world_size, head_index=None): shard_heads, head_offset, head_stride = _resolve_qkv_head_layout(heads, world_size, head_index) payload_shape = (world_size, local_len, 3, shard_heads, hidden_dims) scale_shape = (*payload_shape[:-1], 1) - payload_elems_per_rank, scale_elems_per_rank = _fp8_split_layout(payload_shape, scale_shape, name="fused QKV fp8 split path") + payload_elems_per_rank, scale_elems_per_rank = _quantized_split_layout(payload_shape, scale_shape, name="fused QKV fp8 split path") payload = torch.empty(payload_shape, device=q.device, dtype=torch.float8_e4m3fn) scale = torch.empty(scale_shape, device=q.device, dtype=torch.float32) total_rows = world_size * local_len * shard_heads block_d = _next_power_of_2(hidden_dims) - _qkv_pre_fp8_kernel[(total_rows,)]( + _qkv_pre_quant_kernel[(total_rows,)]( q, k, v, @@ -880,6 +906,46 @@ def qkv_pre_fp8(q, k, v, world_size, head_index=None): payload_elems_per_rank, 0, scale_elems_per_rank, + False, + block_d, + num_warps=4, + ) + return payload, scale, payload_shape, scale_shape + + +def qkv_pre_int8(q, k, v, world_size, head_index=None): + """INT8 version of ``qkv_pre`` with fused layout+quantization.""" + _check_cuda_contiguous(q, k, v) + if q.shape != k.shape or q.shape != v.shape: + raise ValueError(f"q/k/v must have the same shape for qkv fusion, got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}.") + local_len, heads, hidden_dims = q.shape + _check_hidden_dims(hidden_dims, "fused QKV int8 pre path") + + shard_heads, head_offset, head_stride = _resolve_qkv_head_layout(heads, world_size, head_index) + payload_shape = (world_size, local_len, 3, shard_heads, hidden_dims) + scale_shape = (*payload_shape[:-1], 1) + payload_elems_per_rank, scale_elems_per_rank = _quantized_split_layout(payload_shape, scale_shape, name="fused QKV int8 split path") + + payload = torch.empty(payload_shape, device=q.device, dtype=torch.int8) + scale = torch.empty(scale_shape, device=q.device, dtype=torch.float32) + total_rows = world_size * local_len * shard_heads + block_d = _next_power_of_2(hidden_dims) + _qkv_pre_quant_kernel[(total_rows,)]( + q, + k, + v, + payload, + scale, + local_len, + shard_heads, + hidden_dims, + heads, + head_offset, + head_stride, + payload_elems_per_rank, + 0, + scale_elems_per_rank, + True, block_d, num_warps=4, ) @@ -900,6 +966,8 @@ def qkv_post_fp8( q_only=False, block_size=None, head_index=None, + _payload_dtype=torch.float8_e4m3fn, + _quant_name="fp8", ): """FP8 version of ``qkv_post`` for separate payload/scale tensors.""" _check_cuda_contiguous(payload, scale, q_source, k_source, v_source) @@ -920,8 +988,8 @@ def qkv_post_fp8( if shard_heads != expected_shard_heads: raise ValueError(f"qkv shard_heads must be {expected_shard_heads}, got {shard_heads}.") - _check_hidden_dims(hidden_dims, "fused QKV fp8 post path") - payload_elems_per_rank, scale_elems_per_rank = _check_fp8_split(payload, scale, payload_shape, scale_shape, torch.float32, name="split QKV fp8 input") + _check_hidden_dims(hidden_dims, f"fused QKV {_quant_name} post path") + payload_elems_per_rank, scale_elems_per_rank = _check_quantized_split(payload, scale, payload_shape, scale_shape, _payload_dtype, torch.float32, name=f"split QKV {_quant_name} input") block_m = _resolve_block_m(block_size, hidden_dims, default=16) block_d = _next_power_of_2(hidden_dims) @@ -936,7 +1004,7 @@ def qkv_post_fp8( q_rows = q_shape[0] * shard_heads total_rows = max(q_rows, kv_shape[0] * shard_heads) grid = (triton.cdiv(total_rows, block_m),) - _qonly_qkv_post_fp8_kernel[grid]( + _qonly_qkv_post_quant_kernel[grid]( payload, scale, k_source, @@ -970,7 +1038,7 @@ def qkv_post_fp8( out_v = torch.empty_like(out_q) total_rows = shape[0] * shard_heads grid = (triton.cdiv(total_rows, block_m),) - _qkv_post_fp8_kernel[grid]( + _qkv_post_quant_kernel[grid]( payload, scale, q_source, @@ -999,6 +1067,11 @@ def qkv_post_fp8( return out_q, out_k, out_v +def qkv_post_int8(*args, **kwargs): + """INT8 version of ``qkv_post`` for separate payload/scale tensors.""" + return qkv_post_fp8(*args, **kwargs, _payload_dtype=torch.int8, _quant_name="int8") + + def attn_pre_fp8(attn, local_len, world_size, shard_heads, hidden_dims): """FP8 version of ``attn_pre`` with fused layout+quantization.""" _check_cuda_contiguous(attn) @@ -1006,13 +1079,13 @@ def attn_pre_fp8(attn, local_len, world_size, shard_heads, hidden_dims): payload_shape = (world_size, shard_heads, local_len, hidden_dims) scale_shape = (*payload_shape[:-1], 1) - payload_elems_per_rank, scale_elems_per_rank = _fp8_split_layout(payload_shape, scale_shape, name="fused attn fp8 split path") + payload_elems_per_rank, scale_elems_per_rank = _quantized_split_layout(payload_shape, scale_shape, name="fused attn fp8 split path") payload = torch.empty(payload_shape, device=attn.device, dtype=torch.float8_e4m3fn) scale = torch.empty(scale_shape, device=attn.device, dtype=torch.float32) total_rows = world_size * shard_heads * local_len block_d = _next_power_of_2(hidden_dims) - _attn_pre_fp8_kernel[(total_rows,)]( + _attn_pre_quant_kernel[(total_rows,)]( attn, payload, scale, @@ -1022,28 +1095,68 @@ def attn_pre_fp8(attn, local_len, world_size, shard_heads, hidden_dims): payload_elems_per_rank, 0, scale_elems_per_rank, + False, block_d, num_warps=4, ) return payload, scale, payload_shape, scale_shape -def attn_post_fp8(payload, scale, payload_shape, scale_shape, output_dtype, block_size=None): +def attn_pre_int8(attn, local_len, world_size, shard_heads, hidden_dims): + """INT8 version of ``attn_pre`` with fused layout+quantization.""" + _check_cuda_contiguous(attn) + _check_hidden_dims(hidden_dims, "fused attn int8 pre path") + + payload_shape = (world_size, shard_heads, local_len, hidden_dims) + scale_shape = (*payload_shape[:-1], 1) + payload_elems_per_rank, scale_elems_per_rank = _quantized_split_layout(payload_shape, scale_shape, name="fused attn int8 split path") + + payload = torch.empty(payload_shape, device=attn.device, dtype=torch.int8) + scale = torch.empty(scale_shape, device=attn.device, dtype=torch.float32) + total_rows = world_size * shard_heads * local_len + block_d = _next_power_of_2(hidden_dims) + _attn_pre_quant_kernel[(total_rows,)]( + attn, + payload, + scale, + local_len, + shard_heads, + hidden_dims, + payload_elems_per_rank, + 0, + scale_elems_per_rank, + True, + block_d, + num_warps=4, + ) + return payload, scale, payload_shape, scale_shape + + +def attn_post_fp8( + payload, + scale, + payload_shape, + scale_shape, + output_dtype, + block_size=None, + _payload_dtype=torch.float8_e4m3fn, + _quant_name="fp8", +): """FP8 version of ``attn_post`` for separate payload/scale tensors.""" _check_cuda_contiguous(payload, scale) world_size, shard_heads, local_len, hidden_dims = payload_shape if tuple(scale_shape) != (world_size, shard_heads, local_len, 1): raise ValueError("attn scale shape does not match payload shape.") - _check_hidden_dims(hidden_dims, "fused attn fp8 post path") - payload_elems_per_rank, scale_elems_per_rank = _check_fp8_split(payload, scale, payload_shape, scale_shape, torch.float32, name="split attn fp8 input") + _check_hidden_dims(hidden_dims, f"fused attn {_quant_name} post path") + payload_elems_per_rank, scale_elems_per_rank = _check_quantized_split(payload, scale, payload_shape, scale_shape, _payload_dtype, torch.float32, name=f"split attn {_quant_name} input") out = torch.empty((local_len, world_size * shard_heads * hidden_dims), device=payload.device, dtype=output_dtype) total_rows = local_len * world_size * shard_heads block_m = _resolve_block_m(block_size, hidden_dims, default=16) block_d = _next_power_of_2(hidden_dims) grid = (triton.cdiv(total_rows, block_m),) - _attn_post_fp8_kernel[grid]( + _attn_post_quant_kernel[grid]( payload, scale, out, @@ -1059,3 +1172,8 @@ def attn_post_fp8(payload, scale, payload_shape, scale_shape, output_dtype, bloc num_warps=4, ) return out + + +def attn_post_int8(*args, **kwargs): + """INT8 version of ``attn_post`` for separate payload/scale tensors.""" + return attn_post_fp8(*args, **kwargs, _payload_dtype=torch.int8, _quant_name="int8") diff --git a/lightx2v/common/ops/attn/ulysses_prepost.py b/lightx2v/common/ops/attn/ulysses_prepost.py index 6cfba37b8..e33f7235d 100644 --- a/lightx2v/common/ops/attn/ulysses_prepost.py +++ b/lightx2v/common/ops/attn/ulysses_prepost.py @@ -3,12 +3,16 @@ from .kernels.ulysses_layout import ( attn_post, attn_post_fp8, + attn_post_int8, attn_pre, attn_pre_fp8, + attn_pre_int8, qkv_post, qkv_post_fp8, + qkv_post_int8, qkv_pre, qkv_pre_fp8, + qkv_pre_int8, ) from .utils.seq_p import pack_seq_p_tensor, unpack_seq_p_tensor, validate_quant_scheme @@ -126,7 +130,7 @@ def unpack_attn(packed, output_dtype, hidden_dims): class TritonUlyssesPrePost: - """Fused Triton layout and FP8 communication-quantization backend.""" + """Fused Triton layout and FP8/INT8 communication-quantization backend.""" @staticmethod def _validate(q, k, v, quant_scheme): @@ -142,6 +146,9 @@ def pack_qkv(cls, q, k, v, world_size, quant_scheme=None, qkv_fusion=True, head_ if quant_scheme == "fp8": payload, scale, _, _ = qkv_pre_fp8(q, k, v, world_size, head_index=head_index) return ((payload, scale),) + if quant_scheme == "int8": + payload, scale, _, _ = qkv_pre_int8(q, k, v, world_size, head_index=head_index) + return ((payload, scale),) return ((qkv_pre(q, k, v, world_size, head_index=head_index), None),) @staticmethod @@ -168,7 +175,13 @@ def unpack_qkv( if scale is None: return qkv_post(payload, q_source, k_source, v_source, rank, aux_len, qkv_first, q_only=q_only, head_index=head_index) - return qkv_post_fp8( + if payload.dtype == torch.float8_e4m3fn: + post_fn = qkv_post_fp8 + elif payload.dtype == torch.int8: + post_fn = qkv_post_int8 + else: + raise ValueError(f"Unsupported quantized QKV payload dtype: {payload.dtype}.") + return post_fn( payload, scale, payload.shape, @@ -185,11 +198,15 @@ def unpack_qkv( @staticmethod def pack_attn(output, local_len, world_size, shard_heads, hidden_dims, quant_scheme=None): + validate_quant_scheme(quant_scheme) if quant_scheme == "fp4": raise ValueError("prepost_backend='triton' does not support FP4 communication.") if quant_scheme == "fp8": payload, scale, _, _ = attn_pre_fp8(output, local_len, world_size, shard_heads, hidden_dims) return ((payload, scale),) + if quant_scheme == "int8": + payload, scale, _, _ = attn_pre_int8(output, local_len, world_size, shard_heads, hidden_dims) + return ((payload, scale),) return ((attn_pre(output, local_len, world_size, shard_heads, hidden_dims), None),) @staticmethod @@ -197,7 +214,11 @@ def unpack_attn(packed, output_dtype, hidden_dims): payload, scale = packed[0] if scale is None: return attn_post(payload) - return attn_post_fp8(payload, scale, payload.shape, scale.shape, output_dtype) + if payload.dtype == torch.float8_e4m3fn: + return attn_post_fp8(payload, scale, payload.shape, scale.shape, output_dtype) + if payload.dtype == torch.int8: + return attn_post_int8(payload, scale, payload.shape, scale.shape, output_dtype) + raise ValueError(f"Unsupported quantized attention payload dtype: {payload.dtype}.") def create_ulysses_prepost_backend(name): diff --git a/lightx2v/common/ops/attn/utils/seq_p.py b/lightx2v/common/ops/attn/utils/seq_p.py index 1bb609cfd..906c30018 100644 --- a/lightx2v/common/ops/attn/utils/seq_p.py +++ b/lightx2v/common/ops/attn/utils/seq_p.py @@ -2,6 +2,8 @@ from lightx2v.utils.quant_utils import dequant_fp8_vllm, quant_fp8_vllm +INT8_MAX = 127.0 + try: from sageattn3_sparse import dequant_fp4 as dequant_fp4_sage3 from sageattn3_sparse import quant_fp4 as quant_fp4_sage3 @@ -64,14 +66,43 @@ def split_main_aux_output(output, main_len, aux_len, aux_first): def validate_quant_scheme(quant_scheme): - if quant_scheme not in (None, "fp8", "fp4"): - raise ValueError(f"Unknown quant_scheme={quant_scheme!r}; expected None, 'fp8', or 'fp4'.") + if quant_scheme not in (None, "fp8", "fp4", "int8"): + raise ValueError(f"Unknown quant_scheme={quant_scheme!r}; expected None, 'fp8', 'fp4', or 'int8'.") if quant_scheme != "fp4": return if quant_fp4_sage3 is None or dequant_fp4_sage3 is None: raise ImportError("sageattn3_sparse quant_fp4/dequant_fp4 is required for sequence-parallel FP4 communication.") +def quant_int8_per_row(tensor): + """Symmetrically quantize the last dimension to INT8 with FP32 row scales.""" + if not tensor.is_floating_point(): + raise ValueError(f"INT8 communication quantization requires a floating-point tensor, got {tensor.dtype}.") + + shape = tensor.shape + hidden_dims = shape[-1] + rows = tensor.reshape(-1, hidden_dims).float() + amax = rows.abs().amax(dim=-1, keepdim=True) + scale = torch.where(amax > 0, amax / INT8_MAX, torch.ones_like(amax)) + + scaled = rows / scale + rounded = torch.round(scaled) + payload = rounded.clamp(-INT8_MAX, INT8_MAX).to(torch.int8) + return payload.reshape(shape).contiguous(), scale.reshape(*shape[:-1], 1).contiguous() + + +def dequant_int8_per_row(payload, scale, output_dtype): + """Dequantize an INT8 communication payload using FP32 row scales.""" + if payload.dtype != torch.int8: + raise ValueError(f"INT8 communication payload must have dtype torch.int8, got {payload.dtype}.") + if scale.dtype != torch.float32: + raise ValueError(f"INT8 communication scale must have dtype torch.float32, got {scale.dtype}.") + expected_scale_shape = (*payload.shape[:-1], 1) + if tuple(scale.shape) != expected_scale_shape: + raise ValueError(f"INT8 communication scale shape must be {expected_scale_shape}, got {tuple(scale.shape)}.") + return (payload.float() * scale).to(output_dtype) + + def pack_seq_p_tensor(tensor, quant_scheme): tensor = tensor.contiguous() validate_quant_scheme(quant_scheme) @@ -83,6 +114,8 @@ def pack_seq_p_tensor(tensor, quant_scheme): if quant_scheme == "fp8": payload, scale = quant_fp8_vllm(tensor.reshape(-1, hidden_dims)) return payload.reshape(shape).contiguous(), scale.reshape(*shape[:-1], 1).contiguous() + if quant_scheme == "int8": + return quant_int8_per_row(tensor) if hidden_dims % 16 != 0: raise ValueError(f"Sequence-parallel FP4 communication requires hidden_dims divisible by 16, got {hidden_dims}.") @@ -96,6 +129,8 @@ def unpack_seq_p_tensor(packed, output_dtype, hidden_dims): return payload if payload.dtype == torch.float8_e4m3fn: return dequant_fp8_vllm(payload, scale, output_dtype) + if payload.dtype == torch.int8 and payload.shape[-1] == hidden_dims and tuple(scale.shape) == (*payload.shape[:-1], 1): + return dequant_int8_per_row(payload, scale, output_dtype) output_shape = (*payload.shape[:-1], hidden_dims) return dequant_fp4_sage3( payload.reshape(1, 1, -1, hidden_dims // 2), diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py index c6abe71d8..06aca9055 100755 --- a/lightx2v/models/networks/wan/infer/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/transformer_infer.py @@ -62,8 +62,8 @@ def __init__(self, config): self.seq_p_configured_quant_scheme = parallel_config.get("seq_p_quant_scheme") self.seq_p_quant_scheme = self.seq_p_configured_quant_scheme - if self.seq_p_quant_scheme is not None and self.seq_p_quant_scheme not in ("fp8", "fp4"): - raise ValueError(f"Unknown seq_p_quant_scheme={self.seq_p_quant_scheme!r}; expected None, 'fp8', or 'fp4'.") + if self.seq_p_quant_scheme is not None and self.seq_p_quant_scheme not in ("fp8", "fp4", "int8"): + raise ValueError(f"Unknown seq_p_quant_scheme={self.seq_p_quant_scheme!r}; expected None, 'fp8', 'fp4', or 'int8'.") if self.seq_p_quant_scheme is not None and legacy_quant_scheme is not None and self.seq_p_quant_scheme != legacy_quant_scheme: raise ValueError("seq_p_quant_scheme conflicts with legacy seq_p_fp8_comm/seq_p_fp4_comm settings.") if self.seq_p_quant_scheme is None: From 47e699b2d36f32562fc8700d486498ca42124ea1 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 19:26:44 +0000 Subject: [PATCH 2/6] feat: support INT8 sparse attntion optimization of Wan22 --- ..._i2v_distill_int8_dynamic_sparse_attn.json | 47 ++ ..._int8_dynamic_sparse_attn_sp_parallel.json | 56 ++ ..._t2v_distill_int8_dynamic_sparse_attn.json | 47 ++ ..._int8_dynamic_sparse_attn_sp_parallel.json | 56 ++ .../wan/weights/transformer_weights.py | 4 + .../ops/attn/ascend_npu/__init__.py | 1 + .../ascend_npu/npu_dynamic_sparse_attn.py | 497 ++++++++++++++++++ .../attn/ascend_npu/npu_rainfusion_attn.py | 6 +- ...un_wan22_moe_i2v_extreme_dynamic_sparse.sh | 26 + ..._i2v_extreme_dynamic_sparse_sp_parallel.sh | 29 + ...un_wan22_moe_t2v_extreme_dynamic_sparse.sh | 25 + ..._t2v_extreme_dynamic_sparse_sp_parallel.sh | 28 + tools/convert/quant_adapter.py | 352 +++++++++++-- 13 files changed, 1122 insertions(+), 52 deletions(-) create mode 100644 configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json create mode 100644 configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json create mode 100644 configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json create mode 100644 configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json create mode 100644 lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py create mode 100755 scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh create mode 100755 scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh create mode 100755 scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh create mode 100755 scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh mode change 100755 => 100644 tools/convert/quant_adapter.py diff --git a/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json b/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json new file mode 100644 index 000000000..f12349f7e --- /dev/null +++ b/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json @@ -0,0 +1,47 @@ +{ + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "npu_dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "q_block_size": 128, + "kv_block_size": 128, + "backend": "mindie_rf_v2", + "allow_v2_fallback": true, + "v3_inner_precise": 0, + "v2_inner_precise": 0 + }, + "cross_attn_1_type": "npu_flash_attn", + "cross_attn_2_type": "npu_flash_attn", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "int8-npu", + "mxfp8_fuse_enable": false, + "modulate_type": "torch", + "rope_type": "npu_rope", + "layer_norm_type": "npu_layer_norm", + "rms_norm_type": "npu_rms_norm", + "high_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-I2V-A14B_INT8_Sparse_high.safetensors", + "low_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-I2V-A14B_INT8_Sparse_low.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null +} diff --git a/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json b/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json new file mode 100644 index 000000000..6fb28e950 --- /dev/null +++ b/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json @@ -0,0 +1,56 @@ +{ + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "npu_dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "q_block_size": 128, + "kv_block_size": 128, + "backend": "mindie_rf_v2", + "allow_v2_fallback": true, + "v3_inner_precise": 0, + "v2_inner_precise": 0 + }, + "cross_attn_1_type": "npu_flash_attn", + "cross_attn_2_type": "npu_flash_attn", + "sample_guide_scale": [ + 3.5, + 3.5 + ], + "sample_shift": 7.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "int8-npu", + "mxfp8_fuse_enable": false, + "modulate_type": "torch", + "rope_type": "npu_rope", + "layer_norm_type": "npu_layer_norm", + "rms_norm_type": "npu_rms_norm", + "high_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-I2V-A14B_INT8_Sparse_high.safetensors", + "low_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-I2V-A14B_INT8_Sparse_low.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "parallel": { + "seq_p_size": 8, + "seq_p_quant_scheme": "int8", + "seq_p_head_parallel": true, + "seq_p_tensor_fusion": true, + "seq_p_attn_type": "ulysses", + "seq_p_prepost_backend": "torch", + "seq_p_a2a_backend": "torch" + } +} diff --git a/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json b/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json new file mode 100644 index 000000000..c370a0fb1 --- /dev/null +++ b/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json @@ -0,0 +1,47 @@ +{ + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "npu_dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "q_block_size": 128, + "kv_block_size": 128, + "backend": "mindie_rf_v2", + "allow_v2_fallback": true, + "v3_inner_precise": 0, + "v2_inner_precise": 0 + }, + "cross_attn_1_type": "npu_flash_attn", + "cross_attn_2_type": "npu_flash_attn", + "sample_guide_scale": [ + 4.0, + 3.0 + ], + "sample_shift": 5.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "int8-npu", + "mxfp8_fuse_enable": false, + "modulate_type": "torch", + "rope_type": "npu_rope", + "layer_norm_type": "npu_layer_norm", + "rms_norm_type": "npu_rms_norm", + "high_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-T2V-A14B_INT8_Sparse_high.safetensors", + "low_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-T2V-A14B_INT8_Sparse_low.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null +} diff --git a/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json b/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json new file mode 100644 index 000000000..3f81475a1 --- /dev/null +++ b/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json @@ -0,0 +1,56 @@ +{ + "infer_steps": 4, + "target_video_length": 81, + "text_len": 512, + "target_height": 720, + "target_width": 1280, + "self_attn_1_type": "npu_dynamic_sparse_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.9, + "q_block_size": 128, + "kv_block_size": 128, + "backend": "mindie_rf_v2", + "allow_v2_fallback": true, + "v3_inner_precise": 0, + "v2_inner_precise": 0 + }, + "cross_attn_1_type": "npu_flash_attn", + "cross_attn_2_type": "npu_flash_attn", + "sample_guide_scale": [ + 4.0, + 3.0 + ], + "sample_shift": 5.0, + "enable_cfg": false, + "cpu_offload": false, + "t5_cpu_offload": true, + "vae_cpu_offload": false, + "use_image_encoder": false, + "boundary_step_index": 2, + "denoising_step_list": [ + 1000, + 750, + 500, + 250 + ], + "dit_quantized": true, + "dit_quant_scheme": "int8-npu", + "mxfp8_fuse_enable": false, + "modulate_type": "torch", + "rope_type": "npu_rope", + "layer_norm_type": "npu_layer_norm", + "rms_norm_type": "npu_rms_norm", + "high_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-T2V-A14B_INT8_Sparse_high.safetensors", + "low_noise_quantized_ckpt": "Wan2.2-Int8-Sparse/Wan2.2-T2V-A14B_INT8_Sparse_low.safetensors", + "high_noise_original_ckpt": null, + "low_noise_original_ckpt": null, + "parallel": { + "seq_p_size": 8, + "seq_p_quant_scheme": "int8", + "seq_p_head_parallel": true, + "seq_p_tensor_fusion": true, + "seq_p_attn_type": "ulysses", + "seq_p_prepost_backend": "torch", + "seq_p_a2a_backend": "torch" + } +} diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py index 05d06a7ec..0283e67ab 100755 --- a/lightx2v/models/networks/wan/weights/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/transformer_weights.py @@ -517,6 +517,10 @@ def __init__( if "operator" in dynamic_sparse_config: attention_weights_cls.operator = dynamic_sparse_config["operator"] + # NPU dynamic sparse attention setting + if self.config["self_attn_1_type"] == "npu_dynamic_sparse_attn": + attention_weights_cls.configure(self.config.get("dynamic_sparse_attn_setting", {})) + # spas_sage_attn2 setting if self.config["self_attn_1_type"] == "sparge_attn": sparge_config = self.config.get("sparge_attn_setting", {}) diff --git a/lightx2v_platform/ops/attn/ascend_npu/__init__.py b/lightx2v_platform/ops/attn/ascend_npu/__init__.py index b04781439..6e12236be 100644 --- a/lightx2v_platform/ops/attn/ascend_npu/__init__.py +++ b/lightx2v_platform/ops/attn/ascend_npu/__init__.py @@ -1,2 +1,3 @@ +from .npu_dynamic_sparse_attn import * from .npu_flash_attn import * from .npu_rainfusion_attn import * diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py new file mode 100644 index 000000000..38f76d01c --- /dev/null +++ b/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py @@ -0,0 +1,497 @@ +import math +from functools import lru_cache + +import torch +from loguru import logger + +from lightx2v_platform.ops.attn.template import AttnWeightTemplate +from lightx2v_platform.registry_factory import PLATFORM_ATTN_WEIGHT_REGISTER + +_V3_BACKEND = "mindie_rf_v3" +_V2_BACKEND = "mindie_rf_v2" +# The current CANN 8.5.1 image can import the MindIE RF v3 Python wrapper, +# but libopapi.so does not export the aclnnBlockSparseAttention operator it +# calls. Remember the first confirmed capability failure so later layers and +# model instances do not repeatedly enter the unsupported v3 path. +_V3_RUNTIME_UNAVAILABLE = False + + +def block_mean(x, block_size): + """Mean-pool a BNSD tensor along S without padding the tail block.""" + if x.ndim != 4: + raise ValueError(f"block_mean expects a BNSD tensor, but got shape {tuple(x.shape)}.") + if not isinstance(block_size, int) or isinstance(block_size, bool) or block_size <= 0: + raise ValueError(f"block_size must be a positive integer, but got {block_size!r}.") + + sequence_length = x.shape[2] + if sequence_length == 0: + raise ValueError("block_mean does not support an empty sequence.") + + full_block_count, tail_size = divmod(sequence_length, block_size) + pooled = [] + if full_block_count: + full_blocks = x[:, :, : full_block_count * block_size, :] + full_blocks = full_blocks.reshape( + x.shape[0], + x.shape[1], + full_block_count, + block_size, + x.shape[3], + ) + pooled.append(full_blocks.mean(dim=3)) + if tail_size: + pooled.append(x[:, :, full_block_count * block_size :, :].mean(dim=2, keepdim=True)) + + return pooled[0] if len(pooled) == 1 else torch.cat(pooled, dim=2) + + +def _validate_qk(q, k, sparsity_ratio, q_block_size, kv_block_size): + if q.ndim != 4 or k.ndim != 4: + raise ValueError( + "build_dynamic_sparse_mask expects q and k in BNSD layout, " + f"but got q={tuple(q.shape)}, k={tuple(k.shape)}." + ) + if q.shape[0] != k.shape[0]: + raise ValueError(f"q and k batch sizes must match, but got {q.shape[0]} and {k.shape[0]}.") + if q.shape[-1] != k.shape[-1]: + raise ValueError(f"q and k head dimensions must match, but got {q.shape[-1]} and {k.shape[-1]}.") + if q.shape[1] % k.shape[1] != 0: + raise ValueError( + "The query head count must be divisible by the key/value head count for GQA, " + f"but got {q.shape[1]} and {k.shape[1]}." + ) + if q.device != k.device: + raise ValueError(f"q and k must be on the same device, but got {q.device} and {k.device}.") + if q.dtype != k.dtype: + raise ValueError(f"q and k must have the same dtype, but got {q.dtype} and {k.dtype}.") + if not q.is_floating_point() or not k.is_floating_point(): + raise TypeError(f"q and k must be floating-point tensors, but got {q.dtype} and {k.dtype}.") + if q.shape[2] == 0 or k.shape[2] == 0: + raise ValueError("q and k sequence lengths must be greater than zero.") + if not isinstance(sparsity_ratio, (int, float)) or isinstance(sparsity_ratio, bool): + raise TypeError(f"sparsity_ratio must be a number, but got {type(sparsity_ratio).__name__}.") + if not 0 <= sparsity_ratio < 1: + raise ValueError(f"sparsity_ratio must be in [0, 1), but got {sparsity_ratio}.") + for name, size in (("q_block_size", q_block_size), ("kv_block_size", kv_block_size)): + if not isinstance(size, int) or isinstance(size, bool) or size <= 0: + raise ValueError(f"{name} must be a positive integer, but got {size!r}.") + + +def build_dynamic_sparse_mask( + q, + k, + sparsity_ratio, + q_block_size=128, + kv_block_size=128, + return_topk_indices=False, +): + """Build an exact-TopK int8 block mask from current BNSD query and key. + + K is smoothed over its sequence dimension before Q and K are mean-pooled. + The returned tensor has shape ``[B, Hq, ceil(Sq/Bq), ceil(Sk/Bkv)]``. + """ + _validate_qk(q, k, sparsity_ratio, q_block_size, kv_block_size) + + smoothed_k = k - k.mean(dim=2, keepdim=True) + pooled_q = block_mean(q, q_block_size) + pooled_k = block_mean(smoothed_k, kv_block_size) + + query_head_count = pooled_q.shape[1] + key_value_head_count = pooled_k.shape[1] + if query_head_count != key_value_head_count: + pooled_k = pooled_k.repeat_interleave(query_head_count // key_value_head_count, dim=1) + + block_scores = torch.matmul(pooled_q, pooled_k.transpose(-1, -2)) + + kv_block_count = block_scores.shape[-1] + keep_block_count = math.ceil(kv_block_count * (1 - sparsity_ratio)) + keep_block_count = min(max(keep_block_count, 1), kv_block_count) + topk_indices = torch.topk( + block_scores, + k=keep_block_count, + dim=-1, + largest=True, + sorted=False, + ).indices + + block_mask = torch.zeros_like(block_scores, dtype=torch.int8) + block_mask.scatter_(-1, topk_indices, 1) + if return_topk_indices: + return block_mask, topk_indices + return block_mask + + +def topk_indices_to_v2_indices(topk_indices, kv_block_count): + """Convert exact TopK indices to MindIE v2 selectors without a full-K sort.""" + if topk_indices.ndim != 4: + raise ValueError( + "topk_indices_to_v2_indices expects B,H,Q,Kkeep indices, " + f"but got shape {tuple(topk_indices.shape)}." + ) + if topk_indices.shape[0] != 1: + raise ValueError( + "MindIE RainFusion v2 selector conversion currently supports batch size 1, " + f"but got {topk_indices.shape[0]}." + ) + if not isinstance(kv_block_count, int) or isinstance(kv_block_count, bool) or kv_block_count <= 0: + raise ValueError(f"kv_block_count must be a positive integer, but got {kv_block_count!r}.") + keep_block_count = topk_indices.shape[-1] + if keep_block_count == 0 or keep_block_count > kv_block_count: + raise ValueError( + "The number of TopK indices must be in [1, kv_block_count], " + f"but got {keep_block_count} and kv_block_count={kv_block_count}." + ) + + sorted_indices = torch.sort(topk_indices.to(torch.int64), dim=-1).values + padding_size = kv_block_count - keep_block_count + if padding_size: + padding = torch.full( + (*sorted_indices.shape[:-1], padding_size), + -1, + dtype=torch.int64, + device=sorted_indices.device, + ) + sorted_indices = torch.cat((sorted_indices, padding), dim=-1) + + select_idx = sorted_indices[0].transpose(0, 1).contiguous() + select_num_idx = torch.full( + select_idx.shape[:2], + keep_block_count, + dtype=torch.int64, + device=select_idx.device, + ) + return select_idx, select_num_idx + + +def binary_mask_to_v2_indices(block_mask): + """Convert a B=1 binary mask to the selector layout required by MindIE v2.""" + if block_mask.ndim != 4: + raise ValueError( + "binary_mask_to_v2_indices expects a BHQK mask, " + f"but got shape {tuple(block_mask.shape)}." + ) + if block_mask.shape[0] != 1: + raise ValueError( + "MindIE RainFusion v2 selector conversion currently supports batch size 1, " + f"but got {block_mask.shape[0]}." + ) + if block_mask.shape[-1] == 0: + raise ValueError("The block mask must contain at least one KV block.") + + selected = block_mask.to(torch.bool) + kv_block_count = selected.shape[-1] + kv_indices = torch.arange(kv_block_count, dtype=torch.int64, device=selected.device) + kv_indices = kv_indices.view(1, 1, 1, kv_block_count).expand_as(selected) + padded_indices = torch.where(selected, kv_indices, kv_block_count) + padded_indices = torch.sort(padded_indices, dim=-1).values + padded_indices = torch.where(padded_indices == kv_block_count, -1, padded_indices) + + select_idx = padded_indices[0].transpose(0, 1).contiguous() + select_num_idx = selected[0].transpose(0, 1).sum(dim=-1, dtype=torch.int64).contiguous() + return select_idx, select_num_idx + + +def _import_v3_backend(): + from mindiesd.layers.flash_attn.sparse_flash_attn_rf_v3 import rain_fusion_attention_v3 + + if not callable(rain_fusion_attention_v3): + raise ImportError("mindiesd RainFusion v3 does not expose rain_fusion_attention_v3.") + return rain_fusion_attention_v3 + + +def _import_v2_backend(): + from mindiesd.layers.flash_attn.sparse_flash_attn_rf_v2 import rain_fusion_attention + + if not callable(rain_fusion_attention): + raise ImportError("mindiesd RainFusion v2 does not expose rain_fusion_attention.") + return rain_fusion_attention + + +def _is_v3_capability_error(error): + """Match only the known CANN 8.5 missing-BlockSparseAttention failure.""" + message = str(error) + return isinstance(error, RuntimeError) and "aclnnBlockSparseAttention" in message and "not in libopapi.so" in message + + +def _disable_v3_runtime(error): + global _V3_RUNTIME_UNAVAILABLE + + if not _V3_RUNTIME_UNAVAILABLE: + logger.warning( + "MindIE RainFusion v3 cannot run because aclnnBlockSparseAttention " + f"is unavailable; falling back to RainFusion v2. Original error: {error}" + ) + _V3_RUNTIME_UNAVAILABLE = True + # Backend resolution is cached. Clear it so the next layer observes the + # process-wide capability result and resolves directly to RF v2. + _load_mindie_backend.cache_clear() + + +@lru_cache(maxsize=4) +def _load_mindie_backend(backend=_V2_BACKEND, allow_v2_fallback=True): + """Resolve MindIE sparse attention for the current CANN 8.5 runtime. + + RF v2 is the compatibility default because CANN 8.5.1 provides + aclnnRainFusionAttention but not aclnnBlockSparseAttention. RF v3 remains + available as an explicit opt-in for newer runtimes; it may fall back to v2 + only when its API or required ACLNN capability is unavailable. + """ + if backend in ("auto", "v3", _V3_BACKEND): + if _V3_RUNTIME_UNAVAILABLE: + if not allow_v2_fallback: + raise RuntimeError( + "MindIE RainFusion v3 was disabled because " + "aclnnBlockSparseAttention is unavailable and v2 fallback is disabled." + ) + try: + return _V2_BACKEND, _import_v2_backend() + except ImportError as v2_error: + raise RuntimeError( + "MindIE RainFusion v3 is unavailable at runtime and its v2 fallback " + "could not be imported." + ) from v2_error + try: + return _V3_BACKEND, _import_v3_backend() + except ImportError as v3_error: + if not allow_v2_fallback: + raise RuntimeError( + "MindIE RainFusion v3 is unavailable and v2 fallback is disabled." + ) from v3_error + logger.warning("MindIE RainFusion v3 API is unavailable; falling back to RainFusion v2.") + try: + return _V2_BACKEND, _import_v2_backend() + except ImportError as v2_error: + raise RuntimeError( + "Neither MindIE RainFusion v3 nor its v2 fallback is available. " + "Install a compatible MindIE-SD package." + ) from v2_error + if backend in ("v2", _V2_BACKEND): + try: + return _V2_BACKEND, _import_v2_backend() + except ImportError as v2_error: + raise RuntimeError( + "MindIE RainFusion v2 is unavailable. Install a compatible MindIE-SD package." + ) from v2_error + raise ValueError( + f"Unsupported NPU dynamic sparse attention backend {backend!r}; " + f"expected 'auto', '{_V3_BACKEND}', or '{_V2_BACKEND}'." + ) + + +def _validate_inputs(q, k, v): + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError( + "npu_dynamic_sparse_attn expects q, k, and v in [S,H,D] layout, " + f"but got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}." + ) + if k.shape != v.shape: + raise ValueError(f"k and v must have identical shapes, but got {tuple(k.shape)} and {tuple(v.shape)}.") + if q.shape[-1] != k.shape[-1]: + raise ValueError(f"q and k/v head dimensions must match, but got {q.shape[-1]} and {k.shape[-1]}.") + if q.shape[1] % k.shape[1] != 0: + raise ValueError( + "The query head count must be divisible by the key/value head count for GQA, " + f"but got {q.shape[1]} and {k.shape[1]}." + ) + if q.shape[0] == 0 or k.shape[0] == 0: + raise ValueError("q and k/v sequence lengths must be greater than zero.") + if q.device != k.device or q.device != v.device: + raise ValueError(f"q, k, and v must be on the same device, but got {q.device}, {k.device}, and {v.device}.") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError(f"q, k, and v must have the same dtype, but got {q.dtype}, {k.dtype}, and {v.dtype}.") + if not q.is_floating_point(): + raise TypeError(f"q, k, and v must be floating-point tensors, but got {q.dtype}.") + + +def _validate_cu_seqlens(cu_seqlens, sequence_length, name): + if cu_seqlens is None: + return + if torch.is_tensor(cu_seqlens): + if cu_seqlens.numel() != 2: + raise ValueError( + f"{name} must describe exactly one sequence and contain 2 elements, " + f"but got {cu_seqlens.numel()}." + ) + if cu_seqlens.device.type != "cpu": + return + values = cu_seqlens.detach().tolist() + else: + values = list(cu_seqlens) + if values != [0, sequence_length]: + raise ValueError( + f"{name} must describe exactly one sequence [0, {sequence_length}], " + f"but got {values}. Packed batches are not supported." + ) + + +@PLATFORM_ATTN_WEIGHT_REGISTER("npu_dynamic_sparse_attn") +class NpuDynamicSparseAttnWeight(AttnWeightTemplate): + sparsity_ratio = 0.9 + q_block_size = 128 + kv_block_size = 128 + # Keep unqualified configs on RF v2: the target CANN 8.5.1 runtime lacks + # the aclnnBlockSparseAttention symbol required by MindIE RF v3. + backend = _V2_BACKEND + allow_v2_fallback = True + v3_inner_precise = 0 + v2_inner_precise = 0 + + def __init__(self): + self.config = {} + cls = type(self) + self.sparsity_ratio = cls.sparsity_ratio + self.q_block_size = cls.q_block_size + self.kv_block_size = cls.kv_block_size + self.backend = cls.backend + self.allow_v2_fallback = cls.allow_v2_fallback + self.v3_inner_precise = cls.v3_inner_precise + self.v2_inner_precise = cls.v2_inner_precise + + @classmethod + def configure(cls, setting): + setting = setting or {} + q_block_size = setting.get("q_block_size", setting.get("block_size", cls.q_block_size)) + kv_block_size = setting.get("kv_block_size", setting.get("block_size", cls.kv_block_size)) + if q_block_size != 128 or kv_block_size != 128: + raise ValueError( + "MindIE NPU dynamic sparse attention currently requires " + f"q_block_size=kv_block_size=128, but got {q_block_size} and {kv_block_size}." + ) + + sparsity_ratio = setting.get("sparsity_ratio", setting.get("sparsity", cls.sparsity_ratio)) + if not isinstance(sparsity_ratio, (int, float)) or isinstance(sparsity_ratio, bool): + raise TypeError(f"sparsity_ratio must be a number, but got {type(sparsity_ratio).__name__}.") + if not 0 <= sparsity_ratio < 1: + raise ValueError(f"sparsity_ratio must be in [0, 1), but got {sparsity_ratio}.") + + backend = setting.get("backend", cls.backend) + if backend not in ("auto", "v3", "v2", _V3_BACKEND, _V2_BACKEND): + raise ValueError(f"Unsupported NPU dynamic sparse attention backend {backend!r}.") + allow_v2_fallback = setting.get("allow_v2_fallback", cls.allow_v2_fallback) + if not isinstance(allow_v2_fallback, bool): + raise TypeError(f"allow_v2_fallback must be bool, but got {type(allow_v2_fallback).__name__}.") + + v3_inner_precise = setting.get("v3_inner_precise", cls.v3_inner_precise) + v2_inner_precise = setting.get("v2_inner_precise", cls.v2_inner_precise) + for name, value in ( + ("v3_inner_precise", v3_inner_precise), + ("v2_inner_precise", v2_inner_precise), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be an integer, but got {type(value).__name__}.") + cls.sparsity_ratio = sparsity_ratio + cls.q_block_size = q_block_size + cls.kv_block_size = kv_block_size + cls.backend = backend + cls.allow_v2_fallback = allow_v2_fallback + cls.v3_inner_precise = v3_inner_precise + cls.v2_inner_precise = v2_inner_precise + + def apply( + self, + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + _validate_inputs(q, k, v) + query_length, query_head_count, head_dim = q.shape + kv_length, key_value_head_count, _ = k.shape + + if max_seqlen_q is not None and max_seqlen_q != query_length: + raise ValueError(f"max_seqlen_q must equal {query_length}, but got {max_seqlen_q}.") + if max_seqlen_kv is not None and max_seqlen_kv != kv_length: + raise ValueError(f"max_seqlen_kv must equal {kv_length}, but got {max_seqlen_kv}.") + _validate_cu_seqlens(cu_seqlens_q, query_length, "cu_seqlens_q") + _validate_cu_seqlens(cu_seqlens_kv, kv_length, "cu_seqlens_kv") + if kwargs.get("attn_mask") is not None: + raise NotImplementedError("npu_dynamic_sparse_attn does not support an additional attention mask.") + if kwargs.get("causal", False): + raise NotImplementedError("npu_dynamic_sparse_attn currently supports non-causal attention only.") + if kwargs.get("dropout", 0.0) not in (None, 0, 0.0): + raise NotImplementedError("npu_dynamic_sparse_attn does not support attention dropout.") + + softmax_scale = kwargs.get("softmax_scale") + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + + q_bnsd = q.transpose(0, 1).unsqueeze(0).contiguous() + k_bnsd = k.transpose(0, 1).unsqueeze(0).contiguous() + v_bnsd = v.transpose(0, 1).unsqueeze(0).contiguous() + backend_name, backend_fn = _load_mindie_backend(self.backend, self.allow_v2_fallback) + block_mask, topk_indices = build_dynamic_sparse_mask( + q_bnsd, + k_bnsd, + self.sparsity_ratio, + q_block_size=self.q_block_size, + kv_block_size=self.kv_block_size, + return_topk_indices=True, + ) + + if backend_name == _V3_BACKEND: + try: + out = backend_fn( + q_bnsd, + k_bnsd, + v_bnsd, + block_sparse_mask=block_mask, + scale=softmax_scale, + head_num=query_head_count, + num_key_value_heads=key_value_head_count, + input_layout="BNSD", + actual_seq_lengths=[query_length], + actual_seq_lengths_kv=[kv_length], + sparse_size=self.q_block_size, + inner_precise=kwargs.get("v3_inner_precise", self.v3_inner_precise), + ) + except RuntimeError as error: + # CANN 8.5.1 discovers the missing BSA symbol only when the v3 + # operator is first executed. Fall back for that exact + # capability error; propagate shape, OOM, numerical, and other + # kernel failures unchanged so real defects are not hidden. + if not self.allow_v2_fallback or not _is_v3_capability_error(error): + raise + _disable_v3_runtime(error) + backend_name, backend_fn = _load_mindie_backend(self.backend, self.allow_v2_fallback) + + if backend_name == _V2_BACKEND: + if query_head_count != key_value_head_count: + raise NotImplementedError( + "MindIE RainFusion v2 fallback does not support GQA in this integration; " + f"got {query_head_count} query heads and {key_value_head_count} KV heads." + ) + # RF v2 consumes padded block selectors instead of the binary + # block mask accepted by RF v3. The Top-K choice itself is shared. + select_idx, select_num_idx = topk_indices_to_v2_indices( + topk_indices, + kv_block_count=block_mask.shape[-1], + ) + out = backend_fn( + q_bnsd, + k_bnsd, + v_bnsd, + scale=softmax_scale, + head_num=query_head_count, + input_layout="BNSD", + select_idx=select_idx, + select_num_idx=select_num_idx, + blockshape=[self.q_block_size, self.kv_block_size], + actual_seq_lengths=[query_length], + actual_seq_lengths_kv=[kv_length], + inner_precise=kwargs.get("v2_inner_precise", self.v2_inner_precise), + ) + elif backend_name != _V3_BACKEND: + raise RuntimeError(f"Internal error: unrecognized MindIE backend {backend_name!r}.") + + if not torch.is_tensor(out): + raise TypeError(f"MindIE sparse attention must return a tensor, but got {type(out).__name__}.") + expected_shape = (1, query_head_count, query_length, head_dim) + if tuple(out.shape) != expected_shape: + raise RuntimeError( + f"MindIE sparse attention returned shape {tuple(out.shape)}, expected {expected_shape}." + ) + return out.squeeze(0).transpose(0, 1).contiguous().reshape(query_length, query_head_count * head_dim) diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py index f7bad13f1..a61a34c5e 100644 --- a/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py +++ b/lightx2v_platform/ops/attn/ascend_npu/npu_rainfusion_attn.py @@ -3,7 +3,11 @@ from mindiesd.layers.flash_attn.attention_forward import attention_forward _HAS_MINDIESD = True -except ImportError: +# Some CANN/MindIE combinations raise TypeError while probing NPU capabilities +# during this optional import. Treat that as an unavailable optional backend so +# importing the Ascend attention registry can continue; selecting RainFusion +# still fails explicitly in NpuRainfusionOperator.__init__. +except (ImportError, TypeError): mindiesd = None attention_forward = None _HAS_MINDIESD = False diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh new file mode 100755 index 000000000..7a4618e73 --- /dev/null +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" +model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-I2V-A14B}" +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +config_path="${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json" + +export PLATFORM=ascend_npu +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" +export PYTHONPATH="${PYTHONPATH:-}" + +source "${lightx2v_path}/scripts/base/base.sh" + +python -m lightx2v.infer \ + --model_cls wan2.2_moe_distill \ + --task i2v \ + --warmup \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ + --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ + --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_npu.mp4" diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh new file mode 100755 index 000000000..f036d7fa4 --- /dev/null +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" +model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-I2V-A14B}" +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. +# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. +config_path="${CONFIG_PATH:-${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json}" +save_result_path="${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4}" + +export PLATFORM=ascend_npu +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export PYTHONPATH="${PYTHONPATH:-}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun --nproc_per_node=8 -m lightx2v.infer \ + --model_cls wan2.2_moe_distill \ + --task i2v \ + --warmup \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ + --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ + --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ + --save_result_path "${save_result_path}" diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh new file mode 100755 index 000000000..b1bb78121 --- /dev/null +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" +model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-T2V-A14B}" +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +config_path="${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json" + +export PLATFORM=ascend_npu +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" +export PYTHONPATH="${PYTHONPATH:-}" + +source "${lightx2v_path}/scripts/base/base.sh" + +python -m lightx2v.infer \ + --model_cls wan2.2_moe_distill \ + --task t2v \ + --warmup \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_npu.mp4" diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh new file mode 100755 index 000000000..9dcfb3c9e --- /dev/null +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" +model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-T2V-A14B}" +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. +# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. +config_path="${CONFIG_PATH:-${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json}" +save_result_path="${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4}" + +export PLATFORM=ascend_npu +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export PYTHONPATH="${PYTHONPATH:-}" + +source "${lightx2v_path}/scripts/base/base.sh" + +torchrun --nproc_per_node=8 -m lightx2v.infer \ + --model_cls wan2.2_moe_distill \ + --task t2v \ + --warmup \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ + --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ + --save_result_path "${save_result_path}" diff --git a/tools/convert/quant_adapter.py b/tools/convert/quant_adapter.py old mode 100755 new mode 100644 index b8d91025e..fc7073c0e --- a/tools/convert/quant_adapter.py +++ b/tools/convert/quant_adapter.py @@ -1,79 +1,329 @@ import argparse +import gc +import os import sys +from collections import Counter from pathlib import Path import safetensors import torch from safetensors.torch import save_file -sys.path.append(str(Path(__file__).parent.parent.parent)) +NVFP4_BLOCK_SIZE = 16 +NVFP4_SCALE_TILE_M = 128 +NVFP4_SCALE_TILE_K = NVFP4_BLOCK_SIZE * 4 +E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) -quant_path = str(Path(__file__).parent / "quant") -if quant_path not in sys.path: - sys.path.insert(0, quant_path) -from quant import * # noqa: E402 +def _ceil_div(value, divisor): + return (value + divisor - 1) // divisor -from lightx2v.utils.quant_utils import FloatQuantizer # noqa: E402 +def _nvfp4_weight_keys(keys): + key_set = set(keys) + groups = {} + for weight_key in keys: + if not weight_key.endswith(".weight"): + continue + prefix = weight_key.removesuffix(".weight") + companions = { + "scale": f"{prefix}.weight_scale", + "input_global_scale": f"{prefix}.input_global_scale", + "alpha": f"{prefix}.alpha", + } + if all(key in key_set for key in companions.values()): + groups[weight_key] = companions + return groups -def main(): - # 获取脚本所在目录 + +def _decode_e2m1(packed): + if packed.dtype != torch.uint8 or packed.ndim != 2: + raise ValueError(f"NVFP4 weight must be a 2D uint8 tensor, got shape={tuple(packed.shape)}, dtype={packed.dtype}.") + + lookup = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device) + low = lookup[(packed & 0x0F).to(torch.int64)] + high = lookup[((packed >> 4) & 0x0F).to(torch.int64)] + return torch.stack((low, high), dim=-1).reshape(packed.shape[0], packed.shape[1] * 2) + + +def _unswizzle_nvfp4_scales(swizzled, rows, columns, block_size=NVFP4_BLOCK_SIZE): + if columns % block_size: + raise ValueError(f"NVFP4 input dimension must be divisible by {block_size}, got {columns}.") + + row_tiles = _ceil_div(rows, NVFP4_SCALE_TILE_M) + column_tiles = _ceil_div(columns, block_size * 4) + expected_numel = row_tiles * column_tiles * 32 * 4 * 4 + if swizzled.numel() != expected_numel: + raise ValueError(f"Unexpected NVFP4 scale layout: shape={tuple(swizzled.shape)}, numel={swizzled.numel()}, expected={expected_numel} for weight shape=({rows}, {columns}).") + + linear = swizzled.reshape(1, row_tiles, column_tiles, 32, 4, 4).permute(0, 1, 4, 3, 2, 5).reshape(row_tiles * NVFP4_SCALE_TILE_M, column_tiles * 4) + return linear[:rows, : columns // block_size].to(torch.float32) + + +@torch.inference_mode() +def dequantize_nvfp4_weight(packed, swizzled_scale, input_global_scale, alpha): + rows, packed_columns = packed.shape + columns = packed_columns * 2 + + input_global_scale = input_global_scale.to(torch.float32) + alpha = alpha.to(torch.float32) + if input_global_scale.numel() != 1 or alpha.numel() != 1: + raise ValueError("NVFP4 input_global_scale and alpha must both be scalar tensors.") + + # The checkpoint stores alpha = 1 / (input_global_scale * weight_global_scale). + weight_global_scale = torch.reciprocal(input_global_scale * alpha) + if not torch.isfinite(weight_global_scale).item() or weight_global_scale.item() <= 0: + raise ValueError(f"Invalid derived NVFP4 weight_global_scale={weight_global_scale.item()}.") + + values = _decode_e2m1(packed) + scales = _unswizzle_nvfp4_scales(swizzled_scale, rows, columns) + values = values.reshape(rows, columns // NVFP4_BLOCK_SIZE, NVFP4_BLOCK_SIZE) + return (values * (scales / weight_global_scale).unsqueeze(-1)).reshape(rows, columns) + + +@torch.inference_mode() +def quantize_int8_per_output_channel(weight): + if weight.dtype != torch.float32 or weight.ndim != 2: + raise ValueError(f"INT8 source weight must be a 2D float32 tensor, got shape={tuple(weight.shape)}, dtype={weight.dtype}.") + max_value = weight.abs().amax(dim=1, keepdim=True).clamp(min=1e-5) + scale = max_value / 127.0 + quantized = torch.clamp(torch.round(weight / scale), -128, 127).to(torch.int8) + return quantized.contiguous(), scale.to(torch.bfloat16).contiguous() + + +def _target_filename(source_name): + if "NVFP4" in source_name: + return source_name.replace("NVFP4", "INT8") + if "nvfp4" in source_name: + return source_name.replace("nvfp4", "int8") + path = Path(source_name) + return f"{path.stem}_int8{path.suffix}" + + +def _validate_int8_file(path, expected_quantized, expected_keys): + dtype_counts = Counter() + quantized_count = 0 + with safetensors.safe_open(path, framework="pt", device="cpu") as source: + keys = list(source.keys()) + if len(keys) != expected_keys: + raise ValueError(f"{path}: expected {expected_keys} tensors, found {len(keys)}.") + if any(key.endswith((".alpha", ".input_global_scale")) for key in keys): + raise ValueError(f"{path}: found stale NVFP4 alpha/input_global_scale tensors.") + + key_set = set(keys) + for key in keys: + tensor_slice = source.get_slice(key) + dtype_counts[tensor_slice.get_dtype()] += 1 + if tensor_slice.get_dtype() != "I8": + continue + quantized_count += 1 + scale_key = f"{key.removesuffix('.weight')}.weight_scale" + if not key.endswith(".weight") or scale_key not in key_set: + raise ValueError(f"{path}: INT8 tensor {key} has no matching weight_scale.") + scale_slice = source.get_slice(scale_key) + weight_shape = tensor_slice.get_shape() + if scale_slice.get_dtype() != "BF16" or scale_slice.get_shape() != [weight_shape[0], 1]: + raise ValueError(f"{path}: invalid scale for {key}: dtype={scale_slice.get_dtype()}, shape={scale_slice.get_shape()}.") + + if quantized_count != expected_quantized: + raise ValueError(f"{path}: expected {expected_quantized} INT8 weights, found {quantized_count}.") + return dtype_counts + + +def convert_wan_nvfp4_file(source_path, output_path, non_quant_dtype=torch.float32, overwrite=False): + if output_path.exists() and not overwrite: + raise FileExistsError(f"Output already exists: {output_path}. Pass --overwrite to replace it.") + output_path.parent.mkdir(parents=True, exist_ok=True) + + temporary_path = output_path.with_name(f".{output_path.name}.tmp.{os.getpid()}") + converted = {} + try: + with safetensors.safe_open(source_path, framework="pt", device="cpu") as source: + keys = list(source.keys()) + groups = _nvfp4_weight_keys(keys) + if not groups: + raise ValueError(f"{source_path} does not contain LightX2V NVFP4 weight groups.") + + skipped_keys = set() + for companions in groups.values(): + skipped_keys.update(companions.values()) + + total = len(groups) + converted_index = 0 + for key in keys: + if key in skipped_keys: + continue + if key not in groups: + converted[key] = source.get_tensor(key).to(non_quant_dtype).contiguous() + continue + + converted_index += 1 + companions = groups[key] + dequantized = dequantize_nvfp4_weight( + source.get_tensor(key), + source.get_tensor(companions["scale"]), + source.get_tensor(companions["input_global_scale"]), + source.get_tensor(companions["alpha"]), + ) + int8_weight, int8_scale = quantize_int8_per_output_channel(dequantized) + converted[key] = int8_weight + converted[companions["scale"]] = int8_scale + del dequantized, int8_weight, int8_scale + + if converted_index == 1 or converted_index % 25 == 0 or converted_index == total: + print(f"[{source_path.name}] converted {converted_index}/{total} NVFP4 weights", flush=True) + if converted_index % 25 == 0: + gc.collect() + + expected_keys = len(keys) - 2 * total + if len(converted) != expected_keys: + raise ValueError(f"Converted tensor count mismatch: expected {expected_keys}, got {len(converted)}.") + + metadata = { + "source_file": source_path.name, + "source_format": "nvfp4", + "target_format": "int8-npu", + "conversion": "nvfp4-dequant-then-int8-per-output-channel", + } + save_file(converted, str(temporary_path), metadata=metadata) + + if output_path.exists(): + output_path.unlink() + os.replace(temporary_path, output_path) + dtype_counts = _validate_int8_file(output_path, total, expected_keys) + print(f"Saved {output_path} ({output_path.stat().st_size / 1e9:.2f} GB), dtypes={dict(dtype_counts)}", flush=True) + return output_path + finally: + converted.clear() + gc.collect() + if temporary_path.exists(): + temporary_path.unlink() + + +def convert_wan_nvfp4_path(model_path, output_path, non_quant_dtype=torch.float32, overwrite=False): + if model_path.is_file(): + target = output_path + if output_path.exists() and output_path.is_dir(): + target = output_path / _target_filename(model_path.name) + return [convert_wan_nvfp4_file(model_path, target, non_quant_dtype, overwrite)] + + if not model_path.is_dir(): + raise FileNotFoundError(f"Input path does not exist: {model_path}") + if output_path.exists() and not output_path.is_dir(): + raise ValueError(f"Directory input requires a directory output, got file: {output_path}") + + sources = sorted(path for path in model_path.glob("*.safetensors") if "_comfy" not in path.stem) + if not sources: + raise FileNotFoundError(f"No non-ComfyUI safetensors files found under {model_path}.") + + comfy_count = sum(1 for path in model_path.glob("*_comfy.safetensors")) + if comfy_count: + print(f"Skipping {comfy_count} ComfyUI checkpoints because their packed weight layout differs.", flush=True) + + outputs = [] + for source_path in sources: + target = output_path / _target_filename(source_path.name) + outputs.append(convert_wan_nvfp4_file(source_path, target, non_quant_dtype, overwrite)) + return outputs + + +def convert_audio_adapter_fp8(model_path, output_path): + # Keep the original SekoTalk adapter conversion available without importing + # CUDA-specific quantization code in the Wan NVFP4 CPU conversion path. + project_root = Path(__file__).parent.parent.parent + if str(project_root) not in sys.path: + sys.path.append(str(project_root)) + quant_path = str(Path(__file__).parent / "quant") + if quant_path not in sys.path: + sys.path.insert(0, quant_path) + + from lightx2v.utils.quant_utils import FloatQuantizer + + output_path.parent.mkdir(parents=True, exist_ok=True) + state_dict = {} + with safetensors.safe_open(model_path, framework="pt", device="cpu") as source: + for key in source.keys(): + state_dict[key] = source.get_tensor(key) + + converted = {} + for key, tensor in state_dict.items(): + if key.startswith("ca") and ".to" in key and "weight" in key: + print(f"Converting {key} to FP8, dtype: {tensor.dtype}") + weight = tensor.to(torch.float32).cuda() + quantizer = FloatQuantizer("e4m3", True, "per_channel") + weight, weight_scale, _ = quantizer.real_quant_tensor(weight) + converted[key] = weight.to(torch.float8_e4m3fn).cpu() + converted[f"{key}_scale"] = weight_scale.to(torch.float32).cpu() + else: + print(f"Converting {key} to BF16, dtype: {tensor.dtype}") + converted[key] = tensor.to(torch.bfloat16) + + save_file(converted, str(output_path)) + print(f"Quantized adapter saved to: {output_path}") + + +def parse_args(): script_dir = Path(__file__).parent project_root = script_dir.parent.parent - - parser = argparse.ArgumentParser(description="Quantize audio adapter model to FP8") + parser = argparse.ArgumentParser(description="Convert a SekoTalk adapter to FP8 or LightX2V Wan2.2 NVFP4 checkpoints to INT8-NPU.") parser.add_argument( "--model_path", - type=str, - default=str(project_root / "models" / "SekoTalk-Distill" / "audio_adapter_model.safetensors"), - help="Path to input model file", + type=Path, + default=project_root / "models" / "SekoTalk-Distill" / "audio_adapter_model.safetensors", + help="Input adapter file, Wan NVFP4 file, or Wan NVFP4 directory.", ) parser.add_argument( "--output_path", - type=str, - default=str(project_root / "models" / "SekoTalk-Distill-fp8" / "audio_adapter_model_fp8.safetensors"), - help="Path to output quantized model file", + type=Path, + default=project_root / "models" / "SekoTalk-Distill-fp8" / "audio_adapter_model_fp8.safetensors", + help="Output file or directory.", ) - args = parser.parse_args() - - model_path = Path(args.model_path) - output_path = Path(args.output_path) - - output_path.parent.mkdir(parents=True, exist_ok=True) + parser.add_argument( + "--mode", + choices=("auto", "adapter-fp8", "wan-nvfp4-to-int8"), + default="auto", + help="Conversion mode. Auto selects Wan conversion for directory inputs and adapter conversion for file inputs.", + ) + parser.add_argument( + "--non_quant_dtype", + choices=("float32", "bfloat16"), + default="float32", + help="Dtype for non-quantized tensors in Wan INT8 checkpoints. float32 matches existing LightX2V INT8-NPU files.", + ) + parser.add_argument("--overwrite", action="store_true", help="Replace existing output files.") + return parser.parse_args() - state_dict = {} - with safetensors.safe_open(model_path, framework="pt", device="cpu") as f: - for key in f.keys(): - state_dict[key] = f.get_tensor(key) - new_state_dict = {} +def main(): + args = parse_args() + mode = args.mode + if mode == "auto": + mode = "wan-nvfp4-to-int8" if args.model_path.is_dir() else "adapter-fp8" - for key in state_dict.keys(): - if key.startswith("ca") and ".to" in key and "weight" in key: - print(f"Converting {key} to FP8, dtype: {state_dict[key].dtype}") - - ## fp8 - weight = state_dict[key].to(torch.float32).cuda() - w_quantizer = FloatQuantizer("e4m3", True, "per_channel") - weight, weight_scale, _ = w_quantizer.real_quant_tensor(weight) - weight = weight.to(torch.float8_e4m3fn) - weight_scale = weight_scale.to(torch.float32) - - ## QuantWeightMxFP4, QuantWeightMxFP6, QuantWeightMxFP8 for mxfp4,mxfp6,mxfp8 - # weight = state_dict[key].to(torch.bfloat16).cuda() - # quantizer = QuantWeightMxFP4(weight) - # weight, weight_scale, _ = quantizer.weight_quant_func(weight) - - new_state_dict[key] = weight.cpu() - new_state_dict[key + "_scale"] = weight_scale.cpu() - else: - # 不匹配的权重转换为BF16 - print(f"Converting {key} to BF16, dtype: {state_dict[key].dtype}") - new_state_dict[key] = state_dict[key].to(torch.bfloat16) + if mode == "adapter-fp8": + convert_audio_adapter_fp8(args.model_path, args.output_path) + return - save_file(new_state_dict, str(output_path)) - print(f"Quantized model saved to: {output_path}") + non_quant_dtype = torch.float32 if args.non_quant_dtype == "float32" else torch.bfloat16 + outputs = convert_wan_nvfp4_path(args.model_path, args.output_path, non_quant_dtype, args.overwrite) + print(f"Converted {len(outputs)} Wan2.2 checkpoint(s).", flush=True) if __name__ == "__main__": From 51d81f9b636dd8d1f14aaad86ff92a90780951da Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 19:42:36 +0000 Subject: [PATCH 3/6] chore: clean up Ascend sparse integration Align the Ascend extreme scripts with the Wan22 script style, restore quant_adapter.py to its pre-sparse implementation, and include the pre-commit formatting changes. --- .../ascend_npu/npu_dynamic_sparse_attn.py | 102 ++--- ...un_wan22_moe_i2v_extreme_dynamic_sparse.sh | 38 +- ..._i2v_extreme_dynamic_sparse_sp_parallel.sh | 43 +-- ...un_wan22_moe_t2v_extreme_dynamic_sparse.sh | 36 +- ..._t2v_extreme_dynamic_sparse_sp_parallel.sh | 41 +- tools/convert/quant_adapter.py | 352 +++--------------- 6 files changed, 146 insertions(+), 466 deletions(-) mode change 100644 => 100755 tools/convert/quant_adapter.py diff --git a/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py b/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py index 38f76d01c..44f67dd75 100644 --- a/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py +++ b/lightx2v_platform/ops/attn/ascend_npu/npu_dynamic_sparse_attn.py @@ -47,19 +47,13 @@ def block_mean(x, block_size): def _validate_qk(q, k, sparsity_ratio, q_block_size, kv_block_size): if q.ndim != 4 or k.ndim != 4: - raise ValueError( - "build_dynamic_sparse_mask expects q and k in BNSD layout, " - f"but got q={tuple(q.shape)}, k={tuple(k.shape)}." - ) + raise ValueError(f"build_dynamic_sparse_mask expects q and k in BNSD layout, but got q={tuple(q.shape)}, k={tuple(k.shape)}.") if q.shape[0] != k.shape[0]: raise ValueError(f"q and k batch sizes must match, but got {q.shape[0]} and {k.shape[0]}.") if q.shape[-1] != k.shape[-1]: raise ValueError(f"q and k head dimensions must match, but got {q.shape[-1]} and {k.shape[-1]}.") if q.shape[1] % k.shape[1] != 0: - raise ValueError( - "The query head count must be divisible by the key/value head count for GQA, " - f"but got {q.shape[1]} and {k.shape[1]}." - ) + raise ValueError(f"The query head count must be divisible by the key/value head count for GQA, but got {q.shape[1]} and {k.shape[1]}.") if q.device != k.device: raise ValueError(f"q and k must be on the same device, but got {q.device} and {k.device}.") if q.dtype != k.dtype: @@ -124,23 +118,14 @@ def build_dynamic_sparse_mask( def topk_indices_to_v2_indices(topk_indices, kv_block_count): """Convert exact TopK indices to MindIE v2 selectors without a full-K sort.""" if topk_indices.ndim != 4: - raise ValueError( - "topk_indices_to_v2_indices expects B,H,Q,Kkeep indices, " - f"but got shape {tuple(topk_indices.shape)}." - ) + raise ValueError(f"topk_indices_to_v2_indices expects B,H,Q,Kkeep indices, but got shape {tuple(topk_indices.shape)}.") if topk_indices.shape[0] != 1: - raise ValueError( - "MindIE RainFusion v2 selector conversion currently supports batch size 1, " - f"but got {topk_indices.shape[0]}." - ) + raise ValueError(f"MindIE RainFusion v2 selector conversion currently supports batch size 1, but got {topk_indices.shape[0]}.") if not isinstance(kv_block_count, int) or isinstance(kv_block_count, bool) or kv_block_count <= 0: raise ValueError(f"kv_block_count must be a positive integer, but got {kv_block_count!r}.") keep_block_count = topk_indices.shape[-1] if keep_block_count == 0 or keep_block_count > kv_block_count: - raise ValueError( - "The number of TopK indices must be in [1, kv_block_count], " - f"but got {keep_block_count} and kv_block_count={kv_block_count}." - ) + raise ValueError(f"The number of TopK indices must be in [1, kv_block_count], but got {keep_block_count} and kv_block_count={kv_block_count}.") sorted_indices = torch.sort(topk_indices.to(torch.int64), dim=-1).values padding_size = kv_block_count - keep_block_count @@ -166,15 +151,9 @@ def topk_indices_to_v2_indices(topk_indices, kv_block_count): def binary_mask_to_v2_indices(block_mask): """Convert a B=1 binary mask to the selector layout required by MindIE v2.""" if block_mask.ndim != 4: - raise ValueError( - "binary_mask_to_v2_indices expects a BHQK mask, " - f"but got shape {tuple(block_mask.shape)}." - ) + raise ValueError(f"binary_mask_to_v2_indices expects a BHQK mask, but got shape {tuple(block_mask.shape)}.") if block_mask.shape[0] != 1: - raise ValueError( - "MindIE RainFusion v2 selector conversion currently supports batch size 1, " - f"but got {block_mask.shape[0]}." - ) + raise ValueError(f"MindIE RainFusion v2 selector conversion currently supports batch size 1, but got {block_mask.shape[0]}.") if block_mask.shape[-1] == 0: raise ValueError("The block mask must contain at least one KV block.") @@ -217,10 +196,7 @@ def _disable_v3_runtime(error): global _V3_RUNTIME_UNAVAILABLE if not _V3_RUNTIME_UNAVAILABLE: - logger.warning( - "MindIE RainFusion v3 cannot run because aclnnBlockSparseAttention " - f"is unavailable; falling back to RainFusion v2. Original error: {error}" - ) + logger.warning(f"MindIE RainFusion v3 cannot run because aclnnBlockSparseAttention is unavailable; falling back to RainFusion v2. Original error: {error}") _V3_RUNTIME_UNAVAILABLE = True # Backend resolution is cached. Clear it so the next layer observes the # process-wide capability result and resolves directly to RF v2. @@ -239,60 +215,38 @@ def _load_mindie_backend(backend=_V2_BACKEND, allow_v2_fallback=True): if backend in ("auto", "v3", _V3_BACKEND): if _V3_RUNTIME_UNAVAILABLE: if not allow_v2_fallback: - raise RuntimeError( - "MindIE RainFusion v3 was disabled because " - "aclnnBlockSparseAttention is unavailable and v2 fallback is disabled." - ) + raise RuntimeError("MindIE RainFusion v3 was disabled because aclnnBlockSparseAttention is unavailable and v2 fallback is disabled.") try: return _V2_BACKEND, _import_v2_backend() except ImportError as v2_error: - raise RuntimeError( - "MindIE RainFusion v3 is unavailable at runtime and its v2 fallback " - "could not be imported." - ) from v2_error + raise RuntimeError("MindIE RainFusion v3 is unavailable at runtime and its v2 fallback could not be imported.") from v2_error try: return _V3_BACKEND, _import_v3_backend() except ImportError as v3_error: if not allow_v2_fallback: - raise RuntimeError( - "MindIE RainFusion v3 is unavailable and v2 fallback is disabled." - ) from v3_error + raise RuntimeError("MindIE RainFusion v3 is unavailable and v2 fallback is disabled.") from v3_error logger.warning("MindIE RainFusion v3 API is unavailable; falling back to RainFusion v2.") try: return _V2_BACKEND, _import_v2_backend() except ImportError as v2_error: - raise RuntimeError( - "Neither MindIE RainFusion v3 nor its v2 fallback is available. " - "Install a compatible MindIE-SD package." - ) from v2_error + raise RuntimeError("Neither MindIE RainFusion v3 nor its v2 fallback is available. Install a compatible MindIE-SD package.") from v2_error if backend in ("v2", _V2_BACKEND): try: return _V2_BACKEND, _import_v2_backend() except ImportError as v2_error: - raise RuntimeError( - "MindIE RainFusion v2 is unavailable. Install a compatible MindIE-SD package." - ) from v2_error - raise ValueError( - f"Unsupported NPU dynamic sparse attention backend {backend!r}; " - f"expected 'auto', '{_V3_BACKEND}', or '{_V2_BACKEND}'." - ) + raise RuntimeError("MindIE RainFusion v2 is unavailable. Install a compatible MindIE-SD package.") from v2_error + raise ValueError(f"Unsupported NPU dynamic sparse attention backend {backend!r}; expected 'auto', '{_V3_BACKEND}', or '{_V2_BACKEND}'.") def _validate_inputs(q, k, v): if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: - raise ValueError( - "npu_dynamic_sparse_attn expects q, k, and v in [S,H,D] layout, " - f"but got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}." - ) + raise ValueError(f"npu_dynamic_sparse_attn expects q, k, and v in [S,H,D] layout, but got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}.") if k.shape != v.shape: raise ValueError(f"k and v must have identical shapes, but got {tuple(k.shape)} and {tuple(v.shape)}.") if q.shape[-1] != k.shape[-1]: raise ValueError(f"q and k/v head dimensions must match, but got {q.shape[-1]} and {k.shape[-1]}.") if q.shape[1] % k.shape[1] != 0: - raise ValueError( - "The query head count must be divisible by the key/value head count for GQA, " - f"but got {q.shape[1]} and {k.shape[1]}." - ) + raise ValueError(f"The query head count must be divisible by the key/value head count for GQA, but got {q.shape[1]} and {k.shape[1]}.") if q.shape[0] == 0 or k.shape[0] == 0: raise ValueError("q and k/v sequence lengths must be greater than zero.") if q.device != k.device or q.device != v.device: @@ -308,20 +262,14 @@ def _validate_cu_seqlens(cu_seqlens, sequence_length, name): return if torch.is_tensor(cu_seqlens): if cu_seqlens.numel() != 2: - raise ValueError( - f"{name} must describe exactly one sequence and contain 2 elements, " - f"but got {cu_seqlens.numel()}." - ) + raise ValueError(f"{name} must describe exactly one sequence and contain 2 elements, but got {cu_seqlens.numel()}.") if cu_seqlens.device.type != "cpu": return values = cu_seqlens.detach().tolist() else: values = list(cu_seqlens) if values != [0, sequence_length]: - raise ValueError( - f"{name} must describe exactly one sequence [0, {sequence_length}], " - f"but got {values}. Packed batches are not supported." - ) + raise ValueError(f"{name} must describe exactly one sequence [0, {sequence_length}], but got {values}. Packed batches are not supported.") @PLATFORM_ATTN_WEIGHT_REGISTER("npu_dynamic_sparse_attn") @@ -353,10 +301,7 @@ def configure(cls, setting): q_block_size = setting.get("q_block_size", setting.get("block_size", cls.q_block_size)) kv_block_size = setting.get("kv_block_size", setting.get("block_size", cls.kv_block_size)) if q_block_size != 128 or kv_block_size != 128: - raise ValueError( - "MindIE NPU dynamic sparse attention currently requires " - f"q_block_size=kv_block_size=128, but got {q_block_size} and {kv_block_size}." - ) + raise ValueError(f"MindIE NPU dynamic sparse attention currently requires q_block_size=kv_block_size=128, but got {q_block_size} and {kv_block_size}.") sparsity_ratio = setting.get("sparsity_ratio", setting.get("sparsity", cls.sparsity_ratio)) if not isinstance(sparsity_ratio, (int, float)) or isinstance(sparsity_ratio, bool): @@ -460,10 +405,7 @@ def apply( if backend_name == _V2_BACKEND: if query_head_count != key_value_head_count: - raise NotImplementedError( - "MindIE RainFusion v2 fallback does not support GQA in this integration; " - f"got {query_head_count} query heads and {key_value_head_count} KV heads." - ) + raise NotImplementedError(f"MindIE RainFusion v2 fallback does not support GQA in this integration; got {query_head_count} query heads and {key_value_head_count} KV heads.") # RF v2 consumes padded block selectors instead of the binary # block mask accepted by RF v3. The Top-K choice itself is shared. select_idx, select_num_idx = topk_indices_to_v2_indices( @@ -491,7 +433,5 @@ def apply( raise TypeError(f"MindIE sparse attention must return a tensor, but got {type(out).__name__}.") expected_shape = (1, query_head_count, query_length, head_dim) if tuple(out.shape) != expected_shape: - raise RuntimeError( - f"MindIE sparse attention returned shape {tuple(out.shape)}, expected {expected_shape}." - ) + raise RuntimeError(f"MindIE sparse attention returned shape {tuple(out.shape)}, expected {expected_shape}.") return out.squeeze(0).transpose(0, 1).contiguous().reshape(query_length, query_head_count * head_dim) diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh index 7a4618e73..9414d14f4 100755 --- a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse.sh @@ -1,26 +1,24 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" -model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-I2V-A14B}" -# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally -# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. -config_path="${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json" +# set path firstly +lightx2v_path= +model_path= export PLATFORM=ascend_npu -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" -export PYTHONPATH="${PYTHONPATH:-}" +export ASCEND_RT_VISIBLE_DEVICES=0 -source "${lightx2v_path}/scripts/base/base.sh" +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. python -m lightx2v.infer \ - --model_cls wan2.2_moe_distill \ - --task i2v \ - --warmup \ - --model_path "${model_path}" \ - --config_json "${config_path}" \ - --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ - --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ - --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ - --save_result_path "${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_npu.mp4" +--model_cls wan2.2_moe_distill \ +--task i2v \ +--warmup \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn.json \ +--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_npu.mp4 diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh index f036d7fa4..4bb389cd8 100755 --- a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_i2v_extreme_dynamic_sparse_sp_parallel.sh @@ -1,29 +1,26 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" -model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-I2V-A14B}" -# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally -# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. -# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. -# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. -config_path="${CONFIG_PATH:-${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json}" -save_result_path="${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4}" +# set path firstly +lightx2v_path= +model_path= export PLATFORM=ascend_npu -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" -export PYTHONPATH="${PYTHONPATH:-}" +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 -source "${lightx2v_path}/scripts/base/base.sh" +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. +# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. torchrun --nproc_per_node=8 -m lightx2v.infer \ - --model_cls wan2.2_moe_distill \ - --task i2v \ - --warmup \ - --model_path "${model_path}" \ - --config_json "${config_path}" \ - --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ - --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ - --image_path "${lightx2v_path}/assets/inputs/imgs/img_0.jpg" \ - --save_result_path "${save_result_path}" +--model_cls wan2.2_moe_distill \ +--task i2v \ +--warmup \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_i2v_distill_int8_dynamic_sparse_attn_sp_parallel.json \ +--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_i2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4 diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh index b1bb78121..caddef8e5 100755 --- a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse.sh @@ -1,25 +1,23 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" -model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-T2V-A14B}" -# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally -# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. -config_path="${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json" +# set path firstly +lightx2v_path= +model_path= export PLATFORM=ascend_npu -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0}" -export PYTHONPATH="${PYTHONPATH:-}" +export ASCEND_RT_VISIBLE_DEVICES=0 -source "${lightx2v_path}/scripts/base/base.sh" +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. python -m lightx2v.infer \ - --model_cls wan2.2_moe_distill \ - --task t2v \ - --warmup \ - --model_path "${model_path}" \ - --config_json "${config_path}" \ - --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ - --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ - --save_result_path "${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_npu.mp4" +--model_cls wan2.2_moe_distill \ +--task t2v \ +--warmup \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn.json \ +--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_npu.mp4 diff --git a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh index 9dcfb3c9e..d61e9dbac 100755 --- a/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh +++ b/scripts/platforms/ascend_npu/extreme/run_wan22_moe_t2v_extreme_dynamic_sparse_sp_parallel.sh @@ -1,28 +1,25 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -lightx2v_path="${LIGHTX2V_PATH:-$(cd -- "${script_dir}/../../../.." && pwd)}" -model_path="${MODEL_PATH:-/data/wushuo1/models/Wan2.2-T2V-A14B}" -# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally -# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. -# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. -# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. -config_path="${CONFIG_PATH:-${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json}" -save_result_path="${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4}" +# set path firstly +lightx2v_path= +model_path= export PLATFORM=ascend_npu -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" -export PYTHONPATH="${PYTHONPATH:-}" +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 -source "${lightx2v_path}/scripts/base/base.sh" +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh +# CANN 8.5.1 lacks aclnnBlockSparseAttention, so this config intentionally +# keeps the same dynamic Top-K mask algorithm on the supported MindIE RF v2 op. +# INT8 Ulysses communication uses the generic PyTorch pre/post path with HCCL. +# Keep the CUDA-only Triton pre/post backend disabled on Ascend NPU. torchrun --nproc_per_node=8 -m lightx2v.infer \ - --model_cls wan2.2_moe_distill \ - --task t2v \ - --warmup \ - --model_path "${model_path}" \ - --config_json "${config_path}" \ - --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ - --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ - --save_result_path "${save_result_path}" +--model_cls wan2.2_moe_distill \ +--task t2v \ +--warmup \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/extreme/wan_moe_t2v_distill_int8_dynamic_sparse_attn_sp_parallel.json \ +--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v_int8_dynamic_sparse_int8_comm_npu_sp8.mp4 diff --git a/tools/convert/quant_adapter.py b/tools/convert/quant_adapter.py old mode 100644 new mode 100755 index fc7073c0e..b8d91025e --- a/tools/convert/quant_adapter.py +++ b/tools/convert/quant_adapter.py @@ -1,329 +1,79 @@ import argparse -import gc -import os import sys -from collections import Counter from pathlib import Path import safetensors import torch from safetensors.torch import save_file -NVFP4_BLOCK_SIZE = 16 -NVFP4_SCALE_TILE_M = 128 -NVFP4_SCALE_TILE_K = NVFP4_BLOCK_SIZE * 4 -E2M1_VALUES = ( - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, -) +sys.path.append(str(Path(__file__).parent.parent.parent)) +quant_path = str(Path(__file__).parent / "quant") +if quant_path not in sys.path: + sys.path.insert(0, quant_path) -def _ceil_div(value, divisor): - return (value + divisor - 1) // divisor +from quant import * # noqa: E402 +from lightx2v.utils.quant_utils import FloatQuantizer # noqa: E402 -def _nvfp4_weight_keys(keys): - key_set = set(keys) - groups = {} - for weight_key in keys: - if not weight_key.endswith(".weight"): - continue - prefix = weight_key.removesuffix(".weight") - companions = { - "scale": f"{prefix}.weight_scale", - "input_global_scale": f"{prefix}.input_global_scale", - "alpha": f"{prefix}.alpha", - } - if all(key in key_set for key in companions.values()): - groups[weight_key] = companions - return groups - -def _decode_e2m1(packed): - if packed.dtype != torch.uint8 or packed.ndim != 2: - raise ValueError(f"NVFP4 weight must be a 2D uint8 tensor, got shape={tuple(packed.shape)}, dtype={packed.dtype}.") - - lookup = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device) - low = lookup[(packed & 0x0F).to(torch.int64)] - high = lookup[((packed >> 4) & 0x0F).to(torch.int64)] - return torch.stack((low, high), dim=-1).reshape(packed.shape[0], packed.shape[1] * 2) - - -def _unswizzle_nvfp4_scales(swizzled, rows, columns, block_size=NVFP4_BLOCK_SIZE): - if columns % block_size: - raise ValueError(f"NVFP4 input dimension must be divisible by {block_size}, got {columns}.") - - row_tiles = _ceil_div(rows, NVFP4_SCALE_TILE_M) - column_tiles = _ceil_div(columns, block_size * 4) - expected_numel = row_tiles * column_tiles * 32 * 4 * 4 - if swizzled.numel() != expected_numel: - raise ValueError(f"Unexpected NVFP4 scale layout: shape={tuple(swizzled.shape)}, numel={swizzled.numel()}, expected={expected_numel} for weight shape=({rows}, {columns}).") - - linear = swizzled.reshape(1, row_tiles, column_tiles, 32, 4, 4).permute(0, 1, 4, 3, 2, 5).reshape(row_tiles * NVFP4_SCALE_TILE_M, column_tiles * 4) - return linear[:rows, : columns // block_size].to(torch.float32) - - -@torch.inference_mode() -def dequantize_nvfp4_weight(packed, swizzled_scale, input_global_scale, alpha): - rows, packed_columns = packed.shape - columns = packed_columns * 2 - - input_global_scale = input_global_scale.to(torch.float32) - alpha = alpha.to(torch.float32) - if input_global_scale.numel() != 1 or alpha.numel() != 1: - raise ValueError("NVFP4 input_global_scale and alpha must both be scalar tensors.") - - # The checkpoint stores alpha = 1 / (input_global_scale * weight_global_scale). - weight_global_scale = torch.reciprocal(input_global_scale * alpha) - if not torch.isfinite(weight_global_scale).item() or weight_global_scale.item() <= 0: - raise ValueError(f"Invalid derived NVFP4 weight_global_scale={weight_global_scale.item()}.") - - values = _decode_e2m1(packed) - scales = _unswizzle_nvfp4_scales(swizzled_scale, rows, columns) - values = values.reshape(rows, columns // NVFP4_BLOCK_SIZE, NVFP4_BLOCK_SIZE) - return (values * (scales / weight_global_scale).unsqueeze(-1)).reshape(rows, columns) - - -@torch.inference_mode() -def quantize_int8_per_output_channel(weight): - if weight.dtype != torch.float32 or weight.ndim != 2: - raise ValueError(f"INT8 source weight must be a 2D float32 tensor, got shape={tuple(weight.shape)}, dtype={weight.dtype}.") - max_value = weight.abs().amax(dim=1, keepdim=True).clamp(min=1e-5) - scale = max_value / 127.0 - quantized = torch.clamp(torch.round(weight / scale), -128, 127).to(torch.int8) - return quantized.contiguous(), scale.to(torch.bfloat16).contiguous() - - -def _target_filename(source_name): - if "NVFP4" in source_name: - return source_name.replace("NVFP4", "INT8") - if "nvfp4" in source_name: - return source_name.replace("nvfp4", "int8") - path = Path(source_name) - return f"{path.stem}_int8{path.suffix}" - - -def _validate_int8_file(path, expected_quantized, expected_keys): - dtype_counts = Counter() - quantized_count = 0 - with safetensors.safe_open(path, framework="pt", device="cpu") as source: - keys = list(source.keys()) - if len(keys) != expected_keys: - raise ValueError(f"{path}: expected {expected_keys} tensors, found {len(keys)}.") - if any(key.endswith((".alpha", ".input_global_scale")) for key in keys): - raise ValueError(f"{path}: found stale NVFP4 alpha/input_global_scale tensors.") - - key_set = set(keys) - for key in keys: - tensor_slice = source.get_slice(key) - dtype_counts[tensor_slice.get_dtype()] += 1 - if tensor_slice.get_dtype() != "I8": - continue - quantized_count += 1 - scale_key = f"{key.removesuffix('.weight')}.weight_scale" - if not key.endswith(".weight") or scale_key not in key_set: - raise ValueError(f"{path}: INT8 tensor {key} has no matching weight_scale.") - scale_slice = source.get_slice(scale_key) - weight_shape = tensor_slice.get_shape() - if scale_slice.get_dtype() != "BF16" or scale_slice.get_shape() != [weight_shape[0], 1]: - raise ValueError(f"{path}: invalid scale for {key}: dtype={scale_slice.get_dtype()}, shape={scale_slice.get_shape()}.") - - if quantized_count != expected_quantized: - raise ValueError(f"{path}: expected {expected_quantized} INT8 weights, found {quantized_count}.") - return dtype_counts - - -def convert_wan_nvfp4_file(source_path, output_path, non_quant_dtype=torch.float32, overwrite=False): - if output_path.exists() and not overwrite: - raise FileExistsError(f"Output already exists: {output_path}. Pass --overwrite to replace it.") - output_path.parent.mkdir(parents=True, exist_ok=True) - - temporary_path = output_path.with_name(f".{output_path.name}.tmp.{os.getpid()}") - converted = {} - try: - with safetensors.safe_open(source_path, framework="pt", device="cpu") as source: - keys = list(source.keys()) - groups = _nvfp4_weight_keys(keys) - if not groups: - raise ValueError(f"{source_path} does not contain LightX2V NVFP4 weight groups.") - - skipped_keys = set() - for companions in groups.values(): - skipped_keys.update(companions.values()) - - total = len(groups) - converted_index = 0 - for key in keys: - if key in skipped_keys: - continue - if key not in groups: - converted[key] = source.get_tensor(key).to(non_quant_dtype).contiguous() - continue - - converted_index += 1 - companions = groups[key] - dequantized = dequantize_nvfp4_weight( - source.get_tensor(key), - source.get_tensor(companions["scale"]), - source.get_tensor(companions["input_global_scale"]), - source.get_tensor(companions["alpha"]), - ) - int8_weight, int8_scale = quantize_int8_per_output_channel(dequantized) - converted[key] = int8_weight - converted[companions["scale"]] = int8_scale - del dequantized, int8_weight, int8_scale - - if converted_index == 1 or converted_index % 25 == 0 or converted_index == total: - print(f"[{source_path.name}] converted {converted_index}/{total} NVFP4 weights", flush=True) - if converted_index % 25 == 0: - gc.collect() - - expected_keys = len(keys) - 2 * total - if len(converted) != expected_keys: - raise ValueError(f"Converted tensor count mismatch: expected {expected_keys}, got {len(converted)}.") - - metadata = { - "source_file": source_path.name, - "source_format": "nvfp4", - "target_format": "int8-npu", - "conversion": "nvfp4-dequant-then-int8-per-output-channel", - } - save_file(converted, str(temporary_path), metadata=metadata) - - if output_path.exists(): - output_path.unlink() - os.replace(temporary_path, output_path) - dtype_counts = _validate_int8_file(output_path, total, expected_keys) - print(f"Saved {output_path} ({output_path.stat().st_size / 1e9:.2f} GB), dtypes={dict(dtype_counts)}", flush=True) - return output_path - finally: - converted.clear() - gc.collect() - if temporary_path.exists(): - temporary_path.unlink() - - -def convert_wan_nvfp4_path(model_path, output_path, non_quant_dtype=torch.float32, overwrite=False): - if model_path.is_file(): - target = output_path - if output_path.exists() and output_path.is_dir(): - target = output_path / _target_filename(model_path.name) - return [convert_wan_nvfp4_file(model_path, target, non_quant_dtype, overwrite)] - - if not model_path.is_dir(): - raise FileNotFoundError(f"Input path does not exist: {model_path}") - if output_path.exists() and not output_path.is_dir(): - raise ValueError(f"Directory input requires a directory output, got file: {output_path}") - - sources = sorted(path for path in model_path.glob("*.safetensors") if "_comfy" not in path.stem) - if not sources: - raise FileNotFoundError(f"No non-ComfyUI safetensors files found under {model_path}.") - - comfy_count = sum(1 for path in model_path.glob("*_comfy.safetensors")) - if comfy_count: - print(f"Skipping {comfy_count} ComfyUI checkpoints because their packed weight layout differs.", flush=True) - - outputs = [] - for source_path in sources: - target = output_path / _target_filename(source_path.name) - outputs.append(convert_wan_nvfp4_file(source_path, target, non_quant_dtype, overwrite)) - return outputs - - -def convert_audio_adapter_fp8(model_path, output_path): - # Keep the original SekoTalk adapter conversion available without importing - # CUDA-specific quantization code in the Wan NVFP4 CPU conversion path. - project_root = Path(__file__).parent.parent.parent - if str(project_root) not in sys.path: - sys.path.append(str(project_root)) - quant_path = str(Path(__file__).parent / "quant") - if quant_path not in sys.path: - sys.path.insert(0, quant_path) - - from lightx2v.utils.quant_utils import FloatQuantizer - - output_path.parent.mkdir(parents=True, exist_ok=True) - state_dict = {} - with safetensors.safe_open(model_path, framework="pt", device="cpu") as source: - for key in source.keys(): - state_dict[key] = source.get_tensor(key) - - converted = {} - for key, tensor in state_dict.items(): - if key.startswith("ca") and ".to" in key and "weight" in key: - print(f"Converting {key} to FP8, dtype: {tensor.dtype}") - weight = tensor.to(torch.float32).cuda() - quantizer = FloatQuantizer("e4m3", True, "per_channel") - weight, weight_scale, _ = quantizer.real_quant_tensor(weight) - converted[key] = weight.to(torch.float8_e4m3fn).cpu() - converted[f"{key}_scale"] = weight_scale.to(torch.float32).cpu() - else: - print(f"Converting {key} to BF16, dtype: {tensor.dtype}") - converted[key] = tensor.to(torch.bfloat16) - - save_file(converted, str(output_path)) - print(f"Quantized adapter saved to: {output_path}") - - -def parse_args(): +def main(): + # 获取脚本所在目录 script_dir = Path(__file__).parent project_root = script_dir.parent.parent - parser = argparse.ArgumentParser(description="Convert a SekoTalk adapter to FP8 or LightX2V Wan2.2 NVFP4 checkpoints to INT8-NPU.") + + parser = argparse.ArgumentParser(description="Quantize audio adapter model to FP8") parser.add_argument( "--model_path", - type=Path, - default=project_root / "models" / "SekoTalk-Distill" / "audio_adapter_model.safetensors", - help="Input adapter file, Wan NVFP4 file, or Wan NVFP4 directory.", + type=str, + default=str(project_root / "models" / "SekoTalk-Distill" / "audio_adapter_model.safetensors"), + help="Path to input model file", ) parser.add_argument( "--output_path", - type=Path, - default=project_root / "models" / "SekoTalk-Distill-fp8" / "audio_adapter_model_fp8.safetensors", - help="Output file or directory.", - ) - parser.add_argument( - "--mode", - choices=("auto", "adapter-fp8", "wan-nvfp4-to-int8"), - default="auto", - help="Conversion mode. Auto selects Wan conversion for directory inputs and adapter conversion for file inputs.", - ) - parser.add_argument( - "--non_quant_dtype", - choices=("float32", "bfloat16"), - default="float32", - help="Dtype for non-quantized tensors in Wan INT8 checkpoints. float32 matches existing LightX2V INT8-NPU files.", + type=str, + default=str(project_root / "models" / "SekoTalk-Distill-fp8" / "audio_adapter_model_fp8.safetensors"), + help="Path to output quantized model file", ) - parser.add_argument("--overwrite", action="store_true", help="Replace existing output files.") - return parser.parse_args() + args = parser.parse_args() + model_path = Path(args.model_path) + output_path = Path(args.output_path) -def main(): - args = parse_args() - mode = args.mode - if mode == "auto": - mode = "wan-nvfp4-to-int8" if args.model_path.is_dir() else "adapter-fp8" + output_path.parent.mkdir(parents=True, exist_ok=True) + + state_dict = {} + with safetensors.safe_open(model_path, framework="pt", device="cpu") as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) - if mode == "adapter-fp8": - convert_audio_adapter_fp8(args.model_path, args.output_path) - return + new_state_dict = {} + + for key in state_dict.keys(): + if key.startswith("ca") and ".to" in key and "weight" in key: + print(f"Converting {key} to FP8, dtype: {state_dict[key].dtype}") + + ## fp8 + weight = state_dict[key].to(torch.float32).cuda() + w_quantizer = FloatQuantizer("e4m3", True, "per_channel") + weight, weight_scale, _ = w_quantizer.real_quant_tensor(weight) + weight = weight.to(torch.float8_e4m3fn) + weight_scale = weight_scale.to(torch.float32) + + ## QuantWeightMxFP4, QuantWeightMxFP6, QuantWeightMxFP8 for mxfp4,mxfp6,mxfp8 + # weight = state_dict[key].to(torch.bfloat16).cuda() + # quantizer = QuantWeightMxFP4(weight) + # weight, weight_scale, _ = quantizer.weight_quant_func(weight) + + new_state_dict[key] = weight.cpu() + new_state_dict[key + "_scale"] = weight_scale.cpu() + else: + # 不匹配的权重转换为BF16 + print(f"Converting {key} to BF16, dtype: {state_dict[key].dtype}") + new_state_dict[key] = state_dict[key].to(torch.bfloat16) - non_quant_dtype = torch.float32 if args.non_quant_dtype == "float32" else torch.bfloat16 - outputs = convert_wan_nvfp4_path(args.model_path, args.output_path, non_quant_dtype, args.overwrite) - print(f"Converted {len(outputs)} Wan2.2 checkpoint(s).", flush=True) + save_file(new_state_dict, str(output_path)) + print(f"Quantized model saved to: {output_path}") if __name__ == "__main__": From 3ee6b9472f8915628690cf7c019b5e00a7f67b96 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 19:54:26 +0000 Subject: [PATCH 4/6] chore: relocate Ascend RainFusion distill example --- ...i2v_distill_model_4step_cfg_ulysses_rainfusion.json | 6 +++--- ...n_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) rename configs/{distill/wan22 => platforms/ascend_npu/distill}/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json (76%) rename scripts/{wan22 => platforms/ascend_npu}/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh (80%) diff --git a/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json b/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json similarity index 76% rename from configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json rename to configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json index 4ec6e8c01..e25a7b4e9 100644 --- a/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json +++ b/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json @@ -33,10 +33,10 @@ 500, 250 ], - "high_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors", - "low_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors", + "high_noise_original_ckpt": "/data/wushuo1/models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors", + "low_noise_original_ckpt": "/data/wushuo1/models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors", "parallel": { - "seq_p_size": 4, + "seq_p_size": 8, "seq_p_attn_type": "ulysses" } } diff --git a/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh b/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh similarity index 80% rename from scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh rename to scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh index 8ac5ee6a3..9795b222e 100644 --- a/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh +++ b/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh @@ -1,19 +1,19 @@ #!/bin/bash # set path firstly -lightx2v_path=/LightX2V -model_path=/Wan2.2-I2V-A14B +lightx2v_path=/data/wushuo1/LightX2V-sparse +model_path=/data/wushuo1/models/Wan2.2-I2V-A14B -export CUDA_VISIBLE_DEVICES=0,1,2,3 +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 # set environment variables source ${lightx2v_path}/scripts/base/base.sh -torchrun --nproc_per_node=4 -m lightx2v.infer \ +torchrun --nproc_per_node=8 -m lightx2v.infer \ --model_cls wan2.2_moe_distill \ --task i2v \ --model_path $model_path \ ---config_json ${lightx2v_path}/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ From 5d902f295936bf5329b29d7c6e7c550a83ebf0c7 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 19:56:02 +0000 Subject: [PATCH 5/6] Revert "chore: relocate Ascend RainFusion distill example" This reverts commit 3ee6b9472f8915628690cf7c019b5e00a7f67b96. --- ...i2v_distill_model_4step_cfg_ulysses_rainfusion.json | 6 +++--- ...n_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) rename configs/{platforms/ascend_npu/distill => distill/wan22}/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json (76%) rename scripts/{platforms/ascend_npu => wan22}/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh (80%) diff --git a/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json similarity index 76% rename from configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json rename to configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json index e25a7b4e9..4ec6e8c01 100644 --- a/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json +++ b/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json @@ -33,10 +33,10 @@ 500, 250 ], - "high_noise_original_ckpt": "/data/wushuo1/models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors", - "low_noise_original_ckpt": "/data/wushuo1/models/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors", + "high_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_high_noise_lightx2v_4step.safetensors", + "low_noise_original_ckpt": "/Wan2.2-Distill-Models/wan2.2_i2v_A14b_low_noise_lightx2v_4step.safetensors", "parallel": { - "seq_p_size": 8, + "seq_p_size": 4, "seq_p_attn_type": "ulysses" } } diff --git a/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh similarity index 80% rename from scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh rename to scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh index 9795b222e..8ac5ee6a3 100644 --- a/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh +++ b/scripts/wan22/distill/run_wan22_moe_i2v_distill_4step_ulysses_rainfusion.sh @@ -1,19 +1,19 @@ #!/bin/bash # set path firstly -lightx2v_path=/data/wushuo1/LightX2V-sparse -model_path=/data/wushuo1/models/Wan2.2-I2V-A14B +lightx2v_path=/LightX2V +model_path=/Wan2.2-I2V-A14B -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export CUDA_VISIBLE_DEVICES=0,1,2,3 # set environment variables source ${lightx2v_path}/scripts/base/base.sh -torchrun --nproc_per_node=8 -m lightx2v.infer \ +torchrun --nproc_per_node=4 -m lightx2v.infer \ --model_cls wan2.2_moe_distill \ --task i2v \ --model_path $model_path \ ---config_json ${lightx2v_path}/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json \ +--config_json ${lightx2v_path}/configs/distill/wan22/wan_moe_i2v_distill_model_4step_cfg_ulysses_rainfusion.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \ From ebbfeba448a04915914c50407159f5ff6a50b59d Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Mon, 27 Jul 2026 20:02:44 +0000 Subject: [PATCH 6/6] chore: relocate Ascend INT8 distill example --- .../distill}/wan_moe_i2v_distill_int8_4step_ulysses_npu.json | 0 .../distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename configs/{distill/wan22 => platforms/ascend_npu/distill}/wan_moe_i2v_distill_int8_4step_ulysses_npu.json (100%) rename scripts/{wan22 => platforms/ascend_npu}/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh (93%) diff --git a/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json b/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_int8_4step_ulysses_npu.json similarity index 100% rename from configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json rename to configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_int8_4step_ulysses_npu.json diff --git a/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh b/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh similarity index 93% rename from scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh rename to scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh index c397df405..abc444e40 100644 --- a/scripts/wan22/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh +++ b/scripts/platforms/ascend_npu/distill/run_wan22_moe_i2v_distill_int8_4step_ulysses_npu.sh @@ -19,7 +19,7 @@ torchrun --nproc_per_node=2 -m lightx2v.infer \ --model_cls wan2.2_moe_distill \ --task i2v \ --model_path $model_path \ ---config_json ${lightx2v_path}/configs/distill/wan22/wan_moe_i2v_distill_int8_4step_ulysses_npu.json \ +--config_json ${lightx2v_path}/configs/platforms/ascend_npu/distill/wan_moe_i2v_distill_int8_4step_ulysses_npu.json \ --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." \ --negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ --image_path ${lightx2v_path}/assets/inputs/imgs/img_0.jpg \