appiq-solution 1.2.0 → 1.3.1

package/#Tools/SkyReels-V2/skycaptioner_v1/scripts/vllm_struct_caption.py

@@ -0,0 +1,156 @@
+
+import torch
+import decord
+import argparse
+
+import pandas as pd
+import numpy as np
+
+from tqdm import tqdm
+from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer, AutoProcessor
+
+from torch.utils.data import DataLoader
+
+SYSTEM_PROMPT = "I need you to generate a structured and detailed caption for the provided video. The structured output and the requirements for each field are as shown in the following JSON content: {\"subjects\": [{\"appearance\": \"Main subject appearance description\", \"action\": \"Main subject action\", \"expression\": \"Main subject expression  (Only for human/animal categories, empty otherwise)\", \"position\": \"Subject position in the video (Can be relative position to other objects or spatial description)\", \"TYPES\": {\"type\": \"Main category (e.g., Human)\", \"sub_type\": \"Sub-category (e.g., Man)\"}, \"is_main_subject\": true}, {\"appearance\": \"Non-main subject appearance description\", \"action\": \"Non-main subject action\", \"expression\": \"Non-main subject expression (Only for human/animal categories, empty otherwise)\", \"position\": \"Position of non-main subject 1\", \"TYPES\": {\"type\": \"Main category (e.g., Vehicles)\", \"sub_type\": \"Sub-category (e.g., Ship)\"}, \"is_main_subject\": false}], \"shot_type\": \"Shot type(Options: long_shot/full_shot/medium_shot/close_up/extreme_close_up/other)\", \"shot_angle\": \"Camera angle(Options: eye_level/high_angle/low_angle/other)\", \"shot_position\": \"Camera position(Options: front_view/back_view/side_view/over_the_shoulder/overhead_view/point_of_view/aerial_view/overlooking_view/other)\", \"camera_motion\": \"Camera movement description\", \"environment\": \"Video background/environment description\", \"lighting\": \"Lighting information in the video\"}"
+
+
+class VideoTextDataset(torch.utils.data.Dataset):
+    def __init__(self, csv_path, model_path):
+        if isinstance(csv_path, pd.DataFrame):
+            self.meta = csv_path
+        else:
+            self.meta = pd.read_csv(csv_path)
+        self._path = 'path'
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.processor = AutoProcessor.from_pretrained(model_path)
+  
+    def __getitem__(self, index):
+        row = self.meta.iloc[index]
+        path = row[self._path]
+        real_index = self.meta.index[index]
+        vr = decord.VideoReader(path, ctx=decord.cpu(0), width=360, height=420)
+        start = 0
+        end = len(vr)
+        # avg_fps = vr.get_avg_fps()
+        index = self.get_index(end-start, 16, st=start)
+        frames = vr.get_batch(index).asnumpy() # n h w c
+        video_inputs = [torch.from_numpy(frames).permute(0, 3, 1, 2)]
+        conversation = {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "video",
+                            "video": row['path'],
+                            "max_pixels": 360 * 420, # 460800
+                            "fps": 2.0,
+                        },
+                        {   
+                            "type": "text", 
+                            "text": SYSTEM_PROMPT
+                        },
+                    ],
+                }
+                
+        # 生成 user_input
+        user_input = self.processor.apply_chat_template(
+            [conversation],
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        results = dict()
+        inputs = {
+            'prompt': user_input,
+            'multi_modal_data': {'video': video_inputs}
+        }
+        results["index"] = real_index
+        results['input'] = inputs
+        return results
+
+    def __len__(self):
+        return len(self.meta)
+
+    def get_index(self, video_size, num_frames, st=0):
+        seg_size = max(0., float(video_size - 1) / num_frames)
+        max_frame = int(video_size) - 1
+        seq = []
+        # index from 1, must add 1
+        for i in range(num_frames):
+            start = int(np.round(seg_size * i))
+            # end = int(np.round(seg_size * (i + 1)))
+            idx = min(start, max_frame)
+            seq.append(idx+st)
+        return seq
+    
+def result_writer(indices_list: list, result_list: list, meta: pd.DataFrame, column):
+    flat_indices = []
+    for x in zip(indices_list):
+        flat_indices.extend(x)
+    flat_results = []
+    for x in zip(result_list):
+        flat_results.extend(x)
+    
+    flat_indices = np.array(flat_indices)
+    flat_results = np.array(flat_results)
+
+    unique_indices, unique_indices_idx = np.unique(flat_indices, return_index=True)
+    meta.loc[unique_indices, column[0]] = flat_results[unique_indices_idx]
+
+    meta = meta.loc[unique_indices]
+    return meta
+
+
+def worker_init_fn(worker_id):
+    # Set different seed for each worker
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    # Prevent deadlocks by setting timeout
+    torch.set_num_threads(1)
+
+def main():
+    parser = argparse.ArgumentParser(description="SkyCaptioner-V1 vllm batch inference")
+    parser.add_argument("--input_csv", default="./examples/test.csv")
+    parser.add_argument("--out_csv", default="./examples/test_result.csv")
+    parser.add_argument("--bs", type=int, default=4)
+    parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--model_path", required=True, type=str, help="skycaptioner-v1 model path")
+    args = parser.parse_args()
+    
+    dataset = VideoTextDataset(csv_path=args.input_csv, model_path=args.model_path)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.bs,
+        num_workers=4,
+        worker_init_fn=worker_init_fn,
+        persistent_workers=True,
+        timeout=180,
+    )
+
+    sampling_params = SamplingParams(temperature=0.05, max_tokens=2048)
+    
+    llm = LLM(model=args.model_path,
+        gpu_memory_utilization=0.6, 
+        max_model_len=31920,
+        tensor_parallel_size=args.tp)
+    
+    indices_list = []
+    caption_save = []
+    for video_batch in tqdm(dataloader):
+        indices = video_batch["index"]
+        inputs = video_batch["input"]
+        batch_user_inputs = []
+        for prompt, video in zip(inputs['prompt'], inputs['multi_modal_data']['video'][0]):
+            usi={'prompt':prompt, 'multi_modal_data':{'video':video}}
+            batch_user_inputs.append(usi)
+        outputs = llm.generate(batch_user_inputs, sampling_params, use_tqdm=False)
+        struct_outputs = [output.outputs[0].text for output in outputs]
+
+        indices_list.extend(indices.tolist())
+        caption_save.extend(struct_outputs)
+    
+    meta_new = result_writer(indices_list, caption_save, dataset.meta, column=["structural_caption"])
+    meta_new.to_csv(args.out_csv, index=False)
+    print(f'Saved structural_caption to {args.out_csv}')
+
+if __name__ == '__main__':
+    main()

package/#Tools/SkyReels-V2/skyreels_v2_infer/__init__.py

@@ -0,0 +1 @@
+from .pipelines import DiffusionForcingPipeline

package/#Tools/SkyReels-V2/skyreels_v2_infer/distributed/xdit_context_parallel.py

@@ -0,0 +1,286 @@
+import numpy as np
+import torch
+import torch.amp as amp
+from torch.backends.cuda import sdp_kernel
+from xfuser.core.distributed import get_sequence_parallel_rank
+from xfuser.core.distributed import get_sequence_parallel_world_size
+from xfuser.core.distributed import get_sp_group
+from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+
+from ..modules.transformer import sinusoidal_embedding_1d
+
+
+def pad_freqs(original_tensor, target_len):
+    seq_len, s1, s2 = original_tensor.shape
+    pad_size = target_len - seq_len
+    padding_tensor = torch.ones(pad_size, s1, s2, dtype=original_tensor.dtype, device=original_tensor.device)
+    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+    return padded_tensor
+
+
+@amp.autocast("cuda", enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    """
+    x:          [B, L, N, C].
+    grid_sizes: [B, 3].
+    freqs:      [M, C // 2].
+    """
+    s, n, c = x.size(1), x.size(2), x.size(3) // 2
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    grid = [grid_sizes.tolist()] * x.size(0)
+    for i, (f, h, w) in enumerate(grid):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2))
+        freqs_i = torch.cat(
+            [
+                freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+            ],
+            dim=-1,
+        ).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        sp_size = get_sequence_parallel_world_size()
+        sp_rank = get_sequence_parallel_rank()
+        freqs_i = pad_freqs(freqs_i, s * sp_size)
+        s_per_rank = s
+        freqs_i_rank = freqs_i[(sp_rank * s_per_rank) : ((sp_rank + 1) * s_per_rank), :, :]
+        x_i = torch.view_as_real(x_i * freqs_i_rank.cuda()).flatten(2)
+        x_i = torch.cat([x_i, x[i, s:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).float()
+
+
+def broadcast_should_calc(should_calc: bool) -> bool:
+    import torch.distributed as dist
+
+    device = torch.cuda.current_device()
+    int_should_calc = 1 if should_calc else 0
+    tensor = torch.tensor([int_should_calc], device=device, dtype=torch.int8)
+    dist.broadcast(tensor, src=0)
+    should_calc = tensor.item() == 1
+    return should_calc
+
+
+def usp_dit_forward(self, x, t, context, clip_fea=None, y=None, fps=None):
+    """
+    x:              A list of videos each with shape [C, T, H, W].
+    t:              [B].
+    context:        A list of text embeddings each with shape [L, C].
+    """
+    if self.model_type == "i2v":
+        assert clip_fea is not None and y is not None
+    # params
+    device = self.patch_embedding.weight.device
+    if self.freqs.device != device:
+        self.freqs = self.freqs.to(device)
+
+    if y is not None:
+        x = torch.cat([x, y], dim=1)
+
+    # embeddings
+    x = self.patch_embedding(x)
+    grid_sizes = torch.tensor(x.shape[2:], dtype=torch.long)
+    x = x.flatten(2).transpose(1, 2)
+
+    if self.flag_causal_attention:
+        frame_num = grid_sizes[0]
+        height = grid_sizes[1]
+        width = grid_sizes[2]
+        block_num = frame_num // self.num_frame_per_block
+        range_tensor = torch.arange(block_num).view(-1, 1)
+        range_tensor = range_tensor.repeat(1, self.num_frame_per_block).flatten()
+        casual_mask = range_tensor.unsqueeze(0) <= range_tensor.unsqueeze(1)  # f, f
+        casual_mask = casual_mask.view(frame_num, 1, 1, frame_num, 1, 1).to(x.device)
+        casual_mask = casual_mask.repeat(1, height, width, 1, height, width)
+        casual_mask = casual_mask.reshape(frame_num * height * width, frame_num * height * width)
+        self.block_mask = casual_mask.unsqueeze(0).unsqueeze(0)
+
+    # time embeddings
+    with amp.autocast("cuda", dtype=torch.float32):
+        if t.dim() == 2:
+            b, f = t.shape
+            _flag_df = True
+        else:
+            _flag_df = False
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(self.patch_embedding.weight.dtype)
+        )  # b, dim
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim))  # b, 6, dim
+
+        if self.inject_sample_info:
+            fps = torch.tensor(fps, dtype=torch.long, device=device)
+
+            fps_emb = self.fps_embedding(fps).float()
+            if _flag_df:
+                e0 = e0 + self.fps_projection(fps_emb).unflatten(1, (6, self.dim)).repeat(t.shape[1], 1, 1)
+            else:
+                e0 = e0 + self.fps_projection(fps_emb).unflatten(1, (6, self.dim))
+
+        if _flag_df:
+            e = e.view(b, f, 1, 1, self.dim)
+            e0 = e0.view(b, f, 1, 1, 6, self.dim)
+            e = e.repeat(1, 1, grid_sizes[1], grid_sizes[2], 1).flatten(1, 3)
+            e0 = e0.repeat(1, 1, grid_sizes[1], grid_sizes[2], 1, 1).flatten(1, 3)
+            e0 = e0.transpose(1, 2).contiguous()
+
+        assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+    # context
+    context = self.text_embedding(context)
+
+    if clip_fea is not None:
+        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+        context = torch.concat([context_clip, context], dim=1)
+
+    # arguments
+    if e0.ndim == 4:
+        e0 = torch.chunk(e0, get_sequence_parallel_world_size(), dim=2)[get_sequence_parallel_rank()]
+    kwargs = dict(e=e0, grid_sizes=grid_sizes, freqs=self.freqs, context=context, block_mask=self.block_mask)
+
+    if self.enable_teacache:
+        modulated_inp = e0 if self.use_ref_steps else e
+        # teacache
+        if self.cnt % 2 == 0:  # even -> conditon
+            self.is_even = True
+            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
+                should_calc_even = True
+                self.accumulated_rel_l1_distance_even = 0
+            else:
+                rescale_func = np.poly1d(self.coefficients)
+                self.accumulated_rel_l1_distance_even += rescale_func(
+                    ((modulated_inp - self.previous_e0_even).abs().mean() / self.previous_e0_even.abs().mean())
+                    .cpu()
+                    .item()
+                )
+                if self.accumulated_rel_l1_distance_even < self.teacache_thresh:
+                    should_calc_even = False
+                else:
+                    should_calc_even = True
+                    self.accumulated_rel_l1_distance_even = 0
+            self.previous_e0_even = modulated_inp.clone()
+        else:  # odd -> unconditon
+            self.is_even = False
+            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
+                should_calc_odd = True
+                self.accumulated_rel_l1_distance_odd = 0
+            else:
+                rescale_func = np.poly1d(self.coefficients)
+                self.accumulated_rel_l1_distance_odd += rescale_func(
+                    ((modulated_inp - self.previous_e0_odd).abs().mean() / self.previous_e0_odd.abs().mean())
+                    .cpu()
+                    .item()
+                )
+                if self.accumulated_rel_l1_distance_odd < self.teacache_thresh:
+                    should_calc_odd = False
+                else:
+                    should_calc_odd = True
+                    self.accumulated_rel_l1_distance_odd = 0
+            self.previous_e0_odd = modulated_inp.clone()
+
+    x = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
+    if self.enable_teacache:
+        if self.is_even:
+            should_calc_even = broadcast_should_calc(should_calc_even)
+            if not should_calc_even:
+                x += self.previous_residual_even
+            else:
+                ori_x = x.clone()
+                for block in self.blocks:
+                    x = block(x, **kwargs)
+                ori_x.mul_(-1)
+                ori_x.add_(x)
+                self.previous_residual_even = ori_x
+        else:
+            should_calc_odd = broadcast_should_calc(should_calc_odd)
+            if not should_calc_odd:
+                x += self.previous_residual_odd
+            else:
+                ori_x = x.clone()
+                for block in self.blocks:
+                    x = block(x, **kwargs)
+                ori_x.mul_(-1)
+                ori_x.add_(x)
+                self.previous_residual_odd = ori_x
+        self.cnt += 1
+        if self.cnt >= self.num_steps:
+            self.cnt = 0
+    else:
+        # Context Parallel
+        for block in self.blocks:
+            x = block(x, **kwargs)
+
+    # head
+    if e.ndim == 3:
+        e = torch.chunk(e, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
+    x = self.head(x, e)
+    # Context Parallel
+    x = get_sp_group().all_gather(x, dim=1)
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+    return x.float()
+
+
+def usp_attn_forward(self, x, grid_sizes, freqs, block_mask):
+
+    r"""
+    Args:
+        x(Tensor): Shape [B, L, num_heads, C / num_heads]
+        seq_lens(Tensor): Shape [B]
+        grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+        freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+    """
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(torch.bfloat16)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    x = x.to(self.q.weight.dtype)
+    q, k, v = qkv_fn(x)
+
+    if not self._flag_ar_attention:
+        q = rope_apply(q, grid_sizes, freqs)
+        k = rope_apply(k, grid_sizes, freqs)
+    else:
+
+        q = rope_apply(q, grid_sizes, freqs)
+        k = rope_apply(k, grid_sizes, freqs)
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
+        # x = torch.nn.functional.scaled_dot_product_attention(
+        #     q.transpose(1, 2),
+        #     k.transpose(1, 2),
+        #     v.transpose(1, 2),
+        #    ).transpose(1, 2).contiguous()
+        with sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+            x = (
+                torch.nn.functional.scaled_dot_product_attention(
+                    q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=block_mask
+                )
+                .transpose(1, 2)
+                .contiguous()
+            )
+    x = xFuserLongContextAttention()(None, query=half(q), key=half(k), value=half(v), window_size=self.window_size)
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x

package/#Tools/SkyReels-V2/skyreels_v2_infer/modules/__init__.py

@@ -0,0 +1,69 @@
+import gc
+import os
+
+import torch
+from safetensors.torch import load_file
+
+from .clip import CLIPModel
+from .t5 import T5EncoderModel
+from .transformer import WanModel
+from .vae import WanVAE
+
+
+def download_model(model_id):
+    if not os.path.exists(model_id):
+        from huggingface_hub import snapshot_download
+
+        model_id = snapshot_download(repo_id=model_id)
+    return model_id
+
+
+def get_vae(model_path, device="cuda", weight_dtype=torch.float32) -> WanVAE:
+    vae = WanVAE(model_path).to(device).to(weight_dtype)
+    vae.vae.requires_grad_(False)
+    vae.vae.eval()
+    gc.collect()
+    torch.cuda.empty_cache()
+    return vae
+
+
+def get_transformer(model_path, device="cuda", weight_dtype=torch.bfloat16) -> WanModel:
+    config_path = os.path.join(model_path, "config.json")
+    transformer = WanModel.from_config(config_path).to(weight_dtype).to(device)
+
+    for file in os.listdir(model_path):
+        if file.endswith(".safetensors"):
+            file_path = os.path.join(model_path, file)
+            state_dict = load_file(file_path)
+            transformer.load_state_dict(state_dict, strict=False)
+            del state_dict
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    transformer.requires_grad_(False)
+    transformer.eval()
+    gc.collect()
+    torch.cuda.empty_cache()
+    return transformer
+
+
+def get_text_encoder(model_path, device="cuda", weight_dtype=torch.bfloat16) -> T5EncoderModel:
+    t5_model = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth")
+    tokenizer_path = os.path.join(model_path, "google", "umt5-xxl")
+    text_encoder = T5EncoderModel(checkpoint_path=t5_model, tokenizer_path=tokenizer_path).to(device).to(weight_dtype)
+    text_encoder.requires_grad_(False)
+    text_encoder.eval()
+    gc.collect()
+    torch.cuda.empty_cache()
+    return text_encoder
+
+
+def get_image_encoder(model_path, device="cuda", weight_dtype=torch.bfloat16) -> CLIPModel:
+    checkpoint_path = os.path.join(model_path, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
+    tokenizer_path = os.path.join(model_path, "xlm-roberta-large")
+    image_enc = CLIPModel(checkpoint_path, tokenizer_path).to(weight_dtype).to(device)
+    image_enc.requires_grad_(False)
+    image_enc.eval()
+    gc.collect()
+    torch.cuda.empty_cache()
+    return image_enc

package/#Tools/SkyReels-V2/skyreels_v2_infer/modules/attention.py

@@ -0,0 +1,179 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn_interface
+
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+import warnings
+
+__all__ = [
+    "flash_attention",
+    "attention",
+]
+
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == "cuda" and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+
+    q = half(q.flatten(0, 1))
+    q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True)
+
+    # preprocess key, value
+
+    k = half(k.flatten(0, 1))
+    v = half(v.flatten(0, 1))
+    k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True)
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.")
+
+    torch.cuda.nvtx.range_push(f"{list(q.shape)}-{list(k.shape)}-{list(v.shape)}-{q.dtype}-{k.dtype}-{v.dtype}")
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        )[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens])
+            .cumsum(0, dtype=torch.int32)
+            .to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+    torch.cuda.nvtx.range_pop()
+
+    # output
+    return x
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
+            )
+        attn_mask = None
+
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
+        )
+
+        out = out.transpose(1, 2).contiguous()
+        return out