upd

fastvideo/models/hunyuan/modules/models.py

@@ -20,6 +20,7 @@ from .posemb_layers import apply_rotary_emb
 from .token_refiner import SingleTokenRefiner
 
 
+counter=0
 class MMDoubleStreamBlock(nn.Module):
     """
     A multimodal dit block with separate modulation for
@@ -647,31 +648,82 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         img_seq_len = img.shape[1]
 
         freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
-        # --------------------- Pass through DiT blocks ------------------------
+        
-        for _, block in enumerate(self.double_blocks):
+        if False:
-            double_block_args = [img, txt, vec, freqs_cis, text_mask]
+            # --------------------- Pass through DiT blocks ------------------------
+            for _, block in enumerate(self.double_blocks):
+                double_block_args = [img, txt, vec, freqs_cis, text_mask]
 
-            img, txt = block(*double_block_args)
+                img, txt = block(*double_block_args)
 
-        # Merge txt and img to pass through single stream blocks.
+            # Merge txt and img to pass through single stream blocks.
-        x = torch.cat((img, txt), 1)
+            x = torch.cat((img, txt), 1)
-        if output_features:
+            if output_features:
-            features_list = []
+                features_list = []
-        if len(self.single_blocks) > 0:
+            if len(self.single_blocks) > 0:
-            for _, block in enumerate(self.single_blocks):
+                for _, block in enumerate(self.single_blocks):
-                single_block_args = [
+                    single_block_args = [
-                    x,
+                        x,
-                    vec,
+                        vec,
-                    txt_seq_len,
+                        txt_seq_len,
-                    (freqs_cos, freqs_sin),
+                        (freqs_cos, freqs_sin),
-                    text_mask,
+                        text_mask,
-                ]
+                    ]
 
-                x = block(*single_block_args)
+                    x = block(*single_block_args)
-                if output_features and _ % output_features_stride == 0:
+                    if output_features and _ % output_features_stride == 0:
-                    features_list.append(x[:, :img_seq_len, ...])
+                        features_list.append(x[:, :img_seq_len, ...])
 
-        img = x[:, :img_seq_len, ...]
+            img = x[:, :img_seq_len, ...]
+        
+        else:
+            # --------------------- Pass through DiT blocks ------------------------
+            SF, EF = 4, 8
+            SB, EB = 42, 46
+            block_id = 0
+            global counter
+            counter += 1
+
+            for _, block in enumerate(self.double_blocks):
+                if block_id == SF and counter%2 == 1:
+                    self.x_init = img
+                    self.y_init = txt
+                elif block_id == EF and counter%2 == 1:
+                    self.x_end = img
+                    self.y_end = txt
+                
+                if EF == block_id and counter%2 == 0 and counter < 25:
+                    img = img + self.x_end - self.x_init
+                    txt = txt + self.y_end - self.y_init
+                else:
+                    double_block_args = [img, txt, vec, freqs_cis, text_mask]
+                    img, txt = block(*double_block_args)
+                    
+                block_id += 1
+
+            x = torch.cat((img, txt), 1)
+            if output_features:
+                features_list = []
+
+            if len(self.single_blocks) > 0:
+                block_id = len(self.double_blocks)
+                for s_idx, block in enumerate(self.single_blocks):
+                    current_block_id = block_id + s_idx
+                    if current_block_id == SB and counter%2 == 1:
+                        self.x_init_single = x
+                    elif current_block_id == EB and counter%2 == 1:
+                        self.x_end_single = x
+                        
+                    if EB == current_block_id and counter%2 == 0 and counter >= 25:
+                        x = x + (self.x_end_single - self.x_init_single)
+                    else:
+                        single_block_args = [x, vec, txt_seq_len, (freqs_cos, freqs_sin), text_mask]
+                        x = block(*single_block_args)
+                        
+                    if output_features and s_idx % output_features_stride == 0:
+                        features_list.append(x[:, :img_seq_len, ...])
+                
+            img = x[:, :img_seq_len, ...]
 
         # ---------------------------- Final layer ------------------------------
         img = self.final_layer(img,
@@ -723,8 +775,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         }
         counts["attn+mlp"] = counts["double"] + counts["single"]
         return counts
-
+    
-
 #################################################################################
 #                             HunyuanVideo Configs                              #
 #################################################################################

fastvideo/ops/modulate/k_modulate.py

@@ -1,102 +0,0 @@
-import triton
-import triton.language as tl
-
-CONFIG_LIST = [
-    triton.Config({"BLOCK_M": 256, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 128, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 32}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128}, num_stages=2, num_warps=4),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 256}, num_stages=2, num_warps=4),
-]
-
-
-@triton.autotune(
-    configs=CONFIG_LIST,
-    key=["M", "N"],
-)
-@triton.jit
-def _modulate_fwd(
-        x_ptr,  # *Pointer* to first input vector.
-        output_ptr,  # *Pointer* to output vector.
-        scale_ptr,
-        shift_ptr,
-        m_stride,
-        s_stride,
-        M,
-        N,
-        seq_len,
-        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
-        BLOCK_N: tl.constexpr,
-        # NOTE: `constexpr` so it can be used as a shape value.
-):
-    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
-    s_rows = (row_id // seq_len) * BLOCK_M
-    col_id = tl.program_id(axis=1)
-    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
-
-    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
-    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
-    shift_ptrs = shift_ptr + s_rows * s_stride + cols[None, :]
-
-    col_mask = cols[None, :] < N
-    block_mask = (rows[:, None] < M) & col_mask
-    s_block_mask = col_mask
-    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
-    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
-    shift = tl.load(shift_ptrs, mask=s_block_mask, other=0.0)
-
-    output = x * (1 + scale) + shift
-    # Write x + y back to DRAM.
-    tl.store(output_ptr + rows[:, None] * m_stride + cols[None, :], output, mask=block_mask)
-
-
-@triton.autotune(
-    configs=CONFIG_LIST,
-    key=["M", "N"],
-)
-@triton.jit
-def _modulate_bwd(
-        dx_ptr,  # *Pointer* to first input vector.
-        x_ptr,
-        dy_ptr,  # *Pointer* to output vector.
-        scale_ptr,
-        dscale_ptr,
-        m_stride,
-        s_stride,
-        M,
-        N,
-        seq_len,
-        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
-        BLOCK_N: tl.constexpr,
-        # NOTE: `constexpr` so it can be used as a shape value.
-):
-    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
-    s_rows = (row_id // seq_len) * BLOCK_M
-    col_id = tl.program_id(axis=1)
-    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
-
-    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
-    dy_ptrs = dy_ptr + rows[:, None] * m_stride + cols[None, :]
-    dx_ptrs = dx_ptr + rows[:, None] * m_stride + cols[None, :]
-    dscale_ptrs = dscale_ptr + rows[:, None] * m_stride + cols[None, :]
-
-    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
-
-    col_mask = cols[None, :] < N
-    block_mask = (rows[:, None] < M) & col_mask
-    s_block_mask = col_mask
-    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
-    dy = tl.load(dy_ptrs, mask=block_mask, other=0.0)
-    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
-
-    dx = dy * (1 + scale)
-    dscale = dy * x
-    # Write x + y back to DRAM.
-    tl.store(dx_ptrs, dx, mask=block_mask)
-    tl.store(dscale_ptrs, dscale, mask=block_mask)

fastvideo/ops/modulate/modulate.py

@@ -1,63 +0,0 @@
-import torch
-import triton
-
-from fastvideo.ops.modulate.k_modulate import _modulate_fwd, _modulate_bwd
-
-
-class _FusedModulate(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, scale, shift):
-        y = torch.empty_like(x)
-        batch, seq_len, dim = x.shape
-        M = batch * seq_len
-        N = dim
-        x = x.view(-1, dim).contiguous()
-        scale = scale.view(-1, dim).contiguous()
-        shift = shift.view(-1, dim).contiguous()
-
-        def grid(meta):
-            return (
-                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
-                triton.cdiv(dim, meta["BLOCK_N"]),
-            )
-
-        _modulate_fwd[grid](x, y, scale, shift, x.stride(0), scale.stride(0), M, N, seq_len)
-
-        ctx.save_for_backward(x, scale)
-        ctx.batch = batch
-        ctx.seq_len = seq_len
-        ctx.dim = dim
-        return y
-
-    @staticmethod
-    def backward(ctx, dy):  # pragma: no cover  # this is covered, but called directly from C++
-        x, scale = ctx.saved_tensors
-
-        batch, seq_len, dim = ctx.batch, ctx.seq_len, ctx.dim
-        M = batch * seq_len
-        N = dim
-
-        # allocate output
-        dy = dy.contiguous()
-        dx = torch.empty_like(dy)
-        dscale = torch.empty_like(dy)
-        dshift = torch.sum(dy, dim=1)
-
-        def grid(meta):
-            return (
-                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
-                triton.cdiv(dim, meta["BLOCK_N"]),
-            )
-
-        _modulate_bwd[grid](dx, x, dy, scale, dscale, x.stride(0), scale.stride(0), M, N, seq_len)
-
-        dscale = torch.sum(dscale, dim=1)
-        return dx, dscale, dshift
-
-
-def fused_modulate(
-        x: torch.Tensor,
-        scale: torch.Tensor,
-        shift: torch.Tensor,
-) -> torch.Tensor:
-    return _FusedModulate.apply(x, scale, shift)

fastvideo/sample/sample_t2v_cogvideo.py

@@ -1,306 +0,0 @@
-"""
-This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline.
-The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v),
-and video-to-video (v2v), depending on the input data and different weight.
-
-- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
-- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
-- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V
-
-Running the Script:
-To run the script, use the following command with appropriate arguments:
-
-```bash
-$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v"
-```
-
-Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths.
-"""
-
-import argparse
-import logging
-import os
-from typing import Literal, Optional
-
-import torch
-from diffusers import (
-    CogVideoXDPMScheduler,
-    CogVideoXImageToVideoPipeline,
-    CogVideoXPipeline,
-    CogVideoXVideoToVideoPipeline,
-)
-from diffusers.utils import export_to_video, load_image, load_video
-
-
-logging.basicConfig(level=logging.INFO)
-
-# Recommended resolution for each model (width, height)
-RESOLUTION_MAP = {
-    # cogvideox1.5-*
-    "cogvideox1.5-5b-i2v": (1360, 768),
-    "cogvideox1.5-5b": (1360, 768),
-    # cogvideox-*
-    "cogvideox-5b-i2v": (720, 480),
-    "cogvideox-5b": (720, 480),
-    "cogvideox-2b": (720, 480),
-}
-
-
-def generate_video(
-    prompt: str,
-    model_path: str,
-    lora_path: str = None,
-    lora_rank: int = 128,
-    num_frames: int = 81,
-    width: Optional[int] = None,
-    height: Optional[int] = None,
-    output_path: str = "./output.mp4",
-    image_or_video_path: str = "",
-    num_inference_steps: int = 50,
-    guidance_scale: float = 6.0,
-    num_videos_per_prompt: int = 1,
-    dtype: torch.dtype = torch.bfloat16,
-    generate_type: str = Literal[
-        "t2v", "i2v", "v2v"
-    ],  # i2v: image to video, v2v: video to video
-    seed: int = 42,
-    fps: int = 16,
-):
-    """
-    Generates a video based on the given prompt and saves it to the specified path.
-
-    Parameters:
-    - prompt (str): The description of the video to be generated.
-    - model_path (str): The path of the pre-trained model to be used.
-    - lora_path (str): The path of the LoRA weights to be used.
-    - lora_rank (int): The rank of the LoRA weights.
-    - output_path (str): The path where the generated video will be saved.
-    - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
-    - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps.
-    - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V
-    - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V
-    - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
-    - num_videos_per_prompt (int): Number of videos to generate per prompt.
-    - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).
-    - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·
-    - seed (int): The seed for reproducibility.
-    - fps (int): The frames per second for the generated video.
-    """
-
-    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).
-    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
-    # function to use Multi GPUs.
-
-    image = None
-    video = None
-
-    model_name = model_path.split("/")[-1].lower()
-    desired_resolution = RESOLUTION_MAP[model_name]
-    if width is None or height is None:
-        width, height = desired_resolution
-        logging.info(
-            f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m"
-        )
-    elif (width, height) != desired_resolution:
-        if generate_type == "i2v":
-            # For i2v models, use user-defined width and height
-            logging.warning(
-                f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m"
-            )
-        else:
-            # Otherwise, use the recommended width and height
-            logging.warning(
-                f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m"
-            )
-            width, height = desired_resolution
-
-    if generate_type == "i2v":
-        pipe = CogVideoXImageToVideoPipeline.from_pretrained(
-            model_path, torch_dtype=dtype
-        )
-        image = load_image(image=image_or_video_path)
-    elif generate_type == "t2v":
-        pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype)
-    else:
-        pipe = CogVideoXVideoToVideoPipeline.from_pretrained(
-            model_path, torch_dtype=dtype
-        )
-        video = load_video(image_or_video_path)
-
-    # If you're using with lora, add this code
-    if lora_path:
-        pipe.load_lora_weights(
-            lora_path,
-            weight_name="pytorch_lora_weights.safetensors",
-            adapter_name="test_1",
-        )
-        pipe.fuse_lora(lora_scale=1 / lora_rank)
-
-    # 2. Set Scheduler.
-    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
-    # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B.
-    # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V.
-
-    # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-    pipe.scheduler = CogVideoXDPMScheduler.from_config(
-        pipe.scheduler.config, timestep_spacing="trailing"
-    )
-
-    # 3. Enable CPU offload for the model.
-    # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference
-    # and enable to("cuda")
-
-    # pipe.to("cuda")
-    pipe.enable_sequential_cpu_offload()
-    pipe.vae.enable_slicing()
-    pipe.vae.enable_tiling()
-
-    # 4. Generate the video frames based on the prompt.
-    # `num_frames` is the Number of frames to generate.
-    if generate_type == "i2v":
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            image=image,
-            # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480
-            num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
-            num_inference_steps=num_inference_steps,  # Number of inference steps
-            num_frames=num_frames,  # Number of frames to generate
-            use_dynamic_cfg=True,  # This id used for DPM scheduler, for DDIM scheduler, it should be False
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(
-                seed
-            ),  # Set the seed for reproducibility
-        ).frames[0]
-    elif generate_type == "t2v":
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            num_videos_per_prompt=num_videos_per_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            use_dynamic_cfg=True,
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(seed),
-        ).frames[0]
-    else:
-        video_generate = pipe(
-            height=height,
-            width=width,
-            prompt=prompt,
-            video=video,  # The path of the video to be used as the background of the video
-            num_videos_per_prompt=num_videos_per_prompt,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-            use_dynamic_cfg=True,
-            guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(
-                seed
-            ),  # Set the seed for reproducibility
-        ).frames[0]
-    export_to_video(video_generate, output_path, fps=fps)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Generate a video from a text prompt using CogVideoX"
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="The description of the video to be generated",
-    )
-    parser.add_argument(
-        "--image_or_video_path",
-        type=str,
-        default=None,
-        help="The path of the image to be used as the background of the video",
-    )
-    parser.add_argument(
-        "--model_path",
-        type=str,
-        default="THUDM/CogVideoX1.5-5B",
-        help="Path of the pre-trained model use",
-    )
-    parser.add_argument(
-        "--lora_path",
-        type=str,
-        default=None,
-        help="The path of the LoRA weights to be used",
-    )
-    parser.add_argument(
-        "--lora_rank", type=int, default=128, help="The rank of the LoRA weights"
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        default="./output.mp4",
-        help="The path save generated video",
-    )
-    parser.add_argument(
-        "--guidance_scale",
-        type=float,
-        default=6.0,
-        help="The scale for classifier-free guidance",
-    )
-    parser.add_argument(
-        "--num_inference_steps", type=int, default=50, help="Inference steps"
-    )
-    parser.add_argument(
-        "--num_frames",
-        type=int,
-        default=81,
-        help="Number of steps for the inference process",
-    )
-    parser.add_argument(
-        "--width", type=int, default=None, help="The width of the generated video"
-    )
-    parser.add_argument(
-        "--height", type=int, default=None, help="The height of the generated video"
-    )
-    parser.add_argument(
-        "--fps",
-        type=int,
-        default=16,
-        help="The frames per second for the generated video",
-    )
-    parser.add_argument(
-        "--num_videos_per_prompt",
-        type=int,
-        default=1,
-        help="Number of videos to generate per prompt",
-    )
-    parser.add_argument(
-        "--generate_type", type=str, default="t2v", help="The type of video generation"
-    )
-    parser.add_argument(
-        "--dtype", type=str, default="bfloat16", help="The data type for computation"
-    )
-    parser.add_argument(
-        "--seed", type=int, default=42, help="The seed for reproducibility"
-    )
-
-    args = parser.parse_args()
-    dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
-    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
-    generate_video(
-        prompt=args.prompt,
-        model_path=args.model_path,
-        lora_path=args.lora_path,
-        lora_rank=args.lora_rank,
-        output_path=args.output_path,
-        num_frames=args.num_frames,
-        width=args.width,
-        height=args.height,
-        image_or_video_path=args.image_or_video_path,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        num_videos_per_prompt=args.num_videos_per_prompt,
-        dtype=dtype,
-        generate_type=args.generate_type,
-        seed=args.seed,
-        fps=args.fps,
-    )

tests/test_kernels/benchmark.py

@@ -1,297 +0,0 @@
-# Copyright (c) 2023, Tri Dao.
-"""Useful functions for writing test code."""
-
-import torch
-import torch.utils.benchmark as benchmark
-
-
-def get_abs_err(x, y):
-    return (x - y).flatten().abs().max().item()
-
-
-def get_err_ratio(x, y):
-    err = (x - y).flatten().square().mean().sqrt().item()
-    base = x.flatten().square().mean().sqrt().item()
-    return err / base
-
-
-def assert_close(prefix, ref, tri, ratio):
-    msg = f"{prefix} diff: {get_abs_err(ref, tri):.6f} ratio: {get_err_ratio(ref, tri):.6f}"
-    print(msg)
-    assert get_err_ratio(ref, tri) < ratio, msg
-
-
-def benchmark_forward(
-    fn,
-    *inputs,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Forward pass")
-
-    def amp_wrapper(*inputs, **kwinputs):
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            fn(*inputs, **kwinputs)
-
-    t = benchmark.Timer(
-        stmt="fn_amp(*inputs, **kwinputs)",
-        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_backward(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Backward pass")
-    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-        y = fn(*inputs, **kwinputs)
-        if type(y) is tuple:
-            y = y[0]
-    if grad is None:
-        grad = torch.randn_like(y)
-    else:
-        if grad.shape != y.shape:
-            raise RuntimeError("Grad shape does not match output shape")
-
-    def f(*inputs, y, grad):
-        # Set .grad to None to avoid extra operation of gradient accumulation
-        for x in inputs:
-            if isinstance(x, torch.Tensor):
-                x.grad = None
-        y.backward(grad, retain_graph=True)
-
-    t = benchmark.Timer(
-        stmt="f(*inputs, y=y, grad=grad)",
-        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_combined(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    if verbose:
-        print(desc, "- Forward + Backward pass")
-    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-        y = fn(*inputs, **kwinputs)
-        if type(y) is tuple:
-            y = y[0]
-    if grad is None:
-        grad = torch.randn_like(y)
-    else:
-        if grad.shape != y.shape:
-            raise RuntimeError("Grad shape does not match output shape")
-
-    def f(grad, *inputs, **kwinputs):
-        for x in inputs:
-            if isinstance(x, torch.Tensor):
-                x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            y = fn(*inputs, **kwinputs)
-            if type(y) is tuple:
-                y = y[0]
-        y.backward(grad, retain_graph=True)
-
-    t = benchmark.Timer(
-        stmt="f(grad, *inputs, **kwinputs)",
-        globals={
-            "f": f,
-            "fn": fn,
-            "inputs": inputs,
-            "grad": grad,
-            "kwinputs": kwinputs,
-        },
-        num_threads=torch.get_num_threads(),
-    )
-    m = t.timeit(repeats)
-    if verbose:
-        print(m)
-    return t, m
-
-
-def benchmark_fwd_bwd(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    return (
-        benchmark_forward(
-            fn,
-            *inputs,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_backward(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-    )
-
-
-def benchmark_all(
-    fn,
-    *inputs,
-    grad=None,
-    repeats=10,
-    desc="",
-    verbose=True,
-    amp=False,
-    amp_dtype=torch.float16,
-    **kwinputs,
-):
-    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
-    return (
-        benchmark_forward(
-            fn,
-            *inputs,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_backward(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-        benchmark_combined(
-            fn,
-            *inputs,
-            grad=grad,
-            repeats=repeats,
-            desc=desc,
-            verbose=verbose,
-            amp=amp,
-            amp_dtype=amp_dtype,
-            **kwinputs,
-        ),
-    )
-
-
-def pytorch_profiler(
-    fn,
-    *inputs,
-    trace_filename=None,
-    backward=False,
-    amp=False,
-    amp_dtype=torch.float16,
-    cpu=False,
-    verbose=True,
-    **kwinputs,
-):
-    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
-    if backward:
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-            g = torch.randn_like(out)
-    for _ in range(30):  # Warm up
-        if backward:
-            for x in inputs:
-                if isinstance(x, torch.Tensor):
-                    x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-        # Backward should be done outside autocast
-        if backward:
-            out.backward(g, retain_graph=True)
-    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
-        torch.profiler.ProfilerActivity.CUDA
-    ]
-    with torch.profiler.profile(
-        activities=activities,
-        record_shapes=True,
-        # profile_memory=True,
-        with_stack=True,
-    ) as prof:
-        if backward:
-            for x in inputs:
-                if isinstance(x, torch.Tensor):
-                    x.grad = None
-        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
-            out = fn(*inputs, **kwinputs)
-            if type(out) is tuple:
-                out = out[0]
-        if backward:
-            out.backward(g, retain_graph=True)
-    if verbose:
-        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
-        print(prof.key_averages().table(row_limit=50))
-    if trace_filename is not None:
-        prof.export_chrome_trace(trace_filename)
-
-
-def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
-    torch.cuda.empty_cache()
-    torch.cuda.reset_peak_memory_stats()
-    torch.cuda.synchronize()
-    fn(*inputs, **kwinputs)
-    torch.cuda.synchronize()
-    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
-    if verbose:
-        print(f"{desc} max memory: {mem}GB")
-    torch.cuda.empty_cache()
-    return mem

tests/test_kernels/benchmark_modulate.py

@@ -1,83 +0,0 @@
-import torch
-from torch import nn
-
-from benchmark import (
-    benchmark_combined,
-    benchmark_forward,
-    benchmark_backward,
-    assert_close,
-)
-from fastvideo.ops.modulate.modulate import fused_modulate
-
-
-def torch_modulate(x, scale, shift):
-    return x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1)
-
-
-# from flash_attn import flash_attn_func
-
-
-def time_fwd(func, *args, **kwargs):
-    time_fb = benchmark_forward(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-def time_fwd_bwd(func, *args, **kwargs):
-    time_fb = benchmark_combined(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-def time_bwd(func, *args, **kwargs):
-    time_fb = benchmark_backward(func, *args, **kwargs)
-    return time_fb[1].mean
-
-
-device = "cuda"
-dtype = torch.bfloat16
-
-
-batch_sizes = [1, 8, 32]
-seq_lengths = [128, 512, 1024]
-hidden_dims = [768, 1024, 2048]
-
-# methods = (["torch", "triton", "thunderkitten"])
-methods = ["torch", "triton"]
-time_f = {}
-time_b = {}
-time_f_b = {}
-speed_f = {}
-speed_b = {}
-speed_f_b = {}
-for B in batch_sizes:
-    for T in seq_lengths:
-        for D in hidden_dims:
-            config = (B, T, D)
-
-            norm_func = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6)
-            x = torch.randn(B, T, D, device="cuda", requires_grad=True, dtype=dtype)
-            x = norm_func(x.to(torch.float32)).to(dtype)
-            shift = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
-            scale = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
-            # test torch
-            o_ref = torch_modulate(x, scale, shift)
-            o_ref.sum().backward(retain_graph=True)
-            f_b = time_fwd_bwd(torch_modulate, x, scale, shift, verbose=False)
-            time_f_b[config, "torch"] = f_b
-            # test triton
-            o2 = fused_modulate(x, scale, shift)
-            o2.sum().backward(retain_graph=True)
-            f_b = time_fwd_bwd(fused_modulate, x, scale, shift, verbose=False)
-            time_f_b[config, "triton"] = f_b
-            # test if the results are close
-            assert_close("  o", o_ref, o2, 0.005)
-            # time_f_b[config, "thunderkitten"] = f_b
-
-            print(f"### batch size={B}, seq length={T}, B={B}, hidden dim={D} ###")
-            for method in methods:
-                # time_f_b[config, method] = time_f[config, method] + time_b[config, method]
-                print(
-                    f"{method:>50} fwd + bwd:\t {time_f_b[config, method]*1000:>6.4f} ms "
-                )
-
-# with open('flash2_attn_time.plk', 'wb') as fp:
-#     pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL)