Add LongCat T2V (Base, Distillation and Refinement) Support to FastVideo (#883)
Co-authored-by: Shao Duan <shaoxiongduan@gmail.com>
This commit is contained in:
alexzms
committed by
RandNMR73
RandNMR73
co-authored by
Shao Duan
Shao Duan
parent
f2cced218d
commit
9211fa84fc
@@ -30,6 +30,8 @@ env
|
||||
**/build/
|
||||
**.pyc
|
||||
**.txt
|
||||
*.log
|
||||
weights/
|
||||
|
||||
# Distribution / packaging
|
||||
build/
|
||||
|
||||
@@ -2,10 +2,12 @@ from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
|
||||
from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
|
||||
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
|
||||
from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
|
||||
from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
|
||||
from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
|
||||
from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
|
||||
|
||||
__all__ = [
|
||||
"HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
|
||||
"StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig"
|
||||
"StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
|
||||
"LongCatVideoConfig"
|
||||
]
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
LongCat Video DiT configuration for native FastVideo implementation.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
|
||||
from fastvideo.platforms import AttentionBackendEnum
|
||||
|
||||
|
||||
def is_longcat_blocks(n: str, m) -> bool:
|
||||
"""FSDP shard condition for LongCat transformer blocks."""
|
||||
return "blocks" in n and str.isdigit(n.split(".")[-1])
|
||||
|
||||
|
||||
@dataclass
|
||||
class LongCatVideoArchConfig(DiTArchConfig):
|
||||
"""Architecture configuration for native LongCat Video DiT."""
|
||||
|
||||
_fsdp_shard_conditions: list = field(
|
||||
default_factory=lambda: [is_longcat_blocks])
|
||||
|
||||
# Enable torch.compile for transformer blocks (major speedup!)
|
||||
_compile_conditions: list = field(
|
||||
default_factory=lambda: [is_longcat_blocks])
|
||||
|
||||
# Parameter name mapping for weight conversion
|
||||
# Maps original LongCat third_party names -> native FastVideo names
|
||||
param_names_mapping: dict = field(
|
||||
default_factory=lambda: {
|
||||
# Embedders
|
||||
r"^x_embedder\.(.*)$": r"patch_embed.\1",
|
||||
r"^t_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1",
|
||||
r"^t_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\1",
|
||||
r"^y_embedder\.y_proj\.0\.(.*)$": r"caption_embedder.linear_1.\1",
|
||||
r"^y_embedder\.y_proj\.2\.(.*)$": r"caption_embedder.linear_2.\1",
|
||||
|
||||
# Transformer blocks - AdaLN modulation
|
||||
r"^blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$":
|
||||
r"blocks.\1.adaln_linear_1.\2",
|
||||
|
||||
# Transformer blocks - Normalization
|
||||
r"^blocks\.(\d+)\.mod_norm_attn\.(.*)$": r"blocks.\1.norm_attn.\2",
|
||||
r"^blocks\.(\d+)\.mod_norm_ffn\.(.*)$": r"blocks.\1.norm_ffn.\2",
|
||||
r"^blocks\.(\d+)\.pre_crs_attn_norm\.(.*)$":
|
||||
r"blocks.\1.norm_cross.\2",
|
||||
|
||||
# Self-attention: QKV fused -> separate (will need splitting in converter)
|
||||
# Original has attn.qkv.weight -> need to split into to_q, to_k, to_v
|
||||
r"^blocks\.(\d+)\.attn\.qkv\.(.*)$":
|
||||
r"blocks.\1.self_attn.qkv_fused.\2", # Marker for splitting
|
||||
r"^blocks\.(\d+)\.attn\.proj\.(.*)$":
|
||||
r"blocks.\1.self_attn.to_out.\2",
|
||||
r"^blocks\.(\d+)\.attn\.q_norm\.(.*)$":
|
||||
r"blocks.\1.self_attn.q_norm.\2",
|
||||
r"^blocks\.(\d+)\.attn\.k_norm\.(.*)$":
|
||||
r"blocks.\1.self_attn.k_norm.\2",
|
||||
|
||||
# Cross-attention
|
||||
r"^blocks\.(\d+)\.cross_attn\.q_linear\.(.*)$":
|
||||
r"blocks.\1.cross_attn.to_q.\2",
|
||||
r"^blocks\.(\d+)\.cross_attn\.kv_linear\.(.*)$":
|
||||
r"blocks.\1.cross_attn.kv_fused.\2", # Marker for splitting
|
||||
r"^blocks\.(\d+)\.cross_attn\.proj\.(.*)$":
|
||||
r"blocks.\1.cross_attn.to_out.\2",
|
||||
r"^blocks\.(\d+)\.cross_attn\.q_norm\.(.*)$":
|
||||
r"blocks.\1.cross_attn.q_norm.\2",
|
||||
r"^blocks\.(\d+)\.cross_attn\.k_norm\.(.*)$":
|
||||
r"blocks.\1.cross_attn.k_norm.\2",
|
||||
|
||||
# FFN (SwiGLU)
|
||||
r"^blocks\.(\d+)\.ffn\.w1\.(.*)$": r"blocks.\1.ffn.w1.\2", # gate
|
||||
r"^blocks\.(\d+)\.ffn\.w2\.(.*)$": r"blocks.\1.ffn.w2.\2", # down
|
||||
r"^blocks\.(\d+)\.ffn\.w3\.(.*)$": r"blocks.\1.ffn.w3.\2", # up
|
||||
|
||||
# Final layer
|
||||
r"^final_layer\.adaLN_modulation\.1\.(.*)$":
|
||||
r"final_layer.adaln_linear.\1",
|
||||
r"^final_layer\.norm_final\.(.*)$": r"final_layer.norm.\1",
|
||||
r"^final_layer\.linear\.(.*)$": r"final_layer.proj.\1",
|
||||
})
|
||||
|
||||
# Reverse mapping for saving checkpoints: custom -> hf
|
||||
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
|
||||
|
||||
# LoRA parameter name mapping
|
||||
lora_param_names_mapping: dict = field(default_factory=lambda: {})
|
||||
|
||||
# Model architecture parameters
|
||||
hidden_size: int = 4096
|
||||
depth: int = 48 # Number of transformer blocks
|
||||
num_attention_heads: int = 32
|
||||
attention_head_dim: int = 128 # hidden_size / num_attention_heads
|
||||
|
||||
in_channels: int = 16 # Latent space channels
|
||||
out_channels: int = 16
|
||||
num_channels_latents: int = 16
|
||||
|
||||
# Patch embedding
|
||||
patch_size: tuple[int, int,
|
||||
int] = (1, 2, 2) # [T, H, W] - no temporal compression
|
||||
|
||||
# Text/caption embedding
|
||||
caption_channels: int = 4096 # UMT5 d_model
|
||||
|
||||
# Timestep embedding
|
||||
adaln_tembed_dim: int = 512
|
||||
frequency_embedding_size: int = 256
|
||||
|
||||
# FFN
|
||||
mlp_ratio: int = 4
|
||||
|
||||
# Attention backend support
|
||||
_supported_attention_backends: tuple = field(default_factory=lambda: (
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.TORCH_SDPA,
|
||||
))
|
||||
|
||||
# Text padding behavior
|
||||
text_tokens_zero_pad: bool = True
|
||||
|
||||
# Block Sparse Attention (BSA)
|
||||
enable_bsa: bool = False
|
||||
bsa_params: dict | None = field(
|
||||
default_factory=lambda: {
|
||||
"sparsity": 0.9375,
|
||||
"cdf_threshold": None,
|
||||
"chunk_3d_shape_q": [4, 4, 4],
|
||||
"chunk_3d_shape_k": [4, 4, 4],
|
||||
})
|
||||
|
||||
# LoRA exclusions
|
||||
exclude_lora_layers: list[str] = field(default_factory=lambda: [])
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.out_channels = self.out_channels or self.in_channels
|
||||
# Ensure attention_head_dim matches
|
||||
self.attention_head_dim = self.hidden_size // self.num_attention_heads
|
||||
|
||||
|
||||
@dataclass
|
||||
class LongCatVideoConfig(DiTConfig):
|
||||
"""Main configuration for LongCat Video DiT."""
|
||||
|
||||
arch_config: DiTArchConfig = field(default_factory=LongCatVideoArchConfig)
|
||||
|
||||
prefix: str = "longcat"
|
||||
@@ -0,0 +1,355 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
import html
|
||||
|
||||
import ftfy
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
from fastvideo.configs.models import DiTConfig, VAEConfig
|
||||
from fastvideo.configs.models.dits.base import DiTArchConfig
|
||||
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
|
||||
from fastvideo.configs.models.vaes import WanVAEConfig
|
||||
from fastvideo.configs.pipelines.base import PipelineConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class LongCatDiTArchConfig(DiTArchConfig):
|
||||
"""Extended DiTArchConfig with LongCat-specific fields.
|
||||
|
||||
NOTE: This is for Phase 1 wrapper compatibility. For native model (Phase 2),
|
||||
use LongCatVideoConfig from fastvideo.configs.models.dits.longcat instead.
|
||||
"""
|
||||
# LongCat-specific architecture parameters
|
||||
adaln_tembed_dim: int = 512
|
||||
caption_channels: int = 4096
|
||||
depth: int = 48
|
||||
enable_bsa: bool = False
|
||||
enable_flashattn3: bool = False
|
||||
enable_flashattn2: bool = True
|
||||
enable_xformers: bool = False
|
||||
frequency_embedding_size: int = 256
|
||||
in_channels: int = 16
|
||||
mlp_ratio: int = 4
|
||||
num_heads: int = 32
|
||||
out_channels: int = 16
|
||||
text_tokens_zero_pad: bool = True
|
||||
patch_size: list[int] = field(default_factory=lambda: [1, 2, 2])
|
||||
cp_split_hw: list[int] | None = None
|
||||
bsa_params: dict | None = None
|
||||
|
||||
|
||||
def longcat_preprocess_text(prompt: str) -> str:
|
||||
"""Clean and preprocess text like original LongCat implementation.
|
||||
|
||||
This function applies the same text cleaning pipeline as the original
|
||||
LongCat-Video implementation to ensure identical tokenization results.
|
||||
|
||||
Steps:
|
||||
1. basic_clean: Fix unicode issues and unescape HTML entities
|
||||
2. whitespace_clean: Normalize whitespace to single spaces
|
||||
|
||||
Args:
|
||||
prompt: Raw input text prompt
|
||||
|
||||
Returns:
|
||||
Cleaned and normalized text prompt
|
||||
"""
|
||||
# basic_clean: fix unicode and HTML entities
|
||||
text = ftfy.fix_text(prompt)
|
||||
text = html.unescape(html.unescape(text))
|
||||
text = text.strip()
|
||||
|
||||
# whitespace_clean: normalize whitespace
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = text.strip()
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def umt5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
|
||||
"""
|
||||
Postprocess UMT5/T5 encoder outputs to fixed length 512 embeddings.
|
||||
"""
|
||||
mask: torch.Tensor = outputs.attention_mask
|
||||
hidden_state: torch.Tensor = outputs.last_hidden_state
|
||||
seq_lens = mask.gt(0).sum(dim=1).long()
|
||||
assert torch.isnan(hidden_state).sum() == 0
|
||||
prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
|
||||
prompt_embeds_tensor: torch.Tensor = torch.stack([
|
||||
torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
|
||||
for u in prompt_embeds
|
||||
],
|
||||
dim=0)
|
||||
return prompt_embeds_tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class LongCatT2V480PConfig(PipelineConfig):
|
||||
"""Configuration for LongCat pipeline (480p) aligned to LongCat-Video modules.
|
||||
|
||||
Components expected by loaders:
|
||||
- tokenizer: AutoTokenizer
|
||||
- text_encoder: UMT5EncoderModel
|
||||
- transformer: LongCatVideoTransformer3DModel (Phase 1 wrapper)
|
||||
OR LongCatTransformer3DModel (Phase 2 native)
|
||||
- vae: AutoencoderKLWan (Wan VAE, 4x8 compression)
|
||||
- scheduler: FlowMatchEulerDiscreteScheduler
|
||||
"""
|
||||
|
||||
# DiT config with LongCat-specific arch_config
|
||||
# NOTE: For Phase 1 wrapper, uses LongCatDiTArchConfig
|
||||
# For Phase 2 native model, can use LongCatVideoConfig directly
|
||||
dit_config: DiTConfig = field(
|
||||
default_factory=lambda: DiTConfig(arch_config=LongCatDiTArchConfig()))
|
||||
|
||||
# VAE config: Wan VAE with encoder+decoder enabled
|
||||
vae_config: VAEConfig = field(default_factory=WanVAEConfig)
|
||||
vae_tiling: bool = False
|
||||
vae_sp: bool = False
|
||||
|
||||
# Precision defaults
|
||||
dit_precision: str = "bf16"
|
||||
vae_precision: str = "bf16"
|
||||
text_encoder_precisions: tuple[str, ...] = field(
|
||||
default_factory=lambda: ("bf16", ))
|
||||
|
||||
# Text encoding (UMT5 uses T5-like config; postprocess to fixed 512)
|
||||
text_encoder_configs: tuple[T5Config, ...] = field(
|
||||
default_factory=lambda: (T5Config(), ))
|
||||
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
|
||||
default_factory=lambda: (longcat_preprocess_text, ))
|
||||
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
|
||||
...] = field(default_factory=lambda:
|
||||
(umt5_postprocess_text, ))
|
||||
|
||||
# LongCat-specific runtime toggles (consumed by pipeline/stages)
|
||||
enable_kv_cache: bool = True
|
||||
offload_kv_cache: bool = False
|
||||
enable_bsa: bool = False
|
||||
use_distill: bool = False
|
||||
enhance_hf: bool = False
|
||||
# Optional BSA parameter dict (kept for backward/phase-1 compatibility).
|
||||
# `LongCatPipeline.initialize_pipeline()` uses this as a base and then applies
|
||||
# CLI overrides (bsa_sparsity / bsa_chunk_{q,k} / bsa_cdf_threshold).
|
||||
bsa_params: dict | None = None
|
||||
# BSA runtime overrides (preferred over bsa_params if provided via CLI)
|
||||
bsa_sparsity: float | None = None
|
||||
bsa_cdf_threshold: float | None = None
|
||||
bsa_chunk_q: list[int] | None = None
|
||||
bsa_chunk_k: list[int] | None = None
|
||||
t_thresh: float | None = None # refine stage default controlled by sampling args
|
||||
|
||||
# LongCat does not need flow_shift
|
||||
flow_shift: float | None = None
|
||||
dmd_denoising_steps: list[int] | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
# LongCat inference requires vae encoder and decoder
|
||||
self.vae_config.load_encoder = True
|
||||
self.vae_config.load_decoder = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class LongCatT2V704PConfig(LongCatT2V480PConfig):
|
||||
"""Configuration for LongCat pipeline (704p) with BSA enabled by default.
|
||||
|
||||
Uses the same resolution and BSA parameters as original LongCat refinement stage.
|
||||
BSA parameters configured in transformer config.json with chunk_3d_shape=[4,4,4]:
|
||||
- Input: 704×1280×96
|
||||
- VAE (8x): 88×160×96
|
||||
- Patch [1,2,2]: 44×80×96
|
||||
- chunk [4,4,4]: 96%4=0, 44%4=0, 80%4=0 ✅
|
||||
|
||||
This configuration matches the original LongCat refinement stage parameters.
|
||||
"""
|
||||
|
||||
# Enable BSA by default for 704p
|
||||
enable_bsa: bool = True
|
||||
|
||||
|
||||
ASPECT_RATIO_627 = {
|
||||
'0.26': ([320, 1216], 1),
|
||||
'0.31': ([352, 1120], 1),
|
||||
'0.38': ([384, 1024], 1),
|
||||
'0.43': ([416, 960], 1),
|
||||
'0.52': ([448, 864], 1),
|
||||
'0.58': ([480, 832], 1),
|
||||
'0.67': ([512, 768], 1),
|
||||
'0.74': ([544, 736], 1),
|
||||
'0.86': ([576, 672], 1),
|
||||
'0.95': ([608, 640], 1),
|
||||
'1.05': ([640, 608], 1),
|
||||
'1.17': ([672, 576], 1),
|
||||
'1.29': ([704, 544], 1),
|
||||
'1.35': ([736, 544], 1),
|
||||
'1.50': ([768, 512], 1),
|
||||
'1.67': ([800, 480], 1),
|
||||
'1.73': ([832, 480], 1),
|
||||
'2.00': ([896, 448], 1),
|
||||
'2.31': ([960, 416], 1),
|
||||
'2.58': ([992, 384], 1),
|
||||
'2.75': ([1056, 384], 1),
|
||||
'3.09': ([1088, 352], 1),
|
||||
'3.70': ([1184, 320], 1),
|
||||
'3.80': ([1216, 320], 1),
|
||||
'3.90': ([1248, 320], 1),
|
||||
'4.00': ([1280, 320], 1)
|
||||
}
|
||||
|
||||
ASPECT_RATIO_627_F64 = {
|
||||
'0.26': ([320, 1216], 1),
|
||||
'0.38': ([384, 1024], 1),
|
||||
'0.50': ([448, 896], 1),
|
||||
'0.67': ([512, 768], 1),
|
||||
'0.82': ([576, 704], 1),
|
||||
'1.00': ([640, 640], 1),
|
||||
'1.22': ([704, 576], 1),
|
||||
'1.50': ([768, 512], 1),
|
||||
'1.86': ([832, 448], 1),
|
||||
'2.00': ([896, 448], 1),
|
||||
'2.50': ([960, 384], 1),
|
||||
'2.83': ([1088, 384], 1),
|
||||
'3.60': ([1152, 320], 1),
|
||||
'3.80': ([1216, 320], 1),
|
||||
'4.00': ([1280, 320], 1)
|
||||
}
|
||||
|
||||
ASPECT_RATIO_627_F128 = {
|
||||
'0.25': ([256, 1024], 1),
|
||||
'0.38': ([384, 1024], 1),
|
||||
'0.43': ([384, 896], 1),
|
||||
'0.57': ([512, 896], 1),
|
||||
'0.67': ([512, 768], 1),
|
||||
'1.00': ([640, 640], 1),
|
||||
'1.50': ([768, 512], 1),
|
||||
'1.75': ([896, 512], 1),
|
||||
'2.33': ([896, 384], 1),
|
||||
'2.67': ([1024, 384], 1),
|
||||
'4.00': ([1024, 256], 1),
|
||||
}
|
||||
|
||||
ASPECT_RATIO_627_F256 = {
|
||||
'0.25': ([256, 1024], 1),
|
||||
'0.33': ([256, 768], 1),
|
||||
'0.50': ([256, 512], 1),
|
||||
'0.67': ([512, 768], 1),
|
||||
'1.00': ([512, 512], 1),
|
||||
'1.50': ([768, 512], 1),
|
||||
'2.00': ([512, 256], 1),
|
||||
'3.00': ([768, 256], 1),
|
||||
'4.00': ([1024, 256], 1),
|
||||
}
|
||||
|
||||
ASPECT_RATIO_960 = {
|
||||
'0.25': ([480, 1920], 1),
|
||||
'0.29': ([512, 1792], 1),
|
||||
'0.32': ([544, 1696], 1),
|
||||
'0.36': ([576, 1600], 1),
|
||||
'0.40': ([608, 1504], 1),
|
||||
'0.49': ([672, 1376], 1),
|
||||
'0.54': ([704, 1312], 1),
|
||||
'0.59': ([736, 1248], 1),
|
||||
'0.69': ([800, 1152], 1),
|
||||
'0.74': ([832, 1120], 1),
|
||||
'0.82': ([864, 1056], 1),
|
||||
'0.88': ([896, 1024], 1),
|
||||
'0.94': ([928, 992], 1),
|
||||
'1.00': ([960, 960], 1),
|
||||
'1.07': ([992, 928], 1),
|
||||
'1.14': ([1024, 896], 1),
|
||||
'1.22': ([1056, 864], 1),
|
||||
'1.31': ([1088, 832], 1),
|
||||
'1.35': ([1120, 832], 1),
|
||||
'1.44': ([1152, 800], 1),
|
||||
'1.70': ([1248, 736], 1),
|
||||
'2.00': ([1344, 672], 1),
|
||||
'2.05': ([1376, 672], 1),
|
||||
'2.47': ([1504, 608], 1),
|
||||
'2.53': ([1536, 608], 1),
|
||||
'2.83': ([1632, 576], 1),
|
||||
'3.06': ([1664, 544], 1),
|
||||
'3.12': ([1696, 544], 1),
|
||||
'3.62': ([1856, 512], 1),
|
||||
'3.93': ([1888, 480], 1),
|
||||
'4.00': ([1920, 480], 1)
|
||||
}
|
||||
|
||||
ASPECT_RATIO_960_F64 = {
|
||||
'0.22': ([448, 2048], 1),
|
||||
'0.29': ([512, 1792], 1),
|
||||
'0.36': ([576, 1600], 1),
|
||||
'0.45': ([640, 1408], 1),
|
||||
'0.55': ([704, 1280], 1),
|
||||
'0.63': ([768, 1216], 1),
|
||||
'0.76': ([832, 1088], 1),
|
||||
'0.88': ([896, 1024], 1),
|
||||
'1.00': ([960, 960], 1),
|
||||
'1.14': ([1024, 896], 1),
|
||||
'1.31': ([1088, 832], 1),
|
||||
'1.50': ([1152, 768], 1),
|
||||
'1.58': ([1216, 768], 1),
|
||||
'1.82': ([1280, 704], 1),
|
||||
'1.91': ([1344, 704], 1),
|
||||
'2.20': ([1408, 640], 1),
|
||||
'2.30': ([1472, 640], 1),
|
||||
'2.67': ([1536, 576], 1),
|
||||
'2.89': ([1664, 576], 1),
|
||||
'3.62': ([1856, 512], 1),
|
||||
'3.75': ([1920, 512], 1)
|
||||
}
|
||||
|
||||
ASPECT_RATIO_960_F128 = {
|
||||
'0.20': ([384, 1920], 1),
|
||||
'0.27': ([512, 1920], 1),
|
||||
'0.33': ([512, 1536], 1),
|
||||
'0.42': ([640, 1536], 1),
|
||||
'0.50': ([640, 1280], 1),
|
||||
'0.60': ([768, 1280], 1),
|
||||
'0.67': ([768, 1152], 1),
|
||||
'0.78': ([896, 1152], 1),
|
||||
'1.00': ([1024, 1024], 1),
|
||||
'1.29': ([1152, 896], 1),
|
||||
'1.50': ([1152, 768], 1),
|
||||
'1.67': ([1280, 768], 1),
|
||||
'2.00': ([1280, 640], 1),
|
||||
'2.40': ([1536, 640], 1),
|
||||
'3.00': ([1536, 512], 1),
|
||||
'3.75': ([1920, 512], 1),
|
||||
'5.00': ([1920, 384], 1),
|
||||
}
|
||||
|
||||
ASPECT_RATIO_960_F256 = {
|
||||
'0.33': ([512, 1536], 1),
|
||||
'0.60': ([768, 1280], 1),
|
||||
'1.00': ([1024, 1024], 1),
|
||||
'1.67': ([1280, 768], 1),
|
||||
'3.00': ([1536, 512], 1),
|
||||
}
|
||||
|
||||
|
||||
def get_bucket_config(resolution, scale_factor_spatial):
|
||||
if resolution == '480p':
|
||||
if scale_factor_spatial == 16 or scale_factor_spatial == 32:
|
||||
return ASPECT_RATIO_627
|
||||
elif scale_factor_spatial == 64:
|
||||
return ASPECT_RATIO_627_F64
|
||||
elif scale_factor_spatial == 128:
|
||||
return ASPECT_RATIO_627_F128
|
||||
elif scale_factor_spatial == 256:
|
||||
return ASPECT_RATIO_627_F256
|
||||
elif resolution == '720p':
|
||||
if scale_factor_spatial == 16 or scale_factor_spatial == 32:
|
||||
return ASPECT_RATIO_960
|
||||
elif scale_factor_spatial == 64:
|
||||
return ASPECT_RATIO_960_F64
|
||||
elif scale_factor_spatial == 128:
|
||||
return ASPECT_RATIO_960_F128
|
||||
elif scale_factor_spatial == 256:
|
||||
return ASPECT_RATIO_960_F256
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported resolution '{resolution}' or scale_factor_spatial '{scale_factor_spatial}'"
|
||||
)
|
||||
@@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig
|
||||
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
|
||||
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
|
||||
from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
|
||||
from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
|
||||
|
||||
# isort: off
|
||||
from fastvideo.configs.pipelines.wan import (
|
||||
@@ -77,11 +78,14 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
|
||||
lambda id: "stepvideo" in id.lower(),
|
||||
"cosmos":
|
||||
lambda id: "cosmos" in id.lower(),
|
||||
"longcat":
|
||||
lambda id: "longcat" in id.lower(),
|
||||
# Add other pipeline architecture detectors
|
||||
}
|
||||
|
||||
# Fallback configs when exact match isn't found but architecture is detected
|
||||
PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
|
||||
"longcat": LongCatT2V480PConfig,
|
||||
"hunyuan":
|
||||
HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant
|
||||
"matrixgame": MatrixGameI2V480PConfig,
|
||||
|
||||
@@ -3,6 +3,7 @@ from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.utils import StoreBoolean
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -27,6 +28,17 @@ class SamplingParam:
|
||||
keyboard_cond: Any | None = None # Shape: (B, T, K)
|
||||
grid_sizes: Any | None = None # Shape: (3,) [F,H,W]
|
||||
|
||||
# Refine inputs (LongCat 480p->720p upscaling)
|
||||
# Path-based refine (load stage1 video from disk, e.g. MP4)
|
||||
refine_from: str | None = None # Path to stage1 video (480p output from distill)
|
||||
t_thresh: float = 0.5 # Threshold for timestep scheduling in refinement
|
||||
spatial_refine_only: bool = False # If True, only spatial (no temporal doubling)
|
||||
num_cond_frames: int = 0 # Number of conditioning frames
|
||||
# In-memory refine input (for two-stage pipeline where stage1 frames are already in memory)
|
||||
# This mirrors LongCat's demo where a list of frames (e.g. np.ndarray or PIL.Image)
|
||||
# is passed directly to the refinement pipeline instead of reloading from disk.
|
||||
stage1_video: Any | None = None
|
||||
|
||||
# Text inputs
|
||||
prompt: str | list[str] | None = None
|
||||
negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
|
||||
@@ -216,6 +228,31 @@ class SamplingParam:
|
||||
default=SamplingParam.video_path,
|
||||
help="Path to input video for video-to-video generation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--refine-from",
|
||||
type=str,
|
||||
default=SamplingParam.refine_from,
|
||||
help="Path to stage1 video for refinement (LongCat 480p->720p)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--t-thresh",
|
||||
type=float,
|
||||
default=SamplingParam.t_thresh,
|
||||
help=
|
||||
"Threshold for timestep scheduling in refinement (default: 0.5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spatial-refine-only",
|
||||
action=StoreBoolean,
|
||||
default=SamplingParam.spatial_refine_only,
|
||||
help="Only perform spatial super-resolution (no temporal doubling)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-cond-frames",
|
||||
type=int,
|
||||
default=SamplingParam.num_cond_frames,
|
||||
help="Number of conditioning frames for refinement",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--moba-config-path",
|
||||
type=str,
|
||||
|
||||
@@ -319,6 +319,62 @@ class FastVideoArgs:
|
||||
"Path to a text file containing prompts (one per line) for batch processing",
|
||||
)
|
||||
|
||||
# LoRA parameters (inference-time adapter loading)
|
||||
parser.add_argument(
|
||||
"--lora-path",
|
||||
type=str,
|
||||
default=FastVideoArgs.lora_path,
|
||||
help=
|
||||
"Path to a LoRA adapter (directory or HF repo id). If set, LoRA will be applied at inference.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-nickname",
|
||||
type=str,
|
||||
default=FastVideoArgs.lora_nickname,
|
||||
help=
|
||||
"Nickname to refer to the loaded LoRA adapter (useful for swapping).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-target-modules",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=FastVideoArgs.lora_target_modules,
|
||||
help=
|
||||
"Optional list of module name substrings to restrict LoRA injection (e.g. q_proj k_proj v_proj).",
|
||||
)
|
||||
|
||||
# BSA runtime control (LongCat)
|
||||
parser.add_argument(
|
||||
"--enable-bsa",
|
||||
action=StoreBoolean,
|
||||
help=
|
||||
"Enable Block Sparse Attention (BSA) at runtime (overrides config).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bsa-sparsity",
|
||||
type=float,
|
||||
help="BSA sparsity (e.g., 0.9375).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bsa-cdf-threshold",
|
||||
type=float,
|
||||
help="BSA CDF threshold (optional).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bsa-chunk-q",
|
||||
nargs=3,
|
||||
type=int,
|
||||
metavar=("T", "H", "W"),
|
||||
help="BSA chunk_3d_shape_q as three ints, e.g., 4 4 4.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bsa-chunk-k",
|
||||
nargs=3,
|
||||
type=int,
|
||||
metavar=("T", "H", "W"),
|
||||
help="BSA chunk_3d_shape_k as three ints, e.g., 4 4 4.",
|
||||
)
|
||||
|
||||
# STA (Sliding Tile Attention) parameters
|
||||
parser.add_argument(
|
||||
"--STA-mode",
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
3D Rotary Position Embedding (RoPE) for video transformers.
|
||||
|
||||
Reference: https://arxiv.org/pdf/2104.09864.pdf
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange, repeat
|
||||
|
||||
|
||||
def broadcast(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
|
||||
"""
|
||||
Broadcast and concatenate tensors along a dimension.
|
||||
"""
|
||||
num_tensors = len(tensors)
|
||||
shape_lens = set(len(t.shape) for t in tensors)
|
||||
assert len(
|
||||
shape_lens) == 1, "tensors must all have the same number of dimensions"
|
||||
|
||||
shape_len = list(shape_lens)[0]
|
||||
dim = (dim + shape_len) if dim < 0 else dim
|
||||
|
||||
dims = list(zip(*[list(t.shape) for t in tensors], strict=False))
|
||||
expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
|
||||
|
||||
assert all(
|
||||
len(set(t[1])) <= 2 for t in
|
||||
expandable_dims), "invalid dimensions for broadcastable concatenation"
|
||||
|
||||
max_dims = [(t[0], max(t[1])) for t in expandable_dims]
|
||||
expanded_dims = [(t[0], (t[1], ) * num_tensors) for t in max_dims]
|
||||
expanded_dims.insert(dim, (dim, dims[dim]))
|
||||
expandable_shapes = list(zip(*[t[1] for t in expanded_dims], strict=False))
|
||||
tensors = [
|
||||
t[0].expand(*t[1])
|
||||
for t in zip(tensors, expandable_shapes, strict=False)
|
||||
]
|
||||
|
||||
return torch.cat(tensors, dim=dim)
|
||||
|
||||
|
||||
def rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Rotate half the hidden dims of the input.
|
||||
"""
|
||||
x = rearrange(x, "... (d r) -> ... d r", r=2)
|
||||
x1, x2 = x.unbind(dim=-1)
|
||||
x = torch.stack((-x2, x1), dim=-1)
|
||||
return rearrange(x, "... d r -> ... (d r)")
|
||||
|
||||
|
||||
class RotaryPositionalEmbedding3D(nn.Module):
|
||||
"""
|
||||
3D Rotary Positional Embedding for video transformers.
|
||||
|
||||
Splits the head dimension across temporal, height, and width dimensions,
|
||||
computing separate rotary embeddings for each and concatenating them.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
base: float = 10000.0,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
head_dim: Dimension of each attention head
|
||||
base: Base value for exponential frequency
|
||||
"""
|
||||
super().__init__()
|
||||
self.head_dim = head_dim
|
||||
assert self.head_dim % 8 == 0, "head_dim must be a multiple of 8 for 3D RoPE"
|
||||
self.base = base
|
||||
|
||||
# Cache for precomputed frequencies
|
||||
self.freqs_dict: dict[tuple, torch.Tensor] = {}
|
||||
|
||||
def register_grid_size(self, grid_size: tuple[int, int, int]) -> None:
|
||||
"""
|
||||
Precompute and register frequencies for a given grid size.
|
||||
|
||||
Args:
|
||||
grid_size: (T, H, W) tuple of grid dimensions
|
||||
"""
|
||||
if grid_size not in self.freqs_dict:
|
||||
self.freqs_dict[grid_size] = self.precompute_freqs_3d(grid_size)
|
||||
|
||||
def precompute_freqs_3d(self, grid_size: tuple[int, int,
|
||||
int]) -> torch.Tensor:
|
||||
"""
|
||||
Precompute 3D rotary frequencies.
|
||||
|
||||
Args:
|
||||
grid_size: (num_frames, height, width)
|
||||
|
||||
Returns:
|
||||
freqs: [T*H*W, head_dim] tensor of frequencies
|
||||
"""
|
||||
num_frames, height, width = grid_size
|
||||
|
||||
# Split head_dim across 3 dimensions
|
||||
# Temporal gets the remainder to ensure exact division
|
||||
dim_t = self.head_dim - 4 * (self.head_dim // 6)
|
||||
dim_h = 2 * (self.head_dim // 6)
|
||||
dim_w = 2 * (self.head_dim // 6)
|
||||
|
||||
# Compute frequency bands for each dimension
|
||||
freqs_t = 1.0 / (self.base**(
|
||||
torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t))
|
||||
freqs_h = 1.0 / (self.base**(
|
||||
torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h))
|
||||
freqs_w = 1.0 / (self.base**(
|
||||
torch.arange(0, dim_w, 2)[:(dim_w // 2)].float() / dim_w))
|
||||
|
||||
# Create position grids
|
||||
grid_t = torch.arange(num_frames, dtype=torch.float32)
|
||||
grid_h = torch.arange(height, dtype=torch.float32)
|
||||
grid_w = torch.arange(width, dtype=torch.float32)
|
||||
|
||||
# Compute frequencies for each position
|
||||
freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t)
|
||||
freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h)
|
||||
freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w)
|
||||
|
||||
# Duplicate for complex pair representation
|
||||
freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2)
|
||||
freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
|
||||
freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)
|
||||
|
||||
# Broadcast and concatenate across all 3 dimensions
|
||||
freqs = broadcast(
|
||||
[
|
||||
freqs_t[:, None, None, :], # [T, 1, 1, dim_t]
|
||||
freqs_h[None, :, None, :], # [1, H, 1, dim_h]
|
||||
freqs_w[None, None, :, :], # [1, 1, W, dim_w]
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
# Flatten spatial dimensions: [T, H, W, head_dim] -> [T*H*W, head_dim]
|
||||
freqs = rearrange(freqs, "T H W D -> (T H W) D")
|
||||
|
||||
return freqs
|
||||
|
||||
def forward(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
grid_size: tuple[int, int, int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply 3D rotary positional embedding to queries and keys.
|
||||
|
||||
Args:
|
||||
q: Query tensor [B, num_heads, seq_len, head_dim]
|
||||
k: Key tensor [B, num_heads, seq_len, head_dim]
|
||||
grid_size: (T, H, W) tuple of grid dimensions
|
||||
|
||||
Returns:
|
||||
(q_rotated, k_rotated): Rotated query and key tensors
|
||||
"""
|
||||
# Register grid size if not cached
|
||||
if grid_size not in self.freqs_dict:
|
||||
self.register_grid_size(grid_size)
|
||||
|
||||
# Get cached frequencies
|
||||
freqs_cis = self.freqs_dict[grid_size].to(q.device)
|
||||
|
||||
# Cast to float32 for precision
|
||||
q_, k_ = q.float(), k.float()
|
||||
freqs_cis = freqs_cis.float()
|
||||
|
||||
# Compute cos and sin
|
||||
cos = freqs_cis.cos()
|
||||
sin = freqs_cis.sin()
|
||||
|
||||
# Reshape for broadcasting: [1, 1, seq_len, head_dim]
|
||||
cos = rearrange(cos, "n d -> 1 1 n d")
|
||||
sin = rearrange(sin, "n d -> 1 1 n d")
|
||||
|
||||
# Apply rotation
|
||||
q_ = (q_ * cos) + (rotate_half(q_) * sin)
|
||||
k_ = (k_ * cos) + (rotate_half(k_) * sin)
|
||||
|
||||
# Cast back to original dtype
|
||||
return q_.type_as(q), k_.type_as(k)
|
||||
|
||||
|
||||
def apply_rotary_emb_3d(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
rope_module: RotaryPositionalEmbedding3D,
|
||||
grid_size: tuple[int, int, int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Convenience function to apply 3D RoPE.
|
||||
|
||||
Args:
|
||||
q: Query tensor [B, num_heads, seq_len, head_dim]
|
||||
k: Key tensor [B, num_heads, seq_len, head_dim]
|
||||
rope_module: RotaryPositionalEmbedding3D module
|
||||
grid_size: (T, H, W) grid dimensions
|
||||
|
||||
Returns:
|
||||
(q_rotated, k_rotated): Rotated tensors
|
||||
"""
|
||||
return rope_module(q, k, grid_size)
|
||||
@@ -0,0 +1,866 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Native LongCat Video DiT implementation using FastVideo conventions.
|
||||
|
||||
This is a Phase 2 reimplementation that replaces the third_party wrapper
|
||||
with native FastVideo layers for better performance and integration.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
from fastvideo.configs.models.dits import LongCatVideoConfig
|
||||
from fastvideo.layers.linear import ReplicatedLinear
|
||||
from fastvideo.layers.layernorm import RMSNorm, FP32LayerNorm
|
||||
from fastvideo.layers.activation import get_act_fn
|
||||
from fastvideo.layers.rotary_embedding_3d import RotaryPositionalEmbedding3D
|
||||
from fastvideo.attention.layer import DistributedAttention, LocalAttention
|
||||
from fastvideo.models.dits.base import CachableDiT
|
||||
from fastvideo.platforms import AttentionBackendEnum
|
||||
from fastvideo.third_party.longcat_video.block_sparse_attention.bsa_interface import flash_attn_bsa_3d
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Embeddings
|
||||
# ============================================================================
|
||||
|
||||
class PatchEmbed3D(nn.Module):
|
||||
"""
|
||||
3D patch embedding using Conv3d.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: tuple[int, int, int] = (1, 2, 2),
|
||||
in_channels: int = 16,
|
||||
embed_dim: int = 4096,
|
||||
):
|
||||
super().__init__()
|
||||
self.patch_size = patch_size
|
||||
self.in_channels = in_channels
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.proj = nn.Conv3d(
|
||||
in_channels,
|
||||
embed_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, C, T, H, W]
|
||||
Returns:
|
||||
[B, N, C] where N = (T/pt) * (H/ph) * (W/pw)
|
||||
"""
|
||||
# Padding if needed
|
||||
_, _, T, H, W = x.shape
|
||||
if W % self.patch_size[2] != 0:
|
||||
x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
|
||||
if H % self.patch_size[1] != 0:
|
||||
x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
|
||||
if T % self.patch_size[0] != 0:
|
||||
x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - T % self.patch_size[0]))
|
||||
|
||||
x = self.proj(x) # [B, C, T', H', W']
|
||||
x = x.flatten(2).transpose(1, 2) # [B, N, C]
|
||||
return x
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Sinusoidal timestep embedding + MLP projection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
frequency_embedding_size: int = 256,
|
||||
adaln_tembed_dim: int = 512,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
# Use FastVideo's ReplicatedLinear
|
||||
self.linear_1 = ReplicatedLinear(
|
||||
frequency_embedding_size,
|
||||
adaln_tembed_dim,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
self.act = nn.SiLU()
|
||||
self.linear_2 = ReplicatedLinear(
|
||||
adaln_tembed_dim,
|
||||
adaln_tembed_dim,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
"""
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
|
||||
)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
return embedding
|
||||
|
||||
def forward(self, t: torch.Tensor, latent_shape: tuple | None = None) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
t: [B] or [B, T] timesteps
|
||||
latent_shape: (T, H, W) for temporal expansion
|
||||
Returns:
|
||||
[B, T, C]
|
||||
"""
|
||||
# Sinusoidal embedding in FP32
|
||||
t_freq = self.timestep_embedding(t.flatten(), self.frequency_embedding_size)
|
||||
|
||||
# Cast to model dtype before MLP
|
||||
# Handle LoRA wrapper if present
|
||||
linear_layer = self.linear_1.base_layer if hasattr(self.linear_1, 'base_layer') else self.linear_1
|
||||
target_dtype = linear_layer.weight.dtype
|
||||
if t_freq.dtype != target_dtype:
|
||||
t_freq = t_freq.to(target_dtype)
|
||||
|
||||
# MLP projection
|
||||
t_emb, _ = self.linear_1(t_freq)
|
||||
t_emb = self.act(t_emb)
|
||||
t_emb, _ = self.linear_2(t_emb)
|
||||
|
||||
# Reshape if needed
|
||||
if latent_shape is not None and len(t.shape) > 1:
|
||||
B = t.shape[0]
|
||||
T = latent_shape[0]
|
||||
t_emb = t_emb.reshape(B, T, -1)
|
||||
|
||||
return t_emb
|
||||
|
||||
|
||||
class CaptionEmbedder(nn.Module):
|
||||
"""
|
||||
Caption embedding with MLP projection and optional text compaction.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
caption_channels: int = 4096,
|
||||
hidden_size: int = 4096,
|
||||
text_tokens_zero_pad: bool = True,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.text_tokens_zero_pad = text_tokens_zero_pad
|
||||
|
||||
# Two-layer MLP using ReplicatedLinear
|
||||
self.linear_1 = ReplicatedLinear(
|
||||
caption_channels,
|
||||
hidden_size,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
self.act = nn.SiLU()
|
||||
self.linear_2 = ReplicatedLinear(
|
||||
hidden_size,
|
||||
hidden_size,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_attention_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
encoder_hidden_states: [B, N_text, C_text] or [B, 1, N_text, C_text]
|
||||
encoder_attention_mask: [B, N_text] or [B, 1, 1, N_text]
|
||||
Returns:
|
||||
y: [B, N_text, C] - standard padded representation (like other models)
|
||||
"""
|
||||
# Handle extra dimension from wrapper
|
||||
if len(encoder_hidden_states.shape) == 4:
|
||||
encoder_hidden_states = encoder_hidden_states.squeeze(1)
|
||||
|
||||
# Project
|
||||
y, _ = self.linear_1(encoder_hidden_states)
|
||||
y = self.act(y)
|
||||
y, _ = self.linear_2(y) # [B, N_text, C]
|
||||
|
||||
# Handle attention masking - just zero out padded tokens if requested
|
||||
if encoder_attention_mask is not None:
|
||||
# Remove extra dimensions
|
||||
if len(encoder_attention_mask.shape) == 4:
|
||||
encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1)
|
||||
elif len(encoder_attention_mask.shape) == 3:
|
||||
encoder_attention_mask = encoder_attention_mask.squeeze(1)
|
||||
|
||||
# Zero out padded tokens if requested
|
||||
if self.text_tokens_zero_pad:
|
||||
y = y * encoder_attention_mask.unsqueeze(-1)
|
||||
|
||||
# Return standard format [B, N_text, C] - no compaction!
|
||||
return y
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Attention Modules (Placeholders for now)
|
||||
# ============================================================================
|
||||
|
||||
class LongCatSelfAttention(nn.Module):
|
||||
"""
|
||||
Self-attention with 3D RoPE support and optional BSA.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
config: LongCatVideoConfig,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
# Separate Q/K/V projections (not fused like original)
|
||||
self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
|
||||
# Per-head RMS normalization
|
||||
self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
|
||||
self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
|
||||
|
||||
# Output projection
|
||||
self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
|
||||
# 3D RoPE
|
||||
self.rope_3d = RotaryPositionalEmbedding3D(head_dim=self.head_dim)
|
||||
|
||||
# BSA configuration
|
||||
self.enable_bsa = getattr(config, 'enable_bsa', False)
|
||||
self.bsa_params = getattr(config, 'bsa_params', None)
|
||||
|
||||
# FastVideo attention backend (used when BSA is disabled)
|
||||
self.attn = DistributedAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
supported_attention_backends=config._supported_attention_backends,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor, # [B, N, C]
|
||||
latent_shape: tuple, # (T, H, W)
|
||||
**kwargs
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass with 3D RoPE and optional BSA.
|
||||
"""
|
||||
B, N, C = x.shape
|
||||
T, H, W = latent_shape
|
||||
|
||||
# Project to Q/K/V
|
||||
q, _ = self.to_q(x)
|
||||
k, _ = self.to_k(x)
|
||||
v, _ = self.to_v(x)
|
||||
|
||||
# Reshape to heads: [B, N, num_heads, head_dim]
|
||||
q = q.view(B, N, self.num_heads, self.head_dim)
|
||||
k = k.view(B, N, self.num_heads, self.head_dim)
|
||||
v = v.view(B, N, self.num_heads, self.head_dim)
|
||||
|
||||
# Per-head RMS normalization
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
# For RoPE: need [B, num_heads, N, head_dim]
|
||||
q_rope = q.transpose(1, 2)
|
||||
k_rope = k.transpose(1, 2)
|
||||
|
||||
# Apply 3D RoPE
|
||||
q_rope, k_rope = self.rope_3d(q_rope, k_rope, grid_size=latent_shape)
|
||||
|
||||
# Transpose back: [B, N, num_heads, head_dim] or [B, H, N, D] for BSA
|
||||
q = q_rope.transpose(1, 2)
|
||||
k = k_rope.transpose(1, 2)
|
||||
|
||||
# === Attention: BSA or standard ===
|
||||
if self.enable_bsa and T > 1: # Only use BSA for multi-frame videos
|
||||
# BSA expects [B, H, S, D] format
|
||||
q_bsa = q.transpose(1, 2).contiguous() # [B, num_heads, N, head_dim]
|
||||
k_bsa = k.transpose(1, 2).contiguous()
|
||||
v_bsa = v.transpose(1, 2).contiguous()
|
||||
|
||||
# Handle SP split: BSA operates on per-rank spatial dimensions
|
||||
# Replicate LongCat's cp_split_hw logic exactly
|
||||
from fastvideo.distributed.parallel_state import get_sp_world_size
|
||||
sp_size = get_sp_world_size()
|
||||
if sp_size > 1:
|
||||
# Calculate optimal 2D split (same as LongCat's get_optimal_split)
|
||||
factors = []
|
||||
for i in range(1, int(sp_size**0.5) + 1):
|
||||
if sp_size % i == 0:
|
||||
factors.append([i, sp_size // i])
|
||||
cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1]))
|
||||
|
||||
# Split H and W dimensions by their respective factors
|
||||
T_bsa, H_bsa, W_bsa = latent_shape
|
||||
assert H_bsa % cp_split_hw[0] == 0 and W_bsa % cp_split_hw[1] == 0, \
|
||||
f"H {H_bsa} must be divisible by {cp_split_hw[0]}, W {W_bsa} must be divisible by {cp_split_hw[1]}"
|
||||
H_bsa = H_bsa // cp_split_hw[0]
|
||||
W_bsa = W_bsa // cp_split_hw[1]
|
||||
latent_shape_bsa = (T_bsa, H_bsa, W_bsa)
|
||||
else:
|
||||
latent_shape_bsa = latent_shape
|
||||
|
||||
# Call BSA with per-rank latent shape
|
||||
out = flash_attn_bsa_3d(
|
||||
q_bsa, k_bsa, v_bsa,
|
||||
latent_shape_q=latent_shape_bsa,
|
||||
latent_shape_k=latent_shape_bsa,
|
||||
**self.bsa_params
|
||||
) # [B, num_heads, N, head_dim]
|
||||
|
||||
# Transpose back: [B, N, num_heads, head_dim]
|
||||
out = out.transpose(1, 2)
|
||||
else:
|
||||
# Standard attention: [B, N, num_heads, head_dim]
|
||||
out, _ = self.attn(q, k, v)
|
||||
|
||||
# Reshape and project out
|
||||
out = out.reshape(B, N, C)
|
||||
out, _ = self.to_out(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class LongCatCrossAttention(nn.Module):
|
||||
"""
|
||||
Cross-attention for text conditioning (standard implementation like other models).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
config: LongCatVideoConfig,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
|
||||
# Separate Q/K/V projections
|
||||
self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
|
||||
# Per-head RMS normalization
|
||||
self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
|
||||
self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
|
||||
|
||||
# Output projection
|
||||
self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
|
||||
|
||||
# Cross-attention uses LocalAttention (FastVideo standard)
|
||||
self.attn = LocalAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
dropout_rate=0,
|
||||
softmax_scale=None,
|
||||
causal=False,
|
||||
supported_attention_backends=config.arch_config._supported_attention_backends,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor, # [B, N_img, C]
|
||||
context: torch.Tensor, # [B, N_text, C]
|
||||
**kwargs
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass for cross-attention (standard implementation).
|
||||
|
||||
Args:
|
||||
x: Image tokens [B, N_img, C]
|
||||
context: Text tokens [B, N_text, C] (standard padded format)
|
||||
"""
|
||||
B, N_img, C = x.shape
|
||||
|
||||
# Project Q, K, V (standard cross-attention like WanVideo/StepVideo/Cosmos)
|
||||
q, _ = self.to_q(x)
|
||||
k, _ = self.to_k(context)
|
||||
v, _ = self.to_v(context)
|
||||
|
||||
N_text = context.shape[1]
|
||||
|
||||
# Reshape to heads
|
||||
q = q.view(B, N_img, self.num_heads, self.head_dim)
|
||||
k = k.view(B, N_text, self.num_heads, self.head_dim)
|
||||
v = v.view(B, N_text, self.num_heads, self.head_dim)
|
||||
|
||||
# Per-head RMS normalization
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
|
||||
# Run cross-attention using FastVideo's LocalAttention
|
||||
# LocalAttention handles different q and k/v sequence lengths automatically
|
||||
out = self.attn(q, k, v) # [B, N_img, num_heads, head_dim]
|
||||
|
||||
# Reshape and project out
|
||||
out = out.reshape(B, N_img, C)
|
||||
out, _ = self.to_out(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Feed-Forward Network
|
||||
# ============================================================================
|
||||
|
||||
class LongCatSwiGLUFFN(nn.Module):
|
||||
"""
|
||||
SwiGLU feed-forward network using FastVideo's ReplicatedLinear.
|
||||
|
||||
FFN(x) = down(gate(x) * SiLU(up(x)))
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Three projections for SwiGLU (no bias as per original)
|
||||
self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype) # gate
|
||||
self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype) # up
|
||||
self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False, params_dtype=dtype) # down
|
||||
|
||||
self.act = nn.SiLU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass: SiLU(w1(x)) * w3(x) -> w2 (matching original LongCat)
|
||||
"""
|
||||
w1_out, _ = self.w1(x)
|
||||
w3_out, _ = self.w3(x)
|
||||
combined = self.act(w1_out) * w3_out
|
||||
out, _ = self.w2(combined)
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Modulation Utilities
|
||||
# ============================================================================
|
||||
|
||||
def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Apply modulation in FP32 for numerical stability (matching original LongCat).
|
||||
|
||||
shift and scale should already be FP32 from torch.amp.autocast context.
|
||||
"""
|
||||
# Ensure modulation params are FP32 (should be from autocast)
|
||||
assert shift.dtype == torch.float32 and scale.dtype == torch.float32, \
|
||||
f"shift and scale must be FP32, got {shift.dtype} and {scale.dtype}"
|
||||
|
||||
orig_dtype = x.dtype
|
||||
|
||||
# Normalize and modulate in FP32
|
||||
x_norm = norm(x.to(torch.float32))
|
||||
x_mod = x_norm * (scale + 1) + shift
|
||||
|
||||
return x_mod.to(orig_dtype)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Transformer Block
|
||||
# ============================================================================
|
||||
|
||||
class LongCatTransformerBlock(nn.Module):
|
||||
"""
|
||||
Single-stream transformer block with:
|
||||
- AdaLN modulation (FP32)
|
||||
- Self-attention
|
||||
- Cross-attention
|
||||
- SwiGLU FFN
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: int,
|
||||
adaln_tembed_dim: int,
|
||||
config: LongCatVideoConfig,
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.num_heads = num_heads
|
||||
|
||||
# AdaLN modulation (6 parameters: scale/shift for attn & ffn, gate for residual)
|
||||
self.adaln_linear_1 = ReplicatedLinear(
|
||||
adaln_tembed_dim,
|
||||
6 * hidden_size,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
self.adaln_act = nn.SiLU()
|
||||
|
||||
# Normalization layers (CRITICAL: Use LayerNorm not RMSNorm like original!)
|
||||
# Original LongCat uses LayerNorm_FP32 with elementwise_affine=False
|
||||
self.norm_attn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
|
||||
self.norm_ffn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
|
||||
# Cross-attention norm has elementwise_affine=True (has weight and bias)
|
||||
self.norm_cross = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True)
|
||||
|
||||
# Self-attention
|
||||
self.self_attn = LongCatSelfAttention(
|
||||
dim=hidden_size,
|
||||
num_heads=num_heads,
|
||||
config=config,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Cross-attention
|
||||
self.cross_attn = LongCatCrossAttention(
|
||||
dim=hidden_size,
|
||||
num_heads=num_heads,
|
||||
config=config,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# SwiGLU FFN
|
||||
ffn_hidden_dim = int(hidden_size * mlp_ratio * 2 / 3)
|
||||
# Round up to nearest multiple of 256
|
||||
ffn_hidden_dim = 256 * ((ffn_hidden_dim + 255) // 256)
|
||||
|
||||
self.ffn = LongCatSwiGLUFFN(
|
||||
dim=hidden_size,
|
||||
hidden_dim=ffn_hidden_dim,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor, # [B, N, C]
|
||||
context: torch.Tensor, # [B, N_text, C]
|
||||
t: torch.Tensor, # [B, T, C_t]
|
||||
latent_shape: tuple, # (T, H, W)
|
||||
**kwargs
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass with AdaLN modulation.
|
||||
"""
|
||||
B, N, C = x.shape
|
||||
T, H, W = latent_shape
|
||||
x_orig_dtype = x.dtype # Save for later casting
|
||||
|
||||
# === AdaLN Modulation (CRITICAL: FP32 for stability like original) ===
|
||||
# Use autocast to compute modulation params in FP32
|
||||
with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
|
||||
t_mod = self.adaln_act(t)
|
||||
mod_params, _ = self.adaln_linear_1(t_mod)
|
||||
# Ensure FP32 output (needed when LoRA is applied)
|
||||
if mod_params.dtype != torch.float32:
|
||||
mod_params = mod_params.float()
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
|
||||
mod_params.unsqueeze(2).chunk(6, dim=-1) # [B, T, 1, C]
|
||||
|
||||
# === Self-Attention ===
|
||||
x_norm = modulate_fp32(self.norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa)
|
||||
x_norm = x_norm.view(B, N, C)
|
||||
|
||||
attn_out = self.self_attn(x_norm, latent_shape=latent_shape)
|
||||
|
||||
# Residual with gating (CRITICAL: FP32 like original, then cast back)
|
||||
with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
|
||||
x = x + (gate_msa * attn_out.view(B, T, -1, C)).view(B, N, C)
|
||||
x = x.to(x_orig_dtype)
|
||||
|
||||
# === Cross-Attention ===
|
||||
x_norm_cross = self.norm_cross(x)
|
||||
cross_out = self.cross_attn(x_norm_cross, context)
|
||||
x = x + cross_out
|
||||
|
||||
# === FFN ===
|
||||
x_norm_ffn = modulate_fp32(self.norm_ffn, x.view(B, T, -1, C), shift_mlp, scale_mlp)
|
||||
x_norm_ffn = x_norm_ffn.view(B, N, C)
|
||||
|
||||
ffn_out = self.ffn(x_norm_ffn)
|
||||
|
||||
# Residual with gating (CRITICAL: FP32 like original, then cast back)
|
||||
with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
|
||||
x = x + (gate_mlp * ffn_out.view(B, T, -1, C)).view(B, N, C)
|
||||
x = x.to(x_orig_dtype)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Final Layer
|
||||
# ============================================================================
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
Final output projection with AdaLN modulation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
out_channels: int,
|
||||
adaln_tembed_dim: int,
|
||||
patch_size: tuple[int, int, int],
|
||||
dtype: torch.dtype | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# AdaLN for final layer (2 parameters: scale and shift)
|
||||
self.adaln_linear = ReplicatedLinear(
|
||||
adaln_tembed_dim,
|
||||
2 * hidden_size,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
self.adaln_act = nn.SiLU()
|
||||
|
||||
# CRITICAL: Use LayerNorm not RMSNorm! (matches original)
|
||||
self.norm = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
|
||||
|
||||
# Output projection
|
||||
num_patch = patch_size[0] * patch_size[1] * patch_size[2]
|
||||
self.proj = ReplicatedLinear(
|
||||
hidden_size,
|
||||
num_patch * out_channels,
|
||||
bias=True,
|
||||
params_dtype=dtype,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor, # [B, N, C]
|
||||
t: torch.Tensor, # [B, T, C_t]
|
||||
latent_shape: tuple,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Returns: [B, N, out_channels * patch_size^3]
|
||||
"""
|
||||
B, N, C = x.shape
|
||||
T, _, _ = latent_shape
|
||||
|
||||
# AdaLN modulation (FP32 for stability like original)
|
||||
with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
|
||||
t_mod = self.adaln_act(t)
|
||||
mod_params, _ = self.adaln_linear(t_mod)
|
||||
# Ensure FP32 output (needed when LoRA is applied)
|
||||
if mod_params.dtype != torch.float32:
|
||||
mod_params = mod_params.float()
|
||||
shift, scale = mod_params.unsqueeze(2).chunk(2, dim=-1)
|
||||
|
||||
# Modulate
|
||||
x = modulate_fp32(self.norm, x.view(B, T, -1, C), shift, scale)
|
||||
x = x.reshape(B, N, C)
|
||||
|
||||
# Project
|
||||
x, _ = self.proj(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Main Model
|
||||
# ============================================================================
|
||||
|
||||
class LongCatTransformer3DModel(CachableDiT):
|
||||
"""
|
||||
Native LongCat Video Transformer using FastVideo layers.
|
||||
|
||||
This is a Phase 2 implementation that replaces third_party dependencies.
|
||||
"""
|
||||
|
||||
# FSDP sharding: shard at each transformer block
|
||||
_fsdp_shard_conditions = [
|
||||
lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(),
|
||||
]
|
||||
# torch.compile optimization: compile each transformer block for speedup
|
||||
_compile_conditions = [
|
||||
lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(),
|
||||
]
|
||||
|
||||
# Parameter name mapping (for weight conversion)
|
||||
param_names_mapping = {} # Will be defined in config
|
||||
reverse_param_names_mapping = {}
|
||||
lora_param_names_mapping = {}
|
||||
|
||||
# Supported attention backends
|
||||
_supported_attention_backends = (
|
||||
AttentionBackendEnum.FLASH_ATTN,
|
||||
AttentionBackendEnum.TORCH_SDPA,
|
||||
)
|
||||
|
||||
def __init__(self, config: LongCatVideoConfig, hf_config: dict[str, Any]):
|
||||
super().__init__(config=config, hf_config=hf_config)
|
||||
|
||||
# Extract architecture parameters
|
||||
self.hidden_size = config.hidden_size # 4096
|
||||
self.num_attention_heads = config.num_attention_heads # 32
|
||||
self.depth = config.depth # 48
|
||||
self.mlp_ratio = config.mlp_ratio # 4
|
||||
self.in_channels = config.in_channels # 16
|
||||
self.out_channels = config.out_channels # 16
|
||||
self.num_channels_latents = config.in_channels
|
||||
self.patch_size = config.patch_size # [1, 2, 2]
|
||||
|
||||
# Embeddings
|
||||
self.patch_embed = PatchEmbed3D(
|
||||
patch_size=self.patch_size,
|
||||
in_channels=self.in_channels,
|
||||
embed_dim=self.hidden_size,
|
||||
)
|
||||
|
||||
self.time_embedder = TimestepEmbedder(
|
||||
frequency_embedding_size=config.frequency_embedding_size,
|
||||
adaln_tembed_dim=config.adaln_tembed_dim,
|
||||
)
|
||||
|
||||
self.caption_embedder = CaptionEmbedder(
|
||||
caption_channels=config.caption_channels,
|
||||
hidden_size=self.hidden_size,
|
||||
text_tokens_zero_pad=getattr(config, 'text_tokens_zero_pad', True),
|
||||
)
|
||||
|
||||
# Transformer blocks (48 blocks)
|
||||
self.blocks = nn.ModuleList([
|
||||
LongCatTransformerBlock(
|
||||
hidden_size=self.hidden_size,
|
||||
num_heads=self.num_attention_heads,
|
||||
mlp_ratio=self.mlp_ratio,
|
||||
adaln_tembed_dim=config.adaln_tembed_dim,
|
||||
config=config,
|
||||
)
|
||||
for _ in range(self.depth)
|
||||
])
|
||||
|
||||
# Output projection
|
||||
self.final_layer = FinalLayer(
|
||||
hidden_size=self.hidden_size,
|
||||
out_channels=self.out_channels,
|
||||
adaln_tembed_dim=config.adaln_tembed_dim,
|
||||
patch_size=self.patch_size,
|
||||
)
|
||||
|
||||
def enable_bsa(self):
|
||||
"""Enable BSA for all self-attention layers."""
|
||||
for block in self.blocks:
|
||||
block.self_attn.enable_bsa = True
|
||||
|
||||
def disable_bsa(self):
|
||||
"""Disable BSA for all self-attention layers."""
|
||||
for block in self.blocks:
|
||||
block.self_attn.enable_bsa = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor, # [B, C, T, H, W]
|
||||
encoder_hidden_states: torch.Tensor | list[torch.Tensor], # [B, N_text, C_text]
|
||||
timestep: torch.LongTensor, # [B] or [B, T]
|
||||
encoder_attention_mask: torch.Tensor | None = None, # [B, N_text]
|
||||
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
|
||||
guidance: float | None = None, # Unused, for API compatibility
|
||||
**kwargs
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass with FastVideo parameter ordering.
|
||||
|
||||
NOTE: This follows FastVideo convention:
|
||||
(hidden_states, encoder_hidden_states, timestep)
|
||||
"""
|
||||
B, _, T, H, W = hidden_states.shape
|
||||
|
||||
N_t = T // self.patch_size[0]
|
||||
N_h = H // self.patch_size[1]
|
||||
N_w = W // self.patch_size[2]
|
||||
|
||||
# Handle list of encoder outputs (take first one)
|
||||
if isinstance(encoder_hidden_states, list):
|
||||
encoder_hidden_states = encoder_hidden_states[0]
|
||||
|
||||
# 1. Patch embedding
|
||||
x = self.patch_embed(hidden_states) # [B, N, C]
|
||||
|
||||
# 2. Timestep embedding
|
||||
# Expand timestep from [B] to [B, T] if needed
|
||||
if timestep.ndim == 1:
|
||||
timestep = timestep.unsqueeze(1).expand(-1, N_t) # [B, T]
|
||||
|
||||
t = self.time_embedder(timestep.flatten(), latent_shape=(N_t, N_h, N_w))
|
||||
if t.ndim == 2:
|
||||
t = t.reshape(B, N_t, -1) # [B, T, C_t]
|
||||
|
||||
# 3. Caption embedding (standard format, no compaction)
|
||||
context = self.caption_embedder(
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask
|
||||
) # [B, N_text, C]
|
||||
|
||||
# 4. Transformer blocks
|
||||
for i, block in enumerate(self.blocks):
|
||||
x = block(
|
||||
x, context, t,
|
||||
latent_shape=(N_t, N_h, N_w)
|
||||
)
|
||||
|
||||
# 5. Output projection
|
||||
output = self.final_layer(x, t, latent_shape=(N_t, N_h, N_w))
|
||||
|
||||
# Reshape to [B, C_out, T, H, W]
|
||||
output = self.unpatchify(output, N_t, N_h, N_w)
|
||||
|
||||
# Cast to float32 for better accuracy (as per original)
|
||||
output = output.to(torch.float32)
|
||||
|
||||
return output
|
||||
|
||||
def unpatchify(self, x: torch.Tensor, N_t: int, N_h: int, N_w: int) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
x: [B, N, C] where C = T_p * H_p * W_p * C_out
|
||||
Returns:
|
||||
[B, C_out, T, H, W]
|
||||
"""
|
||||
T_p, H_p, W_p = self.patch_size
|
||||
x = rearrange(
|
||||
x,
|
||||
"B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)",
|
||||
N_t=N_t,
|
||||
N_h=N_h,
|
||||
N_w=N_w,
|
||||
T_p=T_p,
|
||||
H_p=H_p,
|
||||
W_p=W_p,
|
||||
C_out=self.out_channels,
|
||||
)
|
||||
return x
|
||||
|
||||
@@ -16,6 +16,7 @@ import torch.nn as nn
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
from torch.distributed import init_device_mesh
|
||||
from transformers import AutoImageProcessor, AutoTokenizer
|
||||
from transformers import UMT5EncoderModel
|
||||
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
|
||||
|
||||
from fastvideo.configs.models import EncoderConfig
|
||||
@@ -405,7 +406,7 @@ class VAELoader(ComponentLoader):
|
||||
target_device = get_local_torch_device()
|
||||
|
||||
with set_default_torch_dtype(PRECISION_TO_TYPE[
|
||||
fastvideo_args.pipeline_config.vae_precision]):
|
||||
fastvideo_args.pipeline_config.vae_precision] if fastvideo_args.pipeline_config.vae_precision else torch.bfloat16):
|
||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
vae = vae_cls(vae_config).to(target_device)
|
||||
|
||||
|
||||
@@ -29,7 +29,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
|
||||
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
|
||||
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
|
||||
"StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
|
||||
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel")
|
||||
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
|
||||
"LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1)
|
||||
"LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2)
|
||||
}
|
||||
|
||||
_IMAGE_TO_VIDEO_DIT_MODELS = {
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""LongCat pipeline module."""
|
||||
|
||||
from fastvideo.pipelines.basic.longcat.longcat_pipeline import LongCatPipeline
|
||||
|
||||
__all__ = ["LongCatPipeline"]
|
||||
@@ -0,0 +1,145 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
LongCat video diffusion pipeline implementation (Phase 1: Wrapper).
|
||||
|
||||
This module contains a wrapper implementation of the LongCat video diffusion pipeline
|
||||
using FastVideo's modular pipeline architecture with the original LongCat modules.
|
||||
"""
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from fastvideo.pipelines.stages import (
|
||||
DecodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage
|
||||
from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage
|
||||
from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class LongCatPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""
|
||||
LongCat video diffusion pipeline with LoRA support.
|
||||
|
||||
Phase 1 implementation using wrapper modules from third_party/longcat_video.
|
||||
This validates the pipeline infrastructure before full FastVideo integration.
|
||||
"""
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder", "tokenizer", "vae", "transformer", "scheduler"
|
||||
]
|
||||
|
||||
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
"""Initialize LongCat-specific components."""
|
||||
|
||||
# Enable BSA (Block Sparse Attention) if configured
|
||||
pipeline_config = fastvideo_args.pipeline_config
|
||||
transformer = self.get_module("transformer", None)
|
||||
if transformer is None:
|
||||
raise RuntimeError(
|
||||
"Transformer module not found during initializing LongCat pipeline."
|
||||
)
|
||||
# If user toggles BSA via CLI/config
|
||||
if pipeline_config.enable_bsa:
|
||||
# Build effective BSA params:
|
||||
# 1) from explicit CLI overrides if provided
|
||||
# 2) else from pipeline_config.bsa_params
|
||||
# 3) else fall back to reasonable defaults
|
||||
bsa_params_cfg = pipeline_config.bsa_params
|
||||
sparsity = pipeline_config.bsa_sparsity
|
||||
cdf_threshold = pipeline_config.bsa_cdf_threshold
|
||||
chunk_q = pipeline_config.bsa_chunk_q
|
||||
chunk_k = pipeline_config.bsa_chunk_k
|
||||
|
||||
effective_bsa_params = dict(bsa_params_cfg) if isinstance(
|
||||
bsa_params_cfg, dict) else {}
|
||||
if sparsity is not None:
|
||||
effective_bsa_params['sparsity'] = sparsity
|
||||
if cdf_threshold is not None:
|
||||
effective_bsa_params['cdf_threshold'] = cdf_threshold
|
||||
if chunk_q is not None:
|
||||
effective_bsa_params['chunk_3d_shape_q'] = chunk_q
|
||||
if chunk_k is not None:
|
||||
effective_bsa_params['chunk_3d_shape_k'] = chunk_k
|
||||
# Provide defaults if still missing
|
||||
effective_bsa_params.setdefault('sparsity', 0.9375)
|
||||
effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4])
|
||||
effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4])
|
||||
|
||||
if hasattr(transformer, 'enable_bsa'):
|
||||
logger.info(
|
||||
"Enabling Block Sparse Attention (BSA) for LongCat transformer"
|
||||
)
|
||||
transformer.enable_bsa()
|
||||
# Propagate params to all attention modules
|
||||
if hasattr(transformer, 'blocks'):
|
||||
try:
|
||||
for blk in transformer.blocks:
|
||||
if hasattr(blk, 'self_attn'):
|
||||
blk.self_attn.bsa_params = effective_bsa_params
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to set BSA params on all blocks: %s", e)
|
||||
logger.info("BSA parameters in effect: %s",
|
||||
effective_bsa_params)
|
||||
else:
|
||||
logger.warning(
|
||||
"BSA is enabled in config but transformer does not support it"
|
||||
)
|
||||
else:
|
||||
# Explicitly disable if present
|
||||
if hasattr(transformer, 'disable_bsa'):
|
||||
transformer.disable_bsa()
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(stage_name="input_validation_stage",
|
||||
stage=InputValidationStage())
|
||||
|
||||
self.add_stage(stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
))
|
||||
|
||||
# Add refine initialization stage (will be skipped if not refining)
|
||||
self.add_stage(stage_name="longcat_refine_init_stage",
|
||||
stage=LongCatRefineInitStage(vae=self.get_module("vae")))
|
||||
|
||||
# First prepare generic timesteps (for non-refine paths)
|
||||
self.add_stage(stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(
|
||||
scheduler=self.get_module("scheduler")))
|
||||
|
||||
# Then override timesteps for refinement (will be a no-op if not refining),
|
||||
# matching LongCat's generate_refine schedule.
|
||||
self.add_stage(stage_name="longcat_refine_timestep_stage",
|
||||
stage=LongCatRefineTimestepStage(
|
||||
scheduler=self.get_module("scheduler")))
|
||||
|
||||
self.add_stage(stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None)))
|
||||
|
||||
self.add_stage(stage_name="denoising_stage",
|
||||
stage=LongCatDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
transformer_2=self.get_module("transformer_2", None),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
vae=self.get_module("vae"),
|
||||
pipeline=self))
|
||||
|
||||
self.add_stage(stage_name="decoding_stage",
|
||||
stage=DecodingStage(vae=self.get_module("vae"),
|
||||
pipeline=self))
|
||||
|
||||
|
||||
EntryClass = LongCatPipeline
|
||||
@@ -287,6 +287,7 @@ class ComposedPipelineBase(ABC):
|
||||
# remove keys that are not pipeline modules
|
||||
model_index.pop("_class_name")
|
||||
model_index.pop("_diffusers_version")
|
||||
model_index.pop("workload_type", None)
|
||||
if "boundary_ratio" in model_index and model_index[
|
||||
"boundary_ratio"] is not None:
|
||||
logger.info(
|
||||
|
||||
@@ -91,6 +91,14 @@ class ForwardBatch:
|
||||
video_path: str | None = None
|
||||
video_latent: torch.Tensor | None = None
|
||||
|
||||
# Refine inputs (LongCat)
|
||||
refine_from: str | None = None
|
||||
t_thresh: float = 0.5
|
||||
spatial_refine_only: bool = False
|
||||
num_cond_frames: int = 0
|
||||
stage1_video: list[
|
||||
PIL.Image.Image] | None = None # Loaded frames from refine_from
|
||||
|
||||
# Primary encoder embeddings
|
||||
prompt_embeds: list[torch.Tensor] = field(default_factory=list)
|
||||
negative_prompt_embeds: list[torch.Tensor] | None = None
|
||||
|
||||
@@ -29,6 +29,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
|
||||
"Cosmos2VideoToWorldPipeline": "cosmos",
|
||||
"MatrixGamePipeline": "matrixgame",
|
||||
"MatrixGameCausalDMDPipeline": "matrixgame",
|
||||
"LongCatPipeline": "longcat",
|
||||
}
|
||||
|
||||
_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
|
||||
|
||||
@@ -187,6 +187,21 @@ class DecodingStage(PipelineStage):
|
||||
# Convert to CPU float32 for compatibility
|
||||
frames = frames.cpu().float()
|
||||
|
||||
# Crop padding if this is a LongCat refinement
|
||||
if hasattr(batch, 'num_cond_frames_added') and hasattr(
|
||||
batch, 'new_frame_size_before_padding'):
|
||||
num_cond_frames_added = batch.num_cond_frames_added
|
||||
new_frame_size = batch.new_frame_size_before_padding
|
||||
if num_cond_frames_added > 0 or frames.shape[2] != new_frame_size:
|
||||
# frames is [B, C, T, H, W], crop temporal dimension
|
||||
frames = frames[:, :,
|
||||
num_cond_frames_added:num_cond_frames_added +
|
||||
new_frame_size, :, :]
|
||||
logger.info(
|
||||
"Cropped LongCat refinement padding: %s:%s, final shape: %s",
|
||||
num_cond_frames_added,
|
||||
num_cond_frames_added + new_frame_size, frames.shape)
|
||||
|
||||
# Update batch with decoded image
|
||||
batch.output = frames
|
||||
|
||||
|
||||
@@ -136,16 +136,26 @@ class LatentPreparationStage(PipelineStage):
|
||||
)
|
||||
# Generate or use provided latents
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape,
|
||||
generator=generator,
|
||||
device=device,
|
||||
dtype=dtype)
|
||||
latents = randn_tensor(
|
||||
shape,
|
||||
generator=generator,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
if hasattr(self.scheduler, "init_noise_sigma"):
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
else:
|
||||
# Pre-initialized latents:
|
||||
# - For LongCat refine (refine_from or stage1_video present), we should not re-scale by init_noise_sigma.
|
||||
# - For other models, keep the original behavior.
|
||||
latents = latents.to(device)
|
||||
is_longcat_refine = (batch.refine_from
|
||||
is not None) or (batch.stage1_video
|
||||
is not None)
|
||||
if (not is_longcat_refine) and hasattr(self.scheduler,
|
||||
"init_noise_sigma"):
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
|
||||
# Scale the initial noise if needed
|
||||
if hasattr(self.scheduler, "init_noise_sigma"):
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
# Update batch with prepared latents
|
||||
batch.latents = latents
|
||||
batch.raw_latent_shape = bcthw_shape
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
LongCat-specific denoising stage implementing CFG-zero optimized guidance.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.denoising import DenoisingStage
|
||||
from fastvideo.forward_context import set_forward_context
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class LongCatDenoisingStage(DenoisingStage):
|
||||
"""
|
||||
LongCat denoising stage with CFG-zero optimized guidance scale.
|
||||
|
||||
Implements:
|
||||
1. Optimized CFG scale from CFG-zero paper
|
||||
2. Negation of noise prediction before scheduler step (flow matching convention)
|
||||
3. Batched CFG computation (unlike standard FastVideo separate passes)
|
||||
"""
|
||||
|
||||
def optimized_scale(self, positive_flat, negative_flat) -> torch.Tensor:
|
||||
"""
|
||||
Calculate optimized scale from CFG-zero paper.
|
||||
|
||||
st_star = (v_cond^T * v_uncond) / ||v_uncond||^2
|
||||
|
||||
Args:
|
||||
positive_flat: Conditional prediction, flattened [B, -1]
|
||||
negative_flat: Unconditional prediction, flattened [B, -1]
|
||||
|
||||
Returns:
|
||||
st_star: Optimized scale [B, 1]
|
||||
"""
|
||||
# Calculate dot product
|
||||
dot_product = torch.sum(positive_flat * negative_flat,
|
||||
dim=1,
|
||||
keepdim=True)
|
||||
# Squared norm of uncondition
|
||||
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
|
||||
# st_star = v_cond^T * v_uncond / ||v_uncond||^2
|
||||
st_star = dot_product / squared_norm
|
||||
return st_star
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> ForwardBatch:
|
||||
"""
|
||||
Run LongCat denoising loop with optimized CFG.
|
||||
|
||||
Args:
|
||||
batch: The current batch information.
|
||||
fastvideo_args: The inference arguments.
|
||||
|
||||
Returns:
|
||||
The batch with denoised latents.
|
||||
"""
|
||||
if not fastvideo_args.model_loaded["transformer"]:
|
||||
from fastvideo.models.model_loader import TransformerLoader
|
||||
loader = TransformerLoader()
|
||||
self.transformer = loader.load(
|
||||
fastvideo_args.model_paths["transformer"], fastvideo_args)
|
||||
pipeline = self.pipeline() if self.pipeline else None
|
||||
if pipeline:
|
||||
pipeline.add_module("transformer", self.transformer)
|
||||
fastvideo_args.model_loaded["transformer"] = True
|
||||
|
||||
# Get transformer dtype
|
||||
if hasattr(self.transformer, 'module'):
|
||||
transformer_dtype = next(self.transformer.module.parameters()).dtype
|
||||
else:
|
||||
transformer_dtype = next(self.transformer.parameters()).dtype
|
||||
|
||||
target_dtype = transformer_dtype
|
||||
autocast_enabled = (target_dtype != torch.float32
|
||||
) and not fastvideo_args.disable_autocast
|
||||
|
||||
# Extract batch parameters
|
||||
latents = batch.latents
|
||||
timesteps = batch.timesteps
|
||||
prompt_embeds = batch.prompt_embeds[0] # LongCat uses single encoder
|
||||
prompt_attention_mask = batch.prompt_attention_mask[
|
||||
0] if batch.prompt_attention_mask else None
|
||||
guidance_scale = batch.guidance_scale
|
||||
do_classifier_free_guidance = batch.do_classifier_free_guidance
|
||||
|
||||
# Get negative prompts if doing CFG
|
||||
if do_classifier_free_guidance:
|
||||
negative_prompt_embeds = batch.negative_prompt_embeds[0]
|
||||
negative_prompt_attention_mask = (batch.negative_attention_mask[0]
|
||||
if batch.negative_attention_mask
|
||||
else None)
|
||||
# Concatenate for batched processing
|
||||
prompt_embeds_combined = torch.cat(
|
||||
[negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
if prompt_attention_mask is not None:
|
||||
prompt_attention_mask_combined = torch.cat(
|
||||
[negative_prompt_attention_mask, prompt_attention_mask],
|
||||
dim=0)
|
||||
else:
|
||||
prompt_attention_mask_combined = None
|
||||
else:
|
||||
prompt_embeds_combined = prompt_embeds
|
||||
prompt_attention_mask_combined = prompt_attention_mask
|
||||
|
||||
# Denoising loop
|
||||
num_inference_steps = len(timesteps)
|
||||
with tqdm(total=num_inference_steps,
|
||||
desc="LongCat Denoising") as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# Expand latents for CFG
|
||||
if do_classifier_free_guidance:
|
||||
latent_model_input = torch.cat([latents] * 2)
|
||||
else:
|
||||
latent_model_input = latents
|
||||
|
||||
latent_model_input = latent_model_input.to(target_dtype)
|
||||
|
||||
# Expand timestep to match batch size
|
||||
timestep = t.expand(
|
||||
latent_model_input.shape[0]).to(target_dtype)
|
||||
|
||||
# Run transformer with context
|
||||
batch.is_cfg_negative = False
|
||||
with set_forward_context(
|
||||
current_timestep=i,
|
||||
attn_metadata=None,
|
||||
forward_batch=batch,
|
||||
), torch.autocast(device_type='cuda',
|
||||
dtype=target_dtype,
|
||||
enabled=autocast_enabled):
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=latent_model_input,
|
||||
encoder_hidden_states=prompt_embeds_combined,
|
||||
timestep=timestep,
|
||||
encoder_attention_mask=prompt_attention_mask_combined,
|
||||
)
|
||||
|
||||
# Apply CFG with optimized scale
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
|
||||
|
||||
B = noise_pred_cond.shape[0]
|
||||
positive = noise_pred_cond.reshape(B, -1)
|
||||
negative = noise_pred_uncond.reshape(B, -1)
|
||||
|
||||
# Calculate optimized scale (CFG-zero)
|
||||
st_star = self.optimized_scale(positive, negative)
|
||||
|
||||
# Reshape for broadcasting
|
||||
st_star = st_star.view(B, 1, 1, 1, 1)
|
||||
|
||||
# Apply optimized CFG formula
|
||||
noise_pred = (
|
||||
noise_pred_uncond * st_star + guidance_scale *
|
||||
(noise_pred_cond - noise_pred_uncond * st_star))
|
||||
|
||||
# CRITICAL: Negate noise prediction for flow matching scheduler
|
||||
noise_pred = -noise_pred
|
||||
|
||||
# Compute previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred,
|
||||
t,
|
||||
latents,
|
||||
return_dict=False)[0]
|
||||
|
||||
progress_bar.update()
|
||||
|
||||
# Update batch with denoised latents
|
||||
batch.latents = latents
|
||||
return batch
|
||||
@@ -0,0 +1,310 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
LongCat refinement initialization stage.
|
||||
|
||||
This stage prepares the latent variables for LongCat's 480p->720p refinement by:
|
||||
1. Loading the stage1 (480p) video
|
||||
2. Upsampling it to 720p resolution
|
||||
3. Encoding it with VAE
|
||||
4. Mixing with noise according to t_thresh
|
||||
"""
|
||||
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.vision_utils import load_video
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.base import PipelineStage
|
||||
from fastvideo.configs.pipelines.longcat import get_bucket_config
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class LongCatRefineInitStage(PipelineStage):
|
||||
"""
|
||||
Stage for initializing LongCat refinement from a stage1 (480p) video.
|
||||
|
||||
This replicates the logic from LongCatVideoPipeline.generate_refine():
|
||||
- Load stage1_video frames
|
||||
- Upsample spatially and temporally
|
||||
- VAE encode and normalize
|
||||
- Mix with noise according to t_thresh
|
||||
"""
|
||||
|
||||
def __init__(self, vae) -> None:
|
||||
super().__init__()
|
||||
self.vae = vae
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> ForwardBatch:
|
||||
"""
|
||||
Initialize latents for refinement.
|
||||
|
||||
Args:
|
||||
batch: The current batch information.
|
||||
fastvideo_args: The inference arguments.
|
||||
|
||||
Returns:
|
||||
The batch with initialized latents for refinement.
|
||||
"""
|
||||
refine_from = batch.refine_from
|
||||
in_memory_stage1 = batch.stage1_video
|
||||
|
||||
# Only run for refinement tasks: either a path (refine_from) or in-memory video is provided
|
||||
if refine_from is None and in_memory_stage1 is None:
|
||||
# Not a refinement task, skip
|
||||
return batch
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 1. Obtain stage1 frames (either from disk or from in-memory input)
|
||||
# ------------------------------------------------------------------
|
||||
if in_memory_stage1 is not None:
|
||||
# User provided stage1 frames directly (e.g., from distilled stage output)
|
||||
if len(in_memory_stage1) == 0:
|
||||
raise ValueError(
|
||||
"stage1_video is empty; expected a non-empty list of frames"
|
||||
)
|
||||
|
||||
if isinstance(in_memory_stage1[0], Image.Image):
|
||||
pil_images = in_memory_stage1
|
||||
else:
|
||||
# Assume numpy arrays or torch tensors with shape [H, W, C]
|
||||
pil_images = [
|
||||
Image.fromarray(np.array(frame))
|
||||
for frame in in_memory_stage1
|
||||
]
|
||||
|
||||
logger.info(
|
||||
"Initializing LongCat refinement from in-memory stage1_video (%s frames)",
|
||||
len(pil_images))
|
||||
else:
|
||||
# Path-based refine: load video from disk (original design)
|
||||
logger.info("Initializing LongCat refinement from file: %s",
|
||||
refine_from)
|
||||
stage1_video_path = Path(refine_from)
|
||||
if not stage1_video_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Stage1 video not found: {refine_from}")
|
||||
|
||||
# Load video frames as PIL Images
|
||||
pil_images, original_fps = load_video(str(stage1_video_path),
|
||||
return_fps=True)
|
||||
logger.info("Loaded stage1 video: %s frames @ %s fps",
|
||||
len(pil_images), original_fps)
|
||||
|
||||
# Store in batch for reference (use PIL images, same as official demo)
|
||||
batch.stage1_video = pil_images
|
||||
|
||||
# Get parameters from batch
|
||||
num_frames = len(pil_images)
|
||||
spatial_refine_only = batch.spatial_refine_only
|
||||
t_thresh = batch.t_thresh
|
||||
num_cond_frames = batch.num_cond_frames if hasattr(
|
||||
batch, 'num_cond_frames') else 0
|
||||
|
||||
# Calculate new frame count (temporal upsampling if not spatial_refine_only)
|
||||
new_num_frames = num_frames if spatial_refine_only else 2 * num_frames
|
||||
logger.info(
|
||||
"Refine mode: %s",
|
||||
'spatial only' if spatial_refine_only else 'spatial + temporal')
|
||||
|
||||
# Update batch.num_frames to reflect the upsampled count
|
||||
batch.num_frames = new_num_frames
|
||||
|
||||
# Use bucket system to select resolution (exactly like LongCat)
|
||||
# Calculate scale_factor_spatial considering SP split
|
||||
sp_size = fastvideo_args.sp_size if fastvideo_args.sp_size > 0 else 1
|
||||
vae_scale_factor_spatial = 8 # VAE spatial downsampling
|
||||
patch_size_spatial = 2 # LongCat patch size
|
||||
bsa_latent_granularity = 4
|
||||
scale_factor_spatial = vae_scale_factor_spatial * patch_size_spatial * bsa_latent_granularity # 64
|
||||
|
||||
# Calculate optimal split like LongCat (cp_split_hw logic)
|
||||
# For sp_size=1: [1,1], max=1
|
||||
# For sp_size=2: [1,2], max=2
|
||||
# For sp_size=4: [2,2], max=2
|
||||
# For sp_size=8: [2,4], max=4
|
||||
if sp_size > 1:
|
||||
# Get optimal 2D split factors (mimic context_parallel_util.get_optimal_split)
|
||||
factors = []
|
||||
for i in range(1, int(sp_size**0.5) + 1):
|
||||
if sp_size % i == 0:
|
||||
factors.append([i, sp_size // i])
|
||||
cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1]))
|
||||
scale_factor_spatial *= max(cp_split_hw)
|
||||
logger.info("SP split: sp_size=%s, cp_split_hw=%s, max_split=%s",
|
||||
sp_size, cp_split_hw, max(cp_split_hw))
|
||||
else:
|
||||
cp_split_hw = [1, 1]
|
||||
|
||||
# Get bucket config and find closest bucket for the input aspect ratio
|
||||
bucket_config = get_bucket_config('720p', scale_factor_spatial)
|
||||
|
||||
# Get input aspect ratio from stage1 video
|
||||
input_height, input_width = pil_images[0].height, pil_images[0].width
|
||||
input_ratio = input_height / input_width
|
||||
|
||||
# Find closest bucket
|
||||
closest_ratio = min(bucket_config.keys(),
|
||||
key=lambda x: abs(float(x) - input_ratio))
|
||||
height, width = bucket_config[closest_ratio][0]
|
||||
|
||||
logger.info("Input aspect ratio: %.2f (%sx%s)", input_ratio,
|
||||
input_width, input_height)
|
||||
logger.info("Matched bucket ratio: %s -> resolution: %sx%s",
|
||||
closest_ratio, width, height)
|
||||
logger.info("Target: %sx%s @ %s frames (sp_size=%s, scale_factor=%s)",
|
||||
width, height, new_num_frames, sp_size,
|
||||
scale_factor_spatial)
|
||||
|
||||
# Override batch height/width with bucket-selected resolution
|
||||
batch.height = height
|
||||
batch.width = width
|
||||
|
||||
# Convert PIL images to tensor [T, C, H, W]
|
||||
stage1_video_tensor = torch.stack([
|
||||
torch.from_numpy(np.array(img)).permute(2, 0, 1) # HWC -> CHW
|
||||
for img in pil_images
|
||||
]).float() # [T, C, H, W]
|
||||
|
||||
device = batch.prompt_embeds[0].device
|
||||
dtype = batch.prompt_embeds[0].dtype
|
||||
stage1_video_tensor = stage1_video_tensor.to(device=device, dtype=dtype)
|
||||
|
||||
# Replicate LongCat's exact preprocessing (lines 1227-1235 in pipeline_longcat_video.py)
|
||||
# First: spatial interpolation to target (height, width) on [T, C, H, W]
|
||||
video_down = F.interpolate(stage1_video_tensor,
|
||||
size=(height, width),
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
|
||||
# Rearrange to [C, T, H, W] and add batch dimension -> [1, C, T, H, W]
|
||||
video_down = video_down.permute(1, 0, 2,
|
||||
3).unsqueeze(0) # [1, C, T, H, W]
|
||||
video_down = video_down / 255.0 # Normalize to [0, 1]
|
||||
|
||||
# Then: temporal+spatial interpolation to (new_num_frames, height, width)
|
||||
video_up = F.interpolate(video_down,
|
||||
size=(new_num_frames, height, width),
|
||||
mode='trilinear',
|
||||
align_corners=True)
|
||||
|
||||
# Rescale to [-1, 1] for VAE
|
||||
video_up = video_up * 2.0 - 1.0
|
||||
|
||||
logger.info("Upsampled video shape: %s", video_up.shape)
|
||||
|
||||
# Padding logic (exactly like LongCat lines 1237-1255)
|
||||
# Only pad temporal dimension to ensure BSA compatibility
|
||||
vae_scale_factor_temporal = 4
|
||||
num_noise_frames = video_up.shape[2] - num_cond_frames
|
||||
|
||||
num_cond_latents = 0
|
||||
num_cond_frames_added = 0
|
||||
if num_cond_frames > 0:
|
||||
num_cond_latents = 1 + math.ceil(
|
||||
(num_cond_frames - 1) / vae_scale_factor_temporal)
|
||||
num_cond_latents = math.ceil(
|
||||
num_cond_latents /
|
||||
bsa_latent_granularity) * bsa_latent_granularity
|
||||
num_cond_frames_added = 1 + (
|
||||
num_cond_latents -
|
||||
1) * vae_scale_factor_temporal - num_cond_frames
|
||||
num_cond_frames = num_cond_frames + num_cond_frames_added
|
||||
|
||||
num_noise_latents = math.ceil(num_noise_frames /
|
||||
vae_scale_factor_temporal)
|
||||
num_noise_latents = math.ceil(
|
||||
num_noise_latents / bsa_latent_granularity) * bsa_latent_granularity
|
||||
num_noise_frames_added = num_noise_latents * vae_scale_factor_temporal - num_noise_frames
|
||||
|
||||
if num_cond_frames_added > 0 or num_noise_frames_added > 0:
|
||||
logger.info(
|
||||
"Padding temporal dimension for BSA: cond_frames+=%s, noise_frames+=%s",
|
||||
num_cond_frames_added, num_noise_frames_added)
|
||||
pad_front = video_up[:, :, 0:1].repeat(1, 1, num_cond_frames_added,
|
||||
1, 1)
|
||||
pad_back = video_up[:, :, -1:].repeat(1, 1, num_noise_frames_added,
|
||||
1, 1)
|
||||
video_up = torch.cat([pad_front, video_up, pad_back], dim=2)
|
||||
logger.info("Padded video shape: %s", video_up.shape)
|
||||
|
||||
# Update batch with actual frame count after padding
|
||||
batch.num_frames = video_up.shape[2]
|
||||
|
||||
# Store padding info for later cropping (CRITICAL for correct output!)
|
||||
batch.num_cond_frames_added = num_cond_frames_added
|
||||
batch.num_noise_frames_added = num_noise_frames_added
|
||||
batch.new_frame_size_before_padding = new_num_frames
|
||||
|
||||
# Store num_cond_latents for denoising stage
|
||||
if num_cond_latents > 0:
|
||||
batch.num_cond_latents = num_cond_latents
|
||||
logger.info("Will use num_cond_latents=%s during denoising",
|
||||
num_cond_latents)
|
||||
|
||||
logger.info("Padding info: cond+=%s, noise+=%s, original=%s",
|
||||
num_cond_frames_added, num_noise_frames_added,
|
||||
new_num_frames)
|
||||
|
||||
# VAE encode
|
||||
logger.info("Encoding stage1 video with VAE...")
|
||||
vae_dtype = next(self.vae.parameters()).dtype
|
||||
vae_device = next(self.vae.parameters()).device
|
||||
video_up = video_up.to(dtype=vae_dtype, device=vae_device)
|
||||
|
||||
with torch.no_grad():
|
||||
latent_dist = self.vae.encode(video_up)
|
||||
# Extract tensor from latent distribution
|
||||
if hasattr(latent_dist, 'latent_dist'):
|
||||
# Nested distribution wrapper
|
||||
latent_up = latent_dist.latent_dist.sample()
|
||||
elif hasattr(latent_dist, 'sample'):
|
||||
# DiagonalGaussianDistribution or similar
|
||||
latent_up = latent_dist.sample()
|
||||
elif hasattr(latent_dist, 'latents'):
|
||||
# Direct latents tensor
|
||||
latent_up = latent_dist.latents
|
||||
else:
|
||||
# Assume it's already a tensor
|
||||
latent_up = latent_dist
|
||||
|
||||
# Normalize latents using VAE config (exactly like LongCat)
|
||||
if hasattr(self.vae.config, 'latents_mean') and hasattr(
|
||||
self.vae.config, 'latents_std'):
|
||||
latents_mean = torch.tensor(self.vae.config.latents_mean).view(
|
||||
1, self.vae.config.z_dim, 1, 1, 1).to(latent_up.device,
|
||||
latent_up.dtype)
|
||||
# LongCat uses: 1.0 / latents_std (equivalent to dividing by latents_std)
|
||||
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
|
||||
1, self.vae.config.z_dim, 1, 1, 1).to(latent_up.device,
|
||||
latent_up.dtype)
|
||||
# LongCat: (latents - mean) * (1/std)
|
||||
latent_up = (latent_up - latents_mean) * latents_std
|
||||
|
||||
logger.info("Encoded latent shape: %s", latent_up.shape)
|
||||
|
||||
# Mix with noise according to t_thresh
|
||||
# latent_up = (1 - t_thresh) * latent_up + t_thresh * noise
|
||||
noise = torch.randn_like(latent_up).contiguous()
|
||||
latent_up = (1 - t_thresh) * latent_up + t_thresh * noise
|
||||
|
||||
logger.info("Applied t_thresh=%s noise mixing", t_thresh)
|
||||
|
||||
# Store in batch
|
||||
batch.latents = latent_up.to(dtype)
|
||||
batch.raw_latent_shape = latent_up.shape
|
||||
|
||||
logger.info("LongCat refinement initialization complete")
|
||||
|
||||
return batch
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
LongCat refinement timestep preparation stage.
|
||||
|
||||
This stage prepares special timesteps for LongCat refinement that start from t_thresh.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.pipelines.stages.base import PipelineStage
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class LongCatRefineTimestepStage(PipelineStage):
|
||||
"""
|
||||
Stage for preparing timesteps specific to LongCat refinement.
|
||||
|
||||
For refinement, we need to start from t_thresh instead of t=1.0, so we:
|
||||
1. Generate normal timesteps for num_inference_steps
|
||||
2. Filter to only keep timesteps < t_thresh * 1000
|
||||
3. Prepend t_thresh * 1000 as the first timestep
|
||||
"""
|
||||
|
||||
def __init__(self, scheduler) -> None:
|
||||
super().__init__()
|
||||
self.scheduler = scheduler
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
) -> ForwardBatch:
|
||||
"""
|
||||
Prepare refinement-specific timesteps.
|
||||
|
||||
Args:
|
||||
batch: The current batch information.
|
||||
fastvideo_args: The inference arguments.
|
||||
|
||||
Returns:
|
||||
The batch with refinement timesteps.
|
||||
"""
|
||||
# Only apply if this is a refinement task
|
||||
# Trigger when either a refine_from path or in-memory stage1_video is provided
|
||||
if batch.refine_from is None and batch.stage1_video is None:
|
||||
return batch
|
||||
|
||||
device = get_local_torch_device()
|
||||
num_inference_steps = batch.num_inference_steps
|
||||
t_thresh = batch.t_thresh
|
||||
|
||||
logger.info("Preparing LongCat refinement timesteps (t_thresh=%s)",
|
||||
t_thresh)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 1) Match LongCatVideoPipeline.get_timesteps_sigmas (non-distill):
|
||||
# sigmas = linspace(1, 0.001, num_inference_steps) on CPU
|
||||
# ------------------------------------------------------------------
|
||||
base_sigmas = torch.linspace(
|
||||
1.0,
|
||||
0.001,
|
||||
num_inference_steps,
|
||||
dtype=torch.float32,
|
||||
device=
|
||||
"cpu", # scheduler.set_timesteps expects CPU-convertible sigmas
|
||||
)
|
||||
# Let the scheduler build its internal timestep schedule from sigmas
|
||||
self.scheduler.set_timesteps(num_inference_steps,
|
||||
sigmas=base_sigmas,
|
||||
device=device)
|
||||
base_timesteps = self.scheduler.timesteps
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 2) Apply t_thresh cropping exactly like generate_refine:
|
||||
# timesteps = [t_thresh*1000] + [t for t in base_timesteps if t < t_thresh*1000]
|
||||
# sigmas = timesteps / 1000 (with trailing zero)
|
||||
# ------------------------------------------------------------------
|
||||
t_thresh_value = t_thresh * 1000.0
|
||||
t_thresh_tensor = torch.tensor(t_thresh_value,
|
||||
dtype=base_timesteps.dtype,
|
||||
device=device)
|
||||
filtered_timesteps = base_timesteps[base_timesteps < t_thresh_tensor]
|
||||
|
||||
timesteps = torch.cat(
|
||||
[t_thresh_tensor.unsqueeze(0), filtered_timesteps])
|
||||
|
||||
# Update scheduler with these custom timesteps and corresponding sigmas
|
||||
self.scheduler.timesteps = timesteps
|
||||
sigmas = torch.cat([timesteps / 1000.0, torch.zeros(1, device=device)])
|
||||
self.scheduler.sigmas = sigmas
|
||||
|
||||
logger.info("Refinement timesteps: %s steps starting from t=%s",
|
||||
len(timesteps), t_thresh)
|
||||
logger.info("First few timesteps: %s", timesteps[:5].tolist())
|
||||
|
||||
# Store in batch so downstream stages (denoising) use the same schedule
|
||||
batch.timesteps = timesteps
|
||||
|
||||
return batch
|
||||
@@ -7,6 +7,8 @@ This module contains implementations of timestep preparation stages for diffusio
|
||||
|
||||
import inspect
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.distributed import get_local_torch_device
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.logger import init_logger
|
||||
@@ -71,7 +73,12 @@ class TimestepPreparationStage(PipelineStage):
|
||||
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
||||
f" timestep schedules. Please check whether you are using the correct scheduler."
|
||||
)
|
||||
scheduler.set_timesteps(timesteps=timesteps,
|
||||
# Convert timesteps to CPU if it's a tensor (for numpy conversion in scheduler)
|
||||
if isinstance(timesteps, torch.Tensor):
|
||||
timesteps_for_scheduler = timesteps.cpu()
|
||||
else:
|
||||
timesteps_for_scheduler = timesteps
|
||||
scheduler.set_timesteps(timesteps=timesteps_for_scheduler,
|
||||
device=device,
|
||||
**extra_set_timesteps_kwargs)
|
||||
timesteps = scheduler.timesteps
|
||||
|
||||
@@ -4,6 +4,7 @@ app = modal.App()
|
||||
|
||||
import os
|
||||
|
||||
model_vol = modal.Volume.from_name("hf-model-weights")
|
||||
image_version = os.getenv("IMAGE_VERSION")
|
||||
image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
|
||||
print(f"Using image: {image_tag}")
|
||||
@@ -74,9 +75,15 @@ def run_vae_tests():
|
||||
def run_transformer_tests():
|
||||
run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs")
|
||||
|
||||
@app.function(gpu="L40S:2", image=image, timeout=2700, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
|
||||
@app.function(
|
||||
gpu="L40S:4",
|
||||
image=image,
|
||||
timeout=2700,
|
||||
secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})],
|
||||
volumes={"/root/data": model_vol}
|
||||
)
|
||||
def run_ssim_tests():
|
||||
run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
|
||||
run_test("export MODEL_PATH='/root/data/weights' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
|
||||
|
||||
@app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
|
||||
def run_training_tests():
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Block-sparse attention kernels for LongCat."""
|
||||
|
||||
|
||||
@@ -0,0 +1,656 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import math
|
||||
|
||||
from .common import _attn_fwd_gating, _attn_bwd_preprocess, configs_gating_preset
|
||||
from .flash_attn_bsa_varlen_mask import (
|
||||
_attn_fwd_bsa_varlen, _attn_fwd_bsa_varlen_align, _attn_bwd_dkdv_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_wrapper, _attn_bwd_dq_bsa_varlen_align_wrapper,
|
||||
configs_fwd_bsa_varlen_preset, configs_fwd_bsa_varlen_align_preset, configs_bwd_dkdv_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_preset, configs_bwd_dq_bsa_varlen_align_preset
|
||||
)
|
||||
|
||||
torch._dynamo.config.cache_size_limit = 32
|
||||
|
||||
def is_cuda():
|
||||
return triton.runtime.driver.active.get_current_target().backend == "cuda"
|
||||
|
||||
def supports_tma():
|
||||
return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
|
||||
|
||||
HAS_TMA_DESC = "nv_tma_desc_type" in dir(tl)
|
||||
|
||||
if HAS_TMA_DESC:
|
||||
print("TMA benchmarks will be running with experimental grid constant TMA descriptor.", )
|
||||
else:
|
||||
print("TMA benchmarks will be running without grid constant TMA descriptor.", )
|
||||
|
||||
|
||||
# TmaAutoTuneHelper used in htyu's PR #5622
|
||||
class TmaAutoTuneHelper:
|
||||
|
||||
# duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498
|
||||
class KernelParamWrapper:
|
||||
|
||||
def __init__(self, desc):
|
||||
self.desc = desc
|
||||
|
||||
def tma_desc_cpu_ptr(self):
|
||||
return self.desc.data_ptr()
|
||||
|
||||
TMA_SIZE = 128
|
||||
|
||||
def __init__(self):
|
||||
self.fill_1d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_1d_tma_descriptor)
|
||||
self.fill_2d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_2d_tma_descriptor)
|
||||
if HAS_TMA_DESC:
|
||||
self.descriptors = {}
|
||||
else:
|
||||
self.cuda_descriptors = {}
|
||||
|
||||
# Call this method outside of the lambda function for grid size
|
||||
def init_tma_descriptor(self, name):
|
||||
if HAS_TMA_DESC:
|
||||
self.descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cpu", dtype=torch.int8)
|
||||
else:
|
||||
self.cuda_descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device="cuda", dtype=torch.int8)
|
||||
|
||||
# Call this method inside the lambda function for grid size
|
||||
def fill_1d_tma_descriptor(self, name, ptr, dim, block_dim, element_size):
|
||||
if HAS_TMA_DESC:
|
||||
desc_x = self.descriptors[name]
|
||||
assert desc_x.data_ptr() % 64 == 0
|
||||
self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, desc_x.data_ptr())
|
||||
else:
|
||||
desc_x = self.cuda_descriptors[name]
|
||||
buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
|
||||
self.fill_1d_tma_descriptor_inner(ptr, dim, block_dim, element_size, buf_x.data_ptr())
|
||||
desc_x.copy_(buf_x, non_blocking=True)
|
||||
|
||||
# Call this method inside the lambda function for grid size
|
||||
def fill_2d_tma_descriptor(self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size):
|
||||
if HAS_TMA_DESC:
|
||||
desc_x = self.descriptors[name]
|
||||
assert desc_x.data_ptr() % 64 == 0
|
||||
self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr())
|
||||
else:
|
||||
desc_x = self.cuda_descriptors[name]
|
||||
buf_x = torch.empty_like(desc_x, device="cpu", pin_memory=True)
|
||||
self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, buf_x.data_ptr())
|
||||
desc_x.copy_(buf_x, non_blocking=True)
|
||||
|
||||
def get_tma_descriptor_kernel_param(self, name):
|
||||
if HAS_TMA_DESC:
|
||||
assert self.descriptors[name] is not None
|
||||
return self.KernelParamWrapper(self.descriptors[name])
|
||||
else:
|
||||
assert self.cuda_descriptors[name] is not None
|
||||
return self.cuda_descriptors[name]
|
||||
|
||||
|
||||
@triton.jit
|
||||
def create_mask_from_indices_kernel(
|
||||
block_indices,
|
||||
block_mask,
|
||||
stride_bz, stride_bh, stride_bm, stride_bs,
|
||||
stride_mz, stride_mh, stride_mm, stride_mn,
|
||||
H,
|
||||
):
|
||||
i_zh, i_m, i_s = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
i_z, i_h = i_zh // H, i_zh % H
|
||||
|
||||
off_b = i_z.to(tl.int64) * stride_bz + i_h.to(tl.int64) * stride_bh + i_m.to(tl.int64) * stride_bm + i_s.to(tl.int64) * stride_bs
|
||||
|
||||
b_i = tl.load(block_indices + off_b)
|
||||
|
||||
off_m = i_z.to(tl.int64) * stride_mz + i_h.to(tl.int64) * stride_mh + i_m.to(tl.int64) * stride_mm + b_i.to(tl.int64) * stride_mn
|
||||
|
||||
b_m = 1
|
||||
tl.store(block_mask + off_m, b_m.to(block_mask.dtype.element_ty))
|
||||
|
||||
def create_mask_from_indices_triton(
|
||||
block_indices,
|
||||
N_cols
|
||||
):
|
||||
B, H, N_rows, S = block_indices.shape
|
||||
block_mask = torch.zeros((B, H, N_rows, N_cols), dtype=torch.bool, device=block_indices.device)
|
||||
create_mask_from_indices_kernel[(B * H, N_rows, S)](
|
||||
block_indices,
|
||||
block_mask,
|
||||
block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
|
||||
block_mask.stride(0), block_mask.stride(1), block_mask.stride(2), block_mask.stride(3),
|
||||
H,
|
||||
)
|
||||
return block_mask
|
||||
|
||||
@torch.compile
|
||||
def create_mask_from_indices_varlen(block_indices, N_cols_mask):
|
||||
|
||||
B, H, M, _ = block_indices.shape
|
||||
device = block_indices.device
|
||||
|
||||
mask = torch.zeros((B, H, M, N_cols_mask), dtype=torch.bool, device=device)
|
||||
|
||||
valid = block_indices < N_cols_mask
|
||||
|
||||
b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(block_indices)
|
||||
h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(block_indices)
|
||||
m_idx = torch.arange(M, device=device)[None, None, :, None].expand_as(block_indices)
|
||||
|
||||
valid_coords = (b_idx[valid], h_idx[valid], m_idx[valid], block_indices[valid])
|
||||
|
||||
mask[valid_coords] = True
|
||||
|
||||
return mask
|
||||
|
||||
@torch.compile
|
||||
def create_indices_k_from_indices_q_varlen(
|
||||
block_indices,
|
||||
N_cols_mask # indicate the number of the last dimension of the bool mask, since this information cannot be determined by block_indices, which may contain invalid elements
|
||||
):
|
||||
block_mask_qk = create_mask_from_indices_varlen(block_indices, N_cols_mask)
|
||||
B, H, M, N = block_mask_qk.shape
|
||||
block_mask_kq = block_mask_qk.permute(0, 1, 3, 2)
|
||||
indices = torch.arange(M, device=block_indices.device).view(1, 1, 1, -1).expand_as(block_mask_kq)
|
||||
block_indices_k = torch.where(block_mask_kq, indices, M)
|
||||
block_indices_k, _ = torch.sort(block_indices_k, dim=-1)
|
||||
|
||||
block_indices_k_lens = (block_indices_k < M).sum(dim=-1)
|
||||
|
||||
return block_indices_k, block_indices_k_lens
|
||||
|
||||
|
||||
@torch.compile
|
||||
def mean_pooling_compression(
|
||||
x: torch.Tensor,
|
||||
block_size: int
|
||||
) -> torch.Tensor:
|
||||
B, H, S = x.shape[:3]
|
||||
num_block = math.ceil(S / block_size)
|
||||
if S % block_size != 0:
|
||||
x = F.pad(x, (0, 0, 0, num_block * block_size - S))
|
||||
x_cmp = x.view(B, H, num_block, block_size, -1).mean(dim=3)
|
||||
return x_cmp
|
||||
|
||||
@torch.compile
|
||||
def cal_score(q, k):
|
||||
k_transposed = k.transpose(-1, -2) # [b, h, d, s_k]
|
||||
score = torch.matmul(q, k_transposed) # [b, h, s_q, s_k]
|
||||
return score
|
||||
|
||||
def cal_score_triton(q, k):
|
||||
B, H, s_q, D = q.shape
|
||||
s_k = k.shape[2]
|
||||
|
||||
score = torch.empty(B, H, s_q, s_k, device=q.device, dtype=q.dtype)
|
||||
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_gating_preset['default']
|
||||
|
||||
grid = lambda args: (triton.cdiv(s_q, args["BLOCK_M"]), B * H, 1)
|
||||
_attn_fwd_gating[grid](
|
||||
q, k, score,
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
||||
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
|
||||
score.stride(0), score.stride(1), score.stride(2), score.stride(3),
|
||||
H, s_q, s_k,
|
||||
HEAD_DIM=D,
|
||||
**kernel_config
|
||||
)
|
||||
return score
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_topk(q, k, sparsity):
|
||||
score = cal_score(q, k)
|
||||
block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_topk_from_score(score, sparsity):
|
||||
num_selected = int((1 - sparsity) * score.shape[-1])
|
||||
block_indices = torch.topk(score, num_selected)[1]
|
||||
|
||||
block_indices_lens = torch.full(
|
||||
(block_indices.shape[0], block_indices.shape[1], block_indices.shape[2]),
|
||||
num_selected,
|
||||
dtype=torch.int32,
|
||||
device=block_indices.device
|
||||
)
|
||||
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_cdf(q, k, cdf_threshold):
|
||||
score = cal_score(q, k)
|
||||
head_dim = q.shape[-1]
|
||||
block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold, 1 / head_dim**0.5)
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_cdf_from_score(score, cdf_threshold, sm_scale):
|
||||
weights = torch.softmax(score * sm_scale, dim=-1)
|
||||
|
||||
B, H, Sq, Sk = weights.shape
|
||||
cdf_threshold = torch.full((H,), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1)
|
||||
weights_sorted = torch.sort(weights, dim=-1, descending=True)
|
||||
cdf = torch.cumsum(weights_sorted.values, dim=-1)
|
||||
num_selected = torch.searchsorted(cdf, cdf_threshold, right=True)
|
||||
|
||||
return weights_sorted.indices, num_selected.squeeze(-1)
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold):
|
||||
score = cal_score(q, k)
|
||||
head_dim = q.shape[-1]
|
||||
block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, 1 / head_dim**0.5)
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
@torch.compile
|
||||
def get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold, sm_scale):
|
||||
weights = torch.softmax(score * sm_scale, dim=-1)
|
||||
|
||||
B, H, Sq, Sk = weights.shape
|
||||
cdf_threshold = torch.full((H,), cdf_threshold, device=weights.device).view(1, H, 1, 1).expand(B, -1, Sq, -1)
|
||||
weights_sorted = torch.sort(weights, dim=-1, descending=True)
|
||||
cdf = torch.cumsum(weights_sorted.values, dim=-1)
|
||||
num_selected = torch.searchsorted(cdf, cdf_threshold, right=True)
|
||||
|
||||
# max(cdf, topk)
|
||||
num_selected_topk = int((1 - sparsity) * score.shape[-1])
|
||||
num_selected[num_selected < num_selected_topk] = num_selected_topk
|
||||
|
||||
return weights_sorted.indices, num_selected.squeeze(-1)
|
||||
|
||||
def get_select_indices(q, k, sparsity, cdf_threshold):
|
||||
if sparsity is not None and cdf_threshold is None:
|
||||
block_indices, block_indices_lens = get_select_indices_topk(q, k, sparsity)
|
||||
elif sparsity is None and cdf_threshold is not None:
|
||||
block_indices, block_indices_lens = get_select_indices_cdf(q, k, cdf_threshold)
|
||||
elif sparsity is not None and cdf_threshold is not None:
|
||||
block_indices, block_indices_lens = get_select_indices_cdf_topk(q, k, sparsity, cdf_threshold)
|
||||
else:
|
||||
raise ValueError
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
def get_select_indices_from_score(score, sparsity, cdf_threshold):
|
||||
if sparsity is not None and cdf_threshold is None:
|
||||
block_indices, block_indices_lens = get_select_indices_topk_from_score(score, sparsity)
|
||||
elif sparsity is None and cdf_threshold is not None:
|
||||
block_indices, block_indices_lens = get_select_indices_cdf_from_score(score, cdf_threshold)
|
||||
elif sparsity is not None and cdf_threshold is not None:
|
||||
block_indices, block_indices_lens = get_select_indices_cdf_topk_from_score(score, sparsity, cdf_threshold)
|
||||
else:
|
||||
raise ValueError
|
||||
return block_indices, block_indices_lens
|
||||
|
||||
def attn_fwd_bsa_varlen_triton(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
sm_scale,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
chunk_size_q,
|
||||
chunk_size_k,
|
||||
sparsity
|
||||
):
|
||||
|
||||
B, H, Seq, D = q.shape
|
||||
|
||||
o = torch.empty_like(q)
|
||||
M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
|
||||
|
||||
grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
|
||||
|
||||
config_key = 'BLOCK_N_LG=64' if chunk_size_k == 64 else 'default'
|
||||
if chunk_size_k > 128:
|
||||
fwd_func = _attn_fwd_bsa_varlen
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_preset[config_key]
|
||||
else:
|
||||
fwd_func = _attn_fwd_bsa_varlen_align
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_fwd_bsa_varlen_align_preset[config_key]
|
||||
|
||||
block_indices = block_indices.contiguous()
|
||||
block_indices_lens = block_indices_lens.contiguous()
|
||||
|
||||
fwd_func[grid](
|
||||
q, k, v, sm_scale, M, o,
|
||||
block_indices, # [B, H, M_COMPRESS, S]
|
||||
block_indices_lens, # [B, H, M_COMPRESS, S_MAX]
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
||||
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
|
||||
v.stride(0), v.stride(1), v.stride(2), v.stride(3),
|
||||
o.stride(0), o.stride(1), o.stride(2), o.stride(3),
|
||||
block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
|
||||
block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2),
|
||||
H, Seq,
|
||||
D,
|
||||
BLOCK_M=chunk_size_q,
|
||||
BLOCK_N_LG=chunk_size_k,
|
||||
SPARSITY=sparsity,
|
||||
**kernel_config
|
||||
)
|
||||
|
||||
LN2 = 0.6931471824645996
|
||||
lse = M * LN2 # convert back to natural units (M is of base 2)
|
||||
|
||||
return o, lse
|
||||
|
||||
def attn_bwd_bsa_varlen_triton(
|
||||
do,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
dq,
|
||||
dk,
|
||||
dv,
|
||||
sm_scale,
|
||||
M,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
chunk_size_q,
|
||||
chunk_size_k,
|
||||
sparsity
|
||||
):
|
||||
RCP_LN2 = 1.4426950408889634
|
||||
M = M * RCP_LN2 # ln -> log2
|
||||
|
||||
do = do.contiguous()
|
||||
# assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
|
||||
|
||||
BATCH, N_HEAD, N_CTX, HEAD_DIM = q.shape
|
||||
N_CTX_KV = k.shape[-2]
|
||||
|
||||
RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) # reciprocal
|
||||
arg_k = k
|
||||
arg_k = arg_k * (sm_scale * RCP_LN2)
|
||||
|
||||
if min(chunk_size_q, chunk_size_k) >= 128:
|
||||
PRE_BLOCK = 128
|
||||
else:
|
||||
PRE_BLOCK = min(chunk_size_q, chunk_size_k)
|
||||
|
||||
assert N_CTX % PRE_BLOCK == 0
|
||||
pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
|
||||
delta = torch.empty_like(M)
|
||||
_attn_bwd_preprocess[pre_grid](
|
||||
o, do,
|
||||
delta,
|
||||
N_CTX,
|
||||
BLOCK_M=PRE_BLOCK,
|
||||
HEAD_DIM=HEAD_DIM
|
||||
)
|
||||
|
||||
block_indices_k, block_indices_k_lens = create_indices_k_from_indices_q_varlen(
|
||||
block_indices=block_indices,
|
||||
N_cols_mask=N_CTX_KV // chunk_size_k
|
||||
)
|
||||
|
||||
block_indices = block_indices.contiguous()
|
||||
block_indices_lens = block_indices_lens.contiguous()
|
||||
block_indices_k = block_indices_k.contiguous()
|
||||
block_indices_k_lens = block_indices_k_lens.contiguous()
|
||||
|
||||
config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default'
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dkdv_bsa_varlen_preset[config_key]
|
||||
|
||||
grid_dkdv = lambda args: (triton.cdiv(arg_k.shape[2], args["BLOCK_N"]), 1, arg_k.shape[0] * arg_k.shape[1])
|
||||
_attn_bwd_dkdv_bsa_varlen_wrapper[grid_dkdv](
|
||||
q, arg_k, v, sm_scale, # softmax scale
|
||||
do,
|
||||
dk, dv,
|
||||
M, # lse (log2)
|
||||
delta,
|
||||
block_indices_k,
|
||||
block_indices_k_lens,
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
||||
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
|
||||
v.stride(0), v.stride(1), v.stride(2), v.stride(3),
|
||||
dk.stride(0), dk.stride(1), dk.stride(2), dk.stride(3),
|
||||
dv.stride(0), dv.stride(1), dv.stride(2), dv.stride(3),
|
||||
do.stride(0), do.stride(1), do.stride(2), do.stride(3),
|
||||
M.stride(0), M.stride(1), M.stride(2),
|
||||
delta.stride(0), delta.stride(1), delta.stride(2),
|
||||
block_indices_k.stride(0), block_indices_k.stride(1), block_indices_k.stride(2), block_indices_k.stride(3),
|
||||
block_indices_k_lens.stride(0), block_indices_k_lens.stride(1), block_indices_k_lens.stride(2),
|
||||
N_HEAD, N_CTX,
|
||||
BLOCK_M=chunk_size_q,
|
||||
BLOCK_N_DQ_LG=chunk_size_k,
|
||||
HEAD_DIM=HEAD_DIM,
|
||||
SPARSITY=sparsity,
|
||||
**kernel_config
|
||||
)
|
||||
|
||||
config_key = 'BLOCK_N_DQ_LG=64' if chunk_size_k == 64 else 'default'
|
||||
if chunk_size_k > 128:
|
||||
bwd_dq_func = _attn_bwd_dq_bsa_varlen_wrapper
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_preset[config_key]
|
||||
else:
|
||||
bwd_dq_func = _attn_bwd_dq_bsa_varlen_align_wrapper
|
||||
kernel_config = {} if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1' else configs_bwd_dq_bsa_varlen_align_preset[config_key]
|
||||
|
||||
grid_dq = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), 1, q.shape[0] * q.shape[1])
|
||||
bwd_dq_func[grid_dq](
|
||||
q, arg_k, v,
|
||||
do,
|
||||
dq,
|
||||
M, # lse (log2)
|
||||
delta,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
||||
k.stride(0), k.stride(1), k.stride(2), k.stride(3),
|
||||
v.stride(0), v.stride(1), v.stride(2), v.stride(3),
|
||||
dq.stride(0), dq.stride(1), dq.stride(2), dq.stride(3),
|
||||
do.stride(0), do.stride(1), do.stride(2), do.stride(3),
|
||||
M.stride(0), M.stride(1), M.stride(2),
|
||||
delta.stride(0), delta.stride(1), delta.stride(2),
|
||||
block_indices.stride(0), block_indices.stride(1), block_indices.stride(2), block_indices.stride(3),
|
||||
block_indices_lens.stride(0), block_indices_lens.stride(1), block_indices_lens.stride(2),
|
||||
N_HEAD, N_CTX,
|
||||
BLOCK_M=chunk_size_q,
|
||||
BLOCK_N_DQ_LG=chunk_size_k,
|
||||
HEAD_DIM=HEAD_DIM,
|
||||
SPARSITY=sparsity,
|
||||
**kernel_config
|
||||
)
|
||||
|
||||
@torch.compile
|
||||
def make_block_indices_varlen_cp_list(block_indices, cp_size, num_blocks_k_full):
|
||||
"""
|
||||
Args:
|
||||
block_indices: [B, H, num_blocks_q_per_cp_rank, num_blocks_k_full]
|
||||
|
||||
Return:
|
||||
a list of [block_indices, block_indices_lens] for k from each cp_rank
|
||||
- each block_indices starts from zero
|
||||
- block_indices_lens indicates the valid number of elements in the last dimension of block_indices
|
||||
"""
|
||||
res = []
|
||||
num_blocks_per_rank = num_blocks_k_full // cp_size
|
||||
for i in range(cp_size):
|
||||
block_indices_tmp = block_indices.clone()
|
||||
min_block_idx = i * num_blocks_per_rank
|
||||
block_indices_tmp -= min_block_idx
|
||||
block_indices_tmp[block_indices_tmp < 0] = num_blocks_per_rank # block_indices_tmp < 0 indicate invalid indices, set them to num_blocks_per_rank in order to sort them to the tail, so that the first N elements of the block_indices indicated by block_indices_lens are valid
|
||||
|
||||
block_indices_tmp, _ = torch.sort(block_indices_tmp, dim=-1)
|
||||
|
||||
block_indices_tmp_lens = (block_indices_tmp < num_blocks_per_rank).sum(dim=-1)
|
||||
|
||||
res.append([block_indices_tmp, block_indices_tmp_lens])
|
||||
|
||||
return res
|
||||
|
||||
@torch.compile
|
||||
def flash_attn_fwd_softmax_lse_correction(
|
||||
softmax_lse: torch.Tensor,
|
||||
softmax_lse_per_step: torch.Tensor,
|
||||
):
|
||||
"""Merge softmax stats of each step in Attention with context parallelism"""
|
||||
max_scale = torch.max(softmax_lse, softmax_lse_per_step)
|
||||
min_scale = torch.min(softmax_lse, softmax_lse_per_step)
|
||||
lse_diff = min_scale - max_scale
|
||||
lse_diff = lse_diff.nan_to_num(nan=0.) # handle cases: tensor(-inf) - tensor(-inf) = tensor(nan); In the current cp implementation, it is possible that lses of 2 cp ranks are both -inf, if no block is selected from both cp ranks. In such cases, the finally corrected lse should remain -inf.
|
||||
new_scale = max_scale + torch.log1p(torch.exp(lse_diff)) # a + ln(1 + e^(b - a)) = ln(e^a) + ln(1 + e^(b - a)) = ln(e^a + e^b)
|
||||
softmax_lse.copy_(new_scale)
|
||||
|
||||
@torch.compile
|
||||
def flash_attn_fwd_out_correction_init(
|
||||
out_init_step: torch.Tensor, # b h s d
|
||||
softmax_lse: torch.Tensor, # b h s
|
||||
softmax_lse_init_step: torch.Tensor,
|
||||
):
|
||||
"""Merge partial outputs of the first step in Attention with context parallelism"""
|
||||
softmax_lse_corrected_exp = torch.exp(softmax_lse_init_step - softmax_lse)
|
||||
softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1)
|
||||
out_corrected = out_init_step * softmax_lse_corrected_exp
|
||||
return out_corrected.to(out_init_step.dtype)
|
||||
|
||||
|
||||
@torch.compile
|
||||
def flash_attn_fwd_out_correction(
|
||||
out: torch.Tensor,
|
||||
out_per_step: torch.Tensor,
|
||||
softmax_lse: torch.Tensor,
|
||||
softmax_lse_per_step: torch.Tensor,
|
||||
):
|
||||
"""Merge partial outputs of each step in Attention with context parallelism"""
|
||||
softmax_lse_corrected_exp = torch.exp(softmax_lse_per_step - softmax_lse)
|
||||
softmax_lse_corrected_exp = softmax_lse_corrected_exp.unsqueeze(-1)
|
||||
out_corrected = out_per_step * softmax_lse_corrected_exp
|
||||
out.add_(out_corrected)
|
||||
|
||||
@torch.compile
|
||||
def topk_sort(score, num_chunks_selected):
|
||||
block_indices = torch.topk(score, num_chunks_selected)[1]
|
||||
block_indices, _ = torch.sort(block_indices, dim=-1)
|
||||
return block_indices
|
||||
|
||||
class _attention_bsa(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, sm_scale, use_tma=False):
|
||||
# shape constraints
|
||||
HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
|
||||
# when v is in float8_e5m2 it is transposed.
|
||||
HEAD_DIM_V = v.shape[-1]
|
||||
assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
|
||||
assert HEAD_DIM_K in {16, 32, 64, 128, 256}
|
||||
|
||||
# ---------------------- gating ----------------------
|
||||
q_cmp = mean_pooling_compression(q, chunk_size_q)
|
||||
k_cmp = mean_pooling_compression(k, chunk_size_k)
|
||||
block_indices, block_indices_lens = get_select_indices(q_cmp, k_cmp, sparsity, cdf_threshold)
|
||||
|
||||
# ---------------------- bsa ----------------------
|
||||
|
||||
o, lse = attn_fwd_bsa_varlen_triton(
|
||||
q, k, v,
|
||||
sm_scale, block_indices, block_indices_lens,
|
||||
chunk_size_q, chunk_size_k,
|
||||
sparsity
|
||||
)
|
||||
|
||||
ctx.save_for_backward(q, k, v, o, lse, block_indices, block_indices_lens)
|
||||
ctx.sm_scale = sm_scale
|
||||
ctx.HEAD_DIM = HEAD_DIM_K
|
||||
ctx.chunk_size_q = chunk_size_q
|
||||
ctx.chunk_size_k = chunk_size_k
|
||||
ctx.use_tma = use_tma
|
||||
ctx.sparsity = sparsity
|
||||
|
||||
return o
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, do):
|
||||
q, k, v, o, lse, block_indices, block_indices_lens = ctx.saved_tensors
|
||||
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
|
||||
attn_bwd_bsa_varlen_triton(
|
||||
do,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
dq,
|
||||
dk,
|
||||
dv,
|
||||
ctx.sm_scale,
|
||||
lse,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
ctx.chunk_size_q,
|
||||
ctx.chunk_size_k,
|
||||
ctx.sparsity
|
||||
)
|
||||
|
||||
return dq, dk, dv, None, None, None, None, None, None
|
||||
|
||||
flash_attn_bsa = _attention_bsa.apply
|
||||
|
||||
def rearrange_THW_to_3d_block(x, Nt, Nh, Nw, t, h, w, D):
|
||||
B, H, _, D = x.shape
|
||||
x = x.view(B, H, Nt, t, Nh, h, Nw, w, D)
|
||||
x = x.permute(0, 1, 2, 4, 6, 3, 5, 7, 8) # B H Nt Nh Nw t h w D
|
||||
return x.contiguous().view(B, H, Nt * Nh * Nw * t * h * w, D)
|
||||
|
||||
def rearrange_3d_block_to_THW(x, Nt, Nh, Nw, t, h, w, D):
|
||||
B, H, _, D = x.shape
|
||||
x = x.view(B, H, Nt, Nh, Nw, t, h, w, D)
|
||||
x = x.permute(0, 1, 2, 5, 3, 6, 4, 7, 8) # B H Nt t Nh h Nw w D
|
||||
return x.contiguous().view(B, H, Nt * t * Nh * h * Nw * w, D)
|
||||
|
||||
def flash_attn_bsa_3d(
|
||||
q: torch.Tensor, # [B, H, Sq, D]
|
||||
k: torch.Tensor, # [B, H, Skv, D]
|
||||
v: torch.Tensor, # [B, H, Skv, D]
|
||||
latent_shape_q,
|
||||
latent_shape_k,
|
||||
# bsa_params
|
||||
sparsity=0.875,
|
||||
cdf_threshold=None,
|
||||
chunk_3d_shape_q=[4, 4, 8],
|
||||
chunk_3d_shape_k=[4, 4, 8],
|
||||
) -> torch.Tensor:
|
||||
_, _, Sq, head_dim_q = q.shape
|
||||
_, _, Sk, head_dim_k = k.shape
|
||||
|
||||
assert head_dim_q == head_dim_k
|
||||
head_dim = head_dim_q
|
||||
|
||||
Tq, Hq, Wq = latent_shape_q
|
||||
Tk, Hk, Wk = latent_shape_k
|
||||
|
||||
assert Tq * Hq * Wq == Sq
|
||||
assert Tk * Hk * Wk == Sk
|
||||
|
||||
tq, hq, wq = chunk_3d_shape_q
|
||||
tk, hk, wk = chunk_3d_shape_k
|
||||
|
||||
assert Tq % tq == 0 and Hq % hq == 0 and Wq % wq == 0
|
||||
assert Tk % tk == 0 and Hk % hk == 0 and Wk % wk == 0
|
||||
|
||||
Ntq = Tq // tq
|
||||
Nhq = Hq // hq
|
||||
Nwq = Wq // wq
|
||||
|
||||
Ntk = Tk // tk
|
||||
Nhk = Hk // hk
|
||||
Nwk = Wk // wk
|
||||
|
||||
q = rearrange_THW_to_3d_block(q, Ntq, Nhq, Nwq, tq, hq, wq, q.shape[-1])
|
||||
k = rearrange_THW_to_3d_block(k, Ntk, Nhk, Nwk, tk, hk, wk, k.shape[-1])
|
||||
v = rearrange_THW_to_3d_block(v, Ntk, Nhk, Nwk, tk, hk, wk, v.shape[-1])
|
||||
|
||||
chunk_size_q = tq * hq * wq
|
||||
chunk_size_k = tk * hk * wk
|
||||
|
||||
output = flash_attn_bsa(q, k, v, chunk_size_q, chunk_size_k, sparsity, cdf_threshold, 1 / head_dim**0.5)
|
||||
|
||||
output = rearrange_3d_block_to_THW(output, Ntq, Nhq, Nwq, tq, hq, wq, output.shape[-1])
|
||||
return output
|
||||
@@ -0,0 +1,111 @@
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import os
|
||||
|
||||
if os.environ.get('TRITON_AUTOTUNE_ENBALE', '0') == '1':
|
||||
autotune = triton.autotune
|
||||
else:
|
||||
def autotune(*args, **kwargs):
|
||||
def decorator(func):
|
||||
return func
|
||||
return decorator
|
||||
|
||||
configs_gating_preset = {
|
||||
'default': {
|
||||
'BLOCK_M': 64,
|
||||
'BLOCK_N': 64,
|
||||
'num_stages': 3,
|
||||
'num_warps': 8,
|
||||
}
|
||||
}
|
||||
|
||||
configs_gating = [
|
||||
triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
|
||||
for BM in [64, 128] \
|
||||
for BN in [32, 64] \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
|
||||
gating_reevaluate_keys = ["M", "N"] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(configs_gating, key=gating_reevaluate_keys)
|
||||
@triton.jit
|
||||
def _attn_fwd_gating(
|
||||
Q, K, Out,
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_oz, stride_oh, stride_om, stride_on,
|
||||
H, M, N,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
|
||||
tl.static_assert(BLOCK_N <= HEAD_DIM)
|
||||
start_m = tl.program_id(0)
|
||||
off_hz = tl.program_id(1)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
|
||||
|
||||
# block pointers
|
||||
Q_block_ptr = tl.make_block_ptr(
|
||||
base=Q + q_offset,
|
||||
shape=(M, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
K_block_ptr = tl.make_block_ptr(
|
||||
base=K + k_offset,
|
||||
shape=(HEAD_DIM, N),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N),
|
||||
order=(0, 1),
|
||||
)
|
||||
O_block_ptr = tl.make_block_ptr(
|
||||
base=Out + o_offset,
|
||||
shape=(M, N),
|
||||
strides=(stride_om, stride_on),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, BLOCK_N),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
# load q: it will stay in SRAM throughout
|
||||
q = tl.load(Q_block_ptr, boundary_check=(0,))
|
||||
for start_n in range(0, N, BLOCK_N):
|
||||
start_n = tl.multiple_of(start_n, BLOCK_N)
|
||||
# -- compute qk ----
|
||||
k = tl.load(K_block_ptr, boundary_check=(1,))
|
||||
qk = tl.dot(q, k)
|
||||
|
||||
tl.store(O_block_ptr, qk.to(Out.type.element_ty), boundary_check=(0, 1))
|
||||
|
||||
K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
|
||||
O_block_ptr = tl.advance(O_block_ptr, (0, BLOCK_N))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _attn_bwd_preprocess(
|
||||
O, DO,
|
||||
Delta, # output
|
||||
N_CTX,
|
||||
BLOCK_M: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr
|
||||
):
|
||||
|
||||
off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
off_hz = tl.program_id(1)
|
||||
off_n = tl.arange(0, HEAD_DIM)
|
||||
# load
|
||||
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
|
||||
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
|
||||
delta = tl.sum(o * do, axis=1)
|
||||
# write-back
|
||||
tl.store(Delta + off_hz * N_CTX + off_m, delta)
|
||||
@@ -0,0 +1,43 @@
|
||||
import torch
|
||||
|
||||
|
||||
def p2p_communicate(
|
||||
rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm
|
||||
):
|
||||
"""Point-to-point communications of KV and dKV in Attention with context parallelism"""
|
||||
send_recv_ops = []
|
||||
|
||||
if batch_p2p_comm: # int(os.getenv("NVTE_BATCH_MHA_P2P_COMM", "0")) or (cp_size == 2) 为啥呢
|
||||
if rank % 2 == 0:
|
||||
send_op = torch.distributed.P2POp(
|
||||
torch.distributed.isend, send_tensor, send_dst, cp_group
|
||||
)
|
||||
recv_op = torch.distributed.P2POp(
|
||||
torch.distributed.irecv, recv_tensor, recv_src, cp_group
|
||||
)
|
||||
send_recv_ops.append(send_op)
|
||||
send_recv_ops.append(recv_op)
|
||||
else:
|
||||
recv_op = torch.distributed.P2POp(
|
||||
torch.distributed.irecv, recv_tensor, recv_src, cp_group
|
||||
)
|
||||
send_op = torch.distributed.P2POp(
|
||||
torch.distributed.isend, send_tensor, send_dst, cp_group
|
||||
)
|
||||
send_recv_ops.append(recv_op)
|
||||
send_recv_ops.append(send_op)
|
||||
send_recv_reqs = torch.distributed.batch_isend_irecv(send_recv_ops)
|
||||
else:
|
||||
if rank % 2 == 0:
|
||||
send_op = torch.distributed.isend(send_tensor, send_dst, cp_group)
|
||||
recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group)
|
||||
send_recv_ops.append(send_op)
|
||||
send_recv_ops.append(recv_op)
|
||||
else:
|
||||
recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group)
|
||||
send_op = torch.distributed.isend(send_tensor, send_dst, cp_group)
|
||||
send_recv_ops.append(recv_op)
|
||||
send_recv_ops.append(send_op)
|
||||
send_recv_reqs = send_recv_ops
|
||||
|
||||
return send_recv_reqs
|
||||
+946
@@ -0,0 +1,946 @@
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import os
|
||||
|
||||
from .common import autotune
|
||||
"""
|
||||
TRITON_REEVALUATE_KEY=1
|
||||
- autotune whenever params in reevaluate keys change
|
||||
- use in benchmark script to fine the best config
|
||||
|
||||
TRITON_AUTOTUNE_ENBALE=1
|
||||
- if set to 0, autotune will not work, and the related params must be passed to the function call.
|
||||
"""
|
||||
|
||||
configs_fwd_bsa_varlen_preset = {
|
||||
'default': {
|
||||
'BLOCK_N': 64,
|
||||
'num_stages': 3,
|
||||
'num_warps': 8,
|
||||
},
|
||||
'BLOCK_N_LG=64': {
|
||||
'BLOCK_N': 64,
|
||||
'num_stages': 3,
|
||||
'num_warps': 4,
|
||||
},
|
||||
}
|
||||
configs_fwd_bsa_varlen = [
|
||||
triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \
|
||||
for BN in [32, 64, 128] \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
|
||||
fwd_bsa_reevaluate_varlen_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(list(configs_fwd_bsa_varlen), key=fwd_bsa_reevaluate_varlen_keys)
|
||||
@triton.jit
|
||||
def _attn_fwd_bsa_varlen(
|
||||
Q, K, V, sm_scale, M, Out,
|
||||
block_indices, # [B, H, M_COMPRESS, S_MAX]
|
||||
block_indices_lens, # [B, H, M_COMPRESS]
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_vz, stride_vh, stride_vn, stride_vk,
|
||||
stride_oz, stride_oh, stride_om, stride_ok,
|
||||
stride_bz, stride_bh, stride_bm, stride_bs,
|
||||
stride_lz, stride_lh, stride_lm,
|
||||
H, N_CTX,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N_LG: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking
|
||||
):
|
||||
|
||||
start_m = tl.program_id(0)
|
||||
off_hz = tl.program_id(1)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
|
||||
o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
|
||||
b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
|
||||
l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
|
||||
|
||||
# block pointers
|
||||
Q_block_ptr = tl.make_block_ptr(
|
||||
base=Q + q_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
V_block_ptr = tl.make_block_ptr(
|
||||
base=V + v_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_vn, stride_vk),
|
||||
offsets=(0, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
KT_block_ptr = tl.make_block_ptr(
|
||||
base=K + k_offset,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N),
|
||||
order=(0, 1),
|
||||
)
|
||||
O_block_ptr = tl.make_block_ptr(
|
||||
base=Out + o_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_om, stride_ok),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
block_indices += b_offset + start_m * stride_bm
|
||||
block_indices_lens += l_offset + start_m * stride_lm
|
||||
# initialize offsets
|
||||
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
|
||||
# initialize pointer to m and l
|
||||
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
|
||||
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
# load scales
|
||||
qk_scale = sm_scale
|
||||
qk_scale *= 1.44269504 # 1/ln2; exp2(x/ln2) == exp2(ln(e^x) / ln2) == exp2(log2(e^x)) == exp(x)
|
||||
# load q: it will stay in SRAM throughout
|
||||
q = tl.load(Q_block_ptr)
|
||||
S = tl.load(block_indices_lens)
|
||||
for i in range(S):
|
||||
block_id = tl.load(block_indices + i * stride_bs).to(tl.int32)
|
||||
lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG
|
||||
lo = tl.multiple_of(lo, BLOCK_N)
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
|
||||
V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0))
|
||||
|
||||
# loop over k, v and update accumulator
|
||||
for start_n in range(lo, hi, BLOCK_N):
|
||||
start_n = tl.multiple_of(start_n, BLOCK_N)
|
||||
# -- compute qk ----
|
||||
kT = tl.load(KT_block_ptr_i)
|
||||
qkT = tl.dot(q, kT)
|
||||
|
||||
m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale)
|
||||
qkT = qkT * qk_scale - m_ij[:, None]
|
||||
p = tl.math.exp2(qkT)
|
||||
|
||||
# -- update m_i and l_i
|
||||
alpha = tl.math.exp2(m_i - m_ij)
|
||||
l_ij = tl.sum(p, 1)
|
||||
# -- update output accumulator --
|
||||
acc = acc * alpha[:, None]
|
||||
# update acc
|
||||
v = tl.load(V_block_ptr_i)
|
||||
acc = tl.dot(p.to(v.dtype), v, acc)
|
||||
# update m_i and l_i
|
||||
# place this at the end of the loop to reduce register pressure: https://github.com/triton-lang/triton/commit/ee6abd9
|
||||
l_i = l_i * alpha + l_ij
|
||||
m_i = m_ij
|
||||
V_block_ptr_i = tl.advance(V_block_ptr_i, (BLOCK_N, 0))
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N))
|
||||
|
||||
# epilogue
|
||||
m_i += tl.math.log2(l_i)
|
||||
acc = acc / l_i[:, None]
|
||||
m_ptrs = M + off_hz * N_CTX + offs_m
|
||||
tl.store(m_ptrs, m_i)
|
||||
tl.store(O_block_ptr, acc.to(Out.type.element_ty))
|
||||
|
||||
configs_fwd_bsa_varlen_align_preset = {
|
||||
'default': {
|
||||
'num_stages': 3,
|
||||
'num_warps': 8,
|
||||
},
|
||||
'BLOCK_N_LG=64': {
|
||||
'num_stages': 3,
|
||||
'num_warps': 4,
|
||||
},
|
||||
}
|
||||
configs_fwd_bsa_varlen_align = [
|
||||
triton.Config({}, num_stages=s, num_warps=w) \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
|
||||
fwd_bsa_reevaluate_varlen_align_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(list(configs_fwd_bsa_varlen_align), key=fwd_bsa_reevaluate_varlen_align_keys)
|
||||
@triton.jit
|
||||
def _attn_fwd_bsa_varlen_align(
|
||||
Q, K, V, sm_scale, M, Out,
|
||||
block_indices, # [B, H, M_COMPRESS, S_MAX]
|
||||
block_indices_lens, # [B, H, M_COMPRESS]
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_vz, stride_vh, stride_vn, stride_vk,
|
||||
stride_oz, stride_oh, stride_om, stride_on,
|
||||
stride_bz, stride_bh, stride_bm, stride_bs,
|
||||
stride_lz, stride_lh, stride_lm,
|
||||
H, N_CTX,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N_LG: tl.constexpr,
|
||||
SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking
|
||||
):
|
||||
|
||||
start_m = tl.program_id(0)
|
||||
off_hz = tl.program_id(1)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
q_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
k_offset = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
|
||||
o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
|
||||
b_offset = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
|
||||
l_offset = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
|
||||
|
||||
# block pointers
|
||||
Q_block_ptr = tl.make_block_ptr(
|
||||
base=Q + q_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
V_block_ptr = tl.make_block_ptr(
|
||||
base=V + v_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_vn, stride_vk),
|
||||
offsets=(0, 0),
|
||||
block_shape=(BLOCK_N_LG, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
KT_block_ptr = tl.make_block_ptr(
|
||||
base=K + k_offset,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N_LG),
|
||||
order=(0, 1),
|
||||
)
|
||||
O_block_ptr = tl.make_block_ptr(
|
||||
base=Out + o_offset,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_om, stride_on),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
block_indices += b_offset + start_m * stride_bm
|
||||
block_indices_lens += l_offset + start_m * stride_lm
|
||||
# initialize offsets
|
||||
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
|
||||
# initialize pointer to m and l
|
||||
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
|
||||
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
|
||||
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
# load scales
|
||||
qk_scale = sm_scale
|
||||
qk_scale *= 1.44269504 # 1/ln2; exp2(x/ln2) == exp2(ln(e^x) / ln2) == exp2(log2(e^x)) == exp(x)；乘1/ln2后，exp2(x/ln2) == exp(x)，exp2速度更快
|
||||
# load q: it will stay in SRAM throughout
|
||||
q = tl.load(Q_block_ptr)
|
||||
S = tl.load(block_indices_lens)
|
||||
for i in range(S):
|
||||
block_id = tl.load(block_indices + i * stride_bs).to(tl.int32)
|
||||
lo = block_id * BLOCK_N_LG
|
||||
lo = tl.multiple_of(lo, BLOCK_N_LG)
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
|
||||
V_block_ptr_i = tl.advance(V_block_ptr, (lo, 0))
|
||||
|
||||
# -- compute qk ----
|
||||
kT = tl.load(KT_block_ptr_i)
|
||||
qkT = tl.dot(q, kT)
|
||||
|
||||
m_ij = tl.maximum(m_i, tl.max(qkT, 1) * qk_scale)
|
||||
qkT = qkT * qk_scale - m_ij[:, None]
|
||||
p = tl.math.exp2(qkT)
|
||||
|
||||
# -- update m_i and l_i
|
||||
alpha = tl.math.exp2(m_i - m_ij)
|
||||
l_ij = tl.sum(p, 1)
|
||||
# -- update output accumulator --
|
||||
acc = acc * alpha[:, None]
|
||||
# update acc
|
||||
v = tl.load(V_block_ptr_i)
|
||||
acc = tl.dot(p.to(v.dtype), v, acc) # 没除se，fa2引入的优化
|
||||
# update m_i and l_i
|
||||
# place this at the end of the loop to reduce register pressure: https://github.com/triton-lang/triton/commit/ee6abd9
|
||||
l_i = l_i * alpha + l_ij # 当前总se
|
||||
m_i = m_ij
|
||||
|
||||
|
||||
# epilogue
|
||||
m_i += tl.math.log2(l_i)
|
||||
acc = acc / l_i[:, None]
|
||||
m_ptrs = M + off_hz * N_CTX + offs_m
|
||||
tl.store(m_ptrs, m_i)
|
||||
tl.store(O_block_ptr, acc.to(Out.type.element_ty))
|
||||
|
||||
|
||||
# The main inner-loop logic for computing dK and dV.
|
||||
@triton.jit
|
||||
def _attn_bwd_dkdv_bsa_varlen(
|
||||
dk, dv,
|
||||
k, v,
|
||||
Q, DO,
|
||||
M, D,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
# shared by Q/K/V/DO.
|
||||
# stride_tok, stride_d,
|
||||
stride_qm, stride_qk,
|
||||
stride_dom, stride_dok,
|
||||
stride_mm,
|
||||
stride_dm,
|
||||
stride_bm,
|
||||
N_CTX,
|
||||
BLOCK_M1: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
):
|
||||
|
||||
QT_block_ptr = tl.make_block_ptr(
|
||||
base=Q,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_qk, stride_qm),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_M1),
|
||||
order=(0, 1),
|
||||
)
|
||||
|
||||
DO_block_ptr = tl.make_block_ptr(
|
||||
base=DO,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dom, stride_dok),
|
||||
offsets=(0, 0),
|
||||
block_shape=(BLOCK_M1, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
S = tl.load(block_indices_lens)
|
||||
for i in range(S):
|
||||
block_id = tl.load(block_indices + i * stride_bm).to(tl.int32)
|
||||
start_m = block_id * BLOCK_M1
|
||||
start_m = tl.multiple_of(start_m, BLOCK_M1)
|
||||
|
||||
QT_block_ptr_i = tl.advance(QT_block_ptr, (0, start_m))
|
||||
DO_block_ptr_i = tl.advance(DO_block_ptr, (start_m, 0))
|
||||
|
||||
qT = tl.load(QT_block_ptr_i)
|
||||
# Load m before computing qk to reduce pipeline stall.
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M1) * stride_mm
|
||||
m = tl.load(M + offs_m)
|
||||
kqT = tl.dot(k, qT)
|
||||
pT = tl.math.exp2(kqT - m[None, :])
|
||||
|
||||
do = tl.load(DO_block_ptr_i)
|
||||
# Compute dV.
|
||||
ppT = pT
|
||||
ppT = ppT.to(v.dtype)
|
||||
dv += tl.dot(ppT, do)
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
offs_d = start_m + tl.arange(0, BLOCK_M1) * stride_dm
|
||||
Di = tl.load(D + offs_d)
|
||||
# Compute dP and dS.
|
||||
dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
|
||||
dsT = pT * (dpT - Di[None, :])
|
||||
dsT = dsT.to(v.dtype)
|
||||
dk += tl.dot(dsT, tl.trans(qT))
|
||||
|
||||
|
||||
return dk, dv
|
||||
|
||||
|
||||
# the main inner-loop logic for computing dQ
|
||||
@triton.jit
|
||||
def _attn_bwd_dq_bsa_varlen(
|
||||
dq,
|
||||
q, do,
|
||||
m, d,
|
||||
K, V,
|
||||
N_CTX,
|
||||
BLOCK_N2: tl.constexpr,
|
||||
BLOCK_N_LG: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
stride_bn,
|
||||
# stride_tok, stride_d,
|
||||
stride_kn, stride_kk,
|
||||
stride_vn, stride_vk,
|
||||
):
|
||||
|
||||
VT_block_ptr = tl.make_block_ptr(
|
||||
base=V,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_vk, stride_vn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N2),
|
||||
order=(0, 1),
|
||||
)
|
||||
KT_block_ptr = tl.make_block_ptr(
|
||||
base=K,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N2),
|
||||
order=(0, 1),
|
||||
)
|
||||
|
||||
S = tl.load(block_indices_lens)
|
||||
for i in range(S):
|
||||
block_id = tl.load(block_indices + i * stride_bn).to(tl.int32)
|
||||
lo, hi = block_id * BLOCK_N_LG, (block_id + 1) * BLOCK_N_LG
|
||||
lo = tl.multiple_of(lo, BLOCK_N2)
|
||||
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
|
||||
VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo))
|
||||
|
||||
for start_n in range(lo, hi, BLOCK_N2):
|
||||
start_n = tl.multiple_of(start_n, BLOCK_N2)
|
||||
|
||||
kT = tl.load(KT_block_ptr_i)
|
||||
vT = tl.load(VT_block_ptr_i)
|
||||
qkT = tl.dot(q, kT)
|
||||
p = tl.math.exp2(qkT - m)
|
||||
# Compute dP and dS.
|
||||
dp = tl.dot(do, vT).to(tl.float32)
|
||||
ds = p * (dp - d)
|
||||
ds = ds.to(kT.dtype) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128
|
||||
# Compute dQ.
|
||||
# NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
|
||||
dq += tl.dot(ds, tl.trans(kT))
|
||||
# Increment pointers.
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr_i, (0, BLOCK_N2))
|
||||
VT_block_ptr_i = tl.advance(VT_block_ptr_i, (0, BLOCK_N2))
|
||||
|
||||
return dq
|
||||
|
||||
# the main inner-loop logic for computing dQ
|
||||
@triton.jit
|
||||
def _attn_bwd_dq_bsa_varlen_align(
|
||||
dq,
|
||||
q, do,
|
||||
m, d,
|
||||
K, V,
|
||||
N_CTX,
|
||||
BLOCK_N_LG: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
stride_bn,
|
||||
stride_kn, stride_kk,
|
||||
stride_vn, stride_vk,
|
||||
):
|
||||
|
||||
VT_block_ptr = tl.make_block_ptr(
|
||||
base=V,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_vk, stride_vn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N_LG),
|
||||
order=(0, 1),
|
||||
)
|
||||
KT_block_ptr = tl.make_block_ptr(
|
||||
base=K,
|
||||
shape=(HEAD_DIM, N_CTX),
|
||||
strides=(stride_kk, stride_kn),
|
||||
offsets=(0, 0),
|
||||
block_shape=(HEAD_DIM, BLOCK_N_LG),
|
||||
order=(0, 1),
|
||||
)
|
||||
|
||||
S = tl.load(block_indices_lens)
|
||||
for i in range(S):
|
||||
block_id = tl.load(block_indices + i * stride_bn).to(tl.int32)
|
||||
lo = block_id * BLOCK_N_LG
|
||||
lo = tl.multiple_of(lo, BLOCK_N_LG)
|
||||
|
||||
KT_block_ptr_i = tl.advance(KT_block_ptr, (0, lo))
|
||||
VT_block_ptr_i = tl.advance(VT_block_ptr, (0, lo))
|
||||
|
||||
kT = tl.load(KT_block_ptr_i)
|
||||
vT = tl.load(VT_block_ptr_i)
|
||||
qkT = tl.dot(q, kT)
|
||||
p = tl.math.exp2(qkT - m)
|
||||
# Compute dP and dS.
|
||||
dp = tl.dot(do, vT).to(tl.float32)
|
||||
ds = p * (dp - d)
|
||||
ds = ds.to(kT.dtype) # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py: Converting ds to q.dtype here reduces register pressure and makes it much faster for BLOCK_HEADDIM=128
|
||||
# Compute dQ.
|
||||
# NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
|
||||
dq += tl.dot(ds, tl.trans(kT))
|
||||
|
||||
|
||||
return dq
|
||||
|
||||
configs_bwd_dkdv_bsa_varlen_preset = {
|
||||
'default': {
|
||||
'BLOCK_N': 128,
|
||||
'num_stages': 2,
|
||||
'num_warps': 8,
|
||||
},
|
||||
'BLOCK_N_DQ_LG=64': {
|
||||
'BLOCK_N': 64,
|
||||
'num_stages': 2,
|
||||
'num_warps': 4,
|
||||
}
|
||||
}
|
||||
configs_bwd_dkdv_bsa_varlen = [
|
||||
triton.Config({'BLOCK_N': BN}, num_stages=s, num_warps=w) \
|
||||
for BN in [32, 64, 128] \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
bwd_dkdv_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(list(configs_bwd_dkdv_bsa_varlen), key=bwd_dkdv_bsa_varlen_reevaluate_keys)
|
||||
@triton.jit
|
||||
def _attn_bwd_dkdv_bsa_varlen_wrapper(
|
||||
Q, K, V, sm_scale, # softmax scale
|
||||
DO,
|
||||
DK, DV,
|
||||
M, # lse (log2)
|
||||
D,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
# stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO.
|
||||
# qkv
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_vz, stride_vh, stride_vn, stride_vk,
|
||||
# dk dv do
|
||||
stride_dkz, stride_dkh, stride_dkn, stride_dkk,
|
||||
stride_dvz, stride_dvh, stride_dvn, stride_dvk,
|
||||
stride_doz, stride_doh, stride_dom, stride_dok,
|
||||
# m, d
|
||||
stride_mz, stride_mh, stride_mm,
|
||||
stride_dz, stride_dh, stride_dm,
|
||||
#
|
||||
stride_bz, stride_bh, stride_bn, stride_bm, # block_indices
|
||||
stride_lz, stride_lh, stride_ln, # block_indices_lens
|
||||
#
|
||||
H, N_CTX,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_N_DQ_LG: tl.constexpr, # logical block size
|
||||
HEAD_DIM: tl.constexpr,
|
||||
SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking
|
||||
):
|
||||
|
||||
tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N == 0)
|
||||
start_n = tl.program_id(0)
|
||||
|
||||
off_hz = tl.program_id(2)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
|
||||
|
||||
off_dk = off_z.to(tl.int64) * stride_dkz + off_h.to(tl.int64) * stride_dkh
|
||||
off_dv = off_z.to(tl.int64) * stride_dvz + off_h.to(tl.int64) * stride_dvh
|
||||
off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
|
||||
|
||||
off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
|
||||
off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
|
||||
|
||||
off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
|
||||
off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
|
||||
|
||||
# offset pointers for batch/head
|
||||
Q += off_q
|
||||
K += off_k
|
||||
V += off_v
|
||||
DO += off_do
|
||||
DK += off_dk
|
||||
DV += off_dv
|
||||
|
||||
M += off_m
|
||||
D += off_d
|
||||
block_indices += off_block_incides
|
||||
block_indices_lens += off_block_incides_lens
|
||||
|
||||
# ---------------------------------------- [DKDV] ----------------------------------------
|
||||
|
||||
dv = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32)
|
||||
dk = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32)
|
||||
|
||||
# load K and V: they stay in SRAM throughout the inner loop.
|
||||
|
||||
K_block_ptr = tl.make_block_ptr(
|
||||
base=K,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_kn, stride_kk),
|
||||
offsets=(start_n * BLOCK_N, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
V_block_ptr = tl.make_block_ptr(
|
||||
base=V,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_vn, stride_vk),
|
||||
offsets=(start_n * BLOCK_N, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DK_block_ptr = tl.make_block_ptr(
|
||||
base=DK,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dkn, stride_dkk),
|
||||
offsets=(start_n * BLOCK_N, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DV_block_ptr = tl.make_block_ptr(
|
||||
base=DV,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dvn, stride_dvk),
|
||||
offsets=(start_n * BLOCK_N, 0),
|
||||
block_shape=(BLOCK_N, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
k = tl.load(K_block_ptr)
|
||||
v = tl.load(V_block_ptr)
|
||||
|
||||
k_compress_idx = start_n * BLOCK_N // BLOCK_N_DQ_LG
|
||||
block_indices_i = block_indices + k_compress_idx * stride_bn
|
||||
block_indices_lens_i = block_indices_lens + k_compress_idx * stride_ln
|
||||
|
||||
dk, dv = _attn_bwd_dkdv_bsa_varlen(
|
||||
dk, dv,
|
||||
k, v,
|
||||
Q, DO,
|
||||
M, D,
|
||||
block_indices_i,
|
||||
block_indices_lens_i,
|
||||
# shared by Q/K/V/DO.
|
||||
stride_qm, stride_qk,
|
||||
stride_dom, stride_dok,
|
||||
stride_mm,
|
||||
stride_dm,
|
||||
#
|
||||
stride_bm,
|
||||
N_CTX,
|
||||
BLOCK_M,
|
||||
HEAD_DIM,
|
||||
)
|
||||
|
||||
# Write back dk
|
||||
dk *= sm_scale # S = scale * QKT; dK = scale * QdST
|
||||
tl.store(DK_block_ptr, dk.to(k.dtype))
|
||||
|
||||
# Write back dv
|
||||
tl.store(DV_block_ptr, dv.to(v.dtype))
|
||||
|
||||
configs_bwd_dq_bsa_varlen_preset = {
|
||||
'default': {
|
||||
'BLOCK_N_DQ': 64,
|
||||
'num_stages': 2,
|
||||
'num_warps': 8,
|
||||
},
|
||||
'BLOCK_N_DQ_LG=64': {
|
||||
'BLOCK_N_DQ': 64,
|
||||
'num_stages': 2,
|
||||
'num_warps': 4,
|
||||
},
|
||||
}
|
||||
configs_bwd_dq_bsa_varlen = [
|
||||
triton.Config({'BLOCK_N_DQ': BN}, num_stages=s, num_warps=w) \
|
||||
for BN in [32, 64, 128] \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
bwd_dq_bsa_varlen_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(list(configs_bwd_dq_bsa_varlen), key=bwd_dq_bsa_varlen_reevaluate_keys)
|
||||
@triton.jit
|
||||
def _attn_bwd_dq_bsa_varlen_wrapper(
|
||||
Q, K, V, # softmax scale
|
||||
DO,
|
||||
DQ,
|
||||
M, # lse (log2)
|
||||
D,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
# stride_z, stride_h, stride_tok, stride_d, # shared by Q/K/V/DO.
|
||||
# qkv
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_vz, stride_vh, stride_vn, stride_vk,
|
||||
# dq do
|
||||
stride_dqz, stride_dqh, stride_dqm, stride_dqk,
|
||||
stride_doz, stride_doh, stride_dom, stride_dok,
|
||||
# m, d
|
||||
stride_mz, stride_mh, stride_mm,
|
||||
stride_dz, stride_dh, stride_dm,
|
||||
#
|
||||
stride_bz, stride_bh, stride_bm, stride_bn, # block_indices
|
||||
stride_lz, stride_lh, stride_lm, # block_indices_lens
|
||||
#
|
||||
H, N_CTX,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N_DQ_LG: tl.constexpr, # logical block size
|
||||
BLOCK_N_DQ: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking
|
||||
):
|
||||
|
||||
tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ == 0)
|
||||
tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0)
|
||||
|
||||
LN2: tl.constexpr = 0.6931471824645996 # = ln(2)
|
||||
|
||||
start_m = tl.program_id(0)
|
||||
|
||||
off_hz = tl.program_id(2)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
|
||||
|
||||
off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh
|
||||
off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
|
||||
|
||||
off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
|
||||
off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
|
||||
|
||||
off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
|
||||
off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
|
||||
|
||||
# offset pointers for batch/head
|
||||
Q += off_q
|
||||
K += off_k
|
||||
V += off_v
|
||||
DO += off_do
|
||||
DQ += off_dq
|
||||
|
||||
M += off_m
|
||||
D += off_d
|
||||
block_indices += off_block_incides
|
||||
block_indices_lens += off_block_incides_lens
|
||||
|
||||
# ---------------------------------------- [DQ] ----------------------------------------
|
||||
|
||||
Q_block_ptr = tl.make_block_ptr(
|
||||
base=Q,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DO_block_ptr = tl.make_block_ptr(
|
||||
base=DO,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dom, stride_dok),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DQ_block_ptr = tl.make_block_ptr(
|
||||
base=DQ,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dqm, stride_dqk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
q = tl.load(Q_block_ptr)
|
||||
do = tl.load(DO_block_ptr)
|
||||
|
||||
dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
|
||||
start_m = start_m * BLOCK_M
|
||||
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm
|
||||
offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm
|
||||
|
||||
m = tl.load(M + offs_m)
|
||||
m = m[:, None]
|
||||
|
||||
d = tl.load(D + offs_d)
|
||||
d = d[:, None]
|
||||
|
||||
block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm
|
||||
block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm
|
||||
|
||||
dq = _attn_bwd_dq_bsa_varlen(
|
||||
dq,
|
||||
q, do,
|
||||
m, d,
|
||||
K, V,
|
||||
N_CTX,
|
||||
BLOCK_N_DQ,
|
||||
BLOCK_N_DQ_LG,
|
||||
HEAD_DIM,
|
||||
block_indices_m,
|
||||
block_indices_lens_m,
|
||||
stride_bn,
|
||||
stride_kn, stride_kk,
|
||||
stride_vn, stride_vk,
|
||||
)
|
||||
|
||||
# Write back dQ.
|
||||
dq *= LN2
|
||||
tl.store(DQ_block_ptr, dq.to(q.dtype))
|
||||
|
||||
configs_bwd_dq_bsa_varlen_align_preset = {
|
||||
'default': {
|
||||
'num_stages': 2,
|
||||
'num_warps': 8,
|
||||
},
|
||||
'BLOCK_N_DQ_LG=64': {
|
||||
'num_stages': 2,
|
||||
'num_warps': 4,
|
||||
},
|
||||
}
|
||||
configs_bwd_dq_bsa_varlen_align = [
|
||||
triton.Config({}, num_stages=s, num_warps=w) \
|
||||
for s in [2, 3, 4, 5] \
|
||||
for w in [4, 8] \
|
||||
]
|
||||
bwd_dq_bsa_varlen_align_reevaluate_keys = ['N_CTX', 'BLOCK_M', 'BLOCK_N_DQ_LG', 'SPARSITY'] if os.environ.get('TRITON_REEVALUATE_KEY', '0') == '1' else []
|
||||
@autotune(list(configs_bwd_dq_bsa_varlen_align), key=bwd_dq_bsa_varlen_align_reevaluate_keys)
|
||||
@triton.jit
|
||||
def _attn_bwd_dq_bsa_varlen_align_wrapper(
|
||||
Q, K, V, # softmax scale
|
||||
DO,
|
||||
DQ,
|
||||
M, # lse (log2)
|
||||
D,
|
||||
block_indices,
|
||||
block_indices_lens,
|
||||
# qkv
|
||||
stride_qz, stride_qh, stride_qm, stride_qk,
|
||||
stride_kz, stride_kh, stride_kn, stride_kk,
|
||||
stride_vz, stride_vh, stride_vn, stride_vk,
|
||||
# dq do
|
||||
stride_dqz, stride_dqh, stride_dqm, stride_dqk,
|
||||
stride_doz, stride_doh, stride_dom, stride_dok,
|
||||
# m, d
|
||||
stride_mz, stride_mh, stride_mm,
|
||||
stride_dz, stride_dh, stride_dm,
|
||||
#
|
||||
stride_bz, stride_bh, stride_bm, stride_bn, # block_indices
|
||||
stride_lz, stride_lh, stride_lm, # block_indices_lens
|
||||
#
|
||||
H, N_CTX,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N_DQ_LG: tl.constexpr, # logical block size
|
||||
HEAD_DIM: tl.constexpr,
|
||||
SPARSITY: tl.constexpr, # not used; just for trigger reevaluate for benchmarking
|
||||
):
|
||||
|
||||
tl.static_assert(BLOCK_N_DQ_LG % BLOCK_N_DQ_LG == 0)
|
||||
tl.static_assert(BLOCK_N_DQ_LG % BLOCK_M == 0)
|
||||
|
||||
LN2: tl.constexpr = 0.6931471824645996 # = ln(2)
|
||||
|
||||
start_m = tl.program_id(0)
|
||||
|
||||
off_hz = tl.program_id(2)
|
||||
off_z = off_hz // H
|
||||
off_h = off_hz % H
|
||||
|
||||
off_q = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
|
||||
off_k = off_z.to(tl.int64) * stride_kz + off_h.to(tl.int64) * stride_kh
|
||||
off_v = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
|
||||
|
||||
off_dq = off_z.to(tl.int64) * stride_dqz + off_h.to(tl.int64) * stride_dqh
|
||||
off_do = off_z.to(tl.int64) * stride_doz + off_h.to(tl.int64) * stride_doh
|
||||
|
||||
off_m = off_z.to(tl.int64) * stride_mz + off_h.to(tl.int64) * stride_mh
|
||||
off_d = off_z.to(tl.int64) * stride_dz + off_h.to(tl.int64) * stride_dh
|
||||
|
||||
off_block_incides = off_z.to(tl.int64) * stride_bz + off_h.to(tl.int64) * stride_bh
|
||||
off_block_incides_lens = off_z.to(tl.int64) * stride_lz + off_h.to(tl.int64) * stride_lh
|
||||
|
||||
# offset pointers for batch/head
|
||||
Q += off_q
|
||||
K += off_k
|
||||
V += off_v
|
||||
DO += off_do
|
||||
DQ += off_dq
|
||||
|
||||
M += off_m
|
||||
D += off_d
|
||||
block_indices += off_block_incides
|
||||
block_indices_lens += off_block_incides_lens
|
||||
|
||||
# ---------------------------------------- [DQ] ----------------------------------------
|
||||
|
||||
Q_block_ptr = tl.make_block_ptr(
|
||||
base=Q,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_qm, stride_qk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DO_block_ptr = tl.make_block_ptr(
|
||||
base=DO,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dom, stride_dok),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
DQ_block_ptr = tl.make_block_ptr(
|
||||
base=DQ,
|
||||
shape=(N_CTX, HEAD_DIM),
|
||||
strides=(stride_dqm, stride_dqk),
|
||||
offsets=(start_m * BLOCK_M, 0),
|
||||
block_shape=(BLOCK_M, HEAD_DIM),
|
||||
order=(1, 0),
|
||||
)
|
||||
|
||||
q = tl.load(Q_block_ptr)
|
||||
do = tl.load(DO_block_ptr)
|
||||
|
||||
dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
|
||||
|
||||
start_m = start_m * BLOCK_M
|
||||
|
||||
offs_m = start_m + tl.arange(0, BLOCK_M) * stride_mm
|
||||
offs_d = start_m + tl.arange(0, BLOCK_M) * stride_dm
|
||||
|
||||
m = tl.load(M + offs_m)
|
||||
m = m[:, None]
|
||||
|
||||
# D (= delta) is pre-divided by ds_scale.
|
||||
d = tl.load(D + offs_d)
|
||||
d = d[:, None]
|
||||
|
||||
block_indices_m = block_indices + (start_m // BLOCK_M) * stride_bm
|
||||
block_indices_lens_m = block_indices_lens + (start_m // BLOCK_M) * stride_lm
|
||||
|
||||
dq = _attn_bwd_dq_bsa_varlen_align(
|
||||
dq,
|
||||
q, do,
|
||||
m, d,
|
||||
K, V,
|
||||
N_CTX,
|
||||
BLOCK_N_DQ_LG,
|
||||
HEAD_DIM,
|
||||
block_indices_m,
|
||||
block_indices_lens_m,
|
||||
stride_bn,
|
||||
stride_kn, stride_kk,
|
||||
stride_vn, stride_vk,
|
||||
)
|
||||
|
||||
# Write back dQ.
|
||||
dq *= LN2
|
||||
tl.store(DQ_block_ptr, dq.to(q.dtype))
|
||||
|
||||
@@ -543,9 +543,18 @@ def maybe_download_lora(model_name_or_path: str,
|
||||
Local path to the model
|
||||
"""
|
||||
|
||||
# If it's already a file path, return it directly
|
||||
if os.path.isfile(model_name_or_path):
|
||||
return model_name_or_path
|
||||
|
||||
local_path = maybe_download_model(model_name_or_path, local_dir, download)
|
||||
weight_name = _best_guess_weight_name(model_name_or_path,
|
||||
file_extension=".safetensors")
|
||||
|
||||
# If weight_name is None, assume local_path is already the full path
|
||||
if weight_name is None:
|
||||
return local_path
|
||||
|
||||
return os.path.join(local_path, weight_name)
|
||||
|
||||
|
||||
|
||||
@@ -74,6 +74,7 @@ dependencies = [
|
||||
# Preprocessing Dependencies
|
||||
"torchcodec==0.5.0",
|
||||
"ray>=2.49.1",
|
||||
"ftfy==6.3.1",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
|
||||
@@ -0,0 +1,708 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Convert LongCat weights to FastVideo native format.
|
||||
|
||||
This script performs a complete conversion from original LongCat weights
|
||||
to FastVideo native implementation in a single step:
|
||||
|
||||
1. Converts transformer weights (with QKV/KV splitting)
|
||||
2. Copies other components (VAE, text encoder, tokenizer, scheduler)
|
||||
3. Converts LoRA weights (cfg_step_lora, refinement_lora)
|
||||
4. Updates config files to point to native model
|
||||
|
||||
Usage:
|
||||
python scripts/checkpoint_conversion/longcat_to_fastvideo.py \
|
||||
--source /path/to/LongCat-Video/weights/LongCat-Video \
|
||||
--output weights/longcat-native \
|
||||
--validate
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file, save_file
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def split_qkv(qkv_weight: torch.Tensor, qkv_bias: torch.Tensor | None = None):
|
||||
"""Split fused QKV projection into separate Q, K, V."""
|
||||
dim = qkv_weight.shape[0] // 3
|
||||
q, k, v = torch.chunk(qkv_weight, 3, dim=0)
|
||||
|
||||
if qkv_bias is not None:
|
||||
q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0)
|
||||
else:
|
||||
q_bias = k_bias = v_bias = None
|
||||
|
||||
return (q, k, v), (q_bias, k_bias, v_bias)
|
||||
|
||||
|
||||
def split_kv(kv_weight: torch.Tensor, kv_bias: torch.Tensor | None = None):
|
||||
"""Split fused KV projection into separate K, V."""
|
||||
dim = kv_weight.shape[0] // 2
|
||||
k, v = torch.chunk(kv_weight, 2, dim=0)
|
||||
|
||||
if kv_bias is not None:
|
||||
k_bias, v_bias = torch.chunk(kv_bias, 2, dim=0)
|
||||
else:
|
||||
k_bias = v_bias = None
|
||||
|
||||
return (k, v), (k_bias, v_bias)
|
||||
|
||||
|
||||
def convert_transformer_weights(source_weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
||||
"""
|
||||
Convert LongCat transformer weights to native FastVideo format.
|
||||
|
||||
Main transformations:
|
||||
1. Split fused QKV projections (self-attention)
|
||||
2. Split fused KV projections (cross-attention)
|
||||
3. Rename parameters according to mapping
|
||||
"""
|
||||
converted = OrderedDict()
|
||||
processed_keys = set()
|
||||
|
||||
print(" Converting transformer weights...")
|
||||
|
||||
for key, value in tqdm(source_weights.items(), desc=" Processing parameters"):
|
||||
if key in processed_keys:
|
||||
continue
|
||||
|
||||
# === Embedders ===
|
||||
if key.startswith("x_embedder."):
|
||||
new_key = key.replace("x_embedder.", "patch_embed.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif key.startswith("t_embedder.mlp.0."):
|
||||
new_key = key.replace("t_embedder.mlp.0.", "time_embedder.linear_1.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif key.startswith("t_embedder.mlp.2."):
|
||||
new_key = key.replace("t_embedder.mlp.2.", "time_embedder.linear_2.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif key.startswith("y_embedder.y_proj.0."):
|
||||
new_key = key.replace("y_embedder.y_proj.0.", "caption_embedder.linear_1.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif key.startswith("y_embedder.y_proj.2."):
|
||||
new_key = key.replace("y_embedder.y_proj.2.", "caption_embedder.linear_2.")
|
||||
converted[new_key] = value
|
||||
|
||||
# === Self-Attention QKV Splitting ===
|
||||
elif ".attn.qkv.weight" in key:
|
||||
block_idx = key.split(".")[1]
|
||||
qkv_weight = value
|
||||
qkv_bias_key = key.replace(".weight", ".bias")
|
||||
qkv_bias = source_weights.get(qkv_bias_key)
|
||||
|
||||
(q, k, v), (q_bias, k_bias, v_bias) = split_qkv(qkv_weight, qkv_bias)
|
||||
|
||||
converted[f"blocks.{block_idx}.self_attn.to_q.weight"] = q
|
||||
converted[f"blocks.{block_idx}.self_attn.to_k.weight"] = k
|
||||
converted[f"blocks.{block_idx}.self_attn.to_v.weight"] = v
|
||||
|
||||
if q_bias is not None:
|
||||
converted[f"blocks.{block_idx}.self_attn.to_q.bias"] = q_bias
|
||||
converted[f"blocks.{block_idx}.self_attn.to_k.bias"] = k_bias
|
||||
converted[f"blocks.{block_idx}.self_attn.to_v.bias"] = v_bias
|
||||
|
||||
processed_keys.add(key)
|
||||
if qkv_bias is not None:
|
||||
processed_keys.add(qkv_bias_key)
|
||||
|
||||
elif ".attn.qkv.bias" in key:
|
||||
continue
|
||||
|
||||
elif ".attn.proj." in key:
|
||||
new_key = key.replace(".attn.proj.", ".self_attn.to_out.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif ".attn.q_norm." in key or ".attn.k_norm." in key:
|
||||
new_key = key.replace(".attn.", ".self_attn.")
|
||||
converted[new_key] = value
|
||||
|
||||
# === Cross-Attention ===
|
||||
elif ".cross_attn.q_linear." in key:
|
||||
new_key = key.replace(".cross_attn.q_linear.", ".cross_attn.to_q.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif ".cross_attn.kv_linear.weight" in key:
|
||||
block_idx = key.split(".")[1]
|
||||
kv_weight = value
|
||||
kv_bias_key = key.replace(".weight", ".bias")
|
||||
kv_bias = source_weights.get(kv_bias_key)
|
||||
|
||||
(k, v), (k_bias, v_bias) = split_kv(kv_weight, kv_bias)
|
||||
|
||||
converted[f"blocks.{block_idx}.cross_attn.to_k.weight"] = k
|
||||
converted[f"blocks.{block_idx}.cross_attn.to_v.weight"] = v
|
||||
|
||||
if k_bias is not None:
|
||||
converted[f"blocks.{block_idx}.cross_attn.to_k.bias"] = k_bias
|
||||
converted[f"blocks.{block_idx}.cross_attn.to_v.bias"] = v_bias
|
||||
|
||||
processed_keys.add(key)
|
||||
if kv_bias is not None:
|
||||
processed_keys.add(kv_bias_key)
|
||||
|
||||
elif ".cross_attn.kv_linear.bias" in key:
|
||||
continue
|
||||
|
||||
elif ".cross_attn.proj." in key:
|
||||
new_key = key.replace(".cross_attn.proj.", ".cross_attn.to_out.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif ".cross_attn.q_norm." in key or ".cross_attn.k_norm." in key:
|
||||
converted[key] = value
|
||||
|
||||
# === Final Layer (must come BEFORE general transformer block patterns) ===
|
||||
elif key.startswith("final_layer.adaLN_modulation.1."):
|
||||
new_key = key.replace("final_layer.adaLN_modulation.1.", "final_layer.adaln_linear.")
|
||||
converted[new_key] = value
|
||||
|
||||
# === Transformer Block AdaLN ===
|
||||
elif ".adaLN_modulation.1." in key:
|
||||
new_key = key.replace(".adaLN_modulation.1.", ".adaln_linear_1.")
|
||||
converted[new_key] = value
|
||||
|
||||
# === Transformer Block Normalization ===
|
||||
elif ".mod_norm_attn." in key or ".mod_norm_ffn." in key:
|
||||
continue
|
||||
|
||||
elif ".pre_crs_attn_norm.weight" in key:
|
||||
new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.")
|
||||
converted[new_key] = value
|
||||
|
||||
elif ".pre_crs_attn_norm.bias" in key:
|
||||
new_key = key.replace(".pre_crs_attn_norm.", ".norm_cross.")
|
||||
converted[new_key] = value
|
||||
|
||||
# === FFN (SwiGLU) ===
|
||||
elif ".ffn.w1." in key or ".ffn.w2." in key or ".ffn.w3." in key:
|
||||
converted[key] = value
|
||||
|
||||
elif key.startswith("final_layer.norm_final."):
|
||||
continue
|
||||
|
||||
elif key.startswith("final_layer.linear."):
|
||||
new_key = key.replace("final_layer.linear.", "final_layer.proj.")
|
||||
converted[new_key] = value
|
||||
|
||||
else:
|
||||
print(f" ⚠️ Unknown key: {key}")
|
||||
converted[key] = value
|
||||
|
||||
return converted
|
||||
|
||||
|
||||
def validate_conversion(original: dict, converted: dict) -> bool:
|
||||
"""Validate that conversion preserved all parameters correctly."""
|
||||
print("\n Validating conversion...")
|
||||
|
||||
orig_count = sum(p.numel() for p in original.values())
|
||||
conv_count = sum(p.numel() for p in converted.values())
|
||||
|
||||
dropped_count = 0
|
||||
for key, value in original.items():
|
||||
if ".mod_norm_attn." in key or ".mod_norm_ffn." in key:
|
||||
dropped_count += value.numel()
|
||||
elif "final_layer.norm_final." in key:
|
||||
dropped_count += value.numel()
|
||||
|
||||
expected_conv_count = orig_count - dropped_count
|
||||
|
||||
print(f" Original parameters: {orig_count:,}")
|
||||
print(f" Converted parameters: {conv_count:,}")
|
||||
print(f" Dropped parameters (norms without params): {dropped_count:,}")
|
||||
|
||||
if conv_count != expected_conv_count:
|
||||
print(f" ⚠️ Parameter count mismatch!")
|
||||
return False
|
||||
|
||||
print(f" ✓ Parameter count matches")
|
||||
|
||||
# Verify QKV/KV splits
|
||||
print("\n Verifying QKV/KV splits...")
|
||||
num_blocks = 48
|
||||
|
||||
for i in range(num_blocks):
|
||||
orig_qkv_weight = original.get(f"blocks.{i}.attn.qkv.weight")
|
||||
if orig_qkv_weight is not None:
|
||||
conv_q = converted[f"blocks.{i}.self_attn.to_q.weight"]
|
||||
conv_k = converted[f"blocks.{i}.self_attn.to_k.weight"]
|
||||
conv_v = converted[f"blocks.{i}.self_attn.to_v.weight"]
|
||||
reconstructed = torch.cat([conv_q, conv_k, conv_v], dim=0)
|
||||
if not torch.allclose(orig_qkv_weight, reconstructed):
|
||||
print(f" ❌ QKV weight mismatch in block {i}")
|
||||
return False
|
||||
|
||||
orig_kv_weight = original.get(f"blocks.{i}.cross_attn.kv_linear.weight")
|
||||
if orig_kv_weight is not None:
|
||||
conv_k = converted[f"blocks.{i}.cross_attn.to_k.weight"]
|
||||
conv_v = converted[f"blocks.{i}.cross_attn.to_v.weight"]
|
||||
reconstructed = torch.cat([conv_k, conv_v], dim=0)
|
||||
if not torch.allclose(orig_kv_weight, reconstructed):
|
||||
print(f" ❌ KV weight mismatch in block {i}")
|
||||
return False
|
||||
|
||||
print(f" ✓ All splits verified successfully")
|
||||
return True
|
||||
|
||||
|
||||
def copy_component(source_dir: Path, output_dir: Path, component: str, mapping: dict = None) -> bool:
|
||||
"""Copy a component directory, optionally with name mapping."""
|
||||
source_name = mapping.get(component, component) if mapping else component
|
||||
source_path = source_dir / source_name
|
||||
|
||||
if source_path.exists():
|
||||
output_path = output_dir / component
|
||||
if output_path.exists():
|
||||
shutil.rmtree(output_path)
|
||||
shutil.copytree(source_path, output_path)
|
||||
print(f" ✓ {component} copied")
|
||||
return True
|
||||
else:
|
||||
print(f" ⚠️ {component} not found, skipping")
|
||||
return False
|
||||
|
||||
|
||||
def create_model_index():
|
||||
"""Create model_index.json for FastVideo native model."""
|
||||
return {
|
||||
"_class_name": "LongCatPipeline",
|
||||
"_diffusers_version": "0.32.0",
|
||||
"workload_type": "video-generation",
|
||||
"tokenizer": ["transformers", "AutoTokenizer"],
|
||||
"text_encoder": ["transformers", "UMT5EncoderModel"],
|
||||
"vae": ["diffusers", "AutoencoderKLWan"],
|
||||
"scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
|
||||
"transformer": ["diffusers", "LongCatTransformer3DModel"] # Native model
|
||||
}
|
||||
|
||||
|
||||
def update_transformer_config(transformer_dir: Path):
|
||||
"""Update transformer config.json to point to native model."""
|
||||
config_path = transformer_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
print(" ⚠️ Transformer config not found, skipping")
|
||||
return
|
||||
|
||||
with open(config_path, 'r') as f:
|
||||
config = json.load(f)
|
||||
|
||||
if '_class_name' in config:
|
||||
old_class = config['_class_name']
|
||||
config['_class_name'] = 'LongCatTransformer3DModel'
|
||||
print(f" Updated _class_name: {old_class} → LongCatTransformer3DModel")
|
||||
else:
|
||||
config['_class_name'] = 'LongCatTransformer3DModel'
|
||||
print(f" Added _class_name: LongCatTransformer3DModel")
|
||||
|
||||
# Fix num_heads -> num_attention_heads for FastVideo compatibility
|
||||
if 'num_heads' in config and 'num_attention_heads' not in config:
|
||||
config['num_attention_heads'] = config.pop('num_heads')
|
||||
print(f" Updated num_heads → num_attention_heads")
|
||||
|
||||
with open(config_path, 'w') as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
print(" ✓ Transformer config updated")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# LoRA Conversion Functions
|
||||
# ============================================================================
|
||||
|
||||
def parse_lora_key(key: str) -> tuple[str, str]:
|
||||
"""Parse LongCat LoRA key into module path and weight type."""
|
||||
if key.startswith("lora___lorahyphen___"):
|
||||
key = key[len("lora___lorahyphen___"):]
|
||||
|
||||
key = key.replace("___lorahyphen___", ".")
|
||||
|
||||
if ".lora_down.weight" in key:
|
||||
return key.replace(".lora_down.weight", ""), "lora_down.weight"
|
||||
elif ".lora_up.weight" in key:
|
||||
return key.replace(".lora_up.weight", ""), "lora_up.weight"
|
||||
elif ".lora_up.blocks." in key:
|
||||
match = re.match(r"(.+)\.lora_up\.blocks\.(\d+)\.weight", key)
|
||||
if match:
|
||||
return match.group(1), f"lora_up.blocks.{match.group(2)}.weight"
|
||||
elif ".alpha_scale" in key:
|
||||
return key.replace(".alpha_scale", ""), "alpha_scale"
|
||||
elif ".lora_alpha" in key:
|
||||
return key.replace(".lora_alpha", ""), "lora_alpha"
|
||||
|
||||
raise ValueError(f"Unknown LoRA key format: {key}")
|
||||
|
||||
|
||||
def map_lora_module(module_path: str) -> list[tuple[str, str]]:
|
||||
"""Map LongCat module path to FastVideo paths. Returns [(path, component)]."""
|
||||
# Self-attention QKV → Q, K, V
|
||||
match = re.match(r"blocks\.(\d+)\.attn\.qkv", module_path)
|
||||
if match:
|
||||
b = match.group(1)
|
||||
return [(f"blocks.{b}.self_attn.to_q", "q"),
|
||||
(f"blocks.{b}.self_attn.to_k", "k"),
|
||||
(f"blocks.{b}.self_attn.to_v", "v")]
|
||||
|
||||
# Self-attention output
|
||||
match = re.match(r"blocks\.(\d+)\.attn\.proj", module_path)
|
||||
if match:
|
||||
return [(f"blocks.{match.group(1)}.self_attn.to_out", "single")]
|
||||
|
||||
# Cross-attention Q
|
||||
match = re.match(r"blocks\.(\d+)\.cross_attn\.q_linear", module_path)
|
||||
if match:
|
||||
return [(f"blocks.{match.group(1)}.cross_attn.to_q", "single")]
|
||||
|
||||
# Cross-attention KV → K, V
|
||||
match = re.match(r"blocks\.(\d+)\.cross_attn\.kv_linear", module_path)
|
||||
if match:
|
||||
b = match.group(1)
|
||||
return [(f"blocks.{b}.cross_attn.to_k", "k"),
|
||||
(f"blocks.{b}.cross_attn.to_v", "v")]
|
||||
|
||||
# FFN
|
||||
match = re.match(r"blocks\.(\d+)\.ffn\.(w[123])", module_path)
|
||||
if match:
|
||||
return [(f"blocks.{match.group(1)}.ffn.{match.group(2)}", "single")]
|
||||
|
||||
# AdaLN modulation
|
||||
match = re.match(r"blocks\.(\d+)\.adaLN_modulation\.1", module_path)
|
||||
if match:
|
||||
return [(f"blocks.{match.group(1)}.adaln_linear_1", "single")]
|
||||
|
||||
# Final layer
|
||||
if module_path == "final_layer.adaLN_modulation.1":
|
||||
return [("final_layer.adaln_linear", "single")]
|
||||
if module_path == "final_layer.linear":
|
||||
return [("final_layer.proj", "single")]
|
||||
|
||||
raise ValueError(f"Unknown LoRA module: {module_path}")
|
||||
|
||||
|
||||
def convert_lora_weights(source_weights: dict[str, torch.Tensor], lora_name: str) -> dict[str, torch.Tensor]:
|
||||
"""Convert LongCat LoRA to FastVideo format."""
|
||||
print(f" Converting {lora_name}...")
|
||||
print(f" Source keys: {len(source_weights)}")
|
||||
|
||||
converted = OrderedDict()
|
||||
|
||||
# Group by module
|
||||
modules = {}
|
||||
for key in source_weights.keys():
|
||||
try:
|
||||
module_path, weight_type = parse_lora_key(key)
|
||||
if module_path not in modules:
|
||||
modules[module_path] = {}
|
||||
modules[module_path][weight_type] = key
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Process each module
|
||||
for module_path, weight_keys in modules.items():
|
||||
try:
|
||||
targets = map_lora_module(module_path)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Get alpha_scale if present (defaults to 1.0 if missing)
|
||||
alpha_scale = 1.0
|
||||
if "alpha_scale" in weight_keys:
|
||||
alpha_scale_tensor = source_weights[weight_keys["alpha_scale"]]
|
||||
alpha_scale = alpha_scale_tensor.item() if alpha_scale_tensor.numel() == 1 else float(alpha_scale_tensor.mean())
|
||||
|
||||
# Handle lora_down (lora_A)
|
||||
if "lora_down.weight" in weight_keys:
|
||||
lora_down = source_weights[weight_keys["lora_down.weight"]]
|
||||
|
||||
if len(targets) == 1:
|
||||
converted[f"{targets[0][0]}.lora_A"] = lora_down
|
||||
# Compute alpha from alpha_scale and rank
|
||||
rank = lora_down.shape[0]
|
||||
alpha = alpha_scale * rank
|
||||
converted[f"{targets[0][0]}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32)
|
||||
else:
|
||||
# Split for fused projections
|
||||
n = len(targets)
|
||||
rank = lora_down.shape[0] // n
|
||||
for i, (path, _) in enumerate(targets):
|
||||
converted[f"{path}.lora_A"] = lora_down[i*rank:(i+1)*rank, :]
|
||||
# Compute alpha from alpha_scale and rank for each split
|
||||
alpha = alpha_scale * rank
|
||||
converted[f"{path}.lora_alpha"] = torch.tensor(alpha, dtype=torch.float32)
|
||||
|
||||
# Handle lora_up (lora_B) - may have multiple blocks
|
||||
lora_up_blocks = []
|
||||
i = 0
|
||||
while f"lora_up.blocks.{i}.weight" in weight_keys:
|
||||
lora_up_blocks.append(source_weights[weight_keys[f"lora_up.blocks.{i}.weight"]])
|
||||
i += 1
|
||||
|
||||
if lora_up_blocks:
|
||||
# Multi-block LoRA: construct block-diagonal lora_B
|
||||
# This is equivalent to the multi-block computation without modifying fastvideo
|
||||
n_blocks = len(lora_up_blocks)
|
||||
out_per_block, rank_per_block = lora_up_blocks[0].shape # e.g., [4096, 128]
|
||||
|
||||
if len(targets) == 1:
|
||||
# Single layer with multi-block: create block-diagonal matrix
|
||||
total_out = out_per_block * n_blocks
|
||||
total_rank = rank_per_block * n_blocks
|
||||
lora_B_blockdiag = torch.zeros(total_out, total_rank, dtype=lora_up_blocks[0].dtype)
|
||||
|
||||
for i in range(n_blocks):
|
||||
lora_B_blockdiag[i*out_per_block:(i+1)*out_per_block,
|
||||
i*rank_per_block:(i+1)*rank_per_block] = lora_up_blocks[i]
|
||||
|
||||
converted[f"{targets[0][0]}.lora_B"] = lora_B_blockdiag
|
||||
# Note: rank for alpha calculation should be total_rank (will be computed from lora_A.shape[0])
|
||||
else:
|
||||
# Multi-block with split targets (e.g., QKV split)
|
||||
# Each target gets one block
|
||||
for i, (path, _) in enumerate(targets):
|
||||
if i < n_blocks:
|
||||
converted[f"{path}.lora_B"] = lora_up_blocks[i]
|
||||
elif "lora_up.weight" in weight_keys:
|
||||
lora_up = source_weights[weight_keys["lora_up.weight"]]
|
||||
# Split if needed
|
||||
if len(targets) == 1:
|
||||
converted[f"{targets[0][0]}.lora_B"] = lora_up
|
||||
else:
|
||||
n = len(targets)
|
||||
out_dim = lora_up.shape[0] // n
|
||||
for i, (path, _) in enumerate(targets):
|
||||
converted[f"{path}.lora_B"] = lora_up[i*out_dim:(i+1)*out_dim, :]
|
||||
else:
|
||||
continue
|
||||
|
||||
print(f" Output keys: {len(converted)} (including lora_alpha)")
|
||||
# Count how many lora_alpha values were added
|
||||
alpha_count = sum(1 for k in converted.keys() if "lora_alpha" in k)
|
||||
print(f" Alpha values saved: {alpha_count}")
|
||||
return converted
|
||||
|
||||
|
||||
def convert_loras(source_dir: Path, output_dir: Path) -> bool:
|
||||
"""Convert all LoRA files in source directory."""
|
||||
lora_source = source_dir / "lora"
|
||||
if not lora_source.exists():
|
||||
print(" No LoRA directory found, skipping")
|
||||
return False
|
||||
|
||||
lora_files = list(lora_source.glob("*.safetensors"))
|
||||
if not lora_files:
|
||||
print(" No LoRA files found, skipping")
|
||||
return False
|
||||
|
||||
print(f" Found {len(lora_files)} LoRA file(s)")
|
||||
|
||||
# Map LoRA filenames to subdirectory names for FastVideo compatibility
|
||||
# Each LoRA gets its own directory under lora/
|
||||
lora_subdir_mapping = {
|
||||
"cfg_step_lora.safetensors": "distilled",
|
||||
"refinement_lora.safetensors": "refinement",
|
||||
}
|
||||
|
||||
for lora_file in lora_files:
|
||||
try:
|
||||
# Determine output subdirectory - use mapping if available, otherwise generic name
|
||||
if lora_file.name in lora_subdir_mapping:
|
||||
lora_subdir_name = lora_subdir_mapping[lora_file.name]
|
||||
else:
|
||||
# For unknown LoRAs, create subdirectory based on filename
|
||||
lora_subdir_name = lora_file.stem
|
||||
|
||||
lora_output = output_dir / "lora" / lora_subdir_name
|
||||
lora_output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load
|
||||
source_weights = load_file(str(lora_file))
|
||||
|
||||
# Convert
|
||||
converted = convert_lora_weights(source_weights, lora_file.stem)
|
||||
|
||||
# Save
|
||||
output_file = lora_output / lora_file.name
|
||||
save_file(converted, str(output_file))
|
||||
|
||||
size_mb = output_file.stat().st_size / (1024**2)
|
||||
print(f" ✓ {lora_file.name} → lora/{lora_subdir_name}/ ({size_mb:.1f} MB)")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ❌ Failed to convert {lora_file.name}: {e}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert LongCat weights to FastVideo native format"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to original LongCat weights (LongCat-Video/weights/LongCat-Video/)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to output directory for native weights",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validate",
|
||||
action="store_true",
|
||||
help="Run validation after conversion",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
source_dir = Path(args.source)
|
||||
output_dir = Path(args.output)
|
||||
|
||||
# Check source directory
|
||||
if not source_dir.exists():
|
||||
print(f"❌ Error: Source directory not found: {source_dir}")
|
||||
return 1
|
||||
|
||||
# Check for dit/transformer directory (original uses 'dit', we output to 'transformer')
|
||||
transformer_source = source_dir / "dit"
|
||||
if not transformer_source.exists():
|
||||
print(f"❌ Error: DiT directory not found in source")
|
||||
return 1
|
||||
|
||||
print("=" * 60)
|
||||
print("LongCat → FastVideo Native Conversion")
|
||||
print("=" * 60)
|
||||
print(f"Source: {source_dir}")
|
||||
print(f"Output: {output_dir}")
|
||||
print()
|
||||
|
||||
# Step 1: Convert transformer weights
|
||||
print("[Step 1/4] Converting transformer weights...")
|
||||
|
||||
# Load source weights
|
||||
shard_files = sorted(glob.glob(str(transformer_source / "*.safetensors")))
|
||||
if not shard_files:
|
||||
print(f"❌ Error: No safetensors files found in {transformer_source}")
|
||||
return 1
|
||||
|
||||
print(f" Found {len(shard_files)} shard(s)")
|
||||
source_weights = {}
|
||||
for shard_file in shard_files:
|
||||
print(f" Loading {Path(shard_file).name}...")
|
||||
source_weights.update(load_file(shard_file))
|
||||
|
||||
print(f" Loaded {len(source_weights)} parameters")
|
||||
|
||||
# Convert
|
||||
converted_weights = convert_transformer_weights(source_weights)
|
||||
print(f" Converted to {len(converted_weights)} parameters")
|
||||
|
||||
# Validate if requested
|
||||
if args.validate:
|
||||
if not validate_conversion(source_weights, converted_weights):
|
||||
print("\n❌ Validation failed!")
|
||||
return 1
|
||||
print("\n✓ Validation passed!")
|
||||
|
||||
# Save
|
||||
transformer_output = output_dir / "transformer"
|
||||
transformer_output.mkdir(parents=True, exist_ok=True)
|
||||
output_file = transformer_output / "model.safetensors"
|
||||
print(f"\n Saving to {output_file}...")
|
||||
save_file(converted_weights, str(output_file))
|
||||
size_gb = output_file.stat().st_size / (1024**3)
|
||||
print(f" ✓ Saved ({size_gb:.2f} GB)")
|
||||
print()
|
||||
|
||||
# Step 2: Copy other components
|
||||
print("[Step 2/5] Copying other components...")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
components = ["vae", "text_encoder", "tokenizer", "scheduler"]
|
||||
for component in components:
|
||||
copy_component(source_dir, output_dir, component)
|
||||
|
||||
print()
|
||||
|
||||
# Step 3: Convert LoRA weights
|
||||
print("[Step 3/5] Converting LoRA weights...")
|
||||
convert_loras(source_dir, output_dir)
|
||||
print()
|
||||
|
||||
# Step 4: Update transformer config
|
||||
print("[Step 4/5] Updating transformer config...")
|
||||
|
||||
# Copy config.json from source if exists
|
||||
source_config = transformer_source / "config.json"
|
||||
output_config = transformer_output / "config.json"
|
||||
if source_config.exists():
|
||||
shutil.copy(source_config, output_config)
|
||||
print(f" Copied config.json")
|
||||
|
||||
update_transformer_config(transformer_output)
|
||||
print()
|
||||
|
||||
# Step 5: Create model_index.json
|
||||
print("[Step 5/5] Creating model_index.json...")
|
||||
model_index_path = output_dir / "model_index.json"
|
||||
with open(model_index_path, 'w') as f:
|
||||
json.dump(create_model_index(), f, indent=2)
|
||||
print(f" ✓ Created {model_index_path}")
|
||||
print()
|
||||
|
||||
print("=" * 60)
|
||||
print("✓ Conversion Complete!")
|
||||
print("=" * 60)
|
||||
print(f"Native weights ready at: {output_dir}")
|
||||
print()
|
||||
print("Converted components:")
|
||||
print(" ✓ Transformer (native FastVideo implementation)")
|
||||
print(" ✓ VAE, text encoder, tokenizer, scheduler")
|
||||
if (output_dir / "lora").exists():
|
||||
lora_dirs = [d for d in (output_dir / "lora").iterdir() if d.is_dir()]
|
||||
if lora_dirs:
|
||||
print(f" ✓ LoRA weights ({len(lora_dirs)} adapters)")
|
||||
for lora_dir in sorted(lora_dirs):
|
||||
print(f" - lora/{lora_dir.name}/")
|
||||
print()
|
||||
print("Next steps:")
|
||||
print()
|
||||
print(" 1. Test basic generation:")
|
||||
print(" from fastvideo import VideoGenerator")
|
||||
print(f" generator = VideoGenerator.from_pretrained('{output_dir}')")
|
||||
print(" video = generator.generate_video(")
|
||||
print(" prompt='A cat playing piano',")
|
||||
print(" num_inference_steps=50")
|
||||
print(" )")
|
||||
print()
|
||||
if (output_dir / "lora" / "distilled").exists():
|
||||
print(" 2. Test distilled generation (16 steps with LoRA):")
|
||||
print(f" generator = VideoGenerator.from_pretrained('{output_dir}',")
|
||||
print(f" lora_path='{output_dir}/lora/distilled',")
|
||||
print(" lora_nickname='distilled')")
|
||||
print(" video = generator.generate_video(")
|
||||
print(" prompt='A cat playing piano',")
|
||||
print(" num_inference_steps=16,")
|
||||
print(" guidance_scale=1.0)")
|
||||
print()
|
||||
print()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit(main())
|
||||
@@ -0,0 +1,241 @@
|
||||
import os
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from safetensors import safe_open
|
||||
|
||||
|
||||
def validate_components(model_path):
|
||||
"""Validate all model components exist."""
|
||||
print("=" * 60)
|
||||
print("VALIDATING COMPONENTS")
|
||||
print("=" * 60)
|
||||
|
||||
components = {
|
||||
"tokenizer": ["special_tokens_map.json", "tokenizer_config.json"],
|
||||
"text_encoder": ["config.json", "model.safetensors.index.json"],
|
||||
"vae": ["config.json", "diffusion_pytorch_model.safetensors"],
|
||||
"scheduler": ["scheduler_config.json"],
|
||||
"transformer": ["config.json", "diffusion_pytorch_model.safetensors.index.json"]
|
||||
}
|
||||
|
||||
all_valid = True
|
||||
for component, required_files in components.items():
|
||||
component_path = os.path.join(model_path, component)
|
||||
print(f"\n{component}:")
|
||||
|
||||
if not os.path.exists(component_path):
|
||||
print(f" ✗ Directory not found")
|
||||
all_valid = False
|
||||
continue
|
||||
|
||||
for req_file in required_files:
|
||||
file_path = os.path.join(component_path, req_file)
|
||||
exists = os.path.exists(file_path)
|
||||
symbol = "✓" if exists else "✗"
|
||||
print(f" {symbol} {req_file}")
|
||||
if not exists:
|
||||
all_valid = False
|
||||
|
||||
return all_valid
|
||||
|
||||
|
||||
def validate_dit_weights(dit_path):
|
||||
"""Validate DiT weights structure."""
|
||||
print("\n" + "=" * 60)
|
||||
print("VALIDATING DiT WEIGHTS")
|
||||
print("=" * 60)
|
||||
|
||||
# Load config
|
||||
config_path = os.path.join(dit_path, "config.json")
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
|
||||
hidden_size = config["hidden_size"]
|
||||
depth = config["depth"]
|
||||
num_heads = config["num_heads"]
|
||||
|
||||
print(f"\nArchitecture:")
|
||||
print(f" - hidden_size: {hidden_size}")
|
||||
print(f" - depth: {depth}")
|
||||
print(f" - num_heads: {num_heads}")
|
||||
|
||||
# Load weight index
|
||||
index_path = os.path.join(dit_path, "diffusion_pytorch_model.safetensors.index.json")
|
||||
with open(index_path) as f:
|
||||
index = json.load(f)
|
||||
|
||||
weight_map = index["weight_map"]
|
||||
all_keys = list(weight_map.keys())
|
||||
|
||||
print(f"\nWeight statistics:")
|
||||
print(f" - Total keys: {len(all_keys)}")
|
||||
print(f" - Total size: {index['metadata']['total_size'] / 1e9:.2f} GB")
|
||||
|
||||
# Check structure
|
||||
embedder_keys = [k for k in all_keys if 'embedder' in k]
|
||||
block_keys = [k for k in all_keys if k.startswith('blocks.')]
|
||||
final_keys = [k for k in all_keys if k.startswith('final_layer.')]
|
||||
|
||||
print(f"\nKey distribution:")
|
||||
print(f" - Embedder layers: {len(embedder_keys)}")
|
||||
print(f" - Transformer blocks: {len(block_keys)}")
|
||||
print(f" - Final layer: {len(final_keys)}")
|
||||
|
||||
# Verify all blocks present
|
||||
block_nums = set()
|
||||
for key in block_keys:
|
||||
if key.startswith('blocks.'):
|
||||
block_num = int(key.split('.')[1])
|
||||
block_nums.add(block_num)
|
||||
|
||||
expected_blocks = set(range(depth))
|
||||
missing_blocks = expected_blocks - block_nums
|
||||
|
||||
if missing_blocks:
|
||||
print(f"\n✗ Missing blocks: {sorted(missing_blocks)}")
|
||||
return False
|
||||
else:
|
||||
print(f"\n✓ All {depth} blocks present (0-{depth-1})")
|
||||
|
||||
# Sample weights
|
||||
first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
|
||||
print(f"\nSampling weights from first shard:")
|
||||
|
||||
with safe_open(first_shard, framework="pt", device="cpu") as f:
|
||||
sample_keys = [k for k in f.keys() if k in all_keys][:5]
|
||||
for key in sample_keys:
|
||||
tensor = f.get_tensor(key)
|
||||
print(f" - {key}")
|
||||
print(f" Shape: {tuple(tensor.shape)}, Dtype: {tensor.dtype}")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def validate_shapes(dit_path):
|
||||
"""Validate weight shapes match expected architecture."""
|
||||
print("\n" + "=" * 60)
|
||||
print("VALIDATING WEIGHT SHAPES")
|
||||
print("=" * 60)
|
||||
|
||||
# Load config
|
||||
config_path = os.path.join(dit_path, "config.json")
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
|
||||
hidden_size = config["hidden_size"]
|
||||
num_heads = config["num_heads"]
|
||||
head_dim = hidden_size // num_heads
|
||||
adaln_dim = config.get("adaln_tembed_dim", 512)
|
||||
mlp_ratio = config.get("mlp_ratio", 4)
|
||||
|
||||
# Calculate FFN hidden_dim using SwiGLU formula from blocks.py
|
||||
# hidden_dim = int(2 * (hidden_size * mlp_ratio) / 3)
|
||||
# rounded to multiple_of=256
|
||||
multiple_of = 256
|
||||
ffn_hidden = int(2 * hidden_size * mlp_ratio / 3)
|
||||
ffn_hidden = multiple_of * ((ffn_hidden + multiple_of - 1) // multiple_of)
|
||||
|
||||
# Expected shapes
|
||||
expected = {
|
||||
"x_embedder.proj.weight": (hidden_size, 16, 1, 2, 2),
|
||||
"x_embedder.proj.bias": (hidden_size,),
|
||||
"t_embedder.mlp.0.weight": (adaln_dim, 256),
|
||||
"t_embedder.mlp.2.weight": (adaln_dim, adaln_dim),
|
||||
"y_embedder.y_proj.0.weight": (hidden_size, 4096),
|
||||
"blocks.0.attn.qkv.weight": (3 * hidden_size, hidden_size),
|
||||
"blocks.0.attn.q_norm.weight": (head_dim,),
|
||||
"blocks.0.attn.proj.weight": (hidden_size, hidden_size),
|
||||
"blocks.0.cross_attn.q_linear.weight": (hidden_size, hidden_size),
|
||||
"blocks.0.cross_attn.kv_linear.weight": (2 * hidden_size, hidden_size),
|
||||
"blocks.0.ffn.w1.weight": (ffn_hidden, hidden_size),
|
||||
"blocks.0.adaLN_modulation.1.weight": (6 * hidden_size, adaln_dim),
|
||||
"final_layer.linear.weight": (64, hidden_size),
|
||||
}
|
||||
|
||||
# Load and check
|
||||
first_shard = os.path.join(dit_path, "diffusion_pytorch_model-00001-of-00006.safetensors")
|
||||
|
||||
all_valid = True
|
||||
with safe_open(first_shard, framework="pt", device="cpu") as f:
|
||||
for key, expected_shape in expected.items():
|
||||
if key in f.keys():
|
||||
tensor = f.get_tensor(key)
|
||||
actual_shape = tuple(tensor.shape)
|
||||
|
||||
if actual_shape == expected_shape:
|
||||
print(f"✓ {key}: {actual_shape}")
|
||||
else:
|
||||
print(f"✗ {key}: expected {expected_shape}, got {actual_shape}")
|
||||
all_valid = False
|
||||
|
||||
return all_valid
|
||||
|
||||
|
||||
def validate_model_index(model_path):
|
||||
"""Validate model_index.json exists and is correct."""
|
||||
print("\n" + "=" * 60)
|
||||
print("VALIDATING MODEL INDEX")
|
||||
print("=" * 60)
|
||||
|
||||
model_index_path = os.path.join(model_path, "model_index.json")
|
||||
|
||||
if not os.path.exists(model_index_path):
|
||||
print("✗ model_index.json not found")
|
||||
return False
|
||||
|
||||
with open(model_index_path) as f:
|
||||
index = json.load(f)
|
||||
|
||||
required_keys = ["_class_name", "workload_type", "tokenizer", "text_encoder",
|
||||
"vae", "scheduler", "transformer"]
|
||||
|
||||
all_valid = True
|
||||
for key in required_keys:
|
||||
if key in index:
|
||||
print(f"✓ {key}: {index[key]}")
|
||||
else:
|
||||
print(f"✗ {key}: missing")
|
||||
all_valid = False
|
||||
|
||||
return all_valid
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Validate LongCat weights for FastVideo")
|
||||
parser.add_argument("--model-path", type=str, required=True,
|
||||
help="Path to LongCat model directory")
|
||||
parser.add_argument("--check-shapes", action="store_true",
|
||||
help="Also validate weight shapes (slower)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"\nValidating: {args.model_path}\n")
|
||||
|
||||
# Run validations
|
||||
components_valid = validate_components(args.model_path)
|
||||
model_index_valid = validate_model_index(args.model_path)
|
||||
dit_valid = validate_dit_weights(os.path.join(args.model_path, "transformer"))
|
||||
|
||||
if args.check_shapes:
|
||||
shapes_valid = validate_shapes(os.path.join(args.model_path, "transformer"))
|
||||
else:
|
||||
shapes_valid = True
|
||||
print("\nSkipping shape validation (use --check-shapes to enable)")
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 60)
|
||||
print("VALIDATION SUMMARY")
|
||||
print("=" * 60)
|
||||
|
||||
all_valid = components_valid and model_index_valid and dit_valid and shapes_valid
|
||||
|
||||
if all_valid:
|
||||
print("✓ All validations passed!")
|
||||
print("✓ Model ready for FastVideo")
|
||||
else:
|
||||
print("✗ Some validations failed")
|
||||
print("✗ Please check errors above")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
Executable
+32
@@ -0,0 +1,32 @@
|
||||
#!/bin/bash
|
||||
|
||||
num_gpus=2
|
||||
|
||||
export FASTVIDEO_ATTENTION_BACKEND=
|
||||
|
||||
# For longcat, we must first convert the official weights to FastVideo native format
|
||||
# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
|
||||
# --source /path/to/LongCat-Video/weights/LongCat-Video
|
||||
# --output weights/longcat-native
|
||||
export MODEL_BASE=weights/longcat-native
|
||||
|
||||
fastvideo generate \
|
||||
--model-path $MODEL_BASE \
|
||||
--sp-size $num_gpus \
|
||||
--tp-size 1 \
|
||||
--num-gpus $num_gpus \
|
||||
--dit-cpu-offload False \
|
||||
--vae-cpu-offload False \
|
||||
--text-encoder-cpu-offload False \
|
||||
--pin-cpu-memory False \
|
||||
--enable-bsa False \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 93 \
|
||||
--num-inference-steps 50 \
|
||||
--fps 15 \
|
||||
--guidance-scale 4.0 \
|
||||
--prompt-txt assets/prompt.txt \
|
||||
--negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
|
||||
--seed 42 \
|
||||
--output-path outputs_video/longcat_480p
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
|
||||
num_gpus=1
|
||||
export FASTVIDEO_ATTENTION_BACKEND=
|
||||
# For longcat, we must first convert the official weights to FastVideo native format
|
||||
# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
|
||||
# --source /path/to/LongCat-Video/weights/LongCat-Video
|
||||
# --output weights/longcat-native
|
||||
export MODEL_BASE=weights/longcat-native
|
||||
|
||||
fastvideo generate \
|
||||
--model-path $MODEL_BASE \
|
||||
--sp-size $num_gpus \
|
||||
--tp-size 1 \
|
||||
--num-gpus $num_gpus \
|
||||
--dit-cpu-offload False \
|
||||
--vae-cpu-offload False \
|
||||
--text-encoder-cpu-offload False \
|
||||
--pin-cpu-memory False \
|
||||
--enable-bsa False \
|
||||
--lora-path "$MODEL_BASE/lora/distilled" \
|
||||
--lora-nickname "distilled" \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 93 \
|
||||
--num-inference-steps 16 \
|
||||
--fps 15 \
|
||||
--guidance-scale 1.0 \
|
||||
--prompt "In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene." \
|
||||
--seed 42 \
|
||||
--output-path outputs_video/longcat_distill
|
||||
@@ -0,0 +1,73 @@
|
||||
#!/bin/bash
|
||||
|
||||
num_gpus=1
|
||||
export FASTVIDEO_ATTENTION_BACKEND=
|
||||
# For longcat, we must first convert the official weights to FastVideo native format
|
||||
# conversion method: python scripts/checkpoint_conversion/longcat_to_fastvideo.py
|
||||
# --source /path/to/LongCat-Video/weights/LongCat-Video
|
||||
# --output weights/longcat-native
|
||||
export MODEL_BASE=weights/longcat-native
|
||||
|
||||
INPUT_VIDEO="outputs_video/longcat_distill/In a realistic photography style, an asian boy around seven or eight years old sits on a park bench,.mp4"
|
||||
REFINE_OUTPUT="outputs_video/longcat_refine_720p"
|
||||
|
||||
# Prompt used for base generation
|
||||
PROMPT="In a realistic photography style, an asian boy around seven or eight years old sits on a park bench, wearing a light yellow T-shirt, denim shorts, and white sneakers. He holds an ice cream cone with vanilla and chocolate flavors, and beside him is a medium-sized golden Labrador. Smiling, the boy offers the ice cream to the dog, who eagerly licks it with its tongue. The sun is shining brightly, and the background features a green lawn and several tall trees, creating a warm and loving scene."
|
||||
|
||||
echo "=========================================="
|
||||
echo "LongCat 480p -> 720p Refinement"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
echo "Input: $INPUT_VIDEO"
|
||||
echo "Output: $REFINE_OUTPUT"
|
||||
echo ""
|
||||
|
||||
# Check if input video exists
|
||||
if [ ! -f "$INPUT_VIDEO" ]; then
|
||||
echo "Error: Input video not found: $INPUT_VIDEO"
|
||||
echo "Please set INPUT_VIDEO to your 480p video path"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "🔧 Configuring refinement (BSA enabled, refinement LoRA)..."
|
||||
echo "✅ Input video: $INPUT_VIDEO"
|
||||
echo "✅ BSA enabled with sparsity=0.875"
|
||||
echo "✅ Refinement LoRA loaded"
|
||||
echo ""
|
||||
|
||||
fastvideo generate \
|
||||
--model-path $MODEL_BASE \
|
||||
--sp-size $num_gpus \
|
||||
--tp-size 1 \
|
||||
--num-gpus $num_gpus \
|
||||
--dit-cpu-offload True \
|
||||
--vae-cpu-offload False \
|
||||
--text-encoder-cpu-offload True \
|
||||
--pin-cpu-memory False \
|
||||
--enable-bsa True \
|
||||
--bsa-sparsity 0.875 \
|
||||
--bsa-chunk-q 4 4 8 \
|
||||
--bsa-chunk-k 4 4 8 \
|
||||
--lora-path "$MODEL_BASE/lora/refinement" \
|
||||
--lora-nickname "refinement" \
|
||||
--refine-from "$INPUT_VIDEO" \
|
||||
--t-thresh 0.5 \
|
||||
--spatial-refine-only False \
|
||||
--num-cond-frames 0 \
|
||||
--height 720 \
|
||||
--width 1280 \
|
||||
--num-inference-steps 50 \
|
||||
--fps 30 \
|
||||
--guidance-scale 1.0 \
|
||||
--prompt "$PROMPT" \
|
||||
--seed 42 \
|
||||
--output-path "$REFINE_OUTPUT"
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "✓ Refinement Complete!"
|
||||
echo "=========================================="
|
||||
echo ""
|
||||
echo "Output directory: $REFINE_OUTPUT"
|
||||
echo ""
|
||||
|
||||
Reference in New Issue
Block a user