Initial ABot-World interactive rollout demo

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/desert_valley.png filter=lfs diff=lfs merge=lfs -text
+examples/example.png filter=lfs diff=lfs merge=lfs -text
+examples/forest_stream.png filter=lfs diff=lfs merge=lfs -text
+examples/mountain_meadow.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,27 @@
 ---
-title: Abot World Interactive
-emoji: 🏃
-colorFrom: pink
-colorTo: indigo
+title: ABot-World Interactive
+emoji: 🌍
+colorFrom: gray
+colorTo: red
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+short_description: Interactive action-conditioned world rollout from one image
+python_version: "3.10"
+startup_duration_timeout: 40m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🌍 ABot-World — Interactive World Rollout
+
+Turn a single image into a navigable world. Provide a first-frame image, describe
+the scene, and choose an action (WASD to move / IJKL to look). The model
+autoregressively rolls out an action-conditioned video of the world.
+
+- **Model:** [`acvlab/ABot-World-0-5B-LF`](https://huggingface.co/acvlab/ABot-World-0-5B-LF) (built on Wan2.2-TI2V-5B)
+- **Code:** [amap-cvlab/ABot-World](https://github.com/amap-cvlab/ABot-World)
+- **Project:** [ABot-World](https://amap-cvlab.github.io/ABot-World/)
+
+This Space adapts the upstream real-time interactive web client into a request/response
+rollout suitable for ZeroGPU: instead of live keyboard streaming, you pick one action
+to hold and a rollout length, and receive a video. Inference is causal / block-wise,
+matching the upstream `CausalInferencePipeline`.

app.py

@@ -0,0 +1,328 @@
+"""
+ABot-World — Interactive Action-Conditioned World Rollout (ZeroGPU Gradio demo)
+
+Given a first-frame image, a scene prompt, and an action (WASD / IJKL), the model
+autoregressively rolls out an action-conditioned video of a navigable world.
+
+Upstream: https://github.com/amap-cvlab/ABot-World
+Model:    https://huggingface.co/acvlab/ABot-World-0-5B-LF  (built on Wan2.2-TI2V-5B)
+"""
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # must precede torch / CUDA-touching imports
+
+import sys
+import time
+import tempfile
+import threading
+from pathlib import Path
+
+import numpy as np
+import torch
+import imageio
+import gradio as gr
+from PIL import Image
+from omegaconf import OmegaConf
+from huggingface_hub import snapshot_download
+
+# ── Repo paths ───────────────────────────────────────────────────────────────
+APP_DIR = Path(__file__).resolve().parent
+if str(APP_DIR) not in sys.path:
+    sys.path.insert(0, str(APP_DIR))
+
+MODEL_ID = "acvlab/ABot-World-0-5B-LF"
+CKPT_DIR = APP_DIR / "checkpoints" / "ABot-World-0-5B-LF"
+
+# ── Stream / rollout configuration ───────────────────────────────────────────
+# 704x1280 is the native training resolution used by the upstream web client.
+# We keep it identical so the model behaves as intended.
+STREAM_HEIGHT = 704
+STREAM_WIDTH = 1280
+VIDEO_FPS = 12
+
+# Actions map to the 8-key one-hot the model was trained on (W A S D I J K L).
+ACTIONS = {
+    "Forward (W)":            {"W": True},
+    "Backward (S)":           {"S": True},
+    "Turn left (A)":          {"A": True},
+    "Turn right (D)":         {"D": True},
+    "Look up (I)":            {"I": True},
+    "Look down (K)":          {"K": True},
+    "Pan left (J)":           {"J": True},
+    "Pan right (L)":          {"L": True},
+    "Forward + turn left":    {"W": True, "A": True},
+    "Forward + turn right":   {"W": True, "D": True},
+    "Idle (no action)":       {},
+}
+
+DEFAULT_PROMPT = (
+    "A realistic outdoor world scene with a navigable path, natural lighting, "
+    "detailed ground texture, and stable forward motion."
+)
+
+# ── Download weights (once, at startup) ──────────────────────────────────────
+print(f"[startup] downloading {MODEL_ID} weights ...", flush=True)
+snapshot_download(
+    repo_id=MODEL_ID,
+    repo_type="model",
+    local_dir=str(CKPT_DIR),
+)
+print("[startup] weights downloaded.", flush=True)
+
+# Point the taehv VAE at the downloaded checkpoints via env (config uses these).
+os.environ["TAEW2_2_CHECKPOINT"] = str(CKPT_DIR / "taew2_2.pth")
+
+# ── Build the causal streaming pipeline (module scope, eager to CUDA) ─────────
+from pipeline import CausalInferencePipeline
+from utils.misc import set_seed
+from utils.wan_wrapper import create_vae_from_config
+
+_CONFIG_PATH = APP_DIR / "configs" / "long_forcing_dmd.yaml"
+_DEFAULT_CFG_PATH = APP_DIR / "configs" / "default_config.yaml"
+
+
+def _build_config():
+    cfg = OmegaConf.merge(
+        OmegaConf.load(str(_DEFAULT_CFG_PATH)),
+        OmegaConf.load(str(_CONFIG_PATH)),
+    )
+    # Rewrite every checkpoint path to the absolute downloaded location.
+    cfg.taew2_2_checkpoint = str(CKPT_DIR / "taew2_2.pth")
+    cfg.lightvae_encoder_checkpoint = str(CKPT_DIR / "Wan2.2_VAE.pth")
+    cfg.model_kwargs.model_name = str(CKPT_DIR)
+    cfg.text_encoder_kwargs.tokenizer_path = str(CKPT_DIR / "google" / "umt5-xxl") + "/"
+    cfg.text_encoder_kwargs.encoder_pth_path = str(CKPT_DIR / "models_t5_umt5-xxl-enc-bf16.pth")
+    cfg.vae_kwargs.pretrained_path = str(CKPT_DIR / "Wan2.2_VAE.pth")
+    cfg.vae_type = "taew2_2"
+    # FP8 GEMM relies on Triton kernels tuned for RTX 5090; disable for portability.
+    cfg.use_fp8_gemm = False
+    return cfg
+
+
+print("[startup] building pipeline ...", flush=True)
+set_seed(42)
+torch.set_grad_enabled(False)
+
+CONFIG = _build_config()
+_vae = create_vae_from_config(CONFIG)
+pipeline = CausalInferencePipeline(CONFIG, device=torch.device("cuda"), vae=_vae)
+
+# Optional Triton norm/rope speedups from the upstream repo. They work on the
+# training GPU; wrap defensively so a kernel incompatibility can't break boot.
+try:
+    from wan.modules.helios_kernels import (
+        replace_all_norms_with_flash_norms,
+        replace_rope_with_flash_rope,
+    )
+    replace_all_norms_with_flash_norms(pipeline.generator.model)
+    replace_rope_with_flash_rope()
+    print("[startup] helios Triton kernels enabled.", flush=True)
+except Exception as e:  # pragma: no cover
+    print(f"[startup] helios kernels disabled ({e!r}); using eager norms/rope.", flush=True)
+
+pipeline = pipeline.to(dtype=torch.bfloat16)
+# Full model on GPU — we have the VRAM budget, no CPU offload.
+pipeline.text_encoder.to(device="cuda")
+pipeline.generator.to(device="cuda")
+pipeline.vae.to(device="cuda")
+if pipeline.encoder is not None:
+    pipeline.encoder.to(device="cuda")
+pipeline.torch_dtype = torch.bfloat16
+print("[startup] pipeline ready.", flush=True)
+
+NUM_FPB = int(getattr(CONFIG, "num_frame_per_block", 3))
+_vae_for_shape = pipeline.encoder if pipeline.encoder is not None else pipeline.vae
+LATENT_CHANNELS = _vae_for_shape.z_dim
+UPSAMPLE = getattr(_vae_for_shape, "upsampling_factor", 16)
+LATENT_H = STREAM_HEIGHT // UPSAMPLE
+LATENT_W = STREAM_WIDTH // UPSAMPLE
+
+_infer_lock = threading.Lock()
+
+
+def _decode_block_to_frames(lat_block):
+    frames = []
+
+    class _Cap:
+        def append_data(self, f):
+            frames.append(f)
+
+    pipeline.decode_block_and_write(lat_block, _Cap())
+    return frames
+
+
+def _estimate_duration(image, prompt, action_label, num_blocks, seed, *args, **kwargs):
+    try:
+        n = int(num_blocks)
+    except Exception:
+        n = 8
+    return min(230, 40 + int(n) * 12)
+
+
+@spaces.GPU(duration=_estimate_duration)
+def rollout(image, prompt, action_label, num_blocks, seed,
+            progress=gr.Progress(track_tqdm=False)):
+    """Roll out an action-conditioned world video from a first-frame image.
+
+    Args:
+        image: first-frame image (PIL / filepath) to start the world from.
+        prompt: text description of the scene and desired dynamics.
+        action_label: which navigation action to hold during the rollout.
+        num_blocks: number of autoregressive latent blocks to generate (video length).
+        seed: RNG seed for reproducibility.
+
+    Returns:
+        Path to the generated MP4 video of the world rollout.
+    """
+    if image is None:
+        raise gr.Error("Please provide a first-frame image to start the world.")
+
+    prompt = (prompt or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
+    keys = ACTIONS.get(action_label, {})
+    num_blocks = max(1, min(40, int(num_blocks)))
+    seed = int(seed)
+
+    device = torch.device("cuda")
+    set_seed(seed)
+
+    # Persist the first-frame image to a path (pipeline reads a file path).
+    if isinstance(image, str):
+        first_frame_path = image
+    else:
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(np.asarray(image))
+        tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        image.convert("RGB").save(tmp_img.name)
+        first_frame_path = tmp_img.name
+
+    with _infer_lock:
+        noise = torch.randn(
+            [1, NUM_FPB, LATENT_CHANNELS, LATENT_H, LATENT_W],
+            device=device, dtype=torch.bfloat16,
+        )
+
+        pipeline.set_prompts([prompt], device=device)
+        # No pre-generated head/left/right/front/back refs -> zero ref_latents
+        # (model ignores them; this matches the upstream fallback path).
+        pipeline.set_ref_latent_mask_from_exists_paths(
+            ref_dir=str(APP_DIR / "__no_such_ref_dir__"), device=device,
+        )
+        pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
+                              device=device, initial_latent=None)
+        pipeline.set_first_frame_latent(
+            first_frame_path, height=STREAM_HEIGHT, width=STREAM_WIDTH, device=device,
+        )
+        pipeline.set_act(keys, height=STREAM_HEIGHT, width=STREAM_WIDTH,
+                         num_frames=NUM_FPB, device=device)
+
+        out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        writer = imageio.get_writer(
+            out_path, fps=VIDEO_FPS, format="FFMPEG", codec="libx264",
+            ffmpeg_params=["-crf", "18", "-preset", "fast", "-pix_fmt", "yuv420p"],
+        )
+        try:
+            t0 = time.perf_counter()
+            for b in range(num_blocks):
+                pipeline.set_act(keys, height=STREAM_HEIGHT, width=STREAM_WIDTH,
+                                 num_frames=NUM_FPB, device=device)
+                lat_block = pipeline.generate_next_block(noise)
+                noise = torch.randn_like(noise)
+                frames = _decode_block_to_frames(lat_block)
+                for f in frames:
+                    writer.append_data(np.asarray(f))
+                progress((b + 1) / num_blocks, desc=f"block {b + 1}/{num_blocks}")
+            elapsed = time.perf_counter() - t0
+            print(f"[rollout] {num_blocks} blocks in {elapsed:.1f}s "
+                  f"(action={keys}, seed={seed})", flush=True)
+        finally:
+            writer.close()
+            # Reset streaming caches for the next request.
+            try:
+                pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
+                                      device=device, initial_latent=None)
+            except Exception:
+                pass
+            try:
+                if hasattr(pipeline.vae, "model") and hasattr(pipeline.vae.model, "clear_cache"):
+                    pipeline.vae.model.clear_cache()
+            except Exception:
+                pass
+
+    return out_path
+
+
+# ── UI ───────���────────────────────────────────────────────────────────────────
+CSS = """
+#col-container { max-width: 1180px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+INTRO = """
+# 🌍 ABot-World — Interactive World Rollout
+Turn a single image into a navigable world. Pick a **first frame**, describe the
+**scene**, choose an **action** to hold, and the model autoregressively rolls out
+an action-conditioned video.
+
+Built on [`acvlab/ABot-World-0-5B-LF`](https://huggingface.co/acvlab/ABot-World-0-5B-LF)
+(Wan2.2-TI2V-5B) · [project page](https://amap-cvlab.github.io/ABot-World/) ·
+[code](https://github.com/amap-cvlab/ABot-World)
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(INTRO)
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_in = gr.Image(label="First frame", type="filepath", height=300)
+                prompt_in = gr.Textbox(
+                    label="Scene prompt", value=DEFAULT_PROMPT, lines=3,
+                )
+                action_in = gr.Dropdown(
+                    label="Action (held during rollout)",
+                    choices=list(ACTIONS.keys()),
+                    value="Forward (W)",
+                )
+                run_btn = gr.Button("Explore the world", variant="primary")
+                with gr.Accordion("Advanced settings", open=False):
+                    blocks_in = gr.Slider(
+                        label="Rollout length (latent blocks)",
+                        minimum=2, maximum=24, step=1, value=8,
+                        info="More blocks = longer video, longer runtime.",
+                    )
+                    seed_in = gr.Number(label="Seed", value=42, precision=0)
+            with gr.Column(scale=1):
+                video_out = gr.Video(label="World rollout", autoplay=True)
+
+        gr.Examples(
+            examples=[
+                ["examples/example.png", DEFAULT_PROMPT, "Forward (W)", 8, 42],
+                ["examples/desert_valley.png",
+                 "A third-person view walking through a desert river valley, "
+                 "sunlit rocky terrain, stable forward motion.",
+                 "Forward (W)", 8, 42],
+                ["examples/mountain_meadow.png",
+                 "A first-person hike along an alpine meadow trail, lush grass, "
+                 "distant peaks, gentle forward motion.",
+                 "Forward + turn right", 8, 7],
+                ["examples/forest_stream.png",
+                 "A first-person walk stepping across stones in a forest stream, "
+                 "dappled light, moss-covered rocks.",
+                 "Turn left (A)", 8, 21],
+            ],
+            inputs=[image_in, prompt_in, action_in, blocks_in, seed_in],
+            outputs=video_out,
+            fn=rollout,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run_btn.click(
+        fn=rollout,
+        inputs=[image_in, prompt_in, action_in, blocks_in, seed_in],
+        outputs=video_out,
+        api_name="rollout",
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(mcp_server=True)

configs/default_config.yaml

@@ -0,0 +1,6 @@
+vae_type: "taew2_2"
+use_fp8_gemm: true
+quant_type: "fp8-per-token"
+
+taew2_2_checkpoint: "checkpoints/ABot-World-0-5B-LF/taew2_2.pth"
+lightvae_encoder_checkpoint: "checkpoints/ABot-World-0-5B-LF/Wan2.2_VAE.pth"

configs/long_forcing_dmd.yaml

@@ -0,0 +1,36 @@
+model_type: ci2v
+
+denoising_step_list:
+  - 1000
+  - 750
+  - 500
+  - 250
+warp_denoising_step: true
+context_noise: 0
+
+image_or_video_shape:
+  - 1
+  - 31
+  - 48
+  - 44
+  - 80
+independent_first_frame: true
+num_frame_per_block: 3
+
+ref_num_slots: 5
+ref_resolution: 512
+
+model_kwargs:
+  model_name: checkpoints/ABot-World-0-5B-LF
+  timestep_shift: 5.0
+  local_attn_size: 21
+  use_relative_rope: true
+  downscale_factor_control_adapter: 16
+
+text_encoder_kwargs:
+  tokenizer_path: checkpoints/ABot-World-0-5B-LF/google/umt5-xxl/
+  encoder_pth_path: checkpoints/ABot-World-0-5B-LF/models_t5_umt5-xxl-enc-bf16.pth
+
+vae_kwargs:
+  vae_type: Wan2.2_VAE
+  pretrained_path: checkpoints/ABot-World-0-5B-LF/Wan2.2_VAE.pth

pipeline/__init__.py

@@ -0,0 +1,5 @@
+from .causal_inference import CausalInferencePipeline
+
+__all__ = [
+    "CausalInferencePipeline",
+]

pipeline/causal_inference.py

@@ -0,0 +1,822 @@
+from typing import List, Optional
+from pathlib import Path
+import math
+import time
+import numpy as np
+import torch
+from torch.nn import functional as Functional
+from PIL import Image, ImageOps
+from einops import repeat
+
+from utils.wan_wrapper import (
+    WanDiffusionWrapper,
+    WanTextEncoder,
+    WanVAEWrapper,
+    model_kwargs_with_relative_rope,
+)
+
+from utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, \
+    move_model_to_device_with_memory_preservation
+import tqdm
+
+# Supported image extensions for reference image loading (case-insensitive)
+SUPPORTED_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif")
+
+
+def _resolve_ref_image_path(directory: Path, name: str) -> Optional[Path]:
+    """Resolve a reference image name to an actual file with any supported extension."""
+    for ext in SUPPORTED_IMAGE_EXTS:
+        for variant in (ext, ext.upper()):
+            candidate = directory / f"{name}{variant}"
+            if candidate.is_file():
+                return candidate
+    return None
+
+
+class CausalInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None,
+            need_vae=True
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        model_kwargs = model_kwargs_with_relative_rope(args)
+        model_kwargs["local_attn_size"] = model_kwargs.get("local_attn_size", getattr(args, "local_attn_size", -1))
+        self.generator = WanDiffusionWrapper(
+            **model_kwargs, is_causal=True, model_type=args.model_type,
+        ) if generator is None else generator
+        self.text_encoder = WanTextEncoder(tokenizer_path=args.text_encoder_kwargs.tokenizer_path,
+                                           encoder_pth_path=args.text_encoder_kwargs.encoder_pth_path) if text_encoder is None else text_encoder
+        if vae is not None and getattr(args, "model_type", "t2v") in ["ci2v", "ti2v"]:
+            # self.encoder = WanVAEWrapper()
+            self.encoder = WanVAEWrapper(**getattr(args, "vae_kwargs", {}))
+        else:
+            self.encoder = None
+        # self.vae = WanVAEWrapper() if vae is None else vae
+        self.vae = WanVAEWrapper(**getattr(args, "vae_kwargs", {})) if vae is None else vae
+
+        # Step 2: Initialize all causal hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.denoising_step_list = torch.tensor(
+            args.denoising_step_list, dtype=torch.long)
+        if args.warp_denoising_step:
+            timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = (
+            args.image_or_video_shape[3] * args.image_or_video_shape[4]
+        ) // (
+            self.generator.model.patch_size[1] * self.generator.model.patch_size[2]
+        )
+
+        self.kv_cache1 = None
+        self.args = args
+        self.device = device
+        self.torch_dtype = torch.bfloat16
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+        self.pyramid_sample_ratio = getattr(args, "pyramid_sample_ratio", None)
+
+        self.conditional_dict = None
+        self.crossattn_cache = None
+
+        # Streaming state
+        self.current_start_frame = 0
+        self.num_input_frames = 0
+        self._stream_block_diffusion_times: List[float] = []
+        self._stream_block_decode_times: List[float] = []
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def _clear_generation_caches(self):
+        self.kv_cache1 = None
+        self.crossattn_cache = None
+        model = getattr(self.generator, "model", None)
+        if model is not None and hasattr(model, "clear_cache"):
+            model.clear_cache()
+
+    def _decode_output_to_video(self, output: torch.Tensor) -> torch.Tensor:
+        self._clear_generation_caches()
+        offload = bool(getattr(self.args, "offload_generator_for_decode", False))
+        generator_device = None
+        if offload:
+            try:
+                generator_device = next(self.generator.parameters()).device
+            except StopIteration:
+                generator_device = None
+            if generator_device is not None and generator_device.type == "cuda":
+                self.generator.to("cpu")
+                torch.cuda.empty_cache()
+
+        try:
+            return self.vae.decode_to_pixel(output, use_cache=False)
+        finally:
+            if (
+                offload
+                and generator_device is not None
+                and generator_device.type == "cuda"
+                and bool(getattr(self.args, "restore_generator_after_decode", False))
+            ):
+                self.generator.to(generator_device)
+
+    def get_condition_split(self, conditional_dict, start, end):
+        # print(f"Getting condition split for frames {start} to {end}",conditional_dict.keys())
+        device = conditional_dict["prompt_embeds"].device
+        act_context = conditional_dict.get("act_context")
+        y = conditional_dict.get("y")
+        clip_fea = conditional_dict.get("clip_fea")
+        first_frame_latents = conditional_dict.get("first_frame_latents")
+        split_conditional_dict = {
+            "prompt_embeds": conditional_dict["prompt_embeds"],
+            "act_context": act_context[:, :, start:end].to(device) if act_context is not None else None,
+            "act_context_scale": conditional_dict.get("act_context_scale", 0.0),
+            "y": y[:, :, start:end].to(device) if y is not None else None,
+            "clip_fea": clip_fea,
+            "first_frame_latents": first_frame_latents,
+            "ref_latents": conditional_dict.get("ref_latents", None),
+            "ref_mask": conditional_dict.get("ref_mask", None),
+        }
+        return split_conditional_dict
+
+
+    def set_prompts(self, text_prompts: List[str], device: torch.device):
+        """
+        刷新你的世界：重新算 conditional_dict，并让 crossattn_cache 在下次 generator 调用时重建。
+        """
+        self.conditional_dict = self.text_encoder(text_prompts=text_prompts, device=device)
+
+        # 强制 cross-attn cache 重新 init（因为 prompt 变了）
+        if hasattr(self, "crossattn_cache") and self.crossattn_cache is not None:
+            for b in range(self.num_transformer_blocks):
+                self.crossattn_cache[b]["is_init"] = False
+
+    def set_ref_latent_mask_from_exists_paths(
+        self,
+        ref_dir: str,
+        device: Optional[torch.device] = None,
+    ):
+        """Set ref_latents and ref_mask from a directory of pre-generated reference images.
+
+        ``ref_dir`` is expected to contain the 5 output images (head, left,
+        right, front, back) in any common image format (.jpg, .jpeg, .png,
+        .webp, .bmp, .tiff). The method only performs the VAE encoding step.
+
+        If any expected file is missing, falls back to zero ref_latents with zero
+        ref_mask (model ignores them).
+        Should be called after set_prompts().
+        """
+        if self.conditional_dict is None:
+            raise RuntimeError("call set_prompts first")
+        dev = device or self.device
+        num_slots = getattr(self.args, "ref_num_slots", 5)
+        ref_resolution = getattr(self.args, "ref_resolution", 512)
+        # Get z_dim and upsampling_factor from encoder (preferred) or vae
+        vae_for_shape = self.encoder if self.encoder is not None else self.vae
+        z_dim = getattr(vae_for_shape, "z_dim", 48)
+        upsampling = getattr(vae_for_shape, "upsampling_factor", 16)
+        ref_h = ref_w = ref_resolution // upsampling
+
+        ref_names = ["head", "left", "right", "front", "back"]
+        ref_dir_path = Path(ref_dir)
+
+        # Resolve each reference name to an actual file with any supported extension
+        resolved_paths = []
+        missing = []
+        for name in ref_names[:num_slots]:
+            found = _resolve_ref_image_path(ref_dir_path, name)
+            if found is not None:
+                resolved_paths.append(found)
+            else:
+                missing.append(name)
+
+        if missing:
+            print(f"[set_ref_latent_mask_from_exists_paths] Missing ref images {missing} "
+                  f"in {ref_dir}, falling back to zero ref_latents.")
+            ref_latents = torch.zeros(
+                [1, num_slots, z_dim, 1, ref_h, ref_w],
+                dtype=torch.bfloat16, device=dev,
+            )
+            ref_mask = torch.zeros([1, num_slots], dtype=torch.float32, device=dev)
+            source = "zeros_missing"
+        else:
+            vae = self.encoder if self.encoder is not None else self.vae
+            vae_device = next(vae.parameters()).device
+
+            latents_list = []
+            for img_path in resolved_paths:
+                img = Image.open(img_path).convert("RGB")
+                img = img.resize((ref_resolution, ref_resolution))
+                image = self.preprocess_image(img, device=dev, torch_dtype=torch.bfloat16)
+                ref_pixel = image.unsqueeze(2)  # [1, C, 1, H, W]
+                with torch.no_grad():
+                    latent = vae.encode_to_latent(ref_pixel.to(vae_device))
+                latent_slot = latent[0, 0].unsqueeze(1).to(device=dev, dtype=torch.bfloat16)
+                latents_list.append(latent_slot)
+
+            ref_latents = torch.stack(latents_list, dim=0).unsqueeze(0)  # [1, K, C_lat, 1, H_lat, W_lat]
+            ref_mask = torch.ones([1, num_slots], dtype=torch.float32, device=dev)
+            source = "exists_paths"
+
+        self.conditional_dict["ref_latents"] = ref_latents
+        self.conditional_dict["ref_mask"] = ref_mask
+        print(f"[set_ref_latent_mask_from_exists_paths] ref_latents={ref_latents.shape}, "
+              f"ref_mask={ref_mask.shape}, "
+              f"source={source}")
+
+    def preprocess_image(self, image, torch_dtype=torch.bfloat16, device=None, pattern="B C H W", min_value=-1, max_value=1):
+        # Transform a PIL.Image to torch.Tensor
+        image = torch.Tensor(np.array(image, dtype=np.float32))
+        image = image.to(dtype=torch_dtype, device=device)
+        image = image * ((max_value - min_value) / 255) + min_value
+        image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {}))
+        return image
+
+    @torch.no_grad()
+    def set_first_frame_latent(
+        self,
+        ref_image_path: str,
+        height: int = 480,
+        width: int = 832,
+        device: Optional[torch.device] = None,
+    ):
+        """
+        设置首帧 clean latent：从路径加载图片，通过 vae 编码得到 latent，
+        存入 self.conditional_dict["first_frame_latents"]。
+        在 generate_next_block 的第一个 block 中会用它替换噪声的首帧。
+        """
+        if self.conditional_dict is None:
+            raise RuntimeError("call set_prompts first")
+        dev = device or self.device
+        if self.encoder is not None:
+            vae = self.encoder
+        else:
+            vae = self.vae
+
+        if not ref_image_path:
+            self.conditional_dict["first_frame_latents"] = None
+        else:
+            input_image = Image.open(ref_image_path).convert("RGB")
+            # Preserve the source aspect ratio: scale to cover the target canvas
+            # then center-crop, instead of non-uniform stretching to (width, height).
+            input_image = ImageOps.fit(
+                input_image,
+                (width, height),
+                method=Image.LANCZOS,
+                centering=(0.5, 0.5),
+            )
+            image = self.preprocess_image(input_image, device=dev, torch_dtype=torch.bfloat16)
+            ref_pixel = image.unsqueeze(2)  # [B, C, 1, H, W]
+            # encode_to_latent returns [B, 1, C_lat, H_lat, W_lat]
+            first_frame_latents = vae.encode_to_latent(ref_pixel.to(next(vae.parameters()).device))
+            self.conditional_dict["first_frame_latents"] = first_frame_latents.to(device=dev, dtype=torch.bfloat16)
+            print(f"set_first_frame_latent: shape={first_frame_latents.shape}")
+
+    def set_act(self, keys_dict, height: int = 480, width: int = 832, num_frames: int = 81, device: Optional[torch.device] = None):
+        """
+        设置 act_context [B, 8]
+        device: 指定设备，默认 self.device
+        """
+        if self.conditional_dict is None:
+            raise RuntimeError("call set_prompts first")
+        key_order = ['W', 'A', 'S', 'D', 'I', 'J', 'K', 'L']
+        action_list = [1 if keys_dict.get(k, False) else 0 for k in key_order]
+        action = torch.tensor(action_list, dtype=torch.float32, device=device).unsqueeze(0)
+        control_action_latents = action[:, None, None, :].repeat(1, height, width, 1).permute([3, 0, 1, 2]).unsqueeze(0).to(device=device, dtype=torch.bfloat16) #f,h,w,8 --> 1,8,f,h,w, ([1, 8, 1, 480, 832])
+        # repeat control_action_latents to num_frames
+        control_action_latents = control_action_latents.repeat_interleave(4, dim=1)  # [1, 8*4, 1, 480, 832]
+        control_action_latents = control_action_latents.repeat(1, 1, num_frames, 1, 1)  # [1, 8*4, F, 480, 832]
+        self.conditional_dict["act_context"] = control_action_latents
+        print(f'action_list, {action_list}')
+
+
+    def reset_stream(self, batch_size: int, dtype, device, initial_latent=None):
+        ref_cache_token_len = self._ref_cache_token_len(getattr(self, 'conditional_dict', None))
+        if self._kv_cache_needs_reinit(ref_cache_token_len):
+            self._initialize_kv_cache(
+                batch_size=batch_size, dtype=dtype, device=device,
+                extra_tokens=ref_cache_token_len,
+            )
+            self._initialize_crossattn_cache(batch_size=batch_size, dtype=dtype, device=device)
+        else:
+            # reset cross-attn cache flags
+            for b in range(self.num_transformer_blocks):
+                self.crossattn_cache[b]["is_init"] = False
+            # reset kv cache indices
+            for b in range(len(self.kv_cache1)):
+                self.kv_cache1[b]["global_end_index"].fill_(0)
+                self.kv_cache1[b]["local_end_index"].fill_(0)
+
+        self.current_start_frame = 0
+        self.num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        self._stream_block_diffusion_times = []
+        self._stream_block_decode_times = []
+
+
+    @torch.no_grad()
+    def generate_next_block(self, noise_block: torch.Tensor):
+        """
+        noise_block: [B, F, C, H, W]，其中 F=1 或 num_frame_per_block
+        return denoised_pred: [B, F, C, H, W] on GPU
+        """
+        t0 = time.perf_counter()
+        B, F, C, H, W = noise_block.shape
+        noisy_input = noise_block
+
+        y = self.conditional_dict.get("y", None)
+        if y is not None:
+            y = y.clone()
+
+        action_context = self.conditional_dict.get("act_context", None)
+        if action_context is not None:
+            action_context = action_context.clone()
+            _, Ca, _, Ha, Wa = action_context.shape
+
+        first_frame_latents = self.conditional_dict.get("first_frame_latents", None)
+        replace_first = (self.current_start_frame - self.num_input_frames == 0) and (first_frame_latents is not None)
+        if replace_first:
+            print(f"replace the first timestep with the clean context. latent shape: {noisy_input.shape}, first_frame_latents shape: {first_frame_latents.shape}")
+            noisy_input[:, 0:1] = first_frame_latents
+
+        #Pyramid denoising
+        if self.pyramid_sample_ratio is not None:
+            assert len(self.pyramid_sample_ratio) == len(self.denoising_step_list)
+            noisy_input = noisy_input.reshape(B, F*C, H, W)
+            noisy_input = Functional.interpolate(noisy_input, scale_factor=self.pyramid_sample_ratio[0], mode='nearest')
+            noisy_input = noisy_input.reshape(B, F, C, int(H*self.pyramid_sample_ratio[0]), int(W*self.pyramid_sample_ratio[0]))
+
+            if y is not None:
+                cur_y = y.reshape(B, C*F, H, W)
+                cur_y = Functional.interpolate(cur_y, scale_factor=self.pyramid_sample_ratio[0], mode='nearest')
+                cur_y = cur_y.reshape(B, C, F, int(H*self.pyramid_sample_ratio[0]), int(W*self.pyramid_sample_ratio[0]))
+                self.conditional_dict["y"] = cur_y.clone()
+
+            if action_context is not None:
+                cur_action_context = action_context.reshape(B, Ca*F, Ha, Wa)
+                cur_action_context = Functional.interpolate(cur_action_context, scale_factor=self.pyramid_sample_ratio[0], mode='nearest')
+                cur_action_context = cur_action_context.reshape(B, Ca, F, int(Ha*self.pyramid_sample_ratio[0]), int(Wa*self.pyramid_sample_ratio[0]))
+                self.conditional_dict["act_context"] = cur_action_context.clone()
+
+
+        for index, current_timestep in enumerate(self.denoising_step_list):
+            timestep = torch.ones([B, F], device=noise_block.device, dtype=torch.int64) * current_timestep
+            if replace_first:
+                timestep[:, 0] = 0
+
+            if self.pyramid_sample_ratio is not None:
+                _, _, _, H_s, W_s = noisy_input.shape
+
+            _, denoised_pred = self.generator(
+                noisy_image_or_video=noisy_input,
+                conditional_dict=self.conditional_dict,
+                timestep=timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=self.current_start_frame * self.frame_seq_length,
+                replace_first_timestep_and_noise_latents=replace_first,
+            )
+
+            if index < len(self.denoising_step_list) - 1:
+                next_timestep = self.denoising_step_list[index + 1]
+                if self.pyramid_sample_ratio is not None:
+                    denoised_pred = denoised_pred.reshape(B, F*C, H_s, W_s)
+                    denoised_pred = Functional.interpolate(denoised_pred, scale_factor=self.pyramid_sample_ratio[index+1], mode='nearest')
+                    denoised_pred = denoised_pred.reshape(B, F, C, int(H_s*self.pyramid_sample_ratio[index+1]), int(W_s*self.pyramid_sample_ratio[index+1]))
+
+                    if y is not None:
+                        cur_y = y.reshape(B, C*F, H, W)
+                        scale = math.prod(self.pyramid_sample_ratio[:index+2])
+                        cur_y = Functional.interpolate(cur_y, scale_factor=scale, mode='nearest')
+                        cur_y = cur_y.reshape(B, C, F, int(H*scale), int(W*scale))
+                        self.conditional_dict["y"] = cur_y
+
+                    if action_context is not None:
+                        cur_action_context = action_context.reshape(B, Ca*F, Ha, Wa)
+                        scale = math.prod(self.pyramid_sample_ratio[:index+2])
+                        cur_action_context = Functional.interpolate(cur_action_context, scale_factor=scale, mode='nearest')
+                        cur_action_context = cur_action_context.reshape(B, Ca, F, int(Ha*scale), int(Wa*scale))
+                        self.conditional_dict["act_context"] = cur_action_context.clone()
+
+                noisy_input = self.scheduler.add_noise(
+                    denoised_pred.flatten(0, 1),
+                    torch.randn_like(denoised_pred.flatten(0, 1)),
+                    next_timestep * torch.ones([B * F], device=noise_block.device, dtype=torch.long),
+                ).unflatten(0, denoised_pred.shape[:2])
+            else:
+                noisy_input = denoised_pred
+
+            if replace_first:
+                noisy_input[:, 0:1] = first_frame_latents
+
+        # 用干净 context 更新 KV cache
+        context_timestep = torch.ones([B, F], device=noise_block.device, dtype=torch.int64) * self.args.context_noise
+        if replace_first:
+            context_timestep = context_timestep.clone()
+            context_timestep[:, 0] = 0
+            noisy_input[:, 0:1] = first_frame_latents
+
+        self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=self.conditional_dict,
+            timestep=context_timestep,
+            kv_cache=self.kv_cache1,
+            crossattn_cache=self.crossattn_cache,
+            current_start=self.current_start_frame * self.frame_seq_length,
+            replace_first_timestep_and_noise_latents=replace_first,
+        )
+
+        self.current_start_frame += F
+
+        torch.cuda.synchronize()
+        block_diffusion_s = time.perf_counter() - t0
+        self._stream_block_diffusion_times.append(block_diffusion_s)
+
+        return noisy_input
+    
+    @torch.no_grad()
+    def decode_block_and_write(self, latents_block: torch.Tensor, writer):
+        # latents_block: [B,F,C,H,W] on GPU
+        t0 = time.perf_counter()
+        vid = self.vae.decode_to_pixel(latents_block, use_cache=True, return_in_cpu=True)
+        torch.cuda.synchronize()
+        block_decode_s = time.perf_counter() - t0
+        self._stream_block_decode_times.append(block_decode_s)
+
+        vid = (vid * 0.5 + 0.5).clamp(0, 1)
+        # print('decode_block_and_write: ', latents_block.shape, vid.shape) # [1, 3, 16, 60, 104]) torch.Size([1, 12, 3, 480, 832])
+        frames = (vid[0].permute(0,2,3,1) * 255.0).clamp(0,255).to(torch.uint8).numpy()
+        for f in frames:
+            writer.append_data(f)
+    
+    def inference(
+            self,
+            noise: torch.Tensor,
+            text_prompts: List[str],
+            conditional_dict=None,
+            initial_latent: Optional[torch.Tensor] = None,
+            return_latents: bool = False,
+            profile: bool = False,
+            low_memory: bool = False,
+            rectified_tf=False
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+                It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame:
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            # default here
+            # self.independent_first_frame: False
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        elif self.independent_first_frame and initial_latent is not None:
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames if self.independent_first_frame else num_frames + num_input_frames
+        # conditional_dict = self.text_encoder(
+        #     text_prompts=text_prompts
+        # )
+
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(self.text_encoder, target_device=gpu,
+                                                          preserved_memory_gb=gpu_memory_preservation)
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Set up profiling if requested
+        if profile:
+            init_start = torch.cuda.Event(enable_timing=True)
+            init_end = torch.cuda.Event(enable_timing=True)
+            diffusion_start = torch.cuda.Event(enable_timing=True)
+            diffusion_end = torch.cuda.Event(enable_timing=True)
+            vae_start = torch.cuda.Event(enable_timing=True)
+            vae_end = torch.cuda.Event(enable_timing=True)
+            block_times = []
+            block_start = torch.cuda.Event(enable_timing=True)
+            block_end = torch.cuda.Event(enable_timing=True)
+            init_start.record()
+
+        # Step 1: Initialize KV cache to all zeros
+        ref_cache_token_len = self._ref_cache_token_len(conditional_dict)
+        if self._kv_cache_needs_reinit(ref_cache_token_len):
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                extra_tokens=ref_cache_token_len,
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache1[block_index]["ref_token_len"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                tmp_conditional_dict = self.get_condition_split(conditional_dict, 0, 1)
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=tmp_conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = \
+                    initial_latent[:, current_start_frame:current_start_frame + self.num_frame_per_block]
+                output[:, current_start_frame:current_start_frame + self.num_frame_per_block] = current_ref_latents
+                tmp_conditional_dict = self.get_condition_split(conditional_dict, current_start_frame,
+                                                                current_start_frame + self.num_frame_per_block)
+
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=tmp_conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += self.num_frame_per_block
+
+        if profile:
+            init_end.record()
+            torch.cuda.synchronize()
+            diffusion_start.record()
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_num_frames in tqdm.tqdm(all_num_frames):
+            if profile:
+                block_start.record()
+
+            if self.independent_first_frame:
+                noise_start_frame = current_start_frame
+                noise_end_frame = current_start_frame + current_num_frames
+            else:
+                noise_start_frame = current_start_frame - num_input_frames
+                noise_end_frame = current_start_frame + current_num_frames - num_input_frames
+            noisy_input = noise[:, noise_start_frame:noise_end_frame]
+            condition_start_frame = current_start_frame if self.independent_first_frame else noise_start_frame
+            condition_end_frame = (
+                current_start_frame + current_num_frames
+                if self.independent_first_frame
+                else noise_end_frame
+            )
+            tmp_conditional_dict = self.get_condition_split(
+                conditional_dict,
+                condition_start_frame,
+                condition_end_frame,
+            )
+            latents = noisy_input
+            replace_first_timestep_and_noise_latents = (
+                current_start_frame == 0
+                and num_input_frames == 0
+                and conditional_dict.get("first_frame_latents") is not None
+            )
+            if replace_first_timestep_and_noise_latents:
+                print(f"replace the first timestep with the clean context. latent shape: {noisy_input.shape}, conditional_dict['first_frame_latents'] shape: {conditional_dict['first_frame_latents'].shape}" )
+                noisy_input[:, 0:1] = conditional_dict["first_frame_latents"] # replace the first timestep with the clean context
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                # print(f"current_timestep: {current_timestep}")
+                # set current timestep
+                timestep = torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if index < len(self.denoising_step_list) - 1:
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=tmp_conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                        replace_first_timestep_and_noise_latents=replace_first_timestep_and_noise_latents,
+                    )
+                    next_timestep = self.denoising_step_list[index + 1]
+                    noisy_input = self.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep * torch.ones(
+                            [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                    ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    # for getting real output
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=tmp_conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                        replace_first_timestep_and_noise_latents=replace_first_timestep_and_noise_latents,
+                    )
+                    noisy_input = denoised_pred
+                if replace_first_timestep_and_noise_latents:
+                    noisy_input[:, 0:1] = conditional_dict["first_frame_latents"]
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = noisy_input
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            context_timestep = torch.ones_like(timestep) * self.args.context_noise
+            if replace_first_timestep_and_noise_latents:
+                context_timestep = context_timestep.clone()
+                context_timestep[:, 0] = 0
+                noisy_input[:, 0:1] = conditional_dict["first_frame_latents"]
+
+            self.generator(
+                noisy_image_or_video=noisy_input,
+                conditional_dict=tmp_conditional_dict,
+                timestep=context_timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=current_start_frame * self.frame_seq_length,
+                replace_first_timestep_and_noise_latents=replace_first_timestep_and_noise_latents,
+            )
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        if profile:
+            # End diffusion timing and synchronize CUDA
+            diffusion_end.record()
+            torch.cuda.synchronize()
+            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
+            init_time = init_start.elapsed_time(init_end)
+            vae_start.record()
+        if rectified_tf:
+            mean = torch.load('laboratory/mean.pt').to(output.device)
+            std = torch.load('laboratory/std.pt').to(output.device)
+            noise = torch.randn_like(output).to(output.device)
+            output -= mean
+        # Step 4: Decode the output
+        video = self._decode_output_to_video(output)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        if profile:
+            # End VAE timing and synchronize CUDA
+            vae_end.record()
+            torch.cuda.synchronize()
+            vae_time = vae_start.elapsed_time(vae_end)
+            total_time = init_time + diffusion_time + vae_time
+
+            print("Profiling results:")
+            print(f"  - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)")
+            print(f"  - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)")
+            for i, block_time in enumerate(block_times):
+                print(
+                    f"    - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)")
+            print(f"  - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)")
+            print(f"  - Total time: {total_time:.2f} ms")
+
+        if return_latents:
+            return video, output
+        else:
+            return video
+
+    def _ref_cache_token_len(self, conditional_dict):
+        if conditional_dict is None:
+            return 0
+        ref_latents = conditional_dict.get("ref_latents", None)
+        if ref_latents is None:
+            return 0
+        in_ch_need = self.generator.model.patch_embedding.in_channels
+        if ref_latents.ndim == 4:
+            ref_latents = ref_latents.unsqueeze(0).unsqueeze(3)
+        elif ref_latents.ndim == 5 and ref_latents.shape[2] == in_ch_need:
+            ref_latents = ref_latents.unsqueeze(3)
+        elif ref_latents.ndim == 5:
+            ref_latents = ref_latents.unsqueeze(0)
+        if ref_latents.ndim != 6:
+            return 0
+        _, num_slots, _, ref_t, ref_h, ref_w = ref_latents.shape
+        patch_t, patch_h, patch_w = self.generator.model.patch_size
+        return (
+            int(num_slots)
+            * (int(ref_t) // int(patch_t))
+            * (int(ref_h) // int(patch_h))
+            * (int(ref_w) // int(patch_w))
+        )
+
+    def _kv_cache_needs_reinit(self, ref_cache_token_len):
+        if self.kv_cache1 is None:
+            return True
+        if getattr(self, "kv_cache_extra_tokens", 0) != int(ref_cache_token_len):
+            return True
+        return any("ref_token_len" not in cache for cache in self.kv_cache1)
+
+    def _initialize_kv_cache(self, batch_size, dtype, device, extra_tokens=0):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache1 = []
+        if self.local_attn_size != -1:
+            # Use the local attention size to compute the KV cache size
+            print("self.local_attn_size:", self.local_attn_size)
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # Use the default KV cache size
+            kv_cache_size = self.args.image_or_video_shape[1] * self.frame_seq_length
+        kv_cache_size = kv_cache_size + int(extra_tokens)
+
+        num_heads = self.generator.model.num_heads
+        head_dim = self.generator.model.dim // num_heads
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "ref_token_len": torch.tensor([0], dtype=torch.long, device=device),
+            })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+        self.kv_cache_extra_tokens = int(extra_tokens)
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        num_heads = self.generator.model.num_heads
+        head_dim = self.generator.model.dim // num_heads
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, num_heads, head_dim], dtype=dtype, device=device),
+                "is_init": False
+            })
+        self.crossattn_cache = crossattn_cache

quantizer/__init__.py

@@ -0,0 +1,49 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, Optional
+
+from .quant import *  # noqa: F401 F403
+from .quant import DynamicDiTQuantizer
+
+DEFAULT_FP8_INCLUDE_PATTERNS = ["blocks"]
+DEFAULT_FP8_EXCLUDE_PATTERNS = []
+
+
+def apply_fp8_quantization(
+    model,
+    quant_type: str = "fp8-per-token",
+    include_patterns: Optional[List[str]] = None,
+    exclude_patterns: Optional[List[str]] = None,
+) -> Dict[str, object]:
+    """Apply DynamicDiTQuantizer to the provided DiT model."""
+    final_include_patterns = (
+        include_patterns if include_patterns is not None else DEFAULT_FP8_INCLUDE_PATTERNS
+    )
+    final_exclude_patterns = (
+        exclude_patterns if exclude_patterns is not None else DEFAULT_FP8_EXCLUDE_PATTERNS
+    )
+
+    quantizer = DynamicDiTQuantizer(
+        quant_type=quant_type,
+        include_patterns=final_include_patterns,
+        exclude_patterns=final_exclude_patterns,
+    )
+    quantizer.convert_linear(model)
+    return {
+        "quantizer": quantizer,
+        "quant_type": quant_type,
+        "include_patterns": final_include_patterns,
+        "exclude_patterns": final_exclude_patterns,
+    }

quantizer/kernels/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

quantizer/kernels/python/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

quantizer/kernels/python/gemm/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+FP8 GEMM kernels with automatic backend selection.
+
+This module automatically selects between Triton (for Linux/CUDA) and
+PyTorch (for Windows/CPU) implementations based on the runtime environment.
+"""
+
+# from angelslim.compressor._platform import use_triton
+
+# Conditional imports based on platform/backend availability
+# if use_triton():
+from .fp8_gemm import fp8_gemm_triton_block
+# else:
+#     # PyTorch fallback implementation
+#     from .fp8_gemm_torch import fp8_gemm_torch_block as fp8_gemm_triton_block
+
+# Also export PyTorch version directly for explicit use
+from .fp8_gemm_torch import fp8_gemm_torch_block
+
+__all__ = ["fp8_gemm_triton_block", "fp8_gemm_torch_block"]

quantizer/kernels/python/gemm/fp8_gemm.py

@@ -0,0 +1,112 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import triton
+import triton.language as tl
+
+# modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
+fp8_gemm_configs = [
+    triton.Config(
+        {"BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": 128},
+        num_stages=num_stages,
+        num_warps=8,
+    )
+    for block_m in [16, 32, 64]
+    for block_n in [32, 64, 128]
+    for num_stages in [3, 4, 5, 6]
+]
+
+
+@triton.autotune(configs=fp8_gemm_configs, key=["N", "K"])
+@triton.jit
+def _fp8_gemm_triton_block_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    a_s_ptr,
+    b_s_ptr,
+    M,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """
+    Performs a matrix multiplication operation on FP8 matrices with scaling factors.
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    k = tl.cdiv(K, BLOCK_SIZE_K)
+    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
+    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
+    a_s_ptrs = a_s_ptr + offs_m * k
+    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for i in range(k):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
+        a_s = tl.load(a_s_ptrs)
+        b_s = tl.load(b_s_ptrs)
+
+        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K
+        b_ptrs += BLOCK_SIZE_K
+        a_s_ptrs += 1
+        b_s_ptrs += 1
+    c = accumulator.to(c_ptr.dtype.element_ty)
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    tl.store(c_ptrs, c, mask=mask)
+
+
+# triton fp8 gemm for fp8 per-block weight & fp8 per-group activation
+# modified from https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
+def fp8_gemm_triton_block(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    out_dtype=torch.bfloat16,
+    bias=None,
+) -> torch.Tensor:
+    """
+    Perform a matrix multiplication using FP8 precision.
+    """
+    assert a.is_contiguous() and b.is_contiguous()
+    assert a_s.is_contiguous() and b_s.is_contiguous()
+    K = a.size(-1)
+    M = a.numel() // K
+    N = b.size(0)
+    c = a.new_empty(*a.size()[:-1], N, dtype=out_dtype)
+
+    def grid(meta):
+        return (
+            triton.cdiv(M, meta["BLOCK_SIZE_M"]),
+            triton.cdiv(N, meta["BLOCK_SIZE_N"]),
+        )
+
+    _fp8_gemm_triton_block_kernel[grid](a, b, c, a_s, b_s, M, N, K)
+
+    if bias is not None:
+        c += bias
+
+    return c

quantizer/kernels/python/gemm/fp8_gemm_torch.py

@@ -0,0 +1,217 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Pure PyTorch implementation of FP8 block-wise GEMM.
+
+This module provides CPU/Windows-compatible implementations that mirror
+the Triton kernel for FP8 GEMM with block-wise quantization.
+"""
+
+from typing import Optional
+
+import torch
+
+
+def fp8_gemm_torch_block(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    bias: Optional[torch.Tensor] = None,
+    block_size: int = 128,
+) -> torch.Tensor:
+    """
+    Pure PyTorch implementation of FP8 GEMM with block-wise quantization.
+
+    Performs a matrix multiplication using FP8 precision with per-block scaling.
+    This implementation dequantizes the inputs and performs standard matmul.
+
+    C = (A * A_scale) @ (B * B_scale).T + bias
+
+    Args:
+        a: Input activation tensor in FP8 format, shape [..., K]
+        a_s: Scale tensor for A, shape [..., K // block_size] or [..., num_k_blocks]
+        b: Weight tensor in FP8 format, shape [N, K]
+        b_s: Scale tensor for B, shape [N // block_size, K // block_size]
+        out_dtype: Output data type (default: bfloat16)
+        bias: Optional bias tensor, shape [N]
+        block_size: Block size used for quantization (default: 128)
+
+    Returns:
+        Output tensor of shape [..., N]
+    """
+    assert a.is_contiguous() and b.is_contiguous()
+    assert a_s.is_contiguous() and b_s.is_contiguous()
+
+    K = a.size(-1)
+    orig_shape = a.shape[:-1]
+    M = a.numel() // K
+    N = b.size(0)
+
+    # Reshape for computation
+    a_2d = a.view(M, K)  # [M, K]
+
+    # Dequantize A: expand scales to match tensor dimensions
+    # a_s shape is typically [M, K//block_size]
+    a_s_2d = a_s.view(M, -1)  # [M, num_k_blocks]
+
+    # Dequantize by expanding scales
+    a_dq = _dequantize_per_group(a_2d, a_s_2d, block_size, K)
+
+    # Dequantize B: b_s is [N//block_size, K//block_size]
+    b_dq = _dequantize_blockwise_2d(b, b_s, block_size)
+
+    # Perform matmul: [M, K] @ [K, N] -> [M, N]
+    c = torch.matmul(a_dq.to(out_dtype), b_dq.to(out_dtype).t())
+
+    # Reshape output
+    c = c.view(*orig_shape, N)
+
+    if bias is not None:
+        c = c + bias
+
+    return c
+
+
+def _dequantize_per_group(
+    x: torch.Tensor,
+    s: torch.Tensor,
+    group_size: int,
+    K: int,
+) -> torch.Tensor:
+    """
+    Dequantize tensor with per-group scales.
+
+    Args:
+        x: Quantized tensor [M, K]
+        s: Scale tensor [M, num_groups]
+        group_size: Size of each group
+        K: Total size of last dimension
+
+    Returns:
+        Dequantized tensor [M, K]
+    """
+    M = x.shape[0]
+    num_groups = s.shape[1]
+
+    x_float = x.to(torch.float32)
+
+    # Expand scales to match K dimension
+    # s: [M, num_groups] -> [M, K]
+    s_expanded = s.unsqueeze(-1).expand(M, num_groups, group_size)
+    s_expanded = s_expanded.reshape(M, num_groups * group_size)
+
+    # Handle case where K is not exactly num_groups * group_size
+    if s_expanded.shape[1] > K:
+        s_expanded = s_expanded[:, :K]
+    elif s_expanded.shape[1] < K:
+        # Pad with last scale value
+        pad_size = K - s_expanded.shape[1]
+        s_expanded = torch.nn.functional.pad(s_expanded, (0, pad_size), mode="replicate")
+
+    return x_float * s_expanded
+
+
+def _dequantize_blockwise_2d(
+    x: torch.Tensor,
+    s: torch.Tensor,
+    block_size: int,
+) -> torch.Tensor:
+    """
+    Dequantize 2D tensor with block-wise scales.
+
+    Args:
+        x: Quantized tensor [N, K]
+        s: Scale tensor [n_blocks, k_blocks]
+        block_size: Block size
+
+    Returns:
+        Dequantized tensor [N, K]
+    """
+    N, K = x.shape
+    n_blocks, k_blocks = s.shape
+
+    x_float = x.to(torch.float32)
+    y = torch.empty_like(x_float)
+
+    for nb in range(n_blocks):
+        n_start = nb * block_size
+        n_end = min(n_start + block_size, N)
+        for kb in range(k_blocks):
+            k_start = kb * block_size
+            k_end = min(k_start + block_size, K)
+
+            scale = s[nb, kb]
+            y[n_start:n_end, k_start:k_end] = x_float[n_start:n_end, k_start:k_end] * scale
+
+    return y
+
+
+def fp8_gemm_torch_simple(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Simplified PyTorch FP8 GEMM using full dequantization.
+
+    This version is simpler but may use more memory for large tensors.
+
+    Args:
+        a: Input activation tensor in FP8 format
+        a_s: Scale tensor for A
+        b: Weight tensor in FP8 format
+        b_s: Scale tensor for B
+        out_dtype: Output data type
+        bias: Optional bias tensor
+
+    Returns:
+        Output tensor
+    """
+    K = a.size(-1)
+    orig_shape = a.shape[:-1]
+    M = a.numel() // K
+    N = b.size(0)
+
+    # Reshape
+    a_2d = a.view(M, K)
+    a_s_2d = a_s.view(M, -1)
+
+    # Simple dequantization: repeat scales to match dimensions
+    block_size = K // a_s_2d.shape[1] if a_s_2d.shape[1] > 0 else K
+
+    # Dequantize A
+    a_dq = a_2d.to(torch.float32)
+    if a_s_2d.shape[1] > 1:
+        a_s_expanded = a_s_2d.repeat_interleave(block_size, dim=1)[:, :K]
+        a_dq = a_dq * a_s_expanded
+    else:
+        a_dq = a_dq * a_s_2d
+
+    # Dequantize B
+    b_dq = _dequantize_blockwise_2d(b, b_s, block_size)
+
+    # Matmul
+    c = torch.matmul(a_dq.to(out_dtype), b_dq.to(out_dtype).t())
+    c = c.view(*orig_shape, N)
+
+    if bias is not None:
+        c = c + bias
+
+    return c

quantizer/kernels/python/mm/triton_kernels.py

@@ -0,0 +1,743 @@
+import torch
+from triton import Config, autotune, cdiv, jit, next_power_of_2
+from triton import language as tl
+
+_ordered_datatypes = [torch.int8, torch.float16, torch.bfloat16, torch.float32]
+
+
+@jit
+def gelu(x):
+    return x * tl.sigmoid(x * 1.702)
+
+
+@jit
+def int8_quantize_kernel(X, OUT, SCALES, HDIM, BLOCK_SIZE: tl.constexpr):
+    row_idx = tl.program_id(0)
+    x_ptr = X + row_idx * HDIM
+    out_ptr = OUT + row_idx * HDIM
+    h_offset = tl.arange(0, BLOCK_SIZE)
+    x = tl.load(x_ptr + h_offset, mask=h_offset < HDIM).to(tl.float32)
+    x_scale = 127.0 / tl.max(tl.abs(x))
+    x_scaled = x * x_scale
+    x_scaled += (0.5 * tl.where(x_scaled >= 0, 1, -1)).to(tl.int8)
+    tl.store(out_ptr + h_offset, x_scaled, mask=h_offset < HDIM)
+    tl.store(SCALES + row_idx, 1 / x_scale)
+
+
+def int8_quantize_triton(x):
+    x_shape_orig = x.shape
+    x = x.view(-1, x_shape_orig[-1])
+    out = torch.empty(x_shape_orig, dtype=torch.int8, device=x.device)
+    scales = torch.empty(x.shape[0], dtype=torch.float32, device=x.device)
+    BLOCK_SIZE = next_power_of_2(x_shape_orig[-1])
+    grid = (x.shape[0],)
+    int8_quantize_kernel[grid](x, out, scales, x_shape_orig[-1], BLOCK_SIZE, num_warps=8)
+    return out.view(x_shape_orig), scales.view(x_shape_orig[:-1])
+
+
+@jit
+def fp8_quantize_kernel(X, OUT, SCALES, HDIM, BLOCK_SIZE: tl.constexpr, FP8_MAX_VAL: tl.constexpr):
+    row_idx = tl.program_id(0)
+    x_ptr = X + row_idx * HDIM
+    out_ptr = OUT + row_idx * HDIM
+    h_offset = tl.arange(0, BLOCK_SIZE)
+    x = tl.load(x_ptr + h_offset, mask=h_offset < HDIM).to(tl.float32)
+    absmax = tl.max(tl.abs(x))
+    eps = 1e-8
+    absmax = tl.maximum(absmax, eps)
+    x_scale = absmax / FP8_MAX_VAL
+    x_scaled = x / x_scale
+    x_scaled = tl.clamp(x_scaled, -FP8_MAX_VAL, FP8_MAX_VAL)
+    tl.store(out_ptr + h_offset, x_scaled, mask=h_offset < HDIM)
+    tl.store(SCALES + row_idx, x_scale)
+
+
+def fp8_quantize_triton(x):
+    x_shape_orig = x.shape
+    x = x.view(-1, x_shape_orig[-1])
+    out_scaled = torch.empty(x_shape_orig, dtype=torch.float32, device=x.device)
+    scales = torch.empty(x.shape[0], dtype=torch.bfloat16, device=x.device)
+    BLOCK_SIZE = next_power_of_2(x_shape_orig[-1])
+    grid = (x.shape[0],)
+    FP8_MAX = 448.0
+    fp8_quantize_kernel[grid](x, out_scaled, scales, x_shape_orig[-1], BLOCK_SIZE, FP8_MAX_VAL=FP8_MAX, num_warps=8)
+    quantized = out_scaled.to(torch.float8_e4m3fn)
+    return quantized.view(x_shape_orig), scales.view(x_shape_orig[:-1])
+
+
+def upcast_if_fp8(a):
+    if "fp8" in str(a):
+        return torch.float16
+    return a
+
+
+def get_higher_dtype(a, b):
+    a = upcast_if_fp8(a)
+    b = upcast_if_fp8(b)
+    if a is b:
+        return a
+
+    assert a in _ordered_datatypes
+    assert b in _ordered_datatypes
+
+    for d in _ordered_datatypes:
+        if a is d:
+            return b
+        if b is d:
+            return a
+
+
+@autotune(
+    configs=[
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+    ],
+    key=["M", "N", "K"],
+)
+@jit
+def int8_gemm_bias_kernel(
+    A,
+    B,
+    BIAS,
+    A_SCALES,
+    B_SCALES,
+    C,
+    M,
+    N,
+    K,  #
+    stride_am,
+    stride_ak,  #
+    stride_bk,
+    stride_bn,  #
+    stride_cm,
+    stride_cn,  #
+    acc_dtype: tl.constexpr,  #
+    fuse_gelu: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,  #
+    GROUP_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    AB_DTYPE: tl.constexpr,  #
+):
+    # matrix multiplication
+    pid = tl.program_id(0)
+    pid_z = tl.program_id(1)
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+    # re-order program ID for better L2 performance
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // (group_size)
+    # do matrix multiplication
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+    rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
+    # pointers
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype)
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        if EVEN_K:
+            a = tl.load(A)
+            b = tl.load(B)
+        else:
+            k_remaining = K - k * (BLOCK_K * SPLIT_K)
+            _0 = tl.zeros((1, 1), dtype=C.dtype.element_ty)
+            a = tl.load(A, mask=rk[None, :] < k_remaining, other=_0)
+            b = tl.load(B, mask=rk[:, None] < k_remaining, other=_0)
+        if AB_DTYPE is not None:
+            a = a.to(AB_DTYPE)
+            b = b.to(AB_DTYPE)
+
+        acc = tl.dot(a, b, acc, out_dtype=acc_dtype, input_precision=None)
+        A += BLOCK_K * SPLIT_K * stride_ak
+        B += BLOCK_K * SPLIT_K * stride_bk
+
+    acc = acc.to(tl.float32)
+    a_scales_ptr = A_SCALES + pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    b_scales_ptr = B_SCALES + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    a_scales = tl.load(a_scales_ptr)  # [BM]
+    b_scales = tl.load(b_scales_ptr)  # [BN]
+    # [BM, BN] * [BM, 1] * [1, BN]
+
+    bias_ptr = BIAS + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    bias = tl.load(bias_ptr)
+    if fuse_gelu:
+        acc = gelu(((acc * a_scales[:, None]) * b_scales[None, :]) + bias[None, :])
+    else:
+        acc = ((acc * a_scales[:, None]) * b_scales[None, :]) + bias[None, :]
+
+    acc = acc.to(C.dtype.element_ty)
+
+    # rematerialize rm and rn to save registers
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+    # handles write-back with reduction-splitting
+    if SPLIT_K == 1:
+        tl.store(C, acc, mask=mask)
+    else:
+        tl.atomic_add(C, acc, mask=mask)
+
+
+# @torch.compiler.disable()
+def int8_gemm_bias_triton(a, b, bias, a_scales, b_scales, fuse_gelu=False, output_dtype=None):
+    device = a.device
+    # handle non-contiguous inputs if necessary
+    a_orig_shape = a.shape
+    a = a.view(-1, a.shape[-1])
+    b = b.t()
+
+    if a.stride(0) > 1 and a.stride(1) > 1:
+        a = a.contiguous()
+    if b.stride(0) > 1 and b.stride(1) > 1:
+        b = b.contiguous()
+    # checks constraints
+    assert a.shape[1] == b.shape[0], f"incompatible dimensions {a.shape} and {b.shape}"
+    M, K = a.shape
+    _, N = b.shape
+    out_shape = a_orig_shape[:-1] + (N,)
+    # common type between a and b
+    ab_dtype = get_higher_dtype(a.dtype, b.dtype)
+
+    # allocates output
+    if output_dtype is None:
+        output_dtype = ab_dtype
+
+    c = torch.empty((M, N), device=device, dtype=output_dtype)
+
+    # Allowed types for acc_type given the types of a and b.
+    supported_acc_dtypes = {
+        torch.float16: (torch.float32, torch.float16),
+        torch.bfloat16: (torch.float32, torch.bfloat16),
+        torch.float32: (torch.float32,),
+        torch.int8: (torch.int32,),
+    }
+
+    acc_dtype = supported_acc_dtypes[ab_dtype][0]
+
+    def to_tl_type(ty):
+        return getattr(tl, str(ty).split(".")[-1])
+
+    acc_dtype = to_tl_type(acc_dtype)
+    ab_dtype = to_tl_type(ab_dtype)
+    output_dtype = to_tl_type(output_dtype)
+
+    # Tensor cores support input with mixed float8 types.
+    if a.dtype in [tl.float8e4nv, tl.float8e5] and b.dtype in [
+        tl.float8e4nv,
+        tl.float8e5,
+    ]:
+        ab_dtype = None
+    # launch kernel
+    grid = lambda META: (  # noqa E731
+        cdiv(M, META["BLOCK_M"]) * cdiv(N, META["BLOCK_N"]),
+        META["SPLIT_K"],
+    )  # noqa E731
+    int8_gemm_bias_kernel[grid](
+        a,
+        b,
+        bias,
+        a_scales,
+        b_scales,
+        c,
+        M,
+        N,
+        K,  #
+        a.stride(0),
+        a.stride(1),  #
+        b.stride(0),
+        b.stride(1),  #
+        c.stride(0),
+        c.stride(1),  #
+        acc_dtype=acc_dtype,  #
+        fuse_gelu=fuse_gelu,
+        GROUP_M=8,
+        EVEN_K=True,
+        AB_DTYPE=ab_dtype,
+    )
+    return c.view(*out_shape)
+
+
+@autotune(
+    configs=[
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+    ],
+    key=["M", "N", "K"],
+)
+@jit
+def int8_gemm_kernel(
+    A,
+    B,
+    A_SCALES,
+    B_SCALES,
+    C,
+    M,
+    N,
+    K,  #
+    stride_am,
+    stride_ak,  #
+    stride_bk,
+    stride_bn,  #
+    stride_cm,
+    stride_cn,  #
+    acc_dtype: tl.constexpr,  #
+    fuse_gelu: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,  #
+    GROUP_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+    AB_DTYPE: tl.constexpr,  #
+):
+    # matrix multiplication
+    pid = tl.program_id(0)
+    pid_z = tl.program_id(1)
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+    # re-order program ID for better L2 performance
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // (group_size)
+    # do matrix multiplication
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+    rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
+    # pointers
+    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype)
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        if EVEN_K:
+            a = tl.load(A)
+            b = tl.load(B)
+        else:
+            k_remaining = K - k * (BLOCK_K * SPLIT_K)
+            _0 = tl.zeros((1, 1), dtype=C.dtype.element_ty)
+            a = tl.load(A, mask=rk[None, :] < k_remaining, other=_0)
+            b = tl.load(B, mask=rk[:, None] < k_remaining, other=_0)
+        if AB_DTYPE is not None:
+            a = a.to(AB_DTYPE)
+            b = b.to(AB_DTYPE)
+
+        acc = tl.dot(a, b, acc, out_dtype=acc_dtype, input_precision=None)
+        A += BLOCK_K * SPLIT_K * stride_ak
+        B += BLOCK_K * SPLIT_K * stride_bk
+
+    acc = acc.to(tl.float32)
+    a_scales_ptr = A_SCALES + pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    b_scales_ptr = B_SCALES + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    a_scales = tl.load(a_scales_ptr)  # [BM]
+    b_scales = tl.load(b_scales_ptr)  # [BN]
+    # [BM, BN] * [BM, 1] * [1, BN]
+    if fuse_gelu:
+        acc = gelu((acc * a_scales[:, None]) * b_scales[None, :])
+    else:
+        acc = (acc * a_scales[:, None]) * b_scales[None, :]
+
+    acc = acc.to(C.dtype.element_ty)
+
+    # rematerialize rm and rn to save registers
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+    # handles write-back with reduction-splitting
+    if SPLIT_K == 1:
+        tl.store(C, acc, mask=mask)
+    else:
+        tl.atomic_add(C, acc, mask=mask)
+
+
+# @torch.compiler.disable()
+def int8_gemm_triton(a, b, a_scales, b_scales, fuse_gelu=False, output_dtype=None):
+    device = a.device
+    # handle non-contiguous inputs if necessary
+    # USE ONLY IN linear layer. NOT GENERAL MATRIX MULTIPLY
+    a_orig_shape = a.shape
+    a = a.view(-1, a.shape[-1])
+    b = b.t()
+
+    if a.stride(0) > 1 and a.stride(1) > 1:
+        a = a.contiguous()
+    if b.stride(0) > 1 and b.stride(1) > 1:
+        b = b.contiguous()
+    # checks constraints
+    assert a.shape[1] == b.shape[0], f"incompatible dimensions {a.shape} and {b.shape}"
+    M, K = a.shape
+    _, N = b.shape
+    out_shape = a_orig_shape[:-1] + (N,)
+    # common type between a and b
+    ab_dtype = get_higher_dtype(a.dtype, b.dtype)
+
+    # allocates output
+    if output_dtype is None:
+        output_dtype = ab_dtype
+
+    c = torch.empty((M, N), device=device, dtype=output_dtype)
+
+    # Allowed types for acc_type given the types of a and b.
+    supported_acc_dtypes = {
+        torch.float16: (torch.float32, torch.float16),
+        torch.bfloat16: (torch.float32, torch.bfloat16),
+        torch.float32: (torch.float32,),
+        torch.int8: (torch.int32,),
+    }
+
+    acc_dtype = supported_acc_dtypes[ab_dtype][0]
+
+    def to_tl_type(ty):
+        return getattr(tl, str(ty).split(".")[-1])
+
+    acc_dtype = to_tl_type(acc_dtype)
+    ab_dtype = to_tl_type(ab_dtype)
+    output_dtype = to_tl_type(output_dtype)
+
+    # Tensor cores support input with mixed float8 types.
+    if a.dtype in [tl.float8e4nv, tl.float8e5] and b.dtype in [
+        tl.float8e4nv,
+        tl.float8e5,
+    ]:
+        ab_dtype = None
+    # launch kernel
+    grid = lambda META: (  # noqa E731
+        cdiv(M, META["BLOCK_M"]) * cdiv(N, META["BLOCK_N"]),
+        META["SPLIT_K"],
+    )  # noqa E731
+    int8_gemm_kernel[grid](
+        a,
+        b,
+        a_scales,
+        b_scales,
+        c,
+        M,
+        N,
+        K,  #
+        a.stride(0),
+        a.stride(1),  #
+        b.stride(0),
+        b.stride(1),  #
+        c.stride(0),
+        c.stride(1),  #
+        acc_dtype=acc_dtype,  #
+        fuse_gelu=fuse_gelu,
+        EVEN_K=True,
+        GROUP_M=8,
+        AB_DTYPE=ab_dtype,
+    )
+    return c.view(*out_shape)
+
+
+@autotune(
+    configs=[
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+    ],
+    key=["M", "N", "K"],
+)
+@jit
+def fp8_gemm_bias_kernel(
+    A,
+    B,
+    BIAS,
+    A_SCALES,
+    B_SCALES,
+    C,
+    M,
+    N,
+    K,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    fuse_gelu: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_z = tl.program_id(1)
+
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // group_size
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+
+    rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
+
+    A_ptr = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B_ptr = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        if EVEN_K:
+            a = tl.load(A_ptr)
+            b = tl.load(B_ptr)
+        else:
+            k_remaining = K - k * (BLOCK_K * SPLIT_K)
+            a = tl.load(A_ptr, mask=rk[None, :] < k_remaining, other=0.0)
+            b = tl.load(B_ptr, mask=rk[:, None] < k_remaining, other=0.0)
+
+        acc = tl.dot(a, b, acc, out_dtype=tl.float32, input_precision=None)
+
+        A_ptr += BLOCK_K * SPLIT_K * stride_ak
+        B_ptr += BLOCK_K * SPLIT_K * stride_bk
+
+    a_scales_ptr = A_SCALES + pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    b_scales_ptr = B_SCALES + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    a_scales = tl.load(a_scales_ptr).to(tl.float32)  # [BM]
+    b_scales = tl.load(b_scales_ptr).to(tl.float32)  # [BN]
+
+    bias_ptr = BIAS + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    bias = tl.load(bias_ptr).to(tl.float32)  # [BN]
+
+    out = (acc * a_scales[:, None]) * b_scales[None, :] + bias[None, :]
+    if fuse_gelu:
+        out = gelu(out)
+
+    out = out.to(C.dtype.element_ty)
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    C_ptr = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+
+    if SPLIT_K == 1:
+        tl.store(C_ptr, out, mask=mask)
+    else:
+        tl.atomic_add(C_ptr, out, mask=mask)
+
+
+def fp8_gemm_bias_triton(a, b, bias, a_scales, b_scales, fuse_gelu=False, output_dtype=None):
+    assert a.is_cuda and b.is_cuda, "This kernel is for CUDA"
+    assert a.dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e4m3fnuz", None)), f"a.dtype={a.dtype} is not FP8 E4M3"
+    assert b.dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e4m3fnuz", None)), f"b.dtype={b.dtype} is not FP8 E4M3"
+
+    a_orig_shape = a.shape
+    a2 = a.view(-1, a.shape[-1])
+    b2 = b.t()
+
+    if a2.stride(0) > 1 and a2.stride(1) > 1:
+        a2 = a2.contiguous()
+    if b2.stride(0) > 1 and b2.stride(1) > 1:
+        b2 = b2.contiguous()
+
+    M, K = a2.shape
+    _, N = b2.shape
+    out_shape = a_orig_shape[:-1] + (N,)
+
+    if output_dtype is None:
+        output_dtype = torch.float16
+
+    c = torch.empty((M, N), device=a.device, dtype=output_dtype)
+
+    grid = lambda META: (cdiv(M, META["BLOCK_M"]) * cdiv(N, META["BLOCK_N"]), META["SPLIT_K"])  # noqa E731
+    even_k = K % 128 == 0
+
+    fp8_gemm_bias_kernel[grid](
+        a2,
+        b2,
+        bias,
+        a_scales,
+        b_scales,
+        c,
+        M,
+        N,
+        K,
+        a2.stride(0),
+        a2.stride(1),
+        b2.stride(0),
+        b2.stride(1),
+        c.stride(0),
+        c.stride(1),
+        fuse_gelu=fuse_gelu,
+        GROUP_M=8,
+        EVEN_K=even_k,
+    )
+    return c.view(*out_shape)
+
+
+@autotune(
+    configs=[
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=3, num_warps=8),
+        Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=8),
+    ],
+    key=["M", "N", "K"],
+)
+@jit
+def fp8_gemm_kernel(
+    A,
+    B,
+    A_SCALES,
+    B_SCALES,
+    C,
+    M,
+    N,
+    K,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    fuse_gelu: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    EVEN_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_z = tl.program_id(1)
+
+    grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+
+    width = GROUP_M * grid_n
+    group_id = pid // width
+    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
+    pid_m = group_id * GROUP_M + (pid % group_size)
+    pid_n = (pid % width) // group_size
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
+    rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
+
+    rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
+
+    A_ptr = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
+    B_ptr = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
+        if EVEN_K:
+            a = tl.load(A_ptr)
+            b = tl.load(B_ptr)
+        else:
+            k_remaining = K - k * (BLOCK_K * SPLIT_K)
+            a = tl.load(A_ptr, mask=rk[None, :] < k_remaining, other=0.0)
+            b = tl.load(B_ptr, mask=rk[:, None] < k_remaining, other=0.0)
+
+        acc = tl.dot(a, b, acc, out_dtype=tl.float32, input_precision=None)
+
+        A_ptr += BLOCK_K * SPLIT_K * stride_ak
+        B_ptr += BLOCK_K * SPLIT_K * stride_bk
+
+    a_scales_ptr = A_SCALES + pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    b_scales_ptr = B_SCALES + pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    a_scales = tl.load(a_scales_ptr).to(tl.float32)  # [BM]
+    b_scales = tl.load(b_scales_ptr).to(tl.float32)  # [BN]
+
+    out = (acc * a_scales[:, None]) * b_scales[None, :]
+    if fuse_gelu:
+        out = gelu(out)
+
+    out = out.to(C.dtype.element_ty)
+
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    C_ptr = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
+    mask = (rm < M)[:, None] & (rn < N)[None, :]
+
+    if SPLIT_K == 1:
+        tl.store(C_ptr, out, mask=mask)
+    else:
+        tl.atomic_add(C_ptr, out, mask=mask)
+
+
+def fp8_gemm_triton(a, b, a_scales, b_scales, fuse_gelu=False, output_dtype=None):
+    assert a.is_cuda and b.is_cuda
+    e4m3_ok = []
+    if hasattr(torch, "float8_e4m3fn"):
+        e4m3_ok.append(torch.float8_e4m3fn)
+    if hasattr(torch, "float8_e4m3fnuz"):
+        e4m3_ok.append(torch.float8_e4m3fnuz)
+    e4m3_ok = tuple(e4m3_ok)
+
+    assert a.dtype in e4m3_ok, f"a.dtype={a.dtype} is not FP8 E4M3"
+    assert b.dtype in e4m3_ok, f"b.dtype={b.dtype} is not FP8 E4M3"
+
+    a_orig_shape = a.shape
+    a2 = a.view(-1, a.shape[-1])
+    b2 = b.t()
+
+    if a2.stride(0) > 1 and a2.stride(1) > 1:
+        a2 = a2.contiguous()
+    if b2.stride(0) > 1 and b2.stride(1) > 1:
+        b2 = b2.contiguous()
+
+    M, K = a2.shape
+    _, N = b2.shape
+    out_shape = a_orig_shape[:-1] + (N,)
+
+    if output_dtype is None:
+        output_dtype = torch.float16
+
+    c = torch.empty((M, N), device=a.device, dtype=output_dtype)
+
+    grid = lambda META: (  # noqa E731
+        cdiv(M, META["BLOCK_M"]) * cdiv(N, META["BLOCK_N"]),
+        META["SPLIT_K"],
+    )  # noqa E731
+    even_k = K % 128 == 0
+
+    fp8_gemm_kernel[grid](
+        a2,
+        b2,
+        a_scales,
+        b_scales,
+        c,
+        M,
+        N,
+        K,
+        a2.stride(0),
+        a2.stride(1),
+        b2.stride(0),
+        b2.stride(1),
+        c.stride(0),
+        c.stride(1),
+        fuse_gelu=fuse_gelu,
+        GROUP_M=8,
+        EVEN_K=even_k,
+    )
+    return c.view(*out_shape)

quantizer/kernels/python/quantizers/__init__.py

@@ -0,0 +1,46 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+FP8 quantization kernels with automatic backend selection.
+
+This module automatically selects between Triton (for Linux/CUDA) and
+PyTorch (for Windows/CPU) implementations based on the runtime environment.
+"""
+
+# from angelslim.compressor._platform import use_triton
+
+# Conditional imports based on platform/backend availability
+# if use_triton():
+from .fp8_per_block import fp8_per_block_quant_triton
+from .fp8_per_token_group import fp8_per_token_group_quant_triton
+# else:
+#     # PyTorch fallback implementations
+#     from .fp8_per_block_torch import (
+#         fp8_per_block_quant_torch as fp8_per_block_quant_triton,
+#     )
+#     from .fp8_per_token_group_torch import (
+#         fp8_per_token_group_quant_torch as fp8_per_token_group_quant_triton,
+#     )
+
+# Also export PyTorch versions directly for explicit use
+from .fp8_per_block_torch import fp8_per_block_quant_torch
+from .fp8_per_token_group_torch import fp8_per_token_group_quant_torch
+
+__all__ = [
+    "fp8_per_token_group_quant_triton",
+    "fp8_per_block_quant_triton",
+    "fp8_per_block_quant_torch",
+    "fp8_per_token_group_quant_torch",
+]

quantizer/kernels/python/quantizers/fp8_per_block.py

@@ -0,0 +1,78 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+# https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
+@triton.jit
+def _fp8_per_block_quant_kernel(x_ptr, y_ptr, s_ptr, M, N, BLOCK_SIZE: tl.constexpr):
+    """Quantizes FP32 tensor to FP8 format using block-wise quantization."""
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    n = tl.cdiv(N, BLOCK_SIZE)
+
+    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs = offs_m[:, None] * N + offs_n[None, :]
+
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+    max_val = tl.max(tl.abs(x))
+    scale = max_val / 448.0
+    scale = tl.where(max_val == 0.0, 1.0, scale)
+    y = x / scale
+    y = y.to(y_ptr.dtype.element_ty)
+
+    tl.store(y_ptr + offs, y, mask=mask)
+    tl.store(s_ptr + pid_m * n + pid_n, scale)
+
+
+# triton implementation
+# for weight quantization on gpu
+def fp8_per_block_quant_triton(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantizes a FP32 2D tensor to FP8 (E4M3FN) using block-wise quantization.
+    For each (block_size x block_size) block:
+        - scale = max(abs(block)) / 448.0 (FP8 E4M3FN max magnitude)
+        - if block is all zeros, use scale = 1.0 to avoid div-by-zero
+        - scale, clamp and cast to FP8
+    Returns:
+        y: Quantized FP8 tensor, same shape as input
+        s: Per-block scales, shape (num_blocks_M, num_blocks_N)
+    """
+    assert x.is_contiguous()
+    assert x.dim() == 2
+
+    M, N = x.size()
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    m_blocks = triton.cdiv(M, block_size)
+    n_blocks = triton.cdiv(N, block_size)
+    s = torch.empty((m_blocks, n_blocks), dtype=torch.float32, device=x.device)
+
+    def grid(meta):
+        return (
+            triton.cdiv(M, meta["BLOCK_SIZE"]),
+            triton.cdiv(N, meta["BLOCK_SIZE"]),
+        )
+
+    _fp8_per_block_quant_kernel[grid](x, y, s, M, N, BLOCK_SIZE=block_size)
+
+    return y, s

quantizer/kernels/python/quantizers/fp8_per_block_torch.py

@@ -0,0 +1,155 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Pure PyTorch implementation of per-block FP8 quantization.
+
+This module provides CPU/Windows-compatible implementations that mirror
+the Triton kernel for FP8 per-block quantization.
+"""
+
+from typing import Tuple
+
+import torch
+
+# FP8 E4M3 max value
+FP8_MAX = 448.0
+
+
+def fp8_per_block_quant_torch(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pure PyTorch implementation of block-wise FP8 quantization.
+
+    Quantizes a FP32 2D tensor to FP8 (E4M3FN) using block-wise quantization.
+    For each (block_size x block_size) block:
+        - scale = max(abs(block)) / 448.0 (FP8 E4M3FN max magnitude)
+        - if block is all zeros, use scale = 1.0 to avoid div-by-zero
+        - scale, clamp and cast to FP8
+
+    Args:
+        x: Input tensor of shape (M, N), must be contiguous
+        block_size: Size of each quantization block (default: 128)
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor):
+            - y: Quantized FP8 tensor, same shape as input
+            - s: Per-block scales, shape (num_blocks_M, num_blocks_N)
+    """
+    assert x.is_contiguous(), "Input tensor must be contiguous"
+    assert x.dim() == 2, "Input tensor must be 2D"
+
+    M, N = x.size()
+    device = x.device
+
+    # Calculate number of blocks
+    m_blocks = (M + block_size - 1) // block_size
+    n_blocks = (N + block_size - 1) // block_size
+
+    # Output tensors
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    s = torch.empty((m_blocks, n_blocks), dtype=torch.float32, device=device)
+
+    # Convert to float32 for computation
+    x_float = x.to(torch.float32)
+
+    # Process each block
+    for mb in range(m_blocks):
+        m_start = mb * block_size
+        m_end = min(m_start + block_size, M)
+        for nb in range(n_blocks):
+            n_start = nb * block_size
+            n_end = min(n_start + block_size, N)
+
+            block = x_float[m_start:m_end, n_start:n_end]
+            max_val = block.abs().amax()
+
+            # Compute scale (guard against zero)
+            scale = max_val / FP8_MAX
+            if scale.item() == 0.0:
+                scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+
+            # Quantize block
+            y_block = (block / scale).to(torch.float8_e4m3fn)
+
+            # Store results
+            y[m_start:m_end, n_start:n_end] = y_block
+            s[mb, nb] = scale
+
+    return y, s
+
+
+def fp8_per_block_quant_torch_fast(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Optimized PyTorch implementation using tensor operations.
+
+    This is faster than the loop-based version for large tensors on GPU.
+
+    Args:
+        x: Input tensor of shape (M, N), must be contiguous
+        block_size: Size of each quantization block (default: 128)
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor)
+    """
+    assert x.is_contiguous(), "Input tensor must be contiguous"
+    assert x.dim() == 2, "Input tensor must be 2D"
+
+    M, N = x.size()
+
+    # Pad tensor to be divisible by block_size
+    pad_m = (block_size - M % block_size) % block_size
+    pad_n = (block_size - N % block_size) % block_size
+
+    if pad_m > 0 or pad_n > 0:
+        x_padded = torch.nn.functional.pad(x, (0, pad_n, 0, pad_m), value=0.0)
+    else:
+        x_padded = x
+
+    M_padded, N_padded = x_padded.size()
+    m_blocks = M_padded // block_size
+    n_blocks = N_padded // block_size
+
+    # Reshape to blocks: [m_blocks, block_size, n_blocks, block_size]
+    x_blocks = x_padded.view(m_blocks, block_size, n_blocks, block_size)
+    x_blocks = x_blocks.permute(0, 2, 1, 3).contiguous()
+
+    # Compute max absolute value per block
+    x_float = x_blocks.to(torch.float32)
+    max_vals = x_float.abs().amax(dim=(2, 3))  # [m_blocks, n_blocks]
+
+    # Compute scales
+    s = max_vals / FP8_MAX
+    s = torch.where(s == 0.0, torch.ones_like(s), s)
+
+    # Quantize
+    s_expanded = s[:, :, None, None]
+    y_blocks = (x_float / s_expanded).to(torch.float8_e4m3fn)
+
+    # Reshape back
+    y_blocks = y_blocks.permute(0, 2, 1, 3).contiguous()
+    y_padded = y_blocks.view(M_padded, N_padded)
+
+    # Remove padding
+    y = y_padded[:M, :N].contiguous()
+
+    # Adjust scale tensor size if needed
+    actual_m_blocks = (M + block_size - 1) // block_size
+    actual_n_blocks = (N + block_size - 1) // block_size
+    s = s[:actual_m_blocks, :actual_n_blocks].contiguous()
+
+    return y, s

quantizer/kernels/python/quantizers/fp8_per_token_group.py

@@ -0,0 +1,170 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+# quant function for per-group fp8 activation
+# https://github.com/sgl-project/sglang/
+# blob/a167fd0bcb9ef4b0f4331a109e40c8cdc770b026/python/sglang/srt/layers/
+# quantization/fp8_kernel.py#L116
+@triton.jit
+def _per_token_group_quant_fp8(
+    y_ptr,
+    y_q_ptr,
+    y_s_ptr,
+    y_stride,
+    N,
+    eps,
+    fp8_min,
+    fp8_max,
+    BLOCK: tl.constexpr,
+):
+    """A Triton-accelerated function for per-token-group quantization."""
+    g_id = tl.program_id(0)
+    y_ptr += g_id * y_stride
+    y_q_ptr += g_id * y_stride
+    y_s_ptr += g_id
+
+    cols = tl.arange(0, BLOCK)
+    mask = cols < N
+
+    y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
+    y_s = _absmax / fp8_max
+    y_s_inv = 1.0 / y_s
+    y_q = tl.clamp(y * y_s_inv, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
+
+    tl.store(y_q_ptr + cols, y_q, mask=mask)
+    tl.store(y_s_ptr, y_s)
+
+
+@triton.jit
+def _per_token_group_quant_fp8_colmajor(
+    y_ptr,
+    y_q_ptr,
+    y_s_ptr,
+    group_size,
+    y_num_columns,
+    y_s_col_stride,
+    eps,
+    fp8_min,
+    fp8_max,
+    BLOCK: tl.constexpr,
+):
+    """A Triton-accelerated function for per-token-group quantization."""
+    g_id = tl.program_id(0)
+    y_ptr += g_id * group_size
+    y_q_ptr += g_id * group_size
+
+    blocks_per_row = y_num_columns // group_size
+    scale_col = g_id % blocks_per_row
+    scale_row = g_id // blocks_per_row
+    y_s_ptr += scale_col * y_s_col_stride + scale_row
+
+    cols = tl.arange(0, BLOCK)
+    mask = cols < group_size
+
+    y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
+    y_s = _absmax / fp8_max
+    y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
+
+    tl.store(y_q_ptr + cols, y_q, mask=mask)
+    tl.store(y_s_ptr, y_s)
+
+
+def fp8_per_token_group_quant_triton(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.float8_e4m3fn,
+    column_major_scales: bool = False,
+    scale_tma_aligned: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Function to perform per-token-group quantization on an input tensor `x`."""
+    assert (
+        x.shape[-1] % group_size == 0
+    ), "the last dimension of `x` cannot be divisible by `group_size`"
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    finfo = torch.finfo(dtype)
+    fp8_max = finfo.max
+    fp8_min = -fp8_max
+
+    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
+    M = x.numel() // group_size
+    N = group_size
+
+    if column_major_scales:
+        if scale_tma_aligned:
+            aligned_size = (x.shape[-2] + 3) // 4 * 4
+            x_s = torch.empty(
+                x.shape[:-2] + (x.shape[-1] // group_size, aligned_size),
+                device=x.device,
+                dtype=torch.float32,
+            ).permute(-1, -2)[: x.shape[-2], :]
+        else:
+            x_s = torch.empty(
+                (x.shape[-1] // group_size,) + x.shape[:-1],
+                device=x.device,
+                dtype=torch.float32,
+            ).permute(-1, -2)
+    else:
+        x_s = torch.empty(
+            x.shape[:-1] + (x.shape[-1] // group_size,),
+            device=x.device,
+            dtype=torch.float32,
+        )
+
+    BLOCK = triton.next_power_of_2(N)
+    num_warps = min(max(BLOCK // 256, 1), 8)
+    num_stages = 1
+
+    if column_major_scales:
+        _per_token_group_quant_fp8_colmajor[(M,)](
+            x,
+            x_q,
+            x_s,
+            group_size,
+            x.shape[1],
+            x_s.stride(1),
+            eps,
+            fp8_min=fp8_min,
+            fp8_max=fp8_max,
+            BLOCK=BLOCK,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+    else:
+        _per_token_group_quant_fp8[(M,)](
+            x,
+            x_q,
+            x_s,
+            group_size,
+            N,
+            eps,
+            fp8_min=fp8_min,
+            fp8_max=fp8_max,
+            BLOCK=BLOCK,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+
+    return x_q, x_s

quantizer/kernels/python/quantizers/fp8_per_token_group_torch.py

@@ -0,0 +1,163 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Pure PyTorch implementation of per-token-group FP8 quantization.
+
+This module provides CPU/Windows-compatible implementations that mirror
+the Triton kernels for FP8 per-token-group quantization.
+"""
+
+from typing import Tuple
+
+import torch
+
+
+def fp8_per_token_group_quant_torch(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.float8_e4m3fn,
+    column_major_scales: bool = False,
+    scale_tma_aligned: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pure PyTorch implementation of per-token-group FP8 quantization.
+
+    Function to perform per-token-group quantization on an input tensor `x`.
+    Each group of `group_size` elements along the last dimension is quantized
+    with its own scale factor.
+
+    Args:
+        x: Input tensor, must be contiguous
+        group_size: Size of each quantization group
+        eps: Small value to avoid division by zero (default: 1e-10)
+        dtype: Target FP8 dtype (default: float8_e4m3fn)
+        column_major_scales: Whether to use column-major scale layout (default: False)
+        scale_tma_aligned: Whether to use TMA-aligned scales (default: False)
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor)
+    """
+    assert (
+        x.shape[-1] % group_size == 0
+    ), f"the last dimension ({x.shape[-1]}) cannot be divisible by group_size ({group_size})"
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    finfo = torch.finfo(dtype)
+    fp8_max = finfo.max
+    fp8_min = -fp8_max
+
+    device = x.device
+    orig_shape = x.shape
+    num_groups = orig_shape[-1] // group_size
+
+    # Reshape for group-wise processing: [..., N] -> [..., num_groups, group_size]
+    x_grouped = x.view(*orig_shape[:-1], num_groups, group_size)
+
+    # Convert to float32 for computation
+    x_float = x_grouped.to(torch.float32)
+
+    # Compute per-group max absolute values
+    absmax = x_float.abs().amax(dim=-1, keepdim=True)  # [..., num_groups, 1]
+    absmax = torch.maximum(absmax, torch.tensor(eps, device=device, dtype=torch.float32))
+
+    # Compute scales
+    scales = absmax / fp8_max  # [..., num_groups, 1]
+
+    # Quantize
+    x_scaled = x_float / scales
+    x_q = x_scaled.clamp(fp8_min, fp8_max).to(dtype)
+
+    # Reshape back to original shape
+    x_q = x_q.view(orig_shape)
+
+    # Prepare scale tensor based on output format
+    scales = scales.squeeze(-1)  # [..., num_groups]
+
+    if column_major_scales:
+        if scale_tma_aligned:
+            # TMA-aligned column-major scales
+            aligned_size = (orig_shape[-2] + 3) // 4 * 4
+            x_s = torch.empty(
+                orig_shape[:-2] + (num_groups, aligned_size),
+                device=device,
+                dtype=torch.float32,
+            )
+            # Transpose and fill
+            scales_transposed = scales.permute(*range(len(orig_shape) - 2), -1, -2).contiguous()
+            x_s[..., : orig_shape[-2]] = scales_transposed
+            x_s = x_s.permute(-1, -2)[: orig_shape[-2], :]
+        else:
+            # Standard column-major scales
+            x_s = scales.permute(*range(len(orig_shape) - 2), -1, -2).contiguous()
+    else:
+        # Row-major scales (default)
+        x_s = scales
+
+    return x_q, x_s.float()
+
+
+def fp8_per_token_group_quant_torch_simple(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.float8_e4m3fn,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Simplified PyTorch implementation for row-major scales only.
+
+    This version is more straightforward and easier to understand.
+
+    Args:
+        x: Input tensor, must be contiguous
+        group_size: Size of each quantization group
+        eps: Small value to avoid division by zero
+        dtype: Target FP8 dtype
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor)
+    """
+    assert x.shape[-1] % group_size == 0
+    assert x.is_contiguous()
+
+    finfo = torch.finfo(dtype)
+    fp8_max = finfo.max
+    fp8_min = -fp8_max
+
+    orig_shape = x.shape
+    num_groups = orig_shape[-1] // group_size
+
+    # Flatten all dimensions except last, then reshape for groups
+    x_flat = x.view(-1, orig_shape[-1])  # [M, N]
+    M = x_flat.shape[0]
+
+    # Reshape to [M, num_groups, group_size]
+    x_grouped = x_flat.view(M, num_groups, group_size).float()
+
+    # Compute per-group absmax
+    absmax = x_grouped.abs().amax(dim=-1, keepdim=True)  # [M, num_groups, 1]
+    absmax = torch.clamp(absmax, min=eps)
+
+    # Compute scales
+    scales = absmax / fp8_max
+
+    # Quantize
+    x_q = (x_grouped / scales).clamp(fp8_min, fp8_max).to(dtype)
+
+    # Reshape back
+    x_q = x_q.view(orig_shape)
+    scales = scales.squeeze(-1).view(*orig_shape[:-1], num_groups)
+
+    return x_q, scales.float()

quantizer/kernels/python/sgl/int8_kernel.py

@@ -0,0 +1,465 @@
+import functools
+import json
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+from functools import lru_cache, partial
+
+import torch
+import triton
+import triton.language as tl
+
+@lru_cache(maxsize=1)
+def is_musa() -> bool:
+    try:
+        import torchada  # noqa: F401
+    except ImportError:
+        return False
+    return hasattr(torch.version, "musa") and torch.version.musa is not None
+
+def get_device_name(device_id: int = 0) -> str:
+    if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
+        return torch.cuda.get_device_name(device_id)
+
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.xpu.get_device_name(device_id)
+
+    if hasattr(torch, "hpu") and torch.hpu.is_available():
+        return torch.hpu.get_device_name(device_id)
+
+    if hasattr(torch, "npu") and torch.npu.is_available():
+        return torch.npu.get_device_name(device_id)
+
+@lru_cache(maxsize=1)
+def is_cuda():
+    return torch.cuda.is_available() and torch.version.cuda is not None
+
+_is_cuda = is_cuda()
+if _is_cuda:
+    # Temporary
+    try:
+        from sgl_kernel import sgl_per_token_group_quant_8bit
+
+        enable_sgl_per_token_group_quant_8bit = True
+    except ImportError:
+        from sgl_kernel import sgl_per_token_group_quant_int8
+
+        enable_sgl_per_token_group_quant_8bit = False
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _per_token_quant_int8(
+    x_ptr,
+    xq_ptr,
+    scale_ptr,
+    x_sum_ptr,
+    stride_x,
+    stride_xq,
+    N,
+    CAL_SUM: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
+    row_id = tl.program_id(0)
+
+    cols = tl.arange(0, BLOCK)
+    mask = cols < N
+
+    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
+    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
+    scale_x = absmax / 127
+    x_q = x * (127 / absmax)
+    x_q = tl.extra.cuda.libdevice.round(x_q).to(tl.int8)
+    if CAL_SUM:
+        x_sum = tl.sum(x, axis=0)
+        tl.store(x_sum_ptr + row_id, x_sum.to(x_sum_ptr.dtype.element_ty))
+
+    tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
+    tl.store(scale_ptr + row_id, scale_x.to(scale_ptr.dtype.element_ty))
+
+
+def per_token_quant_int8(x, scale_dtype=torch.float32, cal_sum=False):
+    M = x.numel() // x.shape[-1]
+    N = x.shape[-1]
+    x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
+    scales = torch.empty(x.shape[:-1] + (1,), device=x.device, dtype=scale_dtype)
+    if cal_sum:
+        x_sum = torch.empty(x.shape[:-1], device=x.device, dtype=x.dtype)
+    else:
+        x_sum = None
+    BLOCK = triton.next_power_of_2(N)
+    # heuristics for number of warps
+    num_warps = min(max(BLOCK // 256, 1), 8)
+
+    assert x.is_contiguous()
+    _per_token_quant_int8[(M,)](
+        x,
+        x_q,
+        scales,
+        x_sum,
+        stride_x=x.stride(-2),
+        stride_xq=x_q.stride(-2),
+        N=N,
+        CAL_SUM=cal_sum,
+        BLOCK=BLOCK,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    if cal_sum:
+        return x_q, scales, x_sum
+    else:
+        return x_q, scales
+
+
+@triton.jit
+def _per_token_group_quant_int8(
+    # Pointers to inputs and output
+    y_ptr,
+    y_q_ptr,
+    y_s_ptr,
+    # Stride of input
+    y_stride,
+    # Columns of input
+    N,
+    # Avoid to divide zero
+    eps,
+    # Information for int8
+    int8_min,
+    int8_max,
+    # Meta-parameters
+    BLOCK: tl.constexpr,
+):
+    """A Triton-accelerated function to perform per-token-group quantization on a
+    tensor.
+
+    This function converts the tensor values into int8 values.
+    """
+    # Map the program id to the row of X and Y it should compute.
+    g_id = tl.program_id(0)
+    y_ptr += g_id * y_stride
+    y_q_ptr += g_id * y_stride
+    y_s_ptr += g_id
+
+    cols = tl.arange(0, BLOCK)  # N <= BLOCK
+    mask = cols < N
+
+    y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    # Quant
+    _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
+    y_s = _absmax / int8_max
+    y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty)
+
+    tl.store(y_q_ptr + cols, y_q, mask=mask)
+    tl.store(y_s_ptr, y_s)
+
+
+def per_token_group_quant_int8(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.int8,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Function to perform per-token-group quantization on an input tensor `x`.
+
+    It converts the tensor values into signed int8 values and returns the
+    quantized tensor along with the scaling factor used for quantization.
+
+    Args:
+        x: The input tensor with ndim >= 2.
+        group_size: The group size used for quantization.
+        eps: The minimum to avoid dividing zero.
+        dtype: The dype of output tensor. Note that only `torch.int8` is supported for now.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the scaling factor for quantization.
+    """
+    assert (
+        x.shape[-1] % group_size == 0
+    ), "the last dimension of `x` cannot be divisible by `group_size`"
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    iinfo = torch.iinfo(dtype)
+    int8_max = iinfo.max
+    int8_min = iinfo.min
+
+    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
+    M = x.numel() // group_size
+    N = group_size
+    x_s = torch.empty(
+        x.shape[:-1] + (x.shape[-1] // group_size,),
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    BLOCK = triton.next_power_of_2(N)
+    # heuristics for number of warps
+    num_warps = min(max(BLOCK // 256, 1), 8)
+    num_stages = 1
+    _per_token_group_quant_int8[(M,)](
+        x,
+        x_q,
+        x_s,
+        group_size,
+        N,
+        eps,
+        int8_min=int8_min,
+        int8_max=int8_max,
+        BLOCK=BLOCK,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+
+    return x_q, x_s
+
+
+def sglang_per_token_group_quant_int8(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.int8,
+    enable_v2: Optional[bool] = None,
+):
+    assert (
+        x.shape[-1] % group_size == 0
+    ), "the last dimension of `x` cannot be divisible by `group_size`"
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    iinfo = torch.iinfo(dtype)
+    int8_max = iinfo.max
+    int8_min = iinfo.min
+
+    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
+    x_s = torch.empty(
+        x.shape[:-1] + (x.shape[-1] // group_size,),
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    # Temporary
+    if enable_sgl_per_token_group_quant_8bit:
+        sgl_per_token_group_quant_8bit(
+            x, x_q, x_s, group_size, eps, int8_min, int8_max, enable_v2=enable_v2
+        )
+    else:
+        assert not enable_v2
+        sgl_per_token_group_quant_int8(x, x_q, x_s, group_size, eps, int8_min, int8_max)
+
+    return x_q, x_s
+
+
+@triton.jit
+def _w8a8_block_int8_matmul(
+    # Pointers to inputs and output
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    # Shape for matmul
+    M,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Stride for inputs and output
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_As_m,
+    stride_As_k,
+    stride_Bs_k,
+    stride_Bs_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    """Triton-accelerated function used to perform linear operations (dot
+    product) on input tensors `A` and `B` with block-wise quantization, and store the result in output
+    tensor `C`.
+    """
+
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+
+    As_ptrs = As + offs_am * stride_As_m
+    offs_bsn = offs_bn // group_n
+    Bs_ptrs = Bs + offs_bsn * stride_Bs_n
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
+
+        accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+
+    if C.dtype.element_ty == tl.bfloat16:
+        c = accumulator.to(tl.bfloat16)
+    elif C.dtype.element_ty == tl.float16:
+        c = accumulator.to(tl.float16)
+    else:
+        c = accumulator.to(tl.float32)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+@functools.lru_cache
+def get_w8a8_block_int8_configs(
+    N: int, K: int, block_n: int, block_k: int
+) -> Optional[Dict[int, Any]]:
+    """
+    Return optimized configurations for the w8a8 block fp8 kernel.
+
+    The return value will be a dictionary that maps an irregular grid of
+    batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
+    kernel on a given batch size bs, the closest batch size in the grid should
+    be picked and the associated configuration chosen to invoke the kernel.
+    """
+
+    # First look up if an optimized configuration is available in the configs
+    # directory
+    device_name = get_device_name().replace(" ", "_")
+    json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json"
+
+    config_file_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name
+    )
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info(
+                "Using configuration from %s for W8A8 Block INT8 kernel.",
+                config_file_path,
+            )
+            # If a configuration has been found, return it
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # If no optimized configuration is available, we will use the default
+    # configuration
+    logger.warning(
+        (
+            "Using default W8A8 Block INT8 kernel config. Performance might be sub-optimal! "
+            "Config file not found at %s"
+        ),
+        config_file_path,
+    )
+    return None
+
+
+def w8a8_block_int8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: List[int],
+    output_dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    """This function performs matrix multiplication with block-wise quantization.
+
+    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+    The output is returned in the specified `output_dtype`.
+
+    Args:
+        A: The input tensor, e.g., activation.
+        B: The input tensor, e.g., weight.
+        As: The per-token-group quantization scale for `A`.
+        Bs: The per-block quantization scale for `B`.
+        block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128].
+        output_dytpe: The dtype of the returned tensor.
+
+    Returns:
+        torch.Tensor: The result of matmul.
+    """
+    assert len(block_size) == 2
+    block_n, block_k = block_size[0], block_size[1]
+
+    assert A.shape[-1] == B.shape[-1]
+    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
+    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
+    M = A.numel() // A.shape[-1]
+
+    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
+    N, K = B.shape
+    assert triton.cdiv(N, block_n) == Bs.shape[0]
+    assert triton.cdiv(K, block_k) == Bs.shape[1]
+
+    C_shape = A.shape[:-1] + (N,)
+    C = A.new_empty(C_shape, dtype=output_dtype)
+
+    configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1])
+    if configs:
+        # If an optimal configuration map has been found, look up the
+        # optimal config
+        config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
+    else:
+        # Default config
+        # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1]
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": block_size[0],
+            "BLOCK_SIZE_K": block_size[1],
+            "GROUP_SIZE_M": 32,
+            "num_warps": 4,
+            "num_stages": 3,
+        }
+
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+
+    _w8a8_block_int8_matmul[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(-2),
+        A.stride(-1),
+        B.stride(1),
+        B.stride(0),
+        C.stride(-2),
+        C.stride(-1),
+        As.stride(-2),
+        As.stride(-1),
+        Bs.stride(1),
+        Bs.stride(0),
+        **config,
+    )
+
+    return C

quantizer/quant/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .ptq import DynamicDiTQuantizer
+
+__all__ = [
+    "DynamicDiTQuantizer",
+]

quantizer/quant/modules/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .linear import FP4DynamicLinear, FP8DynamicLinear, FP8WeightOnlyLinear, INT8DynamicLinear
+
+__all__ = ["FP8DynamicLinear", "FP8WeightOnlyLinear", "FP4DynamicLinear", "INT8DynamicLinear"]

quantizer/quant/modules/linear.py

@@ -0,0 +1,589 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+
+import torch
+from lightx2v_kernel.gemm import (
+    cutlass_scaled_mxfp4_mm,
+    cutlass_scaled_mxfp6_mxfp8_mm,
+    cutlass_scaled_mxfp8_mm,
+    cutlass_scaled_nvfp4_mm,
+)
+
+try:
+    from torchao.quantization.utils import quant_int8_per_token_matmul as torchao_int8_gemm
+    from torchao.quantization.utils import quantize_activation_per_token_absmax as torchao_int8_quant
+except ImportError:
+    try:
+        from torchao.quantization.utils import _quant_int8_per_token_matmul as torchao_int8_gemm
+        from torchao.quantization.utils import _quantize_activation_per_token_absmax as torchao_int8_quant
+    except ImportError:
+        torchao_int8_gemm, torchao_int8_quant = None, None
+
+try:
+    from vllm import _custom_ops as vllm_ops
+except ImportError:
+    vllm_ops = None
+
+try:
+    from ...kernels.python.sgl.int8_kernel import per_token_quant_int8 as sglang_int8_act_quant
+except ImportError:
+    sglang_int8_act_quant = None
+
+try:
+    import sgl_kernel
+except ImportError:
+    sgl_kernel = None
+
+try:
+    from q8_kernels.functional.linear import q8_linear
+except ImportError:
+    q8_linear = None
+
+try:
+    from ...kernels.python.mm.triton_kernels import (
+        int8_gemm_bias_triton,
+        int8_gemm_triton,
+        int8_quantize_triton,
+    )
+except ImportError:
+    int8_gemm_bias_triton, int8_gemm_triton, int8_quantize_triton = None, None, None
+
+from ..quant_func import (
+    fp8_gemm,
+    fp8_per_block_quant,
+    fp8_per_tensor_quant,
+    fp8_per_token_group_quant,
+    fp8_per_token_quant_sgl,
+    fp8_weight_only_gemm,
+    mxfp4_per_tensor_quant,
+    mxfp6_per_tensor_quant,
+    mxfp8_per_tensor_quant,
+    nvfp4_per_tensor_quant,
+)
+
+
+# modified from https://github.com/neuralmagic/AutoFP8/blob/main/auto_fp8/quantize.py
+class FP8DynamicLinear(torch.nn.Module):
+    def __init__(
+        self,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.nn.Parameter,
+        native_fp8_support: bool = False,
+        quant_type: str = "fp8-per-tensor",
+        block_size: int = 128,
+    ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+        self.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+        self.bias = bias
+        self.native_fp8_support = native_fp8_support
+        self.quant_type = quant_type
+        self.block_size = block_size
+        self.profile_enabled = os.environ.get("ANGELSLIM_FP8_PROFILE", "0") == "1"
+
+    @torch.compiler.disable(recursive=True)
+    def forward(self, x):
+        ori_dtype = x.dtype
+        assert ori_dtype in [
+            torch.float32,
+            torch.bfloat16,
+            torch.float16,
+        ], "x.dtype must be float32, bfloat16, or float16"
+
+        if ori_dtype == torch.float32:
+            x = x.to(torch.bfloat16)
+
+        if self.profile_enabled and x.is_cuda:
+            torch.cuda.synchronize(x.device)
+        t0 = time.perf_counter()
+
+        if self.quant_type == "fp8-per-tensor":
+            origin_shape = None
+            qinput, x_scale = fp8_per_tensor_quant(x)
+        elif self.quant_type == "fp8-per-token":
+            origin_shape = None
+            x_2d = x.view(-1, x.shape[-1])
+            qinput, x_scale = fp8_per_token_group_quant(x_2d, x_2d.shape[-1])
+        elif self.quant_type == "fp8-per-token-sgl" and self.native_fp8_support:
+            origin_shape = x.shape
+            x_2d = x.view(-1, x.shape[-1])
+            qinput, x_scale = fp8_per_token_quant_sgl(x_2d)
+        elif self.quant_type == "fp8-per-block" and self.native_fp8_support:
+            origin_shape = x.shape
+            x = x.view(-1, x.shape[-1])
+            qinput, x_scale = fp8_per_token_group_quant(
+                x, group_size=128, column_major_scales=True, scale_tma_aligned=True
+            )
+        elif self.quant_type == "fp8-per-block" and not self.native_fp8_support:
+            origin_shape = x.shape
+            x_2d = x.view(-1, x.shape[-1])
+            qinput, x_scale = fp8_per_block_quant(x_2d, block_size=128)
+        elif self.quant_type == "fp8-per-channel-vllm":
+            if vllm_ops is None:
+                raise ImportError(
+                    "quant_type='fp8-per-channel-vllm' requires vllm._custom_ops, but vllm is not installed"
+                )
+            origin_shape = x.shape if x.dim() == 3 else None
+            x_2d = x.view(-1, x.shape[-1]) if x.dim() == 3 else x
+            qinput, x_scale = vllm_ops.scaled_fp8_quant(
+                x_2d, None, scale_ub=None, use_per_token_if_dynamic=True
+            )
+        else:
+            raise ValueError(f"Invalid quant_type: {self.quant_type}")
+
+        if self.profile_enabled and qinput.is_cuda:
+            torch.cuda.synchronize(qinput.device)
+        t1 = time.perf_counter()
+
+        output = fp8_gemm(
+            A=qinput,
+            A_scale=x_scale,
+            B=self.weight,
+            B_scale=self.weight_scale,
+            bias=self.bias,
+            out_dtype=x.dtype,
+            native_fp8_support=self.native_fp8_support,
+            quant_type=self.quant_type,
+            origin_shape=origin_shape,
+        )
+
+        if self.profile_enabled and output.is_cuda:
+            torch.cuda.synchronize(output.device)
+        t2 = time.perf_counter()
+
+        if self.profile_enabled:
+            qshape = tuple(qinput.shape)
+            print(
+                f"[FP8Linear:{self.quant_type}] quant_ms={(t1 - t0) * 1000:.3f}, "
+                f"gemm_ms={(t2 - t1) * 1000:.3f}, qshape={qshape}"
+            )
+
+        if (
+            self.quant_type in ["fp8-per-token", "fp8-per-token-sgl"]
+            and x.dim() == 3
+            and output.dim() == 2
+        ):
+            output = output.unsqueeze(0)
+
+        # Restore original shape for fp8-per-block with native_fp8_support=False
+        # (native_fp8_support=True case is handled in fp8_gemm_deepgemm_block)
+        if (
+            (
+                (self.quant_type == "fp8-per-block" and not self.native_fp8_support)
+                or self.quant_type == "fp8-per-channel-vllm"
+            )
+            and origin_shape is not None
+            and len(origin_shape) == 3
+            and output.dim() == 2
+        ):
+            output = output.view(origin_shape[0], origin_shape[1], -1)
+
+        return output
+
+
+class FP8WeightOnlyLinear(torch.nn.Module):
+    """
+    FP8 Weight-Only Quantized Linear Layer.
+
+    This layer quantizes only the weights to FP8 while keeping activations
+    in higher precision (bfloat16/float16). This provides a good balance
+    between memory savings and accuracy.
+    """
+
+    def __init__(
+        self,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.nn.Parameter,
+        native_fp8_support: bool = False,  # not used
+        quant_type: str = "fp8-per-tensor-weight-only",
+    ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+        self.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+        self.bias = bias
+        self.native_fp8_support = native_fp8_support  # not used
+        self.quant_type = quant_type
+
+    @torch.compiler.disable(recursive=True)
+    def forward(self, x):
+        ori_dtype = x.dtype
+        assert ori_dtype in [
+            torch.float32,
+            torch.bfloat16,
+            torch.float16,
+        ], "x.dtype must be float32, bfloat16, or float16"
+
+        if ori_dtype == torch.float32:
+            x = x.to(torch.bfloat16)
+
+        # For weight-only quantization, we don't quantize activations
+        # Just use the original activations with quantized weights
+        output = fp8_weight_only_gemm(
+            A=x,  # Keep activations in original precision
+            B=self.weight,
+            B_scale=self.weight_scale,
+            bias=self.bias,
+            out_dtype=x.dtype,
+        )
+
+        return output
+
+
+class INT8DynamicLinear(torch.nn.Module):
+    """
+    INT8 weight-only linear layer with per-channel scales.
+    """
+
+    def __init__(
+        self,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.nn.Parameter,
+        native_fp8_support: bool = False,  # not used
+        quant_type: str = "int8",
+    ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+        self.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+        self.bias = bias
+        self.native_fp8_support = native_fp8_support  # not used
+        self.quant_type = quant_type
+
+    @staticmethod
+    def _is_backend_available(backend: str) -> bool:
+        if backend == "torchao":
+            return torchao_int8_quant is not None and torchao_int8_gemm is not None
+        if backend == "vllm":
+            return vllm_ops is not None and hasattr(torch.ops, "_C")
+        if backend == "triton":
+            return (
+                int8_quantize_triton is not None
+                and int8_gemm_triton is not None
+                and int8_gemm_bias_triton is not None
+            )
+        if backend == "sgl":
+            has_act = (
+                sglang_int8_act_quant is not None
+                or vllm_ops is not None
+                or torchao_int8_quant is not None
+                or int8_quantize_triton is not None
+            )
+            return sgl_kernel is not None and has_act
+        if backend == "q8f":
+            has_act = (
+                vllm_ops is not None
+                or torchao_int8_quant is not None
+                or int8_quantize_triton is not None
+            )
+            return q8_linear is not None and has_act
+        return False
+
+    def _resolve_int8_backend(self) -> str:
+        explicit_backend = {
+            "int8-torchao": "torchao",
+            "int8-vllm": "vllm",
+            "int8-triton": "triton",
+            "int8-sgl": "sgl",
+            "int8-q8f": "q8f",
+        }
+        if self.quant_type in explicit_backend:
+            backend = explicit_backend[self.quant_type]
+            if not self._is_backend_available(backend):
+                raise ImportError(
+                    f"quant_type='{self.quant_type}' requires '{backend}' backend dependencies"
+                )
+            return backend
+
+        # quant_type='int8' uses auto priority: sgl > vllm > torchao > triton
+        for backend in ("sgl", "vllm", "torchao", "triton"):
+            if self._is_backend_available(backend):
+                return backend
+        raise ImportError(
+            "quant_type='int8' requires one of backends [sgl, vllm, torchao, triton], but none is available"
+        )
+
+    def _act_quant_int8_torchao(self, x_2d: torch.Tensor):
+        input_tensor_quant, input_tensor_scale = torchao_int8_quant(x_2d)
+        return input_tensor_quant, input_tensor_scale.float()
+
+    def _act_quant_int8_vllm(self, x_2d: torch.Tensor):
+        input_tensor_quant, input_tensor_scale, _ = vllm_ops.scaled_int8_quant(
+            x_2d, scale=None, azp=None, symmetric=True
+        )
+        return input_tensor_quant, input_tensor_scale.float()
+
+    def _act_quant_int8_triton(self, x_2d: torch.Tensor):
+        input_tensor_quant, input_tensor_scale = int8_quantize_triton(x_2d)
+        return input_tensor_quant, input_tensor_scale.float()
+
+    def _act_quant_int8_sgl(self, x_2d: torch.Tensor):
+        if sglang_int8_act_quant is not None:
+            input_tensor_quant, input_tensor_scale = sglang_int8_act_quant(x_2d)
+            return input_tensor_quant, input_tensor_scale.float()
+        if vllm_ops is not None:
+            return self._act_quant_int8_vllm(x_2d)
+        if torchao_int8_quant is not None:
+            return self._act_quant_int8_torchao(x_2d)
+        if int8_quantize_triton is not None:
+            return self._act_quant_int8_triton(x_2d)
+        raise ImportError("int8-sgl activation quantization requires sglang/vllm/torchao/triton")
+
+    def _act_quant_by_backend(self, x_2d: torch.Tensor, backend: str):
+        if backend == "torchao":
+            return self._act_quant_int8_torchao(x_2d)
+        if backend == "vllm":
+            return self._act_quant_int8_vllm(x_2d)
+        if backend == "triton":
+            return self._act_quant_int8_triton(x_2d)
+        if backend == "sgl":
+            return self._act_quant_int8_sgl(x_2d)
+        if backend == "q8f":
+            if vllm_ops is not None:
+                return self._act_quant_int8_vllm(x_2d)
+            if torchao_int8_quant is not None:
+                return self._act_quant_int8_torchao(x_2d)
+            return self._act_quant_int8_triton(x_2d)
+        raise ValueError(f"Unsupported int8 backend: {backend}")
+
+    def _gemm_int8_torchao(self, qinput, x_scale, out_dtype):
+        output = torchao_int8_gemm(
+            qinput,
+            x_scale,
+            self.weight.t(),
+            self.weight_scale.t().float(),
+            output_dtype=out_dtype,
+        )
+        if self.bias is not None:
+            output.add_(self.bias.to(output.dtype))
+        return output
+
+    def _gemm_int8_vllm(self, qinput, x_scale, out_dtype):
+        shape = (qinput.shape[0], self.weight.shape[0])
+        output = torch.empty(shape, dtype=out_dtype, device=qinput.device, requires_grad=False)
+        torch.ops._C.cutlass_scaled_mm(
+            output,
+            qinput,
+            self.weight.t(),
+            x_scale,
+            self.weight_scale.t(),
+            self.bias,
+        )
+        return output
+
+    def _gemm_int8_triton(self, qinput, x_scale, out_dtype):
+        if self.bias is not None:
+            return int8_gemm_bias_triton(
+                qinput,
+                self.weight,
+                self.bias,
+                x_scale,
+                self.weight_scale,
+                output_dtype=out_dtype,
+            )
+        return int8_gemm_triton(
+            qinput,
+            self.weight,
+            x_scale,
+            self.weight_scale,
+            output_dtype=out_dtype,
+        )
+
+    def _gemm_int8_sgl(self, qinput, x_scale, out_dtype):
+        return sgl_kernel.int8_scaled_mm(
+            qinput,
+            self.weight.t(),
+            x_scale,
+            self.weight_scale.t(),
+            out_dtype,
+            self.bias,
+        )
+
+    def _gemm_int8_q8f(self, qinput, x_scale, out_dtype):
+        bias_fp32 = self.bias.float() if self.bias is not None else None
+        return q8_linear(
+            qinput,
+            self.weight,
+            bias_fp32,
+            x_scale.float(),
+            self.weight_scale,
+            fuse_gelu=False,
+            out_dtype=out_dtype,
+        )
+
+    def _gemm_by_backend(self, qinput, x_scale, out_dtype, backend: str):
+        if backend == "torchao":
+            return self._gemm_int8_torchao(qinput, x_scale, out_dtype)
+        if backend == "vllm":
+            return self._gemm_int8_vllm(qinput, x_scale, out_dtype)
+        if backend == "triton":
+            return self._gemm_int8_triton(qinput, x_scale, out_dtype)
+        if backend == "sgl":
+            return self._gemm_int8_sgl(qinput, x_scale, out_dtype)
+        if backend == "q8f":
+            return self._gemm_int8_q8f(qinput, x_scale, out_dtype)
+        raise ValueError(f"Unsupported int8 backend: {backend}")
+
+    @torch.compiler.disable(recursive=True)
+    def forward(self, x):
+        ori_dtype = x.dtype
+        assert ori_dtype in [
+            torch.float32,
+            torch.bfloat16,
+            torch.float16,
+        ], "x.dtype must be float32, bfloat16, or float16"
+
+        if ori_dtype == torch.float32:
+            x = x.to(torch.bfloat16)
+
+        need_reshape = x.dim() == 3
+        if need_reshape:
+            origin_shape = x.shape
+            x_2d = x.view(-1, x.shape[-1])
+        else:
+            origin_shape = None
+            x_2d = x
+
+        backend = self._resolve_int8_backend()
+        qinput, x_scale = self._act_quant_by_backend(x_2d, backend)
+        output = self._gemm_by_backend(qinput, x_scale, x.dtype, backend)
+
+        if need_reshape and output.dim() == 2:
+            output = output.view(origin_shape[0], origin_shape[1], -1)
+        return output.to(ori_dtype)
+
+
+class FP4DynamicLinear(torch.nn.Module):
+    def __init__(
+        self,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.nn.Parameter,
+        weight_global_scale: torch.Tensor = None,
+        native_fp8_support: bool = False,
+        quant_type: str = "nvfp4",
+        block_size: int = 16,
+    ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+        self.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+        self.bias = bias
+        self.native_fp8_support = native_fp8_support
+        self.quant_type = quant_type
+        self.block_size = block_size
+        self.profile_enabled = os.environ.get("ANGELSLIM_NVFP4_PROFILE", "0") == "1"
+        if weight_global_scale is None:
+            weight_global_scale = torch.tensor(1.0, dtype=torch.float32, device=weight.device)
+        self.weight_global_scale = torch.nn.Parameter(
+            weight_global_scale.to(dtype=torch.float32), requires_grad=False
+        )
+        self.calibrate_x_absmax()
+
+    def calibrate_x_absmax(self):
+        if self.quant_type in ("mxfp4", "mxfp6", "mxfp8"):
+            self.x_absmax = torch.tensor(1.0, dtype=torch.float32, device=self.weight.device)
+            self.input_global_scale = torch.tensor(
+                1.0, dtype=torch.float32, device=self.weight.device
+            )
+        else:
+            self.x_absmax = torch.tensor(5.0, dtype=torch.float32, device=self.weight.device)
+            self.input_global_scale = (2688.0 / self.x_absmax).to(torch.float32)
+        self.alpha = 1.0 / (self.input_global_scale * self.weight_global_scale)
+
+    @torch.compiler.disable(recursive=True)
+    def forward(self, x):
+        ori_dtype = x.dtype
+        assert ori_dtype in [
+            torch.float32,
+            torch.bfloat16,
+            torch.float16,
+        ], "x.dtype must be float32, bfloat16, or float16"
+
+        if ori_dtype == torch.float32:
+            x = x.to(torch.bfloat16)
+
+        need_reshape = x.dim() == 3
+        if need_reshape:
+            origin_shape = x.shape
+            x_2d = x.view(-1, x.shape[-1])
+        else:
+            x_2d = x
+
+        if self.profile_enabled and x_2d.is_cuda:
+            torch.cuda.synchronize(x_2d.device)
+        t0 = time.perf_counter()
+        if self.quant_type == "nvfp4":
+            qinput, x_scale, _ = nvfp4_per_tensor_quant(x_2d, self.input_global_scale)
+            output = cutlass_scaled_nvfp4_mm(
+                qinput,
+                self.weight,
+                x_scale,
+                self.weight_scale,
+                self.alpha,
+                bias=self.bias,
+            )
+        elif self.quant_type == "mxfp4":
+            qinput, x_scale, _ = mxfp4_per_tensor_quant(x_2d)
+            output = cutlass_scaled_mxfp4_mm(
+                qinput,
+                self.weight,
+                x_scale,
+                self.weight_scale,
+                self.alpha,
+                bias=self.bias,
+            )
+        elif self.quant_type == "mxfp8":
+            qinput, x_scale, _ = mxfp8_per_tensor_quant(x_2d)
+            output = cutlass_scaled_mxfp8_mm(
+                qinput,
+                self.weight,
+                x_scale,
+                self.weight_scale,
+                self.alpha,
+                bias=self.bias,
+            )
+        elif self.quant_type == "mxfp6":
+            qinput, x_scale, _ = mxfp8_per_tensor_quant(x_2d)
+            output = cutlass_scaled_mxfp6_mxfp8_mm(
+                qinput,
+                self.weight,
+                x_scale,
+                self.weight_scale,
+                self.alpha,
+                bias=self.bias,
+            )
+        else:
+            raise ValueError(f"Invalid quant_type for FP4DynamicLinear: {self.quant_type}")
+        if self.profile_enabled and x_2d.is_cuda:
+            torch.cuda.synchronize(x_2d.device)
+        t1 = time.perf_counter()
+        if self.profile_enabled and x_2d.is_cuda:
+            torch.cuda.synchronize(x_2d.device)
+        t2 = time.perf_counter()
+
+        if self.profile_enabled:
+            print(
+                f"[NVFP4Linear] quant_ms={(t1 - t0) * 1000:.3f}, "
+                f"gemm_ms={(t2 - t1) * 1000:.3f}, "
+                f"shape=({x_2d.shape[0]}, {x_2d.shape[1]})"
+            )
+
+        if need_reshape:
+            output = output.view(origin_shape[0], origin_shape[1], -1)
+
+        return output.to(ori_dtype)

quantizer/quant/ptq.py

@@ -0,0 +1,272 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import logging
+import re
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import tqdm
+
+from .modules import FP4DynamicLinear, FP8DynamicLinear, FP8WeightOnlyLinear, INT8DynamicLinear
+from .quant_func import (
+    fp8_per_block_quant,
+    fp8_per_channel_quant,
+    fp8_per_tensor_quant,
+    fp8_per_token_group_quant,
+    int8_per_channel_quant,
+    mxfp4_per_tensor_quant,
+    mxfp6_per_tensor_quant,
+    mxfp8_per_tensor_quant,
+    nvfp4_per_tensor_quant,
+)
+from .utils import (
+    QuantType,
+    _ensure_deep_gemm,
+    _ensure_sgl_kernel,
+    cleanup_memory,
+    load_fp8_scales,
+    load_quantized_model,
+    replace_module,
+    save_quantized_model,
+    should_quantize_layer,
+)
+
+__all__ = ["DynamicDiTQuantizer"]
+
+logger = logging.getLogger(__name__)
+
+
+class DynamicDiTQuantizer:
+    """
+    Quantizer for DiT that supports various FP8 quantization strategies.
+    Optimized for efficient initialization and conversion.
+    """
+
+    def __init__(
+        self,
+        quant_type: str = QuantType.FP8_PER_TENSOR,
+        layer_filter: Optional[Callable[[str], bool]] = None,
+        include_patterns: Optional[List[Union[str, re.Pattern]]] = None,
+        exclude_patterns: Optional[List[Union[str, re.Pattern]]] = None,
+        native_fp8_support: Optional[bool] = None,
+    ):
+        QuantType.validate(quant_type)
+        self.fp8_scales_map = {}
+        self.quant_type = quant_type
+        self.include_patterns = include_patterns or [
+            "wrapped_module",
+            "block",
+            "lin",
+            "img",
+            "txt",
+        ]
+        self.exclude_patterns = exclude_patterns or ["embed"]
+
+        # Configure layer filter function
+        self.layer_filter = (
+            layer_filter
+            if layer_filter is not None
+            else lambda name: should_quantize_layer(
+                name, self.include_patterns, self.exclude_patterns
+            )
+        )
+
+        # Auto-detect FP8 native support, fallback to False if not present
+        if native_fp8_support is not None:
+            self.native_fp8_support = native_fp8_support
+        else:
+            self.native_fp8_support = (
+                torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9)
+            )
+
+        self.quantize_linear_module = self._set_quantize_linear_module()
+
+    def _set_quantize_linear_module(self) -> torch.nn.Module:
+        if self.quant_type in (
+            QuantType.INT8,
+            QuantType.INT8_TORCHAO,
+            QuantType.INT8_VLLM,
+            QuantType.INT8_TRITON,
+            QuantType.INT8_SGL,
+            QuantType.INT8_Q8F,
+        ):
+            return INT8DynamicLinear
+        if self.quant_type in (QuantType.NVFP4, QuantType.MXFP4, QuantType.MXFP6, QuantType.MXFP8):
+            return FP4DynamicLinear
+        if "fp8" in self.quant_type:
+            if self.quant_type == QuantType.FP8_PER_TENSOR_WEIGHT_ONLY:
+                return FP8WeightOnlyLinear
+            else:
+                return FP8DynamicLinear
+        raise ValueError(f"Invalid quant_type: {self.quant_type}")
+
+    def _quantize_linear_weight(
+        self, linear: torch.nn.Linear
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+        if self.quant_type == QuantType.FP8_PER_TENSOR:
+            quant_weight, weight_scale = fp8_per_tensor_quant(linear.weight)
+        elif self.quant_type == QuantType.FP8_PER_TOKEN:
+            quant_weight, weight_scale = fp8_per_token_group_quant(
+                linear.weight, linear.weight.shape[-1]
+            )
+            weight_scale = weight_scale.t()
+        elif self.quant_type == QuantType.FP8_PER_TOKEN_SGL:
+            if self.native_fp8_support:
+                _ensure_sgl_kernel()
+            quant_weight, weight_scale = fp8_per_channel_quant(linear.weight)
+        elif self.quant_type == QuantType.FP8_PER_BLOCK:
+            if self.native_fp8_support:
+                _ensure_deep_gemm()
+            quant_weight, weight_scale = fp8_per_block_quant(linear.weight)
+        elif self.quant_type == QuantType.FP8_PER_CHANNEL_VLLM:
+            quant_weight, weight_scale = fp8_per_channel_quant(linear.weight)
+        elif self.quant_type == QuantType.FP8_PER_TENSOR_WEIGHT_ONLY:
+            quant_weight, weight_scale = fp8_per_tensor_quant(linear.weight)
+        elif self.quant_type in (
+            QuantType.INT8,
+            QuantType.INT8_TORCHAO,
+            QuantType.INT8_VLLM,
+            QuantType.INT8_TRITON,
+            QuantType.INT8_SGL,
+            QuantType.INT8_Q8F,
+        ):
+            quant_weight, weight_scale = int8_per_channel_quant(linear.weight)
+        elif self.quant_type == QuantType.NVFP4:
+            quant_weight, weight_scale, weight_global_scale = nvfp4_per_tensor_quant(linear.weight)
+            return quant_weight, weight_scale, weight_global_scale
+        elif self.quant_type == QuantType.MXFP4:
+            quant_weight, weight_scale, weight_global_scale = mxfp4_per_tensor_quant(linear.weight)
+            return quant_weight, weight_scale, weight_global_scale
+        elif self.quant_type == QuantType.MXFP8:
+            quant_weight, weight_scale, weight_global_scale = mxfp8_per_tensor_quant(linear.weight)
+            return quant_weight, weight_scale, weight_global_scale
+        elif self.quant_type == QuantType.MXFP6:
+            # Mixed kernel route expects MXFP8-formatted weights.
+            quant_weight, weight_scale, weight_global_scale = mxfp6_per_tensor_quant(linear.weight)
+            return quant_weight, weight_scale, weight_global_scale
+        else:
+            raise ValueError(f"Invalid quant_type: {self.quant_type}")
+        return quant_weight, weight_scale
+
+    def _convert_linear_with_scale(
+        self, model: torch.nn.Module, scale: Union[torch.Tensor, float]
+    ):
+        model.to(torch.bfloat16)
+        assert scale is not None, "scale is required"
+        self.fp8_scales_map = load_fp8_scales(scale)
+        converted_count = 0
+        for name, module in tqdm.tqdm(list(model.named_modules()), desc="converting linear"):
+            if isinstance(module, torch.nn.Linear) and self.layer_filter(name):
+                # Prefer $name.weight_scale, fallback to "$name" key if needed
+                s = self.fp8_scales_map.get(f"{name}.weight_scale") or self.fp8_scales_map.get(
+                    name
+                )
+                if s is None:
+                    continue
+                # import pdb; pdb.set_trace()
+                weight = module.weight.to(torch.float8_e4m3fn)
+                bias = copy.deepcopy(module.bias) if module.bias is not None else None
+                quant_linear = self.quantize_linear_module(
+                    weight=weight,
+                    weight_scale=s.float(),
+                    bias=bias,
+                    native_fp8_support=self.native_fp8_support,
+                    quant_type=self.quant_type,
+                )
+                replace_module(model, name, quant_linear)
+                del module.weight, module.bias, module
+                converted_count += 1
+        print(f"[INIT] quantized linear layers: {converted_count}")
+        cleanup_memory()
+
+    def _convert_linear(self, model: torch.nn.Module):
+        model.to(torch.bfloat16)
+        named_modules = list(model.named_modules())
+        converted_count = 0
+        for name, module in tqdm.tqdm(named_modules, desc="converting linear"):
+            if isinstance(module, torch.nn.Linear) and self.layer_filter(name):
+                quantized = self._quantize_linear_weight(module)
+                if self.quant_type in (
+                    QuantType.NVFP4,
+                    QuantType.MXFP4,
+                    QuantType.MXFP6,
+                    QuantType.MXFP8,
+                ):
+                    quant_weight, weight_scale, weight_global_scale = quantized
+                else:
+                    quant_weight, weight_scale = quantized
+                self.fp8_scales_map[f"{name}.weight_scale"] = weight_scale
+                bias = copy.deepcopy(module.bias) if module.bias is not None else None
+                if self.quant_type in (
+                    QuantType.NVFP4,
+                    QuantType.MXFP4,
+                    QuantType.MXFP6,
+                    QuantType.MXFP8,
+                ):
+                    self.fp8_scales_map[f"{name}.weight_global_scale"] = weight_global_scale
+                    quant_linear = self.quantize_linear_module(
+                        weight=quant_weight,
+                        weight_scale=weight_scale,
+                        weight_global_scale=weight_global_scale,
+                        bias=bias,
+                        native_fp8_support=self.native_fp8_support,
+                        quant_type=self.quant_type,
+                    )
+                else:
+                    quant_linear = self.quantize_linear_module(
+                        weight=quant_weight,
+                        weight_scale=weight_scale,
+                        bias=bias,
+                        native_fp8_support=self.native_fp8_support,
+                        quant_type=self.quant_type,
+                    )
+                replace_module(model, name, quant_linear)
+                del module.weight, module.bias, module
+                converted_count += 1
+        print(f"[INIT] quantized linear layers: {converted_count}")
+        cleanup_memory()
+
+    def convert_linear(
+        self, model: torch.nn.Module, scale: Optional[Union[torch.Tensor, float]] = None
+    ):
+        if scale is not None:
+            self._convert_linear_with_scale(model, scale)
+        else:
+            self._convert_linear(model)
+
+    def export_quantized_weight(self, model: torch.nn.Module, save_path: str):
+        assert (
+            self.quant_type == QuantType.FP8_PER_TENSOR
+        ), "Currently only FP8_PER_TENSOR is supported for export"
+        self.convert_linear(model)
+        save_quantized_model(model, save_path, self.fp8_scales_map)
+        logger.info(f"Quantized model saved to {save_path}")
+        logger.info(f"Quantized scales saved to {save_path}/fp8_scales.safetensors")
+
+    @staticmethod
+    def load_quantized_model(model_class, save_path: str, device: str = "cpu"):
+        """
+        Load a quantized model from save_path.
+
+        Args:
+            model_class: The model class to instantiate
+            save_path: Path to the saved model directory
+            device: Device to load the model on
+
+        Returns:
+            Loaded model with quantized weights
+        """
+        return load_quantized_model(model_class, save_path, device)

quantizer/quant/quant_func.py

@@ -0,0 +1,520 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+import torch
+
+from ..kernels.python.gemm import fp8_gemm_triton_block
+from lightx2v_kernel.gemm import (
+    cutlass_scaled_mxfp4_mm,
+    cutlass_scaled_mxfp6_mxfp8_mm,
+    cutlass_scaled_mxfp8_mm,
+    cutlass_scaled_nvfp4_mm,
+    scaled_mxfp4_quant,
+    scaled_mxfp6_quant,
+    scaled_mxfp8_quant,
+    scaled_nvfp4_quant,
+)
+from ..kernels.python.quantizers import (
+    fp8_per_block_quant_triton,
+    fp8_per_token_group_quant_triton,
+)
+from .utils import QuantType, _ensure_deep_gemm, _ensure_sgl_kernel
+
+FP8_MAX = float(torch.finfo(torch.float8_e4m3fn).max)
+FP8_MIN = float(torch.finfo(torch.float8_e4m3fn).min)
+
+
+# quant function for per-tensor fp8
+# modified from https://github.com/neuralmagic/AutoFP8/blob/main/auto_fp8/quantize.py
+def fp8_per_tensor_quant(x: torch.Tensor) -> Tuple[torch.Tensor, float]:
+    """Dynamically Quantize a tensor using per-tensor static scaling factor.
+    Args:
+        x: The input tensor.
+    """
+    if x.numel() == 0:
+        # handle empty tensor for empty MoE experts
+        min_val, max_val = (
+            torch.tensor(-16.0, dtype=x.dtype),
+            torch.tensor(16.0, dtype=x.dtype),
+        )
+    else:
+        min_val, max_val = x.aminmax()
+    amax = torch.maximum(min_val.abs(), max_val.abs())
+    scale = FP8_MAX / amax.clamp(min=1e-12)
+    qx = (x * scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    qx = qx.to(torch.float8_e4m3fn)
+    scale = scale.float().reciprocal()
+    return qx, scale
+
+
+# quant function for per-token and per-group fp8
+def fp8_per_token_group_quant(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    dtype: torch.dtype = torch.float8_e4m3fn,
+    column_major_scales: bool = False,
+    scale_tma_aligned: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Per-token-group FP8 quantization with automatic backend selection.
+
+    Uses Triton kernel on GPU when available, otherwise falls back to PyTorch.
+
+    Args:
+        x: Input tensor
+        group_size: Size of each quantization group
+        eps: Small value to avoid division by zero
+        dtype: Target FP8 data type
+        column_major_scales: Whether to use column-major scale layout
+        scale_tma_aligned: Whether to use TMA-aligned scales
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor)
+    """
+    # fp8_per_token_group_quant_triton is automatically selected based on
+    # backend availability (Triton vs PyTorch) via __init__.py
+    return fp8_per_token_group_quant_triton(
+        x,
+        group_size,
+        eps,
+        dtype,
+        column_major_scales,
+        scale_tma_aligned,
+    )
+
+
+def fp8_per_channel_quant(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Per-channel FP8 weight quantization (E4M3 format)
+
+    Args:
+        weight: Original weight tensor with shape [out_features, in_features]
+
+    Returns:
+        weight_quant: Quantized weight [out_features, in_features], dtype=float8_e4m3fn
+        weight_scale: Scale factors [out_features, 1], dtype=float32
+    """
+    abs_max = torch.abs(weight).amax(dim=1, keepdim=True)  # [out_features, 1]
+
+    weight_scale = abs_max / FP8_MAX
+    weight_scale = torch.clamp(weight_scale, min=1e-12)
+
+    weight_scaled = (weight / weight_scale).clamp(min=FP8_MIN, max=FP8_MAX)
+    weight_quant = weight_scaled.to(torch.float8_e4m3fn)
+
+    return weight_quant, weight_scale.float()
+
+
+def int8_per_channel_quant(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Per-channel symmetric INT8 quantization for linear weights.
+
+    Args:
+        weight: Weight tensor with shape [out_features, in_features]
+
+    Returns:
+        weight_quant: Quantized INT8 weight [out_features, in_features]
+        weight_scale: Per-channel scales [out_features, 1], float32
+    """
+    abs_max = torch.abs(weight).amax(dim=1, keepdim=True).clamp(min=1e-5)
+    qmin, qmax = -128, 127
+    weight_scale = (abs_max / qmax).to(torch.float32)
+    weight_quant = torch.clamp(torch.round(weight / weight_scale), qmin, qmax).to(torch.int8)
+    return weight_quant, weight_scale
+
+
+def nvfp4_per_tensor_quant(
+    x: torch.Tensor, input_global_scale: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Per-tensor NVFP4 quantization using LightX2V kernel.
+
+    Returns packed uint8 values and per-block quantized scales.
+    """
+    if x.dim() != 2:
+        raise ValueError(f"nvfp4_per_tensor_quant expects 2D tensor, but got shape={tuple(x.shape)}")
+
+    if input_global_scale is None:
+        x_absmax = torch.abs(x).amax().float().clamp(min=1e-12)
+        input_global_scale = (2688.0 / x_absmax).to(torch.float32)
+    else:
+        input_global_scale = input_global_scale.to(device=x.device, dtype=torch.float32)
+
+    qx, qscale = scaled_nvfp4_quant(x, input_global_scale)
+    return qx, qscale, input_global_scale
+
+
+def mxfp4_per_tensor_quant(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Per-tensor MXFP4 quantization using LightX2V kernel.
+
+    Returns packed uint8 values and per-block quantized scales.
+    """
+    if x.dim() != 2:
+        raise ValueError(f"mxfp4_per_tensor_quant expects 2D tensor, but got shape={tuple(x.shape)}")
+
+    qx, qscale = scaled_mxfp4_quant(x)
+    # Keep a compatible triplet return signature with NVFP4 paths.
+    input_global_scale = torch.tensor(1.0, device=x.device, dtype=torch.float32)
+    return qx, qscale, input_global_scale
+
+
+def mxfp6_per_tensor_quant(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Per-tensor MXFP6 quantization using LightX2V kernel.
+
+    Returns packed uint8 values and per-block quantized scales.
+    """
+    if x.dim() != 2:
+        raise ValueError(f"mxfp6_per_tensor_quant expects 2D tensor, but got shape={tuple(x.shape)}")
+
+    qx, qscale = scaled_mxfp6_quant(x)
+    input_global_scale = torch.tensor(1.0, device=x.device, dtype=torch.float32)
+    return qx, qscale, input_global_scale
+
+
+def mxfp8_per_tensor_quant(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Per-tensor MXFP8 quantization using LightX2V kernel.
+
+    Returns packed uint8 values and per-block quantized scales.
+    """
+    if x.dim() != 2:
+        raise ValueError(f"mxfp8_per_tensor_quant expects 2D tensor, but got shape={tuple(x.shape)}")
+
+    qx, qscale = scaled_mxfp8_quant(x)
+    input_global_scale = torch.tensor(1.0, device=x.device, dtype=torch.float32)
+    return qx, qscale, input_global_scale
+
+
+def fp8_per_token_quant_sgl(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    m, k = x.shape
+    input_tensor_quant = torch.empty(
+        (m, k), dtype=torch.float8_e4m3fn, device="cuda", requires_grad=False
+    )
+    input_tensor_scale = torch.empty(
+        (m, 1), dtype=torch.float32, device="cuda", requires_grad=False
+    )
+    _sgl_kernel = _ensure_sgl_kernel()
+    _sgl_kernel.sgl_per_token_quant_fp8(x, input_tensor_quant, input_tensor_scale)
+    return input_tensor_quant, input_tensor_scale
+
+
+# pure torch implementation of block-wise FP8 quantization on cpu
+def fp8_per_block_quant_torch(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pure Torch implementation of block-wise FP8 (e4m3fn) quantization.
+
+    Args:
+        x (torch.Tensor): 2D tensor (M, N) in float types.
+        block_size (int): Block size for both M and N dimensions.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]:
+            - y: quantized tensor with dtype=torch.float8_e4m3fn and same
+              shape/device as x
+            - s: per-block scale tensor with shape (ceil(M/bs), ceil(N/bs)) on
+              x.device
+    """
+    assert x.is_contiguous()
+    assert x.dim() == 2
+
+    M, N = x.size()
+    device = x.device
+
+    # Output tensors
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    m_blocks = (M + block_size - 1) // block_size
+    n_blocks = (N + block_size - 1) // block_size
+    s = torch.empty((m_blocks, n_blocks), dtype=torch.float32, device=device)
+
+    # Iterate over blocks
+    for mb in range(m_blocks):
+        m_start = mb * block_size
+        m_end = min(m_start + block_size, M)
+        for nb in range(n_blocks):
+            n_start = nb * block_size
+            n_end = min(n_start + block_size, N)
+
+            x_block = x[m_start:m_end, n_start:n_end].to(torch.float32)
+            # Compute per-block scale: max(abs(x))/448.0, guard zero to 1.0
+            max_val = x_block.abs().amax().to(torch.float32)
+            scale = max_val / FP8_MAX
+            if scale.item() == 0.0:
+                scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+
+            # Quantize block
+            y_block = (x_block / scale).to(torch.float8_e4m3fn)
+
+            # Store results
+            y[m_start:m_end, n_start:n_end] = y_block
+            s[mb, nb] = scale
+
+    return y, s
+
+
+# quant function for per-block fp8
+def fp8_per_block_quant(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Per-block FP8 quantization with automatic GPU/CPU dispatch.
+
+    Quantizes a FP32 2D tensor to FP8 (E4M3FN) using block-wise quantization.
+    For each (block_size x block_size) block:
+        - scale = max(abs(block)) / 448.0 (FP8 E4M3FN max magnitude)
+        - if block is all zeros, use scale = 1.0 to avoid div-by-zero
+        - scale, clamp and cast to FP8
+
+    Args:
+        x: Input tensor (2D)
+        block_size: Block size for quantization
+
+    Returns:
+        Tuple of (quantized_tensor, scale_tensor):
+            - y: Quantized FP8 tensor, same shape as input
+            - s: Per-block scales, shape (num_blocks_M, num_blocks_N)
+    """
+    if x.is_cuda:
+        return fp8_per_block_quant_triton(x, block_size)
+    else:
+        return fp8_per_block_quant_torch(x, block_size)
+
+
+# gemm function for per-block quantization fp8 using deepgemm
+def fp8_gemm_deepgemm_block(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    bias: Optional[torch.Tensor] = None,
+    origin_shape: Optional[Tuple[int, ...]] = None,
+) -> torch.Tensor:
+    a_fp8 = (a, a_s)
+    b_fp8 = (b, b_s)
+    out = torch.empty((a.shape[0], b.shape[0]), device=a.device, dtype=torch.bfloat16)
+
+    _deep_gemm = _ensure_deep_gemm()
+    _deep_gemm.fp8_gemm_nt(a_fp8, b_fp8, out)
+
+    if origin_shape is not None:
+        out = out.reshape([*origin_shape[:-1], b.shape[0]])
+    if bias is not None:
+        out += bias
+
+    return out.to(out_dtype)
+
+
+# gemm function for per-token and per-group fp8 quantization using torch native api
+def fp8_gemm_torch_tensor_token(
+    a: torch.Tensor,
+    a_s: torch.Tensor,
+    b: torch.Tensor,
+    b_s: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    need_reshape = a.dim() == 3
+    if need_reshape:
+        batch_size = a.shape[0]
+        A_input = a.reshape(-1, a.shape[-1])
+    else:
+        batch_size = None
+        A_input = a
+
+    output = torch._scaled_mm(
+        A_input,
+        b.t(),
+        out_dtype=out_dtype,
+        scale_a=a_s,
+        scale_b=b_s,
+        bias=bias,
+    )
+    # If output is a tuple, take the first element
+    if isinstance(output, tuple):
+        output = output[0]
+
+    if need_reshape:
+        output = output.reshape(batch_size, output.shape[0] // batch_size, output.shape[1])
+
+    return output
+
+
+def fp8_weight_only_gemm(A, B, B_scale, bias, out_dtype):
+    """Perform FP8 GEMM operation with fallback to standard linear.
+
+    Args:
+        A: Input tensor A.
+        B: Input tensor B.
+        B_scale: Scale factor for tensor B.
+        bias: Optional bias tensor.
+        out_dtype: Output data type.
+        native_fp8_support: Whether to use native FP8 support.
+        quant_type: Quantization type.
+        origin_shape: Original shape for reshaping.
+
+    Returns:
+        torch.Tensor: Result of the GEMM operation.
+    """
+    if A.numel() == 0:
+        return torch.empty(size=(0, B.shape[0]), dtype=out_dtype, device=A.device)
+
+    output = torch.nn.functional.linear(
+        A.to(out_dtype),
+        B.to(out_dtype) * B_scale.to(out_dtype),
+        bias=bias,
+    )
+
+    return output
+
+
+def fp8_gemm_sgl_token(A, A_scale, B, B_scale, out_dtype, bias):
+    """GEMM function for FP8 per-token-sgl quantization using sgl-kernel.
+
+    Args:
+        A: Input activation tensor
+        A_scale: Scale tensor for input activations
+        B: Weight tensor
+        B_scale: Scale tensor for weights
+        out_dtype: Output data type.
+        bias: Optional bias tensor
+
+    Returns:
+        torch.Tensor: Result of the GEMM operation.
+    """
+    _sgl_kernel = _ensure_sgl_kernel()
+    shape = (A.shape[0], B.shape[0])
+    output = torch.empty(shape, dtype=out_dtype, device=A.device, requires_grad=False)
+    output = _sgl_kernel.fp8_scaled_mm(
+        A,
+        B.t(),
+        A_scale,
+        B_scale.float(),
+        out_dtype,
+        bias=bias,
+    )
+
+    return output
+
+
+def fp8_gemm(
+    A: torch.Tensor,
+    A_scale: torch.Tensor,
+    B: torch.Tensor,
+    B_scale: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    out_dtype: torch.dtype,
+    native_fp8_support: bool = False,
+    quant_type: str = QuantType.FP8_PER_TENSOR,
+    origin_shape: Optional[Tuple[int, ...]] = None,
+) -> torch.Tensor:
+    """
+    Unified GEMM function for FP8 quantization.
+
+    Args:
+        A: Input activation tensor
+        A_scale: Scale tensor for input activations
+        B: Weight tensor
+        B_scale: Scale tensor for weights
+        bias: Optional bias tensor
+        out_dtype: Output data type
+        native_fp8_support: Whether to use native FP8 kernels
+        quant_type: Quantization type (QuantType enum values)
+        origin_shape: Original shape for reshaping output
+
+    Returns:
+        Output tensor after GEMM operation
+
+    Raises:
+        ValueError: If unsupported combination of parameters is provided
+    """
+    # Return empty tensor if the input is empty
+    if A.numel() == 0:
+        return torch.empty(size=(0, B.shape[0]), dtype=out_dtype, device=A.device)
+
+    if quant_type == QuantType.FP8_PER_CHANNEL_VLLM:
+        if not hasattr(torch.ops, "_C") or not hasattr(torch.ops._C, "cutlass_scaled_mm"):
+            raise RuntimeError(
+                "quant_type='fp8-per-channel-vllm' requires torch.ops._C.cutlass_scaled_mm"
+            )
+        output = torch.empty(
+            (A.shape[0], B.shape[0]),
+            dtype=out_dtype,
+            device=A.device,
+            requires_grad=False,
+        )
+        weight_scale = B_scale
+        if (
+            weight_scale.dim() == 2
+            and weight_scale.shape[0] == B.shape[0]
+            and weight_scale.shape[1] == 1
+        ):
+            weight_scale = weight_scale.t()
+        torch.ops._C.cutlass_scaled_mm(
+            output,
+            A,
+            B.t(),
+            A_scale,
+            weight_scale,
+            bias,
+        )
+        return output
+
+    # Prioritize native fp8 support if available
+    if native_fp8_support:
+        if quant_type in (QuantType.FP8_PER_TENSOR, QuantType.FP8_PER_TOKEN):
+            # Use torch native fp8 GEMM for per-tensor and per-token fp8 quantization
+            return fp8_gemm_torch_tensor_token(A, A_scale, B, B_scale, out_dtype, bias)
+        elif quant_type == QuantType.FP8_PER_TOKEN_SGL:
+            # Use sgl-kernel for per-token-sgl fp8 quantization
+            return fp8_gemm_sgl_token(A, A_scale, B, B_scale, out_dtype, bias)
+        elif quant_type == QuantType.FP8_PER_BLOCK:
+            # Use deepgemm accelerated blockwise fp8 GEMM
+            return fp8_gemm_deepgemm_block(A, A_scale, B, B_scale, out_dtype, bias, origin_shape)
+    else:
+        if quant_type == QuantType.FP8_PER_BLOCK:
+            # Use triton kernel for blockwise fp8 quantization
+            return fp8_gemm_triton_block(A, A_scale, B, B_scale, out_dtype, bias)
+        elif quant_type == QuantType.FP8_PER_TENSOR:
+            # Fall back to scaled linear for per-tensor fp8 without native fp8 kernel
+            return torch.nn.functional.linear(
+                A.to(out_dtype) * A_scale,
+                B.to(out_dtype) * B_scale.to(out_dtype),
+                bias=bias,
+            )
+
+    # Raise error for unsupported combination of quant_type and native_fp8_support
+    raise ValueError(
+        f"Unsupported combination: "
+        f"\n  quant_type={quant_type},"
+        f"\n  native_fp8_support={native_fp8_support}.\n"
+        "Supported combinations:\n"
+        "  - native_fp8_support=True, "
+        "quant_type in [fp8-per-tensor, fp8-per-token,"
+        " fp8-per-block, fp8-per-token-sgl]\n"
+        "  - native_fp8_support=False, "
+        "quant_type in [fp8-per-tensor, fp8-per-block]"
+    )

quantizer/quant/utils/__init__.py

@@ -0,0 +1,37 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .quant_io import load_fp8_scales, load_quantized_model, save_quantized_model
+from .utils import (
+    QuantType,
+    _compile_pattern,
+    _ensure_deep_gemm,
+    _ensure_sgl_kernel,
+    cleanup_memory,
+    replace_module,
+    should_quantize_layer,
+)
+
+__all__ = [
+    "load_fp8_scales",
+    "load_quantized_model",
+    "save_quantized_model",
+    "QuantType",
+    "_compile_pattern",
+    "_ensure_deep_gemm",
+    "cleanup_memory",
+    "replace_module",
+    "should_quantize_layer",
+    "_ensure_sgl_kernel",
+]

quantizer/quant/utils/quant_io.py

@@ -0,0 +1,149 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Dict, Optional, Union
+
+import torch
+
+__all__ = ["load_fp8_scales", "load_quantized_model", "save_quantized_model"]
+
+
+def load_fp8_scales(
+    quant_scales: Optional[Union[str, Dict[str, torch.Tensor]]]
+) -> Dict[str, torch.Tensor]:
+    """Load FP8 quant scales from dict, file, or dir. Prefer .safetensors."""
+    if quant_scales is None:
+        raise ValueError("quant_scales is required")
+
+    if isinstance(quant_scales, dict) and len(quant_scales) > 0:
+        # Use provided dict
+        return quant_scales
+
+    if isinstance(quant_scales, str):
+        # Check if path is file
+        if os.path.isfile(quant_scales):
+            if quant_scales.endswith(".safetensors"):
+                import safetensors.torch
+
+                print(f"Loaded scale map from {quant_scales}")
+                return safetensors.torch.load_file(quant_scales)
+            else:
+                print(f"Loaded scale map from {quant_scales}")
+                return torch.load(quant_scales)
+        # Check if path is directory
+        if os.path.isdir(quant_scales):
+            safetensors_path = os.path.join(quant_scales, "fp8_scales.safetensors")
+            pth_path = os.path.join(quant_scales, "fp8_scales.pth")
+            if os.path.isfile(safetensors_path):
+                import safetensors.torch
+
+                print(f"Loaded scale map from {safetensors_path}")
+                return safetensors.torch.load_file(safetensors_path)
+            if os.path.isfile(pth_path):
+                print(f"Loaded scale map from {pth_path}")
+                return torch.load(pth_path)
+            raise FileNotFoundError(
+                f"Quant scale file not found: {pth_path} or {safetensors_path}"
+            )
+        raise FileNotFoundError(f"quant_scales path does not exist: {quant_scales}")
+
+    raise ValueError(f"Invalid quant_scales type: {type(quant_scales)}. Only str (path) or dict.")
+
+
+def save_quantized_model(model: torch.nn.Module, save_path: str, fp8_scales_map: Dict):
+    """
+    Save quantized model and scale dict to directory.
+    """
+    import logging
+
+    logger = logging.getLogger(__name__)
+
+    if not os.path.exists(save_path):
+        try:
+            os.makedirs(save_path, exist_ok=True)
+        except Exception as e:
+            raise RuntimeError(f"Cannot create directory for save_path: {save_path}. Error: {e}")
+
+    try:
+        # If Hugging Face style, use save_pretrained
+        if hasattr(model, "save_pretrained"):
+            model.save_pretrained(save_path)
+            logger.info(f"Saved quantized model to {save_path} via save_pretrained")
+        else:
+            # Otherwise, save state_dict with safetensors
+            from safetensors.torch import save_file as safe_save
+
+            model_path = os.path.join(save_path, "model.safetensors")
+            safe_save(model.state_dict(), model_path)
+            logger.info(f"Saved state_dict to {model_path}")
+
+        # Always save scales dict
+        from safetensors.torch import save_file as safe_save
+
+        scale_save_path = os.path.join(save_path, "fp8_scales.safetensors")
+        safe_save(fp8_scales_map, scale_save_path)
+        logger.info(f"Saved scales map to {scale_save_path}")
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to save model and scales map to {save_path}. Error: {e}")
+
+
+def load_quantized_model(model_class, save_path: str, device: str = "cpu"):
+    """
+    Load quantized model from directory.
+    """
+    import logging
+
+    logger = logging.getLogger(__name__)
+
+    try:
+        # Try Hugging Face style first
+        if hasattr(model_class, "from_pretrained"):
+            model = model_class.from_pretrained(save_path)
+            logger.info(f"Loaded Hugging Face model from {save_path}")
+            return model
+    except Exception as e:
+        logger.warning(f"Failed to load as Hugging Face model: {e}")
+
+    try:
+        # Try safetensors file first
+        model_path = os.path.join(save_path, "model.safetensors")
+        if os.path.exists(model_path):
+            from safetensors.torch import load_file as safe_load
+
+            state_dict = safe_load(model_path, device=device)
+            model = model_class()
+            model.load_state_dict(state_dict)
+            model.to(device)
+            logger.info(f"Loaded model from {model_path} (safetensors)")
+            return model
+
+        # Try pytorch .bin next
+        model_path = os.path.join(save_path, "pytorch_model.bin")
+        if os.path.exists(model_path):
+            state_dict = torch.load(model_path, map_location=device)
+            model = model_class()
+            model.load_state_dict(state_dict)
+            model.to(device)
+            logger.info(f"Loaded model from {model_path} (pytorch)")
+            return model
+        else:
+            raise FileNotFoundError(
+                f"Model file not found at {save_path}. "
+                "Expected 'model.safetensors' or 'pytorch_model.bin'"
+            )
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model from {save_path}. Error: {e}")

quantizer/quant/utils/utils.py

@@ -0,0 +1,219 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import gc
+import re
+from typing import List, Optional, Union
+
+import torch
+
+__all__ = [
+    "replace_module",
+    "cleanup_memory",
+    "should_quantize_layer",
+    "_compile_pattern",
+    "_ensure_deep_gemm",
+    "_ensure_sgl_kernel",
+    "QuantType",
+]
+
+
+class QuantType:
+    FP8_PER_TENSOR = "fp8-per-tensor"
+    FP8_PER_TOKEN = "fp8-per-token"
+    FP8_PER_BLOCK = "fp8-per-block"
+    FP8_PER_CHANNEL_VLLM = "fp8-per-channel-vllm"
+    FP8_PER_TENSOR_WEIGHT_ONLY = "fp8-per-tensor-weight-only"
+    FP8_PER_TOKEN_SGL = "fp8-per-token-sgl"
+    INT8 = "int8"
+    INT8_TORCHAO = "int8-torchao"
+    INT8_VLLM = "int8-vllm"
+    INT8_TRITON = "int8-triton"
+    INT8_SGL = "int8-sgl"
+    INT8_Q8F = "int8-q8f"
+    NVFP4 = "nvfp4"
+    MXFP4 = "mxfp4"
+    MXFP6 = "mxfp6"
+    MXFP8 = "mxfp8"
+    VALID_TYPES = [
+        FP8_PER_TENSOR,
+        FP8_PER_TOKEN,
+        FP8_PER_BLOCK,
+        FP8_PER_CHANNEL_VLLM,
+        FP8_PER_TENSOR_WEIGHT_ONLY,
+        FP8_PER_TOKEN_SGL,
+        INT8,
+        INT8_TORCHAO,
+        INT8_VLLM,
+        INT8_TRITON,
+        INT8_SGL,
+        INT8_Q8F,
+        NVFP4,
+        MXFP4,
+        MXFP6,
+        MXFP8,
+    ]
+
+    @classmethod
+    def validate(cls, quant_type: str):
+        if quant_type not in cls.VALID_TYPES:
+            raise ValueError(f"Invalid quant_type: {quant_type}. Valid types: {cls.VALID_TYPES}")
+
+
+def replace_module(model: torch.nn.Module, name: str, new_module: torch.nn.Module):
+    """
+    Replace a submodule in the model with a new module by name.
+    """
+    if "." in name:
+        parent_name = name.rsplit(".", 1)[0]
+        child_name = name[len(parent_name) + 1 :]
+        parent = model.get_submodule(parent_name)
+    else:
+        parent_name = ""
+        parent = model
+        child_name = name
+
+    setattr(parent, child_name, new_module)
+
+
+def cleanup_memory():
+    """
+    Run garbage collection and clear CUDA memory cache.
+    """
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def _compile_pattern(pattern: Union[str, re.Pattern], case_sensitive: bool = False) -> re.Pattern:
+    """
+    Compile a pattern (string or pre-compiled pattern) into a regex pattern object.
+
+    Args:
+        pattern: String pattern or already-compiled regex pattern.
+        case_sensitive: Whether the match is case sensitive.
+
+    Returns:
+        Compiled regex pattern object.
+    """
+    if isinstance(pattern, str):
+        # If the string contains special regex characters, treat as regex.
+        if any(char in pattern for char in ".*+?^${}[]|()\\"):
+            flags = 0 if case_sensitive else re.IGNORECASE
+            return re.compile(pattern, flags)
+        else:
+            # Escape regular string so it matches literally.
+            escaped = re.escape(pattern)
+            flags = 0 if case_sensitive else re.IGNORECASE
+            return re.compile(escaped, flags)
+    else:
+        # Already a compiled pattern.
+        return pattern
+
+
+def should_quantize_layer(
+    layer_name: str,
+    include_patterns: Optional[List[Union[str, re.Pattern]]] = None,
+    exclude_patterns: Optional[List[Union[str, re.Pattern]]] = None,
+    case_sensitive: bool = False,
+) -> bool:
+    """
+    Decide whether a layer should be quantized based on inclusion/exclusion patterns.
+
+    Args:
+        layer_name: Name of the layer.
+        include_patterns: List of patterns (str or regex) to include.
+        exclude_patterns: List of patterns (str or regex) to exclude.
+        case_sensitive: Whether patterns are matched case sensitively.
+
+    Returns:
+        bool: Whether this layer should be quantized.
+
+    Note:
+        String patterns are auto-detected as regex if they include special chars:
+        - If contains any . * + ? ^ $ { } [ ] | ( ) \\ -> treated as regex
+        - Otherwise, escaped and matched literally
+    """
+    if include_patterns is None:
+        include_patterns = []
+    if exclude_patterns is None:
+        exclude_patterns = []
+
+    # Check exclusion patterns
+    for pattern in exclude_patterns:
+        compiled_pattern = _compile_pattern(pattern, case_sensitive)
+        if compiled_pattern.search(layer_name):
+            return False
+
+    # If no include patterns, default is to include all layers
+    if not include_patterns:
+        return True
+
+    # Check inclusion patterns
+    for pattern in include_patterns:
+        compiled_pattern = _compile_pattern(pattern, case_sensitive)
+        if compiled_pattern.search(layer_name):
+            return True
+
+    return False
+
+
+_deep_gemm_cached = None
+
+
+def _ensure_deep_gemm():
+    """
+    Lazy, safe import of deep_gemm with process-level caching. Returns the module
+    if available, otherwise raises a clear error.
+    """
+    global _deep_gemm_cached
+    if _deep_gemm_cached is not None:
+        return _deep_gemm_cached
+    try:
+        import deep_gemm
+
+        _deep_gemm_cached = deep_gemm
+        return _deep_gemm_cached
+    except ImportError as e:
+        raise ImportError(
+            (
+                "deep_gemm is required for 'fp8-per-block' quantization with "
+                "native_fp8_support, but was not found. Please install deep_gemm first."
+            )
+        ) from e
+
+
+_sgl_kernel_cached = None
+
+
+def _ensure_sgl_kernel():
+    """
+    Lazy, safe import of sgl_kernel with process-level caching. Returns the module
+    if available, otherwise raises a clear error.
+    """
+    global _sgl_kernel_cached
+    if _sgl_kernel_cached is not None:
+        return _sgl_kernel_cached
+    try:
+        import sgl_kernel
+
+        _sgl_kernel_cached = sgl_kernel
+        return _sgl_kernel_cached
+    except ImportError as e:
+        raise ImportError(
+            (
+                "sgl_kernel is required for 'fp8-per-token-sgl' quantization with "
+                "native_fp8_support, but was not found. Please install sgl_kernel first"
+            )
+        ) from e

requirements.txt

@@ -0,0 +1,16 @@
+diffusers==0.31.0
+transformers==4.57.3
+tokenizers==0.22.1
+accelerate==1.12.0
+safetensors
+opencv-python-headless
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
+easydict==1.13
+ftfy==6.3.1
+numpy==2.2.6
+omegaconf==2.3.0
+einops==0.8.1
+av==13.1.0
+sentencepiece==0.2.1
+torchvision

utils/cuda_kernel_buckets.py

@@ -0,0 +1,227 @@
+"""
+将 PyTorch CUDA profiler 事件按内核/算子名称粗分为若干类，便于估算
+memcpy、GEMM/Linear、Attention、Norm 等在 GPU 时间中的占比。
+
+说明：
+- 分类基于名称子串启发式，不同 CUDA / Triton / cuBLAS 版本下内核名会有差异；
+- 若需更细粒度，请配合 prof.export_chrome_trace() 在 chrome://tracing 或 Perfetto 中查看。
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import DefaultDict, Dict, Iterable, List, Tuple
+
+# 顺序有意义：先匹配更具体的类别
+_BUCKET_ORDER: List[str] = [
+    "memcpy_sync",
+    "attention",
+    "gemm_linear",
+    "conv",
+    "norm",
+    "elementwise_reduce",
+    "other",
+]
+
+
+def classify_cuda_name(name: str) -> str:
+    """根据内核或 aten 算子名称归入粗分类。"""
+    n = (name or "").lower()
+
+    # 设备同步、拷贝、显存设置（含部分 launch 开销）
+    if any(
+        s in n
+        for s in (
+            "memcpy",
+            "memset",
+            "memgetinfo",
+            "cudaget",
+            "cudaevent",
+            "cudaoccupancy",
+            "cuda stream",
+            "cudastream",
+            "cudadevicesynchronize",
+            "cuda devicesynchronize",
+            "device synchronize",
+            "eventrecord",
+            "eventsynchronize",
+        )
+    ):
+        return "memcpy_sync"
+
+    # FlashAttention / SDPA / 各类 fused attention
+    if any(
+        s in n
+        for s in (
+            "flash_attn",
+            "flash-attn",
+            "scaled_dot_product",
+            "sdpa",
+            "efficient_attention",
+            "mem_eff_attention",
+            "fused_attention",
+            "fmha",
+            "mha_default",
+            "multi_head_attention",
+            "contrib_attn",
+            "attention_forward",
+            "attention_backward",
+            "softmax",
+        )
+    ):
+        # 独立 softmax 小核也可能与 attention 同桶；若需拆开可把 softmax 挪到 elementwise
+        return "attention"
+
+    # MatMul / Linear / GEMM（含 cutlass、cublas、triton、fp8）
+    if any(
+        s in n
+        for s in (
+            "gemm",
+            "cublas",
+            "cutlass",
+            "matmul",
+            "mat_mul",
+            "aten::linear",
+            "aten::mm",
+            "aten::bmm",
+            "aten::addmm",
+            "aten::matmul",
+            "scaled_mm",
+            "fp8",
+            "wmma",
+            "mma_sync",
+            "triton",
+            "dot",
+        )
+    ):
+        return "gemm_linear"
+
+    if any(s in n for s in ("conv", "cudnn", "depthwise", "convolution")):
+        return "conv"
+
+    if any(
+        s in n
+        for s in (
+            "layernorm",
+            "layer_norm",
+            "rms_norm",
+            "group_norm",
+            "flash_norm",
+            "aten::layer_norm",
+            "aten::group_norm",
+            "aten::native_layer_norm",
+        )
+    ):
+        return "norm"
+
+    if any(
+        s in n
+        for s in (
+            "elementwise",
+            "vectorized",
+            "unary",
+            "binary",
+            "reduce",
+            "reduction",
+            "activation",
+            "silu",
+            "gelu",
+            "swiglu",
+            "relu",
+            "aten::add",
+            "aten::mul",
+            "aten::div",
+            "aten::pow",
+            "aten::sqrt",
+            "aten::rsqrt",
+        )
+    ):
+        return "elementwise_reduce"
+
+    return "other"
+
+
+def _event_cuda_time_us(event) -> float:
+    for attr in (
+        "cuda_time_total",
+        "self_cuda_time_total",
+        "self_cuda_time_total_us",
+    ):
+        v = getattr(event, attr, None)
+        if v is not None:
+            return float(v)
+    return 0.0
+
+
+def aggregate_from_profiler(prof) -> Tuple[Dict[str, float], List[Tuple[str, float]]]:
+    """
+    从 torch.profiler.profile 实例聚合：
+    - 返回 (bucket -> 微秒总和, 按耗时排序的 (name, us) 列表)
+    使用 prof.events() 以尽量接近 CUDA 内核级名称。
+    """
+    bucket_us: DefaultDict[str, float] = defaultdict(float)
+    per_name: DefaultDict[str, float] = defaultdict(float)
+
+    try:
+        events = prof.events()
+    except Exception:
+        events = []
+
+    for e in events:
+        us = _event_cuda_time_us(e)
+        if us <= 0:
+            continue
+        name = getattr(e, "name", "") or ""
+        b = classify_cuda_name(name)
+        bucket_us[b] += us
+        per_name[name] += us
+
+    # 若 events() 为空，回退到 key_averages（多为 aten 级）
+    if not bucket_us and not per_name:
+        try:
+            for avg in prof.key_averages():
+                us = float(getattr(avg, "cuda_time_total", 0) or getattr(avg, "self_cuda_time_total", 0) or 0)
+                if us <= 0:
+                    continue
+                name = getattr(avg, "key", "") or ""
+                b = classify_cuda_name(name)
+                bucket_us[b] += us
+                per_name[name] += us
+        except Exception:
+            pass
+
+    sorted_names = sorted(per_name.items(), key=lambda x: -x[1])
+    out = {k: bucket_us.get(k, 0.0) for k in _BUCKET_ORDER}
+    for k, v in bucket_us.items():
+        if k not in out:
+            out[k] = v
+    return out, sorted_names
+
+
+def format_bucket_report(
+    bucket_us: Dict[str, float],
+    top_names: Iterable[Tuple[str, float]],
+    top_k: int = 25,
+) -> str:
+    total = sum(bucket_us.values()) or 1e-9
+    lines: List[str] = []
+    lines.append("=== CUDA 时间按粗分类（微秒 / 占比）===")
+    for k in _BUCKET_ORDER:
+        if k in bucket_us and bucket_us[k] > 0:
+            us = bucket_us[k]
+            lines.append(f"  {k:22s}  {us:12.1f} us  ({100.0 * us / total:5.1f}%)")
+    # 其它未在顺序表中的桶
+    for k, us in sorted(bucket_us.items(), key=lambda x: -x[1]):
+        if k in _BUCKET_ORDER:
+            continue
+        if us > 0:
+            lines.append(f"  {k:22s}  {us:12.1f} us  ({100.0 * us / total:5.1f}%)")
+    lines.append(f"  {'TOTAL':22s}  {total:12.1f} us")
+    lines.append("")
+    lines.append(f"=== 耗时最高的 {top_k} 个事件名（微秒）===")
+    for i, (name, us) in enumerate(top_names):
+        if i >= top_k:
+            break
+        short = name if len(name) <= 120 else name[:117] + "..."
+        lines.append(f"  {us:10.1f}  {short}")
+    return "\n".join(lines)

utils/memory.py

@@ -0,0 +1,135 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import torch
+
+
+cpu = torch.device('cpu')
+gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(module: torch.nn.Module, **kwargs):
+        original_class = module.__class__
+        module.__dict__['forge_backup_original_class'] = original_class
+
+        def hacked_get_attr(self, name: str):
+            if '_parameters' in self.__dict__:
+                _parameters = self.__dict__['_parameters']
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if '_buffers' in self.__dict__:
+                _buffers = self.__dict__['_buffers']
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
+            '__getattr__': hacked_get_attr,
+        })
+
+        return
+
+    @staticmethod
+    def _uninstall_module(module: torch.nn.Module):
+        if 'forge_backup_original_class' in module.__dict__:
+            module.__class__ = module.__dict__.pop('forge_backup_original_class')
+        return
+
+    @staticmethod
+    def install_model(model: torch.nn.Module, **kwargs):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(model: torch.nn.Module):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+    if hasattr(model, 'scale_shift_table'):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, 'weight'):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(device=None):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats['active_bytes.all.current']
+    bytes_reserved = memory_stats['reserved_bytes.all.current']
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024 ** 3)
+
+
+def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(*args):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f'Unloaded {m.__class__.__name__} as complete.')
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(model, target_device, unload=True):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
+
+    gpu_complete_modules.append(model)
+    return

utils/misc.py

@@ -0,0 +1,39 @@
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed: int, deterministic: bool = False):
+    """
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+    Args:
+        seed (`int`):
+            The seed to set.
+        deterministic (`bool`, *optional*, defaults to `False`):
+            Whether to use deterministic algorithms where available. Can slow down training.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.use_deterministic_algorithms(True)
+
+
+def merge_dict_list(dict_list):
+    if len(dict_list) == 1:
+        return dict_list[0]
+
+    merged_dict = {}
+    for k, v in dict_list[0].items():
+        if isinstance(v, torch.Tensor):
+            if v.ndim == 0:
+                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)
+            else:
+                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)
+        else:
+            # for non-tensor values, we just copy the value from the first item
+            merged_dict[k] = v
+    return merged_dict

utils/scheduler.py

@@ -0,0 +1,194 @@
+from abc import abstractmethod, ABC
+import torch
+
+
+class SchedulerInterface(ABC):
+    """
+    Base class for diffusion noise schedule.
+    """
+    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule
+
+    @abstractmethod
+    def add_noise(
+        self, clean_latent: torch.Tensor,
+        noise: torch.Tensor, timestep: torch.Tensor
+    ):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        pass
+
+    def convert_x0_to_noise(
+        self, x0: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's x0 prediction to noise predidction.
+        x0: the predicted clean data with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = x0.dtype
+        x0, xt, alphas_cumprod = map(
+            lambda x: x.double().to(x0.device), [x0, xt,
+                                                 self.alphas_cumprod]
+        )
+
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        noise_pred = (xt - alpha_prod_t **
+                      (0.5) * x0) / beta_prod_t ** (0.5)
+        return noise_pred.to(original_dtype)
+
+    def convert_noise_to_x0(
+        self, noise: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's noise prediction to x0 predidction.
+        noise: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = noise.dtype
+        noise, xt, alphas_cumprod = map(
+            lambda x: x.double().to(noise.device), [noise, xt,
+                                                    self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (xt - beta_prod_t **
+                   (0.5) * noise) / alpha_prod_t ** (0.5)
+        return x0_pred.to(original_dtype)
+
+    def convert_velocity_to_x0(
+        self, velocity: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's velocity prediction to x0 predidction.
+        velocity: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        v = sqrt(alpha_t) * noise - sqrt(beta_t) x0
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)
+        given v, x_t, we have
+        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v
+        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56
+        """
+        # use higher precision for calculations
+        original_dtype = velocity.dtype
+        velocity, xt, alphas_cumprod = map(
+            lambda x: x.double().to(velocity.device), [velocity, xt,
+                                                       self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity
+        return x0_pred.to(original_dtype)
+
+
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
+                          num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (
+                self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        """
+        Input:
+            - timestep: the timestep with shape [B*T]
+        Output: the corresponding weighting [B*T]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights

utils/wan_wrapper.py

@@ -0,0 +1,871 @@
+import types
+from pathlib import Path
+from typing import List, Optional
+import torch
+from torch import nn
+from safetensors.torch import load_file as load_safetensors_file
+
+from utils.scheduler import SchedulerInterface, FlowMatchScheduler
+import os
+
+
+def _wan_models_path(*parts) -> str:
+    """Resolve wan_models path relative to project root (works with symlink and any cwd)."""
+    root = Path(__file__).resolve().parent.parent
+    return str((root / "wan_models").joinpath(Path(*parts)))
+
+
+def _resolve_wan_path(path: str) -> str:
+    """If path starts with wan_models/, resolve to absolute path (project root); else return as-is."""
+    if path and path.startswith("wan_models/"):
+        return _wan_models_path(path[len("wan_models/"):])
+    return path
+
+
+def _resolve_wan_path_with_dir(path: str, wan_models_dir: Optional[str] = None) -> str:
+    """Resolve path: if wan_models_dir is set and path starts with wan_models/, use wan_models_dir as base; else _resolve_wan_path."""
+    if not path:
+        return path
+    if wan_models_dir and path.startswith("wan_models/"):
+        return os.path.join(wan_models_dir, path[len("wan_models/"):])
+    return _resolve_wan_path(path)
+
+
+def model_kwargs_with_relative_rope(args, default: bool = False) -> dict:
+    """Merge top-level use_relative_rope into model_kwargs with a stable default."""
+    raw_model_kwargs = getattr(args, "model_kwargs", {}) or {}
+    model_kwargs = dict(raw_model_kwargs)
+    if "use_relative_rope" not in model_kwargs:
+        try:
+            model_kwargs["use_relative_rope"] = bool(getattr(args, "use_relative_rope"))
+        except Exception:
+            model_kwargs["use_relative_rope"] = bool(default)
+    return model_kwargs
+
+from wan.modules.tokenizers import HuggingfaceTokenizer
+from wan.modules.model import WanModel
+from wan.modules.t5 import umt5_xxl
+from wan.modules.causal_model import CausalWanModel
+
+class WanTextEncoder(torch.nn.Module):
+    def __init__(
+        self,
+        tokenizer_path="wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/",
+        encoder_pth_path="wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
+    ) -> None:
+        super().__init__()
+        # tokenizer_path = _resolve_wan_path_with_dir(tokenizer_path, wan_models_dir)
+        # encoder_pth_path = _resolve_wan_path_with_dir(encoder_pth_path, wan_models_dir)
+
+        self.text_encoder = umt5_xxl(
+            encoder_only=True,
+            return_tokenizer=False,
+            dtype=torch.bfloat16,
+            device=torch.device('cpu')
+        ).eval().requires_grad_(False)
+        state_dict = torch.load(encoder_pth_path,
+                                map_location='cpu', weights_only=False)
+        self.text_encoder.load_state_dict(state_dict)
+        del state_dict
+
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path, seq_len=512, clean='whitespace')
+
+    @property
+    def device(self):
+        return next(self.text_encoder.parameters()).device
+
+    def forward(self, text_prompts: List[str], device: torch.device = None) -> dict:
+        ids, mask = self.tokenizer(
+            text_prompts, return_mask=True, add_special_tokens=True)
+        # When DynamicSwapInstaller is active, self.device returns cpu because
+        # parameters are swapped to GPU only during forward.  Use the explicitly
+        # passed device (the intended execution device) when available.
+        target_device = device if device is not None else self.device
+        ids = ids.to(target_device)
+        mask = mask.to(target_device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.text_encoder(ids, mask)
+
+        for u, v in zip(context, seq_lens):
+            u[v:] = 0.0  # set padding to 0.0
+
+        return {
+            "prompt_embeds": context
+        }
+
+
+class WanVAEWrapper(torch.nn.Module):
+    def __init__(
+        self,
+        pretrained_path=None,
+        z_dim=48,
+        vae_type="Wan2.2_VAE",
+        wan_models_dir=None,
+    ):
+        super().__init__()
+        if vae_type != "Wan2.2_VAE":
+            raise ValueError(f"Unsupported vae_type={vae_type!r}; only 'Wan2.2_VAE' is supported.")
+        from wan.modules.vae2_2 import _video_vae
+        self.mean = torch.tensor([
+                -0.2289, -0.0052, -0.1323, -0.2339, -0.2799,  0.0174,  0.1838,  0.1557,
+                -0.1382,  0.0542,  0.2813,  0.0891,  0.1570, -0.0098,  0.0375, -0.1825,
+                -0.2246, -0.1207, -0.0698,  0.5109,  0.2665, -0.2108, -0.2158,  0.2502,
+                -0.2055, -0.0322,  0.1109,  0.1567, -0.0729,  0.0899, -0.2799, -0.1230,
+                -0.0313, -0.1649,  0.0117,  0.0723, -0.2839, -0.2083, -0.0520,  0.3748,
+                0.0152,  0.1957,  0.1433, -0.2944,  0.3573, -0.0548, -0.1681, -0.0667,
+        ], dtype=torch.float32)
+
+        self.std = torch.tensor([
+                0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,
+                0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,
+                0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,
+                0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,
+                0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,
+                0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744,
+        ], dtype=torch.float32)
+        self.scale = [self.mean, 1.0 / self.std]
+        self.upsampling_factor = 16
+
+        z_dim = 48
+        self.z_dim = z_dim
+        self.model = _video_vae(pretrained_path=pretrained_path,
+            z_dim=z_dim,).eval().requires_grad_(False)
+
+    def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None):
+        # Initialize Gaussian noise
+        generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)
+        noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)
+        noise = noise.to(dtype=torch_dtype, device=device)
+        return noise
+    
+    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
+        # pixel: [batch_size, num_channels, num_frames, height, width]
+        device, dtype = pixel.device, pixel.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        output = [
+            self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+            for u in pixel
+        ]
+        output = torch.stack(output, dim=0)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False, return_in_cpu: bool = False) -> torch.Tensor:
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        if use_cache:
+            decode_function = self.model.cached_decode
+        else:
+            decode_function = self.model.decode
+
+        output = []
+        for u in zs:
+            decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+            if return_in_cpu:
+                decoded = decoded.cpu()
+            output.append(decoded)
+        output = torch.stack(output, dim=0)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+
+class TAEW2_2VAEWrapper(torch.nn.Module):
+    """
+    VAE wrapper using TAEHV (TAEW2.2) for faster decoding.
+    Requires: pip install taehv (or install from https://github.com/madebyollin/taehv)
+    Checkpoint: taew2_2.pth (download from taehv releases)
+    """
+    def __init__(self, checkpoint_path: str = "taew2_2.pth", dtype=torch.float16):
+        super().__init__()
+        try:
+            from wan.modules.taehv import TAEHV, StreamingTAEHV
+        except ImportError as e:
+            raise ImportError(
+                "taehv is required for TAEW2.2 VAE. Install with: pip install taehv"
+            ) from e
+        self.taehv = TAEHV(checkpoint_path).to(dtype).eval().requires_grad_(False)
+        self.taehv = StreamingTAEHV(self.taehv)
+        self.dtype = dtype
+        # For compatibility with pipeline.vae.model.clear_cache()
+        self.model = _TAEW2_2ModelRef(self)
+
+    def warmup_first_frame(self, first_frame_latent: torch.Tensor):
+        """Warm up the streaming decoder's MemBlock memory with the first-frame latent.
+
+        The TAeW2.2 MemBlocks use zero-initialized past context for the first frame,
+        causing blur.  By feeding the first-frame latent as a warmup pass (output
+        discarded), subsequent decodes benefit from real temporal context.
+
+        Args:
+            first_frame_latent: [B, 1, C, H, W] latent of the first frame.
+        """
+        if first_frame_latent is None:
+            return
+        # Reset decoder state, then feed first frame as warmup
+        self.taehv.reset()
+        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self.dtype):
+            # Feed the first-frame latent to populate MemBlock memory;
+            # the output (startup frames) is discarded.
+            _ = self.taehv.decode(first_frame_latent)
+
+    def decode_to_pixel(
+        self,
+        latent: torch.Tensor,
+        use_cache: bool = False,
+        return_in_cpu: bool = False
+    ) -> torch.Tensor:
+        # latent: [B, F, C, H, W] = [B, T, C, H, W] (same as TAEHV's NTCHW)
+        # use_cache=True -> parallel=False for lower memory (streaming)
+        parallel = not use_cache
+        with torch.autocast(device_type="cuda", dtype=self.dtype):
+            # out = self.taehv.decode_video(
+            #     latent, parallel=parallel, show_progress_bar=False
+            # )
+            out = self.taehv.decode(latent)
+        # TAEHV returns [0, 1], convert to [-1, 1] to match WanVAEWrapper
+        out = out.mul(2).sub(1).clamp(-1, 1).float()
+        if return_in_cpu:
+            out = out.cpu()
+        return out
+
+
+class _TAEW2_2ModelRef:
+    """Dummy ref for clear_cache compatibility; delegates to StreamingTAEHV.reset()."""
+
+    def __init__(self, parent):
+        self._parent = parent
+
+    def clear_cache(self):
+        self._parent.taehv.reset()
+
+
+class MGLightVAEWrapper(torch.nn.Module):
+    """VAE wrapper using MG-LightVAE (pruned Wan2.2 VAE) for faster decoding.
+
+    Wraps the ``Wan2_2_VAE`` class which supports different pruning rates.
+    The encoder uses the full (unpruned) Wan2.2 VAE teacher, while the decoder
+    uses the pruned student model.
+
+    Args:
+        vae_pth: Path to the pruned LightVAE checkpoint (student decoder).
+        lightvae_pruning_rate: Pruning rate for the decoder (e.g. 0.5, 0.75).
+        lightvae_encoder_vae_pth: Path to the full Wan2.2 VAE checkpoint
+            (teacher encoder). Required for mg_lightvae.
+        dtype: Data type for the VAE model.
+        device: Device to load the VAE on.
+    """
+
+    def __init__(
+        self,
+        vae_pth: str,
+        lightvae_pruning_rate: float = 0.75,
+        lightvae_encoder_vae_pth: str | None = None,
+        dtype=torch.float,
+        device="cpu",
+    ):
+        super().__init__()
+        from wan.modules.vae2_2 import Wan2_2_VAE
+
+        self._vae = Wan2_2_VAE(
+            z_dim=48,
+            c_dim=160,
+            vae_pth=vae_pth,
+            dtype=dtype,
+            device=device,
+            vae_type="mg_lightvae",
+            lightvae_pruning_rate=lightvae_pruning_rate,
+            lightvae_encoder_vae_pth=lightvae_encoder_vae_pth,
+        )
+        # Register model (pruned decoder) and encoder_model (teacher encoder)
+        # as submodules so that .to(), .eval(), .requires_grad_() propagate.
+        self.model = self._vae.model
+        if self._vae.encoder_model is not None:
+            self.encoder_model = self._vae.encoder_model
+        else:
+            self.encoder_model = None
+
+        self.z_dim = 48
+        self.upsampling_factor = 16
+        self.mean = self._vae.scale[0]
+        self.std = 1.0 / self._vae.scale[1]
+
+        # Initialize streaming cache attributes (_feat_map, _conv_idx, etc.)
+        # so cached_decode() can be called before any explicit clear_cache().
+        self.model.clear_cache()
+        if self.encoder_model is not None:
+            self.encoder_model.clear_cache()
+
+    def generate_noise(self, shape, seed=None, rand_device="cpu",
+                        rand_torch_dtype=torch.float32, device=None, torch_dtype=None):
+        generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)
+        noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)
+        noise = noise.to(dtype=torch_dtype, device=device)
+        return noise
+
+    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
+        # pixel: [batch_size, num_channels, num_frames, height, width]
+        device, dtype = pixel.device, pixel.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        encode_model = self.encoder_model if self.encoder_model is not None else self.model
+        output = [
+            encode_model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+            for u in pixel
+        ]
+        output = torch.stack(output, dim=0)
+        # from [B, C, F, H, W] to [B, F, C, H, W]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False,
+                        return_in_cpu: bool = False) -> torch.Tensor:
+        # from [B, F, C, H, W] to [B, C, F, H, W]
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        if use_cache:
+            decode_function = self.model.cached_decode
+        else:
+            decode_function = self.model.decode
+
+        output = []
+        for u in zs:
+            decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+            if return_in_cpu:
+                decoded = decoded.cpu()
+            output.append(decoded)
+        output = torch.stack(output, dim=0)
+        # from [B, C, F, H, W] to [B, F, C, H, W]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+
+def create_vae_from_config(config) -> Optional[torch.nn.Module]:
+    """Create a VAE wrapper based on the unified ``vae_type`` config field.
+
+    Supported vae_type values:
+        - "wan2.2":         Standard Wan2.2 VAE (returns None, pipeline creates WanVAEWrapper)
+        - "taew2_2":        TAeW2.2/taehv fast streaming decoder
+        - "mg_lightvae":    MG-LightVAE pruned decoder (pruning rate 0.5)
+        - "mg_lightvae_v2": MG-LightVAE v2 pruned decoder (pruning rate 0.75)
+
+    If ``vae_type`` is not set, defaults to "taew2_2".
+
+    Returns:
+        A VAE wrapper instance, or None for wan2.2 (let pipeline create
+        WanVAEWrapper from vae_kwargs).
+    """
+    vae_type = getattr(config, "vae_type", None)
+
+    if vae_type is None:
+        vae_type = "taew2_2"
+    else:
+        vae_type = str(vae_type).strip().lower()
+
+    if vae_type == "wan2.2":
+        return None  # pipeline creates WanVAEWrapper(**vae_kwargs) internally
+
+    if vae_type == "taew2_2":
+        ckpt = os.environ.get("TAEW2_2_CHECKPOINT") or getattr(
+            config, "taew2_2_checkpoint", "taew2_2.pth"
+        )
+        return TAEW2_2VAEWrapper(checkpoint_path=ckpt).eval()
+
+    if vae_type in ("mg_lightvae", "mg_lightvae_v2"):
+        pruning_map = {"mg_lightvae": 0.5, "mg_lightvae_v2": 0.75}
+        # Explicit pruning rate overrides the default mapping
+        explicit_rate = getattr(config, "lightvae_pruning_rate", None)
+        if explicit_rate is not None:
+            pruning_rate = float(explicit_rate)
+        else:
+            pruning_rate = pruning_map[vae_type]
+
+        # Select checkpoint based on vae_type
+        ckpt_map = {
+            "mg_lightvae": "lightvae_checkpoint",
+            "mg_lightvae_v2": "lightvae_v2_checkpoint",
+        }
+        vae_ckpt = getattr(config, ckpt_map[vae_type], None)
+        if vae_ckpt is None:
+            raise ValueError(
+                f"vae_type={vae_type!r} requires '{ckpt_map[vae_type]}' config field "
+                f"(path to MG-LightVAE .pth file)."
+            )
+
+        # Encoder checkpoint: explicit config, or fall back to vae_kwargs.pretrained_path
+        encoder_ckpt = getattr(config, "lightvae_encoder_checkpoint", None)
+        if encoder_ckpt is None:
+            vae_kwargs = getattr(config, "vae_kwargs", {}) or {}
+            if isinstance(vae_kwargs, dict):
+                encoder_ckpt = vae_kwargs.get("pretrained_path")
+            else:
+                encoder_ckpt = getattr(vae_kwargs, "pretrained_path", None)
+        if encoder_ckpt is None:
+            raise ValueError(
+                f"vae_type={vae_type!r} requires 'lightvae_encoder_checkpoint' config field "
+                f"(path to full Wan2.2_VAE.pth for teacher encoder), "
+                f"or 'vae_kwargs.pretrained_path' must be set."
+            )
+
+        return MGLightVAEWrapper(
+            vae_pth=vae_ckpt,
+            lightvae_pruning_rate=pruning_rate,
+            lightvae_encoder_vae_pth=encoder_ckpt,
+        )
+
+    raise ValueError(
+        f"Unsupported vae_type={vae_type!r}. "
+        f"Choose from: wan2.2, taew2_2, mg_lightvae, mg_lightvae_v2."
+    )
+
+
+class WanDiffusionWrapper(torch.nn.Module):
+    @staticmethod
+    def _materialize_meta_tensors(module: torch.nn.Module, device: torch.device = torch.device("cpu")):
+        materialized_names = []
+
+        def _materialize_recursive(mod: torch.nn.Module, prefix: str = ""):
+            for name, param in list(mod.named_parameters(recurse=False)):
+                if getattr(param, "is_meta", False):
+                    new_param = torch.nn.Parameter(
+                        torch.empty(tuple(param.shape), dtype=param.dtype, device=device),
+                        requires_grad=param.requires_grad,
+                    )
+                    setattr(mod, name, new_param)
+                    materialized_names.append(prefix + name)
+
+            for name, buf in list(mod.named_buffers(recurse=False)):
+                if getattr(buf, "is_meta", False):
+                    setattr(mod, name, torch.empty(tuple(buf.shape), dtype=buf.dtype, device=device))
+                    materialized_names.append(prefix + name)
+
+            for child_name, child in mod.named_children():
+                _materialize_recursive(child, prefix + child_name + ".")
+
+        _materialize_recursive(module)
+        return materialized_names
+
+    @staticmethod
+    def _normalize_model_state_dict_keys(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        key_prefixes = (
+            "generator.model._fsdp_wrapped_module.",
+            "generator.model.",
+            "model._fsdp_wrapped_module.",
+            "model.",
+            "_fsdp_wrapped_module.",
+            "module.",
+        )
+        keys = list(state_dict.keys())
+        for prefix in key_prefixes:
+            if keys and all(k.startswith(prefix) for k in keys):
+                return {k[len(prefix):]: v for k, v in state_dict.items()}
+        return state_dict
+
+    def _load_model_safetensors(self, model_safetensors_path: str) -> None:
+        if model_safetensors_path.startswith("oss://"):
+            raise ValueError(
+                "model_safetensors_path must be a local mounted path on AI-Hub, "
+                f"got {model_safetensors_path}"
+            )
+
+        model_safetensors_path = _resolve_wan_path(model_safetensors_path)
+        print(f"[WanDiffusionWrapper] Loading model safetensors from {model_safetensors_path}")
+        state_dict = load_safetensors_file(model_safetensors_path, device="cpu")
+        state_dict = self._normalize_model_state_dict_keys(state_dict)
+
+        model_keys = set(self.model.state_dict().keys())
+        matched_keys = model_keys.intersection(state_dict.keys())
+        if not matched_keys:
+            sample_keys = list(state_dict.keys())[:10]
+            raise ValueError(
+                "No safetensors keys matched the Wan model state_dict. "
+                f"First loaded keys: {sample_keys}"
+            )
+
+        match_ratio = len(matched_keys) / max(1, len(state_dict))
+        if match_ratio < 0.5:
+            sample_unexpected = [k for k in state_dict.keys() if k not in model_keys][:10]
+            raise ValueError(
+                f"Only {len(matched_keys)}/{len(state_dict)} safetensors keys match the Wan model "
+                f"state_dict after prefix normalization. Sample unexpected keys: {sample_unexpected}"
+            )
+
+        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
+        if missing:
+            print(
+                f"[WanDiffusionWrapper] model_safetensors missing {len(missing)} keys "
+                f"(showing first 20): {missing[:20]}"
+            )
+        if unexpected:
+            print(
+                f"[WanDiffusionWrapper] model_safetensors unexpected {len(unexpected)} keys "
+                f"(showing first 20): {unexpected[:20]}"
+            )
+        print(
+            f"[WanDiffusionWrapper] Loaded model safetensors with {len(matched_keys)} "
+            f"matched keys from {model_safetensors_path}"
+        )
+
+    def __init__(
+            self,
+            model_name="Wan2.1-T2V-1.3B",
+            timestep_shift=8.0,
+            is_causal=False,
+            local_attn_size=-1,
+            sink_size=0,
+            subfolder=None,
+            model_type='t2v',
+            num_frame_per_block=3,
+            model_safetensors_path: Optional[str] = None,
+            **model_init_kwargs,
+    ):
+        super().__init__()
+
+        self.model_type = model_type
+        use_relative_rope = bool(model_init_kwargs.pop("use_relative_rope", False))
+        if is_causal:
+            model_init_kwargs["use_relative_rope"] = use_relative_rope
+            self.model = CausalWanModel.from_pretrained(
+                model_name, local_attn_size=local_attn_size, sink_size=sink_size, model_type=model_type, num_frame_per_block=num_frame_per_block,
+                **model_init_kwargs)
+        else:
+            if use_relative_rope:
+                print("[WanDiffusionWrapper] use_relative_rope is ignored for non-causal WanModel.")
+            self.model = WanModel.from_pretrained(model_name, model_type=model_type, **model_init_kwargs)
+        materialized = self._materialize_meta_tensors(self.model, device=torch.device("cpu"))
+        if materialized:
+            print(f"[WanDiffusionWrapper] Materialized {len(materialized)} meta tensors on CPU.")
+        if model_safetensors_path:
+            self._load_model_safetensors(model_safetensors_path)
+        self.model.eval()
+
+        # For non-causal diffusion, all frames share the same timestep
+        self.uniform_timestep = not is_causal
+
+        self.scheduler = FlowMatchScheduler(
+            shift=timestep_shift, sigma_min=0.0, extra_one_step=True
+        )
+        self.scheduler.set_timesteps(1000, training=True)
+
+        self.seq_len = None  # [1, 21, 16, 60, 104]
+        self.post_init()
+
+    def enable_gradient_checkpointing(self) -> None:
+        self.model.enable_gradient_checkpointing()
+
+    def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None:
+        # NOTE: This is hard coded for WAN2.1-T2V-1.3B for now!!!!!!!!!!!!!!!!!!!!
+        self._cls_pred_branch = nn.Sequential(
+            # Input: [B, 384, 21, 60, 104]
+            nn.LayerNorm(atten_dim * 3 + time_embed_dim),
+            nn.Linear(atten_dim * 3 + time_embed_dim, 1536),
+            nn.SiLU(),
+            nn.Linear(atten_dim, num_class)
+        )
+        self._cls_pred_branch.requires_grad_(True)
+        num_registers = 3
+        self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim)
+        self._register_tokens.requires_grad_(True)
+
+        gan_ca_blocks = []
+        for _ in range(num_registers):
+            block = GanAttentionBlock()
+            gan_ca_blocks.append(block)
+        self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks)
+        self._gan_ca_blocks.requires_grad_(True)
+        # self.has_cls_branch = True
+
+    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert flow matching's prediction to x0 prediction.
+        flow_pred: the prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = noise - x0
+        x_t = (1-sigma_t) * x0 + sigma_t * noise
+        we have x0 = x_t - sigma_t * pred
+        see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e
+        """
+        # use higher precision for calculations
+        original_dtype = flow_pred.dtype
+        flow_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(flow_pred.device), [flow_pred, xt,
+                                                        self.scheduler.sigmas,
+                                                        self.scheduler.timesteps]
+        )
+
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        x0_pred = xt - sigma_t * flow_pred
+        return x0_pred.to(original_dtype)
+
+    @staticmethod
+    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert x0 prediction to flow matching's prediction.
+        x0_pred: the x0 prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = (x_t - x_0) / sigma_t
+        """
+        # use higher precision for calculations
+        original_dtype = x0_pred.dtype
+        x0_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
+                                                      scheduler.sigmas,
+                                                      scheduler.timesteps]
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        flow_pred = (xt - x0_pred) / sigma_t
+        return flow_pred.to(original_dtype)
+
+    @staticmethod
+    def _history_x_to_model_format(history_x):
+        if history_x is None:
+            return None
+
+        if torch.is_tensor(history_x):
+            if history_x.ndim != 5:
+                raise ValueError(
+                    f"history_x must be [B,F,C,H,W] when passed as a tensor, got {history_x.shape}"
+                )
+            return [u.permute(1, 0, 2, 3).contiguous() for u in history_x]
+
+        return history_x
+
+    @staticmethod
+    def _history_condition_to_model_format(value):
+        if value is None:
+            return None
+
+        if torch.is_tensor(value):
+            if value.ndim < 3:
+                return value
+            return [u.contiguous() for u in value]
+
+        return value
+
+    def forward(
+        self,
+        noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+
+        classify_mode: Optional[bool] = False, # DF
+        concat_time_embeddings: Optional[bool] = False, #DF
+        clean_x: Optional[torch.Tensor] = None, # TF
+        aug_t: Optional[torch.Tensor] = None, # for TF clean GT, if it's also noisy and needs denoising by the model, aug_t is its timestep
+
+        cache_start: Optional[int] = None,
+        updating_cache: Optional[bool] = False,
+        replace_first_timestep_and_noise_latents: Optional[bool] = False,
+        history_x: Optional[torch.Tensor] = None,
+        history_y: Optional[torch.Tensor] = None,
+        history_act_context: Optional[torch.Tensor] = None,
+        history_y_action: Optional[torch.Tensor] = None,
+        noisy_start_frame: int = 0,
+
+    ) -> torch.Tensor:
+        prompt_embeds = conditional_dict["prompt_embeds"]
+        act_context = conditional_dict.get("act_context", None)
+        act_context_scale = conditional_dict.get("act_context_scale", 1.0)
+        clip_fea = conditional_dict.get("clip_fea", None)
+        y = conditional_dict.get("y", None)
+        y_action = conditional_dict.get("y_action", None)
+        ref_latents = conditional_dict.get("ref_latents", None)
+        ref_mask = conditional_dict.get("ref_mask", None)
+        # first_frame_latents = conditional_dict.get("first_frame_latents", None)
+
+        raw_timestep = timestep
+        b, f, c, h, w = noisy_image_or_video.shape
+
+        if replace_first_timestep_and_noise_latents:
+            # Wan2.2 5B uses the first latent frame as a clean condition. Keep
+            # per-frame timesteps for score models so only frame 0 is forced to t=0.
+            if raw_timestep.dim() == 2:
+                input_timestep = raw_timestep.clone()
+                input_timestep[:, 0] = 0
+            elif raw_timestep.dim() == 1 and raw_timestep.shape[0] == f:
+                input_timestep = raw_timestep.unsqueeze(0).repeat(b, 1)
+                input_timestep[:, 0] = 0
+            elif raw_timestep.dim() == 1 and raw_timestep.shape[0] == b:
+                input_timestep = raw_timestep[:, None].repeat(1, f)
+                input_timestep[:, 0] = 0
+            else:
+                input_timestep = raw_timestep.reshape(-1)[0].view(1, 1).repeat(b, f)
+                input_timestep[:, 0] = 0
+        elif self.uniform_timestep:
+            # [B, F] -> [B] for legacy non-causal uniform score models.
+            input_timestep = raw_timestep[:, 0]
+        else:
+            input_timestep = raw_timestep
+
+        logits = None
+        if history_x is not None:
+            history_kwargs = {
+                "history_x": self._history_x_to_model_format(history_x),
+                "noisy_start_frame": int(noisy_start_frame),
+            }
+            history_y = self._history_condition_to_model_format(history_y)
+            history_act_context = self._history_condition_to_model_format(history_act_context)
+            history_y_action = self._history_condition_to_model_format(history_y_action)
+
+            if history_y is not None:
+                history_kwargs["history_y"] = history_y
+            if history_act_context is not None:
+                history_kwargs["history_act_context"] = history_act_context
+            if history_y_action is not None:
+                history_kwargs["history_y_action"] = history_y_action
+
+            model_out = self.model(
+                noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                t=input_timestep,
+                context=prompt_embeds,
+                seq_len=self.seq_len,
+                kv_cache=None,
+                crossattn_cache=crossattn_cache,
+                current_start=0 if current_start is None else current_start,
+                cache_start=0 if cache_start is None else cache_start,
+                act_context=act_context,
+                y_action=y_action,
+                act_context_scale=act_context_scale,
+                clip_fea=clip_fea,
+                y=y,
+                ref_latents=ref_latents,
+                ref_mask=ref_mask,
+                **history_kwargs,
+            )
+            if isinstance(model_out, tuple):
+                flow_pred = model_out[0]
+            else:
+                flow_pred = model_out
+            flow_pred = flow_pred.permute(0, 2, 1, 3, 4)
+        # X0 prediction
+        elif kv_cache is not None:
+            kwargs = {}
+            if updating_cache:
+                kwargs["updating_cache"] = updating_cache
+            flow_pred = self.model(
+                noisy_image_or_video.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W],
+                t=input_timestep, context=prompt_embeds,
+                seq_len=self.seq_len,
+                kv_cache=kv_cache,
+                crossattn_cache=crossattn_cache,
+                current_start=current_start,
+                cache_start=cache_start,
+                act_context=act_context,
+                act_context_scale=act_context_scale,
+                clip_fea=clip_fea,
+                y=y,
+                ref_latents=ref_latents,
+                ref_mask=ref_mask,
+                **kwargs,
+            ).permute(0, 2, 1, 3, 4)
+        else:
+            if clean_x is not None:
+                # teacher forcing
+                flow_pred = self.model(
+                    noisy_image_or_video.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W]
+                    t=input_timestep, context=prompt_embeds,
+                    seq_len=self.seq_len,
+                    clean_x=clean_x.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W]
+                    aug_t=aug_t,
+                    act_context=act_context,
+                    act_context_scale=act_context_scale,
+                    clip_fea=clip_fea,
+                    y=y,
+                    ref_latents=ref_latents,
+                    ref_mask=ref_mask,
+                ).permute(0, 2, 1, 3, 4)
+            else:
+                # diffusion forcing or bidirectional
+                if classify_mode:
+                    flow_pred, logits = self.model(
+                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                        t=input_timestep, context=prompt_embeds,
+                        seq_len=self.seq_len,
+                        classify_mode=True,
+                        register_tokens=self._register_tokens,
+                        cls_pred_branch=self._cls_pred_branch,
+                        gan_ca_blocks=self._gan_ca_blocks,
+                        concat_time_embeddings=concat_time_embeddings,
+                        act_context=act_context,
+                        act_context_scale=act_context_scale,
+                        clip_fea=clip_fea,
+                        y=y,
+                        ref_latents=ref_latents,
+                        ref_mask=ref_mask,
+                    )
+                    flow_pred = flow_pred.permute(0, 2, 1, 3, 4)
+                else:
+                    flow_pred = self.model(
+                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                        t=input_timestep, context=prompt_embeds,
+                        seq_len=self.seq_len,
+                        act_context=act_context,
+                        act_context_scale=act_context_scale,
+                        clip_fea=clip_fea,
+                        y=y,
+                        ref_latents=ref_latents,
+                        ref_mask=ref_mask,
+                    ).permute(0, 2, 1, 3, 4)
+
+        pred_x0 = self._convert_flow_pred_to_x0(
+            flow_pred=flow_pred.flatten(0, 1),
+            xt=noisy_image_or_video.flatten(0, 1),
+            timestep=timestep.flatten(0, 1)
+        ).unflatten(0, flow_pred.shape[:2])
+
+        if logits is not None:
+            return flow_pred, pred_x0, logits
+
+        return flow_pred, pred_x0
+
+    def get_scheduler(self) -> SchedulerInterface:
+        """
+        Update the current scheduler with the interface's static method
+        """
+        scheduler = self.scheduler
+        scheduler.convert_x0_to_noise = types.MethodType(
+            SchedulerInterface.convert_x0_to_noise, scheduler)
+        scheduler.convert_noise_to_x0 = types.MethodType(
+            SchedulerInterface.convert_noise_to_x0, scheduler)
+        scheduler.convert_velocity_to_x0 = types.MethodType(
+            SchedulerInterface.convert_velocity_to_x0, scheduler)
+        self.scheduler = scheduler
+        return scheduler
+
+    def post_init(self):
+        """
+        A few custom initialization steps that should be called after the object is created.
+        Currently, the only one we have is to bind a few methods to scheduler.
+        We can gradually add more methods here if needed.
+        """
+        self.get_scheduler()

wan/__init__.py

@@ -0,0 +1,2 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from . import modules

wan/modules/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae2_2 import Wan2_2_VAE
+
+__all__ = [
+    'Wan2_2_VAE',
+    'WanModel',
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+    'HuggingfaceTokenizer',
+    'flash_attention',
+]

wan/modules/attention.py

@@ -0,0 +1,476 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn_interface
+
+    def is_hopper_gpu():
+        """flashattn-hopper 仅支持 Hopper (H100)，不支持 Blackwell，故只检测 Hopper。"""
+        if not torch.cuda.is_available():
+            return False
+        device_name = torch.cuda.get_device_name(0).lower()
+        return "h100" in device_name or "hopper" in device_name
+    FLASH_ATTN_3_AVAILABLE = is_hopper_gpu()
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except ModuleNotFoundError:
+    SAGE_ATTN_AVAILABLE = False
+
+try:
+    from sageattn3 import sageattn3_blackwell
+    SAGE_ATTN_3_BLACKWELL_AVAILABLE = True
+except ModuleNotFoundError:
+    SAGE_ATTN_3_BLACKWELL_AVAILABLE = False
+
+from .sla_attn import get_block_map
+from .sla_kernel import _attention
+
+
+# FLASH_ATTN_3_AVAILABLE = False
+
+import warnings
+
+# 优先级从高到低，模块加载时按此顺序选取第一个可用的作为默认后端
+# ATTN_BACKEND_PRIORITY = ['flash_attn_3', 'flash_attn_2', 'sdpa']
+ATTN_BACKEND_PRIORITY = ['sla_triton', 'sageattn', 'sageattn3', 'flash_attn_3', 'flash_attn_2', 'sdpa']
+# ATTN_BACKEND_PRIORITY = ['flash_attn_3', 'flash_attn_2', 'sdpa']
+print(f"Attention backend priority: {ATTN_BACKEND_PRIORITY}")
+print(f"SAGE_ATTN_3_BLACKWELL_AVAILABLE: {SAGE_ATTN_3_BLACKWELL_AVAILABLE}")
+print(f"SAGE_ATTN_AVAILABLE: {SAGE_ATTN_AVAILABLE}")
+print(f"FLASH_ATTN_3_AVAILABLE: {FLASH_ATTN_3_AVAILABLE}")
+print(f"FLASH_ATTN_2_AVAILABLE: {FLASH_ATTN_2_AVAILABLE}")
+
+
+def _resolve_default_attn_backend():
+    """根据优先级数组和当前环境可用性，解析出默认的 attention 后端。"""
+    availability = {
+        'sageattn3': SAGE_ATTN_3_BLACKWELL_AVAILABLE,
+        'sageattn': SAGE_ATTN_AVAILABLE,
+        'flash_attn_3': FLASH_ATTN_3_AVAILABLE,
+        'flash_attn_2': FLASH_ATTN_2_AVAILABLE,
+        'sla_triton': False,
+        'sdpa': True,  # 始终可用
+    }
+    for name in ATTN_BACKEND_PRIORITY:
+        if availability.get(name, False):
+            return name
+    return 'sdpa'
+
+
+def _is_power_of_two(x):
+    return x > 0 and (x & (x - 1)) == 0
+
+
+def _is_backend_available(backend):
+    availability = {
+        'sageattn3': SAGE_ATTN_3_BLACKWELL_AVAILABLE,
+        'sageattn': SAGE_ATTN_AVAILABLE,
+        'flash_attn_3': FLASH_ATTN_3_AVAILABLE,
+        'flash_attn_2': FLASH_ATTN_2_AVAILABLE,
+        'sla_triton': True,
+        'sdpa': True,
+    }
+    return availability.get(backend, False)
+
+
+def _can_use_backend(
+    backend,
+    q,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    causal=False,
+    window_size=(-1, -1),
+):
+    # Layout in this module: [B, L, H, D]
+    b, _, _, d = q.shape
+    has_varlen = q_lens is not None or k_lens is not None
+    has_window = window_size != (-1, -1)
+    has_dropout = dropout_p > 0
+
+    if backend == 'sla_triton':
+        # Current SLA Triton path does not consume varlen/causal/window/dropout args,
+        # and Triton kernel in sla_attn.py requires power-of-two D.
+        return (
+            not has_varlen
+            and not has_window
+            and not causal
+            and not has_dropout
+            and b == 1
+            and _is_power_of_two(d)
+        )
+    if backend in ('sageattn', 'sageattn3'):
+        # Sage kernels in this file do not support varlen/window/dropout controls.
+        return not has_varlen and not has_window and not has_dropout
+    if backend == 'flash_attn_3':
+        # flash-attn3 path here does not support dropout/window_size.
+        return not has_window and not has_dropout
+    if backend == 'flash_attn_2':
+        return True
+    if backend == 'sdpa':
+        return True
+    return False
+
+
+def _resolve_auto_backend(
+    q,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.0,
+    causal=False,
+    window_size=(-1, -1),
+):
+    candidates = [DEFAULT_ATTN_BACKEND] + [b for b in ATTN_BACKEND_PRIORITY if b != DEFAULT_ATTN_BACKEND]
+    for backend in candidates:
+        if not _is_backend_available(backend):
+            continue
+        if _can_use_backend(
+            backend=backend,
+            q=q,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            causal=causal,
+            window_size=window_size,
+        ):
+            return backend
+    return 'sdpa'
+
+
+