longcat: vendor the upstream source instead of cloning it (v0.2.52)

`longcat_video` — the pipeline, DiT, VAE, scheduler and INT8 quantisation helpers
— lives inside the upstream repo and is not on PyPI, so resolve_source_dir looked
for a checkout at ~/.coderai/LongCat-Video or a path baked at build time. Neither
exists in a published image, which means anyone pulling it had to clone from
GitHub before LongCat would run at all: a third-party network dependency in the
middle of someone else's install, on whatever `main` happened to be that day.

The 26 source files are now vendored under third_party/, MIT, pinned to commit
6b3f4b85. 496 KB, no weights, no wheels, nothing compiled. The resolver prefers
them over every fallback, and an explicit longcat_source still wins for anyone
running their own checkout.

They ship automatically: the image does COPY . /opt/coderai/app, and third_party
is not in .dockerignore — asserted by a test, since that is the sort of thing a
later ignore rule breaks silently.

The ISOLATED VENV is unaffected and still builds on first use from
requirements-longcat.txt. Only the source is vendored; torch 2.6+cu124 and the
rest are installed, not carried.

While confirming the fp8 question: upstream's quantization.py implements INT8
only — QuantizedLinear holds weight_int8 + weight_scale, and there is no fp8 code
anywhere in the package. The fp8 checkpoints on HF carry no loader either, and
szwagros/LongCat-Video-Avatar-1.5-fp8's own card is a verbatim copy of upstream's
and documents `--use_int8`, never fp8. So a VRAM-resident DiT goes through INT8,
not fp8.

88 tests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.2.51"
+__version__ = "0.2.52"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/longcat_worker.py

@@ -182,6 +182,13 @@ def resolve_source_dir(config: dict = None) -> Path:
     explicit = os.environ.get("CODERAI_LONGCAT_SRC")
     if explicit:
         return Path(os.path.expanduser(explicit))
+    # Vendored with coderai: third_party/longcat_video, MIT, pinned to a commit we
+    # tested against. Preferred over every fallback below so a fresh install — an
+    # image someone else pulls, a container, a pod — has the package already and
+    # never needs to clone an upstream repo. See third_party/README.md.
+    _vendored = Path(__file__).resolve().parents[2] / "third_party"
+    if (_vendored / "longcat_video").is_dir():
+        return _vendored
     if (_BAKED_SRC / "longcat_video").is_dir():
         return _BAKED_SRC
     return Path(os.path.expanduser("~/.coderai/LongCat-Video"))

tests/test_longcat_vendored_source.py

+"""The LongCat source is vendored, not cloned.
+
+`longcat_video` (the pipeline, DiT, VAE, scheduler and INT8 quantisation helpers)
+lives inside the upstream repo and is not on PyPI. Resolving it by cloning from
+GitHub at install time would mean an image someone else pulls cannot run LongCat
+without network access to a third party — so the 26 MIT-licensed source files are
+vendored under third_party/ and pinned to a commit.
+"""
+import pathlib
+import re
+
+import pytest
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+VENDOR = ROOT / "third_party"
+
+
+def test_the_package_is_present_in_the_repo():
+    """Upstream ships no top-level __init__.py — it is a namespace package — so the
+    modules that are actually imported are what to check for."""
+    pkg = VENDOR / "longcat_video"
+    assert pkg.is_dir()
+    assert (pkg / "pipeline_longcat_video.py").is_file()
+    assert (pkg / "pipeline_longcat_video_avatar.py").is_file()
+    # the INT8 loader coderai calls for a quantised DiT
+    assert (pkg / "modules" / "quantization.py").is_file()
+
+
+def test_the_licence_travels_with_it():
+    lic = VENDOR / "LONGCAT-VIDEO-LICENSE"
+    assert lic.is_file(), "vendored MIT code must carry its licence"
+    text = lic.read_text(encoding="utf-8")
+    assert "MIT License" in text
+    assert "Meituan" in text
+
+
+def test_the_provenance_pins_a_commit():
+    """'whatever main happened to be' is not a version."""
+    readme = (VENDOR / "README.md").read_text(encoding="utf-8")
+    assert "github.com/meituan-longcat/LongCat-Video" in readme
+    assert re.search(r"\b[0-9a-f]{40}\b", readme), "no commit hash recorded"
+
+
+def test_nothing_compiled_or_heavy_was_vendored():
+    """Source only: no weights, no wheels, no .so — it must stay reviewable and
+    small enough to live in the repo."""
+    files = [p for p in (VENDOR / "longcat_video").rglob("*") if p.is_file()]
+    assert files, "nothing vendored"
+    non_py = [p for p in files if p.suffix != ".py"]
+    assert not non_py, f"non-source files vendored: {non_py[:5]}"
+    total = sum(p.stat().st_size for p in files)
+    assert total < 5 * 1024 * 1024, f"{total/1e6:.1f} MB is too much to vendor"
+
+
+def test_the_resolver_prefers_the_vendored_copy():
+    lw = pytest.importorskip("codai.api.longcat_worker")
+    got = lw.resolve_source_dir()
+    assert (got / "longcat_video").is_dir(), f"{got} has no longcat_video"
+    assert got == VENDOR, "the vendored copy should win over clone/baked fallbacks"
+
+
+def test_an_explicit_setting_still_wins_over_the_vendored_copy(tmp_path):
+    """Someone running their own checkout must still be able to point at it."""
+    lw = pytest.importorskip("codai.api.longcat_worker")
+    (tmp_path / "longcat_video").mkdir()
+    assert lw.resolve_source_dir({"longcat_source": str(tmp_path)}) == tmp_path
+
+
+def test_it_is_not_excluded_from_the_image():
+    """COPY . /opt/coderai/app ships it — unless .dockerignore drops it."""
+    ignore = (ROOT / ".dockerignore").read_text(encoding="utf-8")
+    assert not any(l.strip().rstrip("/") == "third_party"
+                   for l in ignore.splitlines()), "third_party is dockerignored"

third_party/LONGCAT-VIDEO-LICENSE

+MIT License
+
+Copyright (c) 2025 Meituan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file

third_party/README.md

+# Vendored third-party source
+
+## `longcat_video/`
+
+Upstream: <https://github.com/meituan-longcat/LongCat-Video>
+Commit:   `6b3f4b8582a8bc3f20f795735f5383716c4ba794`
+Licence:  MIT (see `LONGCAT-VIDEO-LICENSE`) — © 2025 Meituan
+
+`LongCatVideoPipeline` and the DiT, VAE, scheduler and INT8 quantisation helpers
+live in a package *inside* the upstream repo and are not published on PyPI. They
+are vendored here, rather than cloned at install time, so that:
+
+* an image someone else pulls already has them — no network, no git, and no
+  "download the upstream sources" step in anyone's setup;
+* the version is pinned to a commit we tested against instead of whatever `main`
+  happens to be on the day of the install;
+* the isolated Python 3.10 venv still installs its own dependencies from
+  `requirements-longcat.txt` (torch 2.6+cu124, transformers 4.41, …) — only the
+  26 source files are vendored, no weights and no wheels.
+
+496 KB, 26 Python files, nothing compiled. To refresh it, copy `longcat_video/`
+from a clean checkout of the commit you intend to pin and update the hash above.

third_party/longcat_video/audio_process/__init__.py

+from .wav2vec2 import Wav2Vec2ModelWrapper
+from transformers import Wav2Vec2FeatureExtractor, WhisperModel, AutoFeatureExtractor
+
+def get_audio_encoder(checkpoint_path, model_type="avatar-v1.0"):
+    if model_type == "avatar-v1.0":
+        model = Wav2Vec2ModelWrapper(checkpoint_path)
+        model.feature_extractor._freeze_parameters()
+        return model
+    if model_type == "avatar-v1.5":
+        model = WhisperModel.from_pretrained(checkpoint_path).eval()
+        model.requires_grad_(False)
+        return model
+
+def get_audio_feature_extractor(checkpoint_path, model_type="avatar-v1.0"):
+    if model_type == "avatar-v1.0":
+        return Wav2Vec2FeatureExtractor(checkpoint_path, local_files_only=True)
+    if model_type == "avatar-v1.5":
+        return AutoFeatureExtractor.from_pretrained(checkpoint_path)

third_party/longcat_video/audio_process/torch_utils.py

+import os
+import json
+import librosa
+import binascii
+import imageio
+import subprocess
+import numpy as np
+import os.path as osp
+from tqdm import tqdm
+import pyloudnorm as pyln
+from einops import rearrange
+import scipy.signal as ss
+
+import torch
+import torch.nn.functional as F
+import torchvision
+
+from ..context_parallel import context_parallel_util
+
+
+def torch_gc():
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
+
+def linear_interpolation(features, seq_len):
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+
+
+@torch.compile
+def calculate_x_ref_attn_map(noise_q, ref_k, ref_target_masks, attn_bias=None):
+
+    # compute cross-reference attention maps between query features and reference key features.
+    ref_k = ref_k.to(noise_q.dtype).to(noise_q.device)
+    scale = 1.0 / noise_q.shape[-1] ** 0.5
+    noise_q = noise_q * scale
+    noise_q = noise_q.transpose(1, 2)
+    ref_k = ref_k.transpose(1, 2)
+    attn = noise_q @ ref_k.transpose(-2, -1)
+
+    if attn_bias is not None:
+        attn = attn + attn_bias
+
+    x_ref_attn_map_source = attn.softmax(-1)
+
+    x_ref_attn_maps = []
+    ref_target_masks = ref_target_masks.to(noise_q.dtype)
+    x_ref_attn_map_source = x_ref_attn_map_source.to(noise_q.dtype)
+
+    for _, ref_target_mask in enumerate(ref_target_masks):
+        torch_gc()
+        ref_target_mask = ref_target_mask[None, None, None, ...]
+        x_ref_attn_map = x_ref_attn_map_source.clone()
+        x_ref_attn_map = x_ref_attn_map * ref_target_mask
+        x_ref_attn_map = x_ref_attn_map.sum(-1) / ref_target_mask.sum() 
+        x_ref_attn_map = x_ref_attn_map.permute(0, 2, 1) 
+        x_ref_attn_map = x_ref_attn_map.mean(-1) 
+        
+        x_ref_attn_maps.append(x_ref_attn_map)
+    
+    del attn
+    del x_ref_attn_map_source
+    torch_gc()
+
+    return torch.concat(x_ref_attn_maps, dim=0)
+
+
+def get_attn_map_with_target(noise_q, key, shape, ref_target_masks=None, split_num=2, cp_split_hw=None):
+    
+    N_t, N_h, N_w = shape
+    x_seqlens = N_h * N_w
+    if cp_split_hw[0] * cp_split_hw[1] > 1:
+        cp_size = context_parallel_util.get_cp_size()
+        (split_h, split_w) = (cp_size, 1) if cp_split_hw is None else cp_split_hw
+
+        assert N_h % split_h == 0, N_w % split_w == 0
+
+        N_h_ = N_h // split_h
+        N_w_ = N_w // split_w
+        x_seqlens = N_h_ * N_w_
+    
+    ref_k = key[:, :x_seqlens] # reference image size after cp split
+    noise_q = noise_q.contiguous()
+
+    # gather reference image key features
+    if cp_split_hw[0] * cp_split_hw[1] > 1:
+        _, _, H, _ = ref_k.shape
+        ref_k = ref_k.permute(0, 2, 1, 3)
+        ref_k = rearrange(ref_k, "b h m k -> b (h m) k")
+        ref_k = context_parallel_util.gather_cp_2d(ref_k, shape=(H, N_h, N_w), split_hw=cp_split_hw) 
+        ref_k = rearrange(ref_k, "b (h m) k -> b h m k", h=H) 
+        ref_k = ref_k.permute(0, 2, 1, 3)
+
+    _, seq_lens, heads, _ = noise_q.shape
+    class_num, _ = ref_target_masks.shape
+    x_ref_attn_maps = torch.zeros(class_num, seq_lens).to(noise_q.device).to(noise_q.dtype)
+
+    split_chunk = heads // split_num
+    
+    # calculate attn map within each group and take the mean
+    for i in range(split_num):
+        x_ref_attn_maps_perhead = calculate_x_ref_attn_map(noise_q[:, :, i*split_chunk:(i+1)*split_chunk, :], ref_k[:, :, i*split_chunk:(i+1)*split_chunk, :], ref_target_masks)
+        x_ref_attn_maps += x_ref_attn_maps_perhead
+    
+    return x_ref_attn_maps / split_num
+
+
+def rand_name(length=8, suffix=''):
+    name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
+    if suffix:
+        if not suffix.startswith('.'):
+            suffix = '.' + suffix
+        name += suffix
+    return name
+
+def cache_video(tensor,
+                save_file=None,
+                fps=30,
+                suffix='.mp4',
+                nrow=8,
+                normalize=True,
+                value_range=(-1, 1),
+                retry=5):
+    
+    # cache file
+    cache_file = osp.join('/tmp', rand_name(
+        suffix=suffix)) if save_file is None else save_file
+
+    # save to cache
+    error = None
+    for _ in range(retry):
+       
+        # preprocess
+        tensor = tensor.clamp(min(value_range), max(value_range))
+        tensor = torch.stack([
+                torchvision.utils.make_grid(
+                    u, nrow=nrow, normalize=normalize, value_range=value_range)
+                for u in tensor.unbind(2)
+            ],
+                                 dim=1).permute(1, 2, 3, 0)
+        tensor = (tensor * 255).type(torch.uint8).cpu()
+
+        # write video
+        writer = imageio.get_writer(cache_file, fps=fps, codec='libx264', quality=10, ffmpeg_params=["-crf", "10"])
+        for frame in tensor.numpy():
+            writer.append_data(frame)
+        writer.close()
+        return cache_file
+
+def get_audio_duration(audio_path):
+    cmd = [
+        "ffprobe",
+        "-v", "quiet",
+        "-print_format", "json",
+        "-show_entries", "format=duration",
+        audio_path,
+    ]
+    out = subprocess.check_output(cmd)
+    info = json.loads(out)
+    return float(info["format"]["duration"])
+
+
+def save_video_ffmpeg(gen_video_samples, save_path, audio_path, fps=25, quality=5, high_quality_save=False):
+
+    def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None):
+        writer = imageio.get_writer(
+            save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params
+        )
+        for frame in tqdm(frames, desc="Saving video"):
+            frame = np.array(frame)
+            writer.append_data(frame)
+        writer.close()
+    save_path_tmp = save_path + "-temp.mp4"
+    
+    os.makedirs(os.path.dirname(save_path_tmp), exist_ok=True)
+    video_audio =  gen_video_samples.cpu().numpy()
+    video_audio = np.clip(video_audio, 0, 255).astype(np.uint8)
+    save_video(video_audio, save_path_tmp, fps=fps, quality=quality)
+
+    # crop audio according to video length
+    T, _, _, _ = gen_video_samples.shape
+    duration = T / fps
+    save_path_crop_audio = save_path + "-cropaudio.wav"
+    final_command = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        audio_path,
+        "-t",
+        f'{duration}',
+        save_path_crop_audio,
+    ]
+    subprocess.run(final_command, check=True)
+
+    # crop video according to audio length
+    crop_audio_duration = get_audio_duration(save_path_crop_audio)
+    save_path_crop_tmp = save_path + "-cropvideo.mp4"
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i", save_path_tmp,
+        "-t", f"{crop_audio_duration}",
+        "-c:v", "copy",
+        "-c:a", "copy",
+        save_path_crop_tmp,
+    ]
+    subprocess.run(cmd, check=True)
+
+    # generate video with audio
+    save_path = save_path + ".mp4"
+    if high_quality_save:
+        final_command = [
+            "ffmpeg",
+            "-y",
+            "-i", save_path_crop_tmp,
+            "-i", save_path_crop_audio,
+            "-c:v", "libx264",
+            "-crf", "0",
+            "-preset", "veryslow", 
+            "-c:a", "aac",
+            "-shortest",
+            save_path,
+        ]
+        subprocess.run(final_command, check=True)
+    else:
+        final_command = [
+            "ffmpeg",
+            "-y",
+            "-i",
+            save_path_crop_tmp,
+            "-i",
+            save_path_crop_audio,
+            "-c:v",
+            "libx264",
+            "-c:a",
+            "aac",
+            "-shortest",
+            save_path,
+        ]
+        subprocess.run(final_command, check=True)
+        
+    os.remove(save_path_tmp)
+    os.remove(save_path_crop_tmp)
+    os.remove(save_path_crop_audio)
\ No newline at end of file

third_party/longcat_video/audio_process/wav2vec2.py

+import os
+import copy
+import logging
+
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2Config
+from transformers import Wav2Vec2Model as Wav2Vec2Model_base
+from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2SamePadLayer, Wav2Vec2PositionalConvEmbedding
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput
+
+from .torch_utils import linear_interpolation
+
+
+def _Wav2Vec2PositionalConvEmbedding_init_hack_(self, config):
+        super(Wav2Vec2PositionalConvEmbedding, self).__init__()
+        self.conv = nn.Conv1d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=config.num_conv_pos_embeddings,
+            padding=config.num_conv_pos_embeddings // 2,
+            groups=config.num_conv_pos_embedding_groups,
+        )
+
+        weight_norm = nn.utils.weight_norm
+        if hasattr(nn.utils.parametrizations, "weight_norm"):
+            weight_norm = nn.utils.parametrizations.weight_norm
+        self.conv = weight_norm(self.conv, name="weight", dim=2)
+
+        self.padding = Wav2Vec2SamePadLayer(config.num_conv_pos_embeddings)
+        self.activation = ACT2FN[config.feat_extract_activation]
+
+Wav2Vec2PositionalConvEmbedding.__init__ = _Wav2Vec2PositionalConvEmbedding_init_hack_
+
+
+class Wav2Vec2ModelWrapper(nn.Module):
+    def __init__(self, config_path, device='cuda', prefix='wav2vec2.'):
+        super(Wav2Vec2ModelWrapper, self).__init__()
+
+        config, model_kwargs = Wav2Vec2Config.from_pretrained(
+                config_path,
+                return_unused_kwargs=True,
+                force_download=False,
+                local_files_only=True,
+            )
+
+        model_path = os.path.join(config_path, 'pytorch_model.bin')
+        state_dict = torch.load(model_path, map_location=device)
+
+        config.name_or_path = config_path
+        config = copy.deepcopy(config)  # We do not want to modify the config inplace in from_pretrained.
+        config = Wav2Vec2Mode._autoset_attn_implementation(config, use_flash_attention_2=False)
+
+        # init model
+        with torch.device('meta'):
+            model = Wav2Vec2Mode(config)
+
+        # load checkpoint
+        logging.info(f'loading {model_path}')
+        if prefix is not None:
+            state_dict = {i.replace(prefix, ''):state_dict[i] for i in state_dict}
+        
+        model.tie_weights()
+        m, u = model.load_state_dict(state_dict, assign=True, strict=False)
+            
+        model.tie_weights()
+        model.eval()
+
+        self.model = model
+    
+    @property
+    def feature_extractor(self):
+        return self.model.feature_extractor
+
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return self.model(
+            input_values,
+            seq_len,
+            attention_mask=attention_mask,
+            mask_time_indices=mask_time_indices,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+    
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return self.model.feature_extract(
+            input_values,
+            seq_len
+        )
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        return self.model.encode(
+            extract_features,
+            attention_mask=attention_mask,
+            mask_time_indices=mask_time_indices,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+# the implementation of Wav2Vec2Model is borrowed from
+# https://github.com/huggingface/transformers/blob/HEAD/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+# initialize our encoder with the pre-trained wav2vec 2.0 weights.
+class Wav2Vec2Mode(Wav2Vec2Model_base):
+    def __init__(self, config: Wav2Vec2Config):
+        config.attn_implementation = "eager"
+        super().__init__(config)
+
+
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config._attn_implementation = "eager"
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return extract_features
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+            
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

third_party/longcat_video/block_sparse_attention/common.py

+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)

third_party/longcat_video/block_sparse_attention/communicate.py

+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

third_party/longcat_video/context_parallel/context_parallel_util.py

+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+
+from einops import rearrange
+
+
+dp_size = dp_group = cp_group = cp_stream = dp_ranks = cp_ranks = dp_rank = None
+cp_size: int = 1
+cp_rank: int = 0
+
+
+def init_context_parallel(context_parallel_size: int = 1,
+                          global_rank: int = 0,
+                          world_size: int = 1,):
+
+    global dp_size, cp_size, dp_group, cp_group, dp_ranks, cp_ranks, dp_rank, cp_rank
+
+    if world_size % context_parallel_size != 0:
+        raise RuntimeError(f'world_size {world_size} must be multiple of context_parallel_size {context_parallel_size}')
+
+    cp_size = context_parallel_size
+    dp_size = world_size//context_parallel_size
+    print(f'[rank {global_rank}] init_device_mesh [dp_size x cp_size]: [{dp_size} x {cp_size}]')
+
+    mesh_2d = init_device_mesh("cuda", (dp_size, cp_size), mesh_dim_names=("dp", "cp"))
+    print(f'[rank {global_rank}] mesh_2d: {mesh_2d}')
+
+    dp_group = mesh_2d.get_group(mesh_dim="dp")
+    cp_group = mesh_2d.get_group(mesh_dim="cp")
+    dp_ranks = torch.distributed.get_process_group_ranks(dp_group)
+    cp_ranks = torch.distributed.get_process_group_ranks(cp_group)
+    dp_rank = dist.get_rank(group=dp_group)
+    cp_rank = dist.get_rank(group=cp_group)
+
+    curr_global_rank = torch.distributed.get_rank()
+    print(f'[rank {curr_global_rank}] [dp_rank, cp_rank]: [{dp_rank}, {cp_rank}],  dp_ranks: {dp_ranks}, cp_ranks: {cp_ranks}')
+
+
+def get_cp_size():
+    global cp_size
+    return cp_size
+
+
+def get_dp_size():
+    global dp_size
+    return dp_size
+
+
+def get_cp_stream():
+    global cp_stream
+    if cp_stream == None:
+        cp_stream = torch.cuda.Stream()
+    return cp_stream
+
+
+def get_dp_group():
+    global dp_group
+    return dp_group
+
+
+def get_cp_group():
+    global cp_group
+    return cp_group
+
+
+def get_dp_rank():
+    global dp_rank
+    return dp_rank
+
+
+def get_cp_rank():
+    global cp_rank
+    return cp_rank
+
+
+def get_cp_rank_list():
+    global cp_ranks
+    if cp_ranks == None:
+        cp_ranks = torch.distributed.get_process_group_ranks(cp_group)
+    return cp_ranks
+
+
+def cp_broadcast(tensor, cp_index=0):
+    global dp_group
+    global cp_group
+    cp_ranks = get_cp_rank_list()
+    torch.distributed.broadcast(tensor, cp_ranks[cp_index], group=cp_group)
+
+
+def split_tensor_in_cp_2d(input, dim_hw, split_hw):
+
+    global cp_size
+    
+    dim_h, dim_w = dim_hw
+    split_h, split_w = split_hw
+    
+    assert cp_size == split_h * split_w
+
+    seq_size_h = input.shape[dim_h]
+    seq_size_w = input.shape[dim_w]
+
+    if seq_size_h % split_h != 0:
+        raise RuntimeError(f'seq_size_h {seq_size_h} in dim_h {dim_h} must be multiple of split_h {split_h}!!!')
+    if seq_size_w % split_w != 0:
+        raise RuntimeError(f'seq_size_w {seq_size_w} in dim_w {dim_w} must be multiple of split_w {split_w}!!!')
+
+    split_seq_size_h = seq_size_h // split_h
+    split_seq_size_w = seq_size_w // split_w
+
+    tensor_splits_h = input.split(split_seq_size_h, dim=dim_h)
+    tensor_splits = []
+    for tensor_split_h in tensor_splits_h:
+        tensor_splits_hw = tensor_split_h.split(split_seq_size_w, dim=dim_w)
+        tensor_splits.extend(tensor_splits_hw)
+
+    cp_rank = get_cp_rank()
+
+    split_tensor = tensor_splits[cp_rank]
+
+    return split_tensor
+
+
+class GatherFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input, process_group, seq_dim, frames):
+        ctx.cp_group = process_group
+        ctx.seq_dim = seq_dim
+        ctx.frames = frames
+        ctx.cp_size = get_cp_size()
+        input = rearrange(input, "B (T S) C -> B T S C", T=frames)
+        with torch.no_grad():
+            input = input.contiguous()
+            output_tensors = [torch.zeros_like(input) for _ in range(ctx.cp_size)]
+            dist.all_gather(output_tensors, input, group=ctx.cp_group)
+            output_tensor = torch.cat(output_tensors, dim=seq_dim)
+        output_tensor = rearrange(output_tensor, "B T S C -> B (T S) C", T=frames)
+        return output_tensor
+
+
+class GatherFunction2D(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input, process_group, seq_dim_hw, shape, split_hw):
+        ctx.cp_group = process_group
+        ctx.seq_dim_hw = seq_dim_hw
+        ctx.split_hw = split_hw
+        ctx.shape = shape
+        ctx.cp_size = get_cp_size()
+
+        T, H, W = shape
+        dim_h, dim_w = seq_dim_hw
+        split_h, split_w = split_hw
+        assert H % split_h == 0, W % split_w == 0
+        assert T * (H // split_h) * (W // split_w) == input.shape[1]
+        input = rearrange(input, "B (T H W) C -> B T H W C", T=T, H=H // split_h, W=W // split_w)
+
+        with torch.no_grad():
+            input = input.contiguous()
+            output_tensors = [torch.zeros_like(input) for _ in range(ctx.cp_size)]
+            dist.all_gather(output_tensors, input, group=ctx.cp_group)
+            output_tensors_hs = []
+            assert ctx.cp_size % split_w == 0
+            for i in range(0, ctx.cp_size // split_w):
+                output_tensors_hs.append(
+                    torch.cat(output_tensors[i * split_w : (i + 1) * split_w], dim=dim_w)
+                )
+            output_tensor = torch.cat(output_tensors_hs, dim=dim_h)
+
+        output_tensor = rearrange(output_tensor, "B T H W C -> B (T H W) C")
+
+        return output_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        T, H, W = ctx.shape
+        with torch.no_grad():
+            grad_output = grad_output * ctx.cp_size
+            grad_output = rearrange(grad_output, "B (T H W) C -> B T H W C", T=T, H=H, W=W)
+            grad_input = split_tensor_in_cp_2d(grad_output, ctx.seq_dim_hw, ctx.split_hw)
+            grad_input = rearrange(grad_input, "B T H W C -> B (T H W) C")
+            
+        return grad_input, None, None, None, None
+
+
+class SplitFunction2D(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input, process_group, seq_dim_hw, split_hw):
+        ctx.cp_group = process_group
+        ctx.seq_dim_hw = seq_dim_hw
+        ctx.split_hw = split_hw
+        ctx.cp_size = get_cp_size()
+        output_tensor = split_tensor_in_cp_2d(input, ctx.seq_dim_hw, split_hw)
+
+        return output_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        with torch.no_grad():
+            grad_output = grad_output / ctx.cp_size
+            output_tensors = [torch.zeros_like(grad_output) for _ in range(ctx.cp_size)]
+            dist.all_gather(output_tensors, grad_output, group=ctx.cp_group)
+            split_h, split_w = ctx.split_hw
+            dim_h, dim_w = ctx.seq_dim_hw
+            output_tensors_hs = []
+            assert ctx.cp_size % split_w == 0
+            for i in range(0, ctx.cp_size // split_w):
+                output_tensors_hs.append(
+                    torch.cat(output_tensors[i * split_w : (i + 1) * split_w], dim=dim_w)
+                )
+            grad_input = torch.cat(output_tensors_hs, dim=dim_h)
+
+        return grad_input, None, None, None
+
+
+def gather_cp(input, frames):
+    cp_process_group = get_cp_group()
+    output_tensor = GatherFunction.apply(input, cp_process_group, 2, frames)
+
+    return output_tensor
+
+def gather_cp_2d(input, shape, split_hw):
+    cp_process_group = get_cp_group()
+    output_tensor = GatherFunction2D.apply(input, cp_process_group, (2, 3), shape, split_hw)
+
+    return output_tensor
+
+
+def split_cp_2d(input, seq_dim_hw, split_hw):
+    cp_process_group = get_cp_group()
+    output_tensor = SplitFunction2D.apply(input, cp_process_group, seq_dim_hw, split_hw)
+
+    return output_tensor
+
+
+def get_optimal_split(size):
+    factors = []
+    for i in range(1, int(size**0.5) + 1):
+        if size % i == 0:
+            factors.append([i, size // i])
+    return min(factors, key=lambda x: abs(x[0] - x[1]))
\ No newline at end of file

third_party/longcat_video/context_parallel/ulysses_wrapper.py

+import torch
+import torch.distributed as dist
+
+from ..context_parallel import context_parallel_util
+
+
+def all_to_all(tensor, scatter_idx, gather_idx, group=None, gather=True):
+    """Perform all-to-all communication on a tensor.
+
+    Args:
+        tensor (torch.Tensor): Input tensor for all-to-all communication
+        scatter_idx (int): Dimension to scatter, will split along this dimension and then scatter to all processes
+        gather_idx (int): Dimension to gather, will gather from all processes and then concatenate along this dimension
+        group (ProcessGroup, optional): Process group to use for communication
+
+    Returns:
+        torch.Tensor
+    """
+    if not dist.is_initialized():
+        return tensor
+
+    world_size = dist.get_world_size(group)
+    ulysses_rank = context_parallel_util.get_cp_rank()
+    if world_size == 1:
+        return tensor
+
+    if scatter_idx == gather_idx:
+        raise ValueError("scatter_idx and gather_idx must be different")
+
+    def chunk_tensor(tensor, scatter_idx):
+        t_shape = list(tensor.shape)
+        if t_shape[scatter_idx] % world_size != 0:
+            raise ValueError(f"Dimension {scatter_idx} must be divisible by world size {world_size}")
+        chunk_size = t_shape[scatter_idx] // world_size
+        new_shape = list()
+        for i in range(len(t_shape)):
+            if i != scatter_idx:
+                new_shape.append(t_shape[i])
+            else:
+                new_shape.extend([world_size, chunk_size])
+        tensor = tensor.reshape(*new_shape)
+        # move scatter_idx to front
+        tensor = tensor.permute(scatter_idx, *[i for i in range(len(new_shape)) if i != scatter_idx]).contiguous()
+        return tensor
+
+    # chunk tensor for all_to_all
+    tensor = chunk_tensor(tensor, scatter_idx)
+
+    # Perform all2all
+    output = torch.empty_like(tensor)
+    dist.all_to_all_single(output, tensor, group=group)
+
+    # output: e.g., [world_size, B, chunked_H, chunked_S, D] if scatter_idx == 1, gather_idx == 2 -> [B, chunked_H, S, D]
+    def reorder_tensor(tensor, gather_idx):
+        t_shape = list(tensor.shape)
+        world_size = t_shape[0]
+        # insert front to gather_idx + 1
+        permute_idx = list()
+        for i in range(1, len(t_shape)):
+            if i != gather_idx + 1:
+                permute_idx.append(i)
+            else:
+                permute_idx.extend([0, i])
+        tensor = tensor.permute(*permute_idx).contiguous() # permute(1,2,0,3) W B CH CS D -> B CH W CS D
+
+        # reshape tensor
+        new_shape = list()
+        if gather:
+            for i in range(1, len(t_shape)): # B CH CS D
+                if i != gather_idx + 1:
+                    new_shape.append(t_shape[i])
+                else:
+                    new_shape.append(world_size * t_shape[i]) # B CH W*CS D
+
+            tensor = tensor.reshape(*new_shape)
+        else:
+            tensor = tensor[:,ulysses_rank] # W B CS CH D -> B CS W CH D
+
+        return tensor
+
+    output = reorder_tensor(output, gather_idx)
+
+    return output
+
+
+@torch.compiler.disable
+def ulysses_a2a_in(query, key, value):
+    if context_parallel_util.get_cp_size() == 1:
+        return query, key, value
+
+    # [B, H, S/N, D] -> [B, H/N, S, D]
+    query = all_to_all(query, scatter_idx=1, gather_idx=2, group=context_parallel_util.get_cp_group())
+    key = all_to_all(key, scatter_idx=1, gather_idx=2, group=context_parallel_util.get_cp_group())
+    value = all_to_all(value, scatter_idx=1, gather_idx=2, group=context_parallel_util.get_cp_group())
+    return query, key, value
+
+
+@torch.compiler.disable
+def ulysses_a2a_out(output):
+    if context_parallel_util.get_cp_size() == 1:
+        return output
+
+    # [B, H/N, S, D] -> [B, H, S/N, D]
+    output = all_to_all(output, scatter_idx=2, gather_idx=1, group=context_parallel_util.get_cp_group())
+    return output
+
+
+def ulysses_wrapper(func):
+    def wrapper(self, query, key, value, shape):
+        # Apply ulysses_a2a_in before the function call, gather sequence and split head
+        query, key, value = ulysses_a2a_in(query, key, value)
+        output = func(self, query, key, value, shape)        
+        output = ulysses_a2a_out(output)
+        return output
+
+    return wrapper