update seko ar rope

configs/seko_talk/ar/seko_talk_ar_kv_dist.json

@@ -17,11 +17,11 @@
     "offload_granularity": "block",
     "use_31_block": true,
     "dit_quantized": true,
-    "dit_quantized_ckpt": "/SekoTalk-Distill-AR/converted_fp8.safetensors",
+    "dit_quantized_ckpt": "/data/nvme5/gushiqiao/models/SekoTalk-Distill-AR-0703/seko_ar_new_rope_fp8_4steps.safetensors",
     "dit_quant_scheme": "fp8-sgl",
     "adapter_quantized": true,
     "adapter_quant_scheme": "fp8",
-    "adapter_model_path": "/SekoTalk-Distill-AR/-audio_adapter_fp8.pt",
+    "adapter_model_path": "/data/nvme5/gushiqiao/models/SekoTalk-Distill-AR-0703/semi_causal_audio_adapter_v2.2_8k_fp8.pth",
     "audio_feature_dim": 1024,
     "audio_projection_dim": 1024,
     "audio_num_tokens": 32,
@@ -31,15 +31,15 @@
     "audio_feat_window_neighbor_frame": 0,
     "audio_feat_fps": 50,
     "parallel": {
-        "seq_p_size": 8,
+        "seq_p_size": 4,
         "seq_p_attn_type": "ulysses"
     },
     "ar_config": {
-        "num_frame_per_chunk": 1,
+        "num_frame_per_chunk": 3,
         "step_kv_cache": true,
         "sink_size": 2,
         "local_attn_size": 21,
         "kv_offload": false,
-        "async_vae_decode": true
+        "async_vae_decode": false
     }
 }

lightx2v/models/networks/wan/infer/audio/transformer_infer.py

@@ -167,25 +167,65 @@ def _spatial_freqs_for_rank(self, freqs, h, w, local_per_frame, world_size=1, ra
             spatial_freqs = torch.cat([spatial_freqs, pad], dim=0)
         return torch.chunk(spatial_freqs, world_size, dim=0)[rank][:local_per_frame]
 
-    def _rope_freqs_for_cache_range(self, freqs, h, w, world_size, rank, token_start, token_end, ref_tokens, local_per_frame):
+    def _cache_positions_for_range(self, token_start, token_end, device, global_end=None, sink_tokens=0):
+        token_idx = torch.arange(token_start, token_end, device=device, dtype=torch.long)
+        if global_end is None:
+            return token_idx
+
+        sink_tokens = int(sink_tokens)
+        recent_local_tokens = max(0, token_end - sink_tokens)
+        recent_global_start = int(global_end) - recent_local_tokens
+        recent_idx = recent_global_start + token_idx - sink_tokens
+        return torch.where(token_idx < sink_tokens, token_idx, recent_idx)
+
+    def _rope_freqs_for_cache_range(
+        self,
+        freqs,
+        h,
+        w,
+        world_size,
+        rank,
+        token_start,
+        token_end,
+        ref_tokens,
+        local_per_frame,
+        global_end=None,
+        sink_tokens=0,
+    ):
         c = self.head_dim // 2
         temporal_dim = c - 2 * (c // 3)
         freqs_split = freqs.split([temporal_dim, c // 3, c // 3], dim=1)
         spatial_freqs = self._spatial_freqs_for_rank(freqs, h, w, local_per_frame, world_size, rank)
 
-        token_idx = torch.arange(token_start, token_end, device=freqs.device, dtype=torch.long)
-        is_ref = token_idx < ref_tokens
-        gen_idx = torch.clamp(token_idx - ref_tokens, min=0)
-        frame_idx = gen_idx // local_per_frame
-        ref_spatial_idx = token_idx % local_per_frame
+        position_idx = self._cache_positions_for_range(token_start, token_end, freqs.device, global_end=global_end, sink_tokens=sink_tokens)
+        is_ref = position_idx < ref_tokens
+        gen_idx = torch.clamp(position_idx - ref_tokens, min=0)
+        ref_frames = ref_tokens // local_per_frame
+        ref_frame_idx = position_idx // local_per_frame
+        gen_frame_idx = ref_frames + gen_idx // local_per_frame
+        frame_idx = torch.where(is_ref, ref_frame_idx, gen_frame_idx)
+        ref_spatial_idx = position_idx % local_per_frame
         gen_spatial_idx = gen_idx % local_per_frame
         spatial_idx = torch.where(is_ref, ref_spatial_idx, gen_spatial_idx)
 
         temporal_freqs = freqs_split[0][frame_idx]
-        temporal_freqs = torch.where(is_ref.unsqueeze(-1), torch.ones_like(temporal_freqs), temporal_freqs)
         return torch.cat([temporal_freqs, spatial_freqs[spatial_idx]], dim=-1).unsqueeze(1)
 
-    def _apply_rope_with_cache_range(self, x, freqs, h, w, world_size, rank, token_start, token_end, ref_tokens, local_per_frame):
+    def _apply_rope_with_cache_range(
+        self,
+        x,
+        freqs,
+        h,
+        w,
+        world_size,
+        rank,
+        token_start,
+        token_end,
+        ref_tokens,
+        local_per_frame,
+        global_end=None,
+        sink_tokens=0,
+    ):
         orig_dtype = x.dtype
         if self.config.get("causal_rope_type", "triton") == "triton":
             return apply_audio_cache_rope(
@@ -198,8 +238,22 @@ def _apply_rope_with_cache_range(self, x, freqs, h, w, world_size, rank, token_s
                 local_per_frame=local_per_frame,
                 world_size=world_size,
                 rank=rank,
+                global_end=global_end,
+                sink_tokens=sink_tokens,
             ).to(orig_dtype)
-        pos_freqs = self._rope_freqs_for_cache_range(freqs, h, w, world_size, rank, token_start, token_end, ref_tokens, local_per_frame)
+        pos_freqs = self._rope_freqs_for_cache_range(
+            freqs,
+            h,
+            w,
+            world_size,
+            rank,
+            token_start,
+            token_end,
+            ref_tokens,
+            local_per_frame,
+            global_end=global_end,
+            sink_tokens=sink_tokens,
+        )
         n = x.size(1)
         x_c = torch.view_as_complex(x.float().reshape(x.size(0), n, -1, 2))
         out = torch.view_as_real(x_c * pos_freqs.to(torch.complex64)).flatten(2)
@@ -296,7 +350,7 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
                 h1 = h0 + shard_heads
                 kv_cache.store_kv(k_rope[:, h0:h1], v[:, h0:h1], local_start_idx, local_end_idx, self.block_idx)
             else:
-                start_frame = segment_idx * frames
+                start_frame = self.kv_cache_manager.ref_num_frames + segment_idx * frames
                 q_rope, k_rope = self._apply_rope_sp(q, k, grid_sizes, freqs, start_frame)
                 use_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False)
                 use_fp4_comm = self.config["parallel"].get("seq_p_fp4_comm", False)
@@ -350,8 +404,34 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
             attn_v = kv_cache.v_cache(self.block_idx, attn_start, local_end_idx)
             if self.config.get("ar_config", {}).get("kv_quant", {}).get("calibrate", False):
                 kv_cache.capture_attn(self.block_idx, attn_start, local_end_idx)
-            q = self._apply_rope_with_cache_range(q, freqs, h, w, sp_world_size, sp_rank, local_start_idx, local_end_idx, local_ref_tokens, local_per_frame)
-            attn_k = self._apply_rope_with_cache_range(attn_k, freqs, h, w, sp_world_size, sp_rank, attn_start, local_end_idx, local_ref_tokens, local_per_frame)
+            q = self._apply_rope_with_cache_range(
+                q,
+                freqs,
+                h,
+                w,
+                sp_world_size,
+                sp_rank,
+                local_start_idx,
+                local_end_idx,
+                local_ref_tokens,
+                local_per_frame,
+                global_end=current_end,
+                sink_tokens=sink_tokens,
+            )
+            attn_k = self._apply_rope_with_cache_range(
+                attn_k,
+                freqs,
+                h,
+                w,
+                sp_world_size,
+                sp_rank,
+                attn_start,
+                local_end_idx,
+                local_ref_tokens,
+                local_per_frame,
+                global_end=current_end,
+                sink_tokens=sink_tokens,
+            )
             if isinstance(attn_k, tuple):
                 k_lens = torch.empty_like(seq_lens).fill_(attn_k[0].size(0))
             else:

lightx2v/models/networks/wan/infer/triton_ops.py

@@ -343,10 +343,13 @@ def _audio_cache_rope_kernel(
     h: tl.constexpr,
     w: tl.constexpr,
     token_start,
+    global_end,
     ref_tokens: tl.constexpr,
+    sink_tokens: tl.constexpr,
     local_per_frame: tl.constexpr,
     world_size: tl.constexpr,
     rank: tl.constexpr,
+    seq_len: tl.constexpr,
     num_heads: tl.constexpr,
     head_dim: tl.constexpr,
     t_dim: tl.constexpr,
@@ -356,7 +359,12 @@ def _audio_cache_rope_kernel(
     row_idx = tl.program_id(0)
     token_head_idx = row_idx // num_heads
     head_idx = row_idx - token_head_idx * num_heads
-    token_idx = token_start + token_head_idx
+    local_token_idx = token_start + token_head_idx
+    range_end = token_start + seq_len
+    recent_local_tokens = tl.maximum(range_end - sink_tokens, 0)
+    recent_global_start = global_end - recent_local_tokens
+    recent_token_idx = recent_global_start + local_token_idx - sink_tokens
+    token_idx = tl.where(local_token_idx < sink_tokens, local_token_idx, recent_token_idx)
 
     half_head_dim: tl.constexpr = head_dim // 2
     col_offsets = tl.arange(0, BLOCK_SIZE)
@@ -372,7 +380,10 @@ def _audio_cache_rope_kernel(
     spatial_chunk: tl.constexpr = padded_hw // world_size
     is_ref = token_idx < ref_tokens
     gen_idx = tl.maximum(token_idx - ref_tokens, 0)
-    frame_idx = gen_idx // local_per_frame
+    ref_frames: tl.constexpr = ref_tokens // local_per_frame
+    ref_frame_idx = token_idx // local_per_frame
+    gen_frame_idx = ref_frames + gen_idx // local_per_frame
+    frame_idx = tl.where(is_ref, ref_frame_idx, gen_frame_idx)
     ref_spatial_idx = token_idx % local_per_frame
     gen_spatial_idx = gen_idx % local_per_frame
     local_spatial_idx = tl.where(is_ref, ref_spatial_idx, gen_spatial_idx)
@@ -386,7 +397,7 @@ def _audio_cache_rope_kernel(
     freq_row = tl.where(is_temporal, frame_idx, tl.where(is_h, y, z))
     freq_col = col_offsets
     freq_offset = (freq_row * half_head_dim + freq_col) * 2
-    load_mask = mask & (~is_temporal | ~is_ref) & (is_temporal | spatial_valid)
+    load_mask = mask & (is_temporal | spatial_valid)
     cos_vals = tl.load(freqs_real_ptr + freq_offset, mask=load_mask, other=1.0)
     sin_vals = tl.load(freqs_real_ptr + freq_offset + 1, mask=load_mask, other=0.0)
 
@@ -407,6 +418,8 @@ def apply_audio_cache_rope(
     local_per_frame: int,
     world_size: int = 1,
     rank: int = 0,
+    global_end: int | None = None,
+    sink_tokens: int = 0,
 ) -> torch.Tensor:
     seq_len, num_heads, head_dim = x.shape
     c = head_dim // 2
@@ -416,17 +429,22 @@ def apply_audio_cache_rope(
     out = torch.empty_like(x_contig)
     freqs_real = torch.view_as_real(freqs.to(device=x.device))
     block_size = triton.next_power_of_2(c)
+    if global_end is None:
+        global_end = token_start + seq_len
     _audio_cache_rope_kernel[(seq_len * num_heads,)](
         out,
         x_contig,
         freqs_real,
         h,
         w,
         int(token_start),
+        int(global_end),
         int(ref_tokens),
+        int(sink_tokens),
         int(local_per_frame),
         int(world_size),
         int(rank),
+        int(seq_len),
         num_heads,
         head_dim,
         t_dim,

scripts/seko_talk/ar/run_seko_talk_ar_01_base.sh

@@ -1,14 +1,14 @@
 #!/bin/bash
 
-lightx2v_path=/data/nvme4/gushiqiao/new/LightX2V
-model_path=/data/nvme5/gushiqiao/models/SekoTalk-Distill-AR/
+lightx2v_path=/data/nvme5/gushiqiao/codes/LightX2V
+model_path=/data/nvme5/gushiqiao/models/SekoTalk-Distill-AR-0703/
 
-export CUDA_VISIBLE_DEVICES=0
+export CUDA_VISIBLE_DEVICES=0,1,2,3
 
 # set environment variables
 source ${lightx2v_path}/scripts/base/base.sh
 
-python -m lightx2v.infer \
+torchrun --nproc-per-node 4 -m lightx2v.infer \
 --model_cls seko_talk_ar \
 --task rs2v \
 --model_path $model_path \
@@ -17,5 +17,5 @@ python -m lightx2v.infer \
 --negative_prompt "low quality,blurry,pixelated,low resolution,noise,artifacts,poor lighting, overexposed, underexposed, distorted, unnatural, deformed, weird,scared,anatomy,mutated, wrong proportions, extra limbs,floating objects, disconnected, gravity-defying, impossible shadows,wrong lighting,non-existent reflections,inconsistent perspective, repetitive, monotonous, monotonous, generic, watermark, ugly, high contrast, bad photo, font, username, error, logo, words, letters, digits, autograph, trademark, name, twisted face, (poorly drawn hands, malformed hands, missing fingers, unnatural hand positions, blur hand, multiple fingers, multiple arms), static, naked, artifacts, oversaturated" \
 --image_path "/data/nvme4/gushiqiao/new/example/1_素材图.png" \
 --audio_path "/data/nvme4/gushiqiao/new/example/1_素材图.mp3" \
---save_result_path ${lightx2v_path}/save_results/output_lightx2v_seko_talk_ar_sp2.mp4 \
+--save_result_path ${lightx2v_path}/save_results/output_lightx2v_seko_talk_ar_sp4.mp4 \
 --seed 0