longcat: context-parallel multi-GPU and the video-longcat pod image (v0.2.36)

Step 7, the last of the server-side work. Two things that are more than flags.

CONTEXT PARALLELISM is N processes under torchrun, not N threads, and every rank must
enter the same pipeline call or the collectives deadlock. So:
- NCCL and init_context_parallel() run BEFORE the model loads — that call is what makes
  the DiT split its spatial dims, so load_pipeline cannot come first;
- rank 0 owns the HTTP socket and BROADCASTS each job before entering the pipeline; the
  other ranks sit in cp_worker_loop waiting for one and discard their copy of the result.
  A rank-0-only server that simply called generate() would hang on the first collective;
- the group is told to shut down when the server stops, or the non-zero ranks outlive it;
- cp_size 1 (the default) initialises no distributed group at all;
- asking for more ranks than there are visible GPUs is refused up front, because NCCL's
  own failure for that names nothing useful.

THE POD IMAGE could not use the generic builder: Dockerfile.capability is python:3.12 on
the cu128 torch index, and LongCat needs 3.10 on cu124 — two interpreters in one image,
the light core running coderai and a 3.10 venv running the model. So video-longcat
carries its own Dockerfile (the escape hatch the engines profile uses), builds the venv
from the SAME requirements-longcat.txt a local install uses, asserts 3.10 / torch 2.6 /
transformers 4.41 at build time rather than discovering a different resolution on a
rented pod, clones the pipeline source (it is not on PyPI) and checks it imports, and
bakes no weights.

model_capability() now routes a `backend: longcat` entry to that image. Left alone it
would land on the plain 'video' image, which cannot run LongCat at all — a pod that
boots, passes its health check, accepts the request and then has nothing to generate
with, exactly the trap the OCR profile documents for Surya. An explicit `capability`
override still wins.

24 new tests, 212 across the seven LongCat suites.
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.35"
+__version__ = "0.2.36"
 __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

@@ -271,6 +271,21 @@ def ensure_built(config: dict = None) -> Path:
 
 # ── service lifecycle ─────────────────────────────────────────────────────────
 
+def _visible_gpu_count(env: dict = None) -> int:
+    """How many CUDA devices this service will see, or 0 when it cannot be told.
+
+    Asking for more context-parallel ranks than there are GPUs fails inside NCCL with a
+    message about nothing in particular, so it is worth catching here."""
+    sel = str((env or os.environ).get("CUDA_VISIBLE_DEVICES") or "").strip()
+    if sel:
+        return len([x for x in sel.split(",") if x.strip() != ""])
+    try:
+        from codai.models.gpu_query import visible_gpu_memory
+        return len(visible_gpu_memory() or [])
+    except Exception:
+        return 0
+
+
 def _free_port() -> int:
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     s.bind(("127.0.0.1", 0))
@@ -364,11 +379,32 @@ def ensure_service(model_path: str, config: dict = None,
             # The upstream pipeline exposes only the diffusers CPU-offload hooks; the
             # manager's other strategies have no equivalent here.
             offload = "model"
-        cmd = [str(py), str(_SERVICE_SCRIPT),
+        # Context parallelism: upstream splits the DiT's spatial dims across GPUs under
+        # torchrun, so cp_size > 1 means N processes, not N threads. Only rank 0 binds
+        # the port; the others take their jobs over a broadcast (see the service).
+        cp = 0
+        try:
+            cp = int(config.get("cp_size") or 0)
+        except (TypeError, ValueError):
+            cp = 0
+        if cp > 1:
+            visible = _visible_gpu_count(env)
+            if visible and cp > visible:
+                raise RuntimeError(
+                    f"cp_size is {cp} but only {visible} GPU(s) are visible to this "
+                    f"service — lower it, or widen longcat.gpu / gpu_device")
+            launcher = [str(py), "-m", "torch.distributed.run",
+                        f"--nproc_per_node={cp}", "--nnodes=1",
+                        "--rdzv-backend=c10d", "--rdzv-endpoint=127.0.0.1:0"]
+        else:
+            launcher = [str(py)]
+        cmd = launcher + [str(_SERVICE_SCRIPT),
                "--model", model_path,
                "--source", str(source),
                "--host", "127.0.0.1", "--port", str(port),
                "--dtype", str(config.get("dtype") or "bfloat16")]
+        if cp > 1:
+            cmd += ["--context-parallel-size", str(cp)]
         if offload:
             cmd += ["--offload", offload]
         if sec is not None and str(getattr(sec, "attention", "") or "").strip():

codai/api/runpod_worker.py

@@ -208,6 +208,9 @@ PUBLISHED_CAPABILITY_IMAGES = (
     "stt-nemo",      # NeMo Canary: large, and its own torch
     "stt-crisper",   # CrisperWhisper: transformers 4.46 AND torch 2.5
     "ocr-paddle",    # paddlepaddle-gpu: brings its own CUDA runtime
+    "video-longcat", # LongCat-Video: Python 3.10 + torch 2.6+cu124 + transformers 4.41,
+                     #   so TWO interpreters in one image (see its own Dockerfile). The
+                     #   plain 'video' image cannot serve it at all.
     # Text: a coderai pod that serves LLMs. Unlike a vLLM/llama.cpp pod it can be
     # SENT a local LoRA adapter, which is the only way a local text adapter runs
     # on RunPod. Build and push it with packaging/runpod before using it.
@@ -298,6 +301,14 @@ def model_capability(entry: dict, include_text: bool = False) -> str:
     explicit = str(entry.get("capability") or "").strip().lower()
     if explicit:
         return explicit
+    # A LongCat-Video model is a video model, but the 'video' image cannot serve it: its
+    # pipeline is not in diffusers and it needs its own Python 3.10 venv. Without this
+    # it would be placed on an image that boots, passes its health check, accepts the
+    # request and then has nothing to run it with — the same trap the OCR profile
+    # documents for Surya.
+    if str(entry.get("backend") or "").strip().lower() in (
+            "longcat", "longcat-video", "longcat_video"):
+        return "video-longcat"
     types = entry.get("model_types") or [entry.get("model_type") or ""]
     for mt in types:
         if mt in _LLM_MODEL_TYPES and not include_text:

docs/longcat-video.md

@@ -126,13 +126,38 @@ returns — shorter, with a warning naming why — rather than failing. Bounded
 Set `max_instances` on the entry to have video requests queue at the front as text does;
 without it, video passes through unqueued as it always has.
 
+## Multi-GPU
+
+`cp_size` on the model entry splits the DiT's spatial dimensions across that many GPUs.
+It is **N processes under torchrun**, not N threads: NCCL is initialised and
+`init_context_parallel()` called *before* the model loads (it is what makes the DiT split
+at all), and every rank must enter the same pipeline call or the collectives deadlock. So
+rank 0 owns the HTTP socket and broadcasts each job to the others, which wait in a loop
+for one. Asking for more ranks than there are visible GPUs is refused up front, because
+NCCL's own failure for that names nothing useful.
+
+`cp_size: 1` — the default — initialises no distributed group at all.
+
 ## Remote and scaled
 
 `longcat.service_url` (or `CODERAI_LONGCAT_SERVICE_URL`) points at a service running
 elsewhere — another host, a container, a rented pod — and coderai proxies to it without
 building or downloading anything. `distribute` fans independent generations across
-engines and nodes. A `video-longcat` RunPod capability image needs a custom Dockerfile
-(the generic one is python:3.12 + cu128, wrong on both axes) and is not built yet.
+engines and nodes.
+
+For rented GPUs there is a **`video-longcat`** capability image. The generic capability
+builder is python:3.12 on the cu128 torch index, so it cannot express this image at all —
+LongCat needs 3.10 on cu124, meaning two interpreters in one image: the light core runs
+coderai, a 3.10 venv runs the model. It therefore carries its own Dockerfile (the same
+escape hatch the `engines` profile uses), builds the venv from
+`requirements-longcat.txt` and asserts 3.10 / torch 2.6 / transformers 4.41 at build
+time, clones the pipeline source and checks it imports, and bakes **no weights**.
+
+A model with `backend: longcat` is routed to that image automatically — placing it on the
+plain `video` image would give a pod that boots, passes its health check, accepts the
+request and has nothing to generate with. Build it with
+`packaging/runpod/build_capability_image.sh video-longcat`, which runs the container and
+polls `/healthz` for 120 s before declaring success.
 
 ## Variants and the avatar families
 
@@ -179,7 +204,5 @@ set on the model entry.
 ## Not done yet
 
 - GGUF variants (they need the ComfyUI gguf loader)
-- context-parallel multi-GPU (`cp_size` is accepted but not yet acted on)
-- the `video-longcat` pod image
 - LoRA/QLoRA training on top of it (feasible — upstream's own refinement expert is a LoRA)
 - **no generation has been verified against real weights**, avatar included

packaging/runpod/Dockerfile.capability-video-longcat

+# CoderAI LongCat-Video pod.
+#
+# LongCat pins Python 3.10 + torch 2.6.0+cu124 + transformers 4.41 + numpy 1.26 - four
+# conflicts with coderai's own 3.13 / 2.11 / 5.x / 2.4 - so it runs in an isolated venv
+# behind a managed subprocess. On a pod that means TWO interpreters in one image: the
+# light core (python:3.12, no torch) runs coderai, and a 3.10 venv runs the model.
+#
+# The generic Dockerfile.capability cannot express that - it is one interpreter on the
+# cu128 torch index - which is why this profile carries its own Dockerfile, the same
+# escape hatch the engines profile uses.
+#
+# `longcat_video` is a package INSIDE the upstream repo, not on PyPI, so the source is
+# cloned here. Weights are NOT baked: ~83 GB across variants, downloaded at runtime into
+# the pod's volume, like every other model.
+
+ARG CORE_BASE=coderai-capability-base-light:latest
+FROM ${CORE_BASE}
+
+ARG PROFILE=video-longcat
+ARG IMAGE_TAG=coderai-video-longcat:latest
+ARG LONGCAT_REF=main
+
+# Python 3.10 for the model's venv: the light core carries only its own 3.12.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+      software-properties-common ca-certificates gnupg \
+ && add-apt-repository -y ppa:deadsnakes/ppa \
+ && apt-get update && apt-get install -y --no-install-recommends \
+      python3.10 python3.10-venv python3.10-dev \
+ && rm -rf /var/lib/apt/lists/*
+
+# The venv, from the same requirements file a local install uses - one source of truth.
+COPY requirements-longcat.txt /tmp/requirements-longcat.txt
+RUN python3.10 -m venv /opt/coderai/longcat_venv \
+ && /opt/coderai/longcat_venv/bin/python -m pip install -U pip \
+ && /opt/coderai/longcat_venv/bin/python -m pip install -r /tmp/requirements-longcat.txt \
+ && /opt/coderai/longcat_venv/bin/python -c "import sys,torch,transformers; \
+      assert sys.version_info[:2]==(3,10), sys.version; \
+      assert torch.__version__.startswith('2.6'), torch.__version__; \
+      assert transformers.__version__.startswith('4.41'), transformers.__version__" \
+ && rm -rf /root/.cache/pip /tmp/requirements-longcat.txt
+
+# The pipeline code. Shallow, no history, no weights, no demo assets.
+RUN git clone --depth 1 --branch "${LONGCAT_REF}" \
+      https://github.com/meituan-longcat/LongCat-Video /opt/coderai/LongCat-Video \
+ && rm -rf /opt/coderai/LongCat-Video/.git /opt/coderai/LongCat-Video/assets \
+ && /opt/coderai/longcat_venv/bin/python -c \
+      "import sys; sys.path.insert(0,'/opt/coderai/LongCat-Video'); \
+       import longcat_video.pipeline_longcat_video"
+
+COPY . /opt/coderai/app
+WORKDIR /opt/coderai/app
+RUN python -c "import codai.api.app" \
+ && python -c "import importlib.util as u; \
+      s=u.spec_from_file_location('c','tools/longcat_common.py'); \
+      m=u.module_from_spec(s); s.loader.exec_module(m); \
+      assert m.source_problems('/opt/coderai/LongCat-Video') == []"
+
+ENV CODERAI_IMAGE_TAG=${IMAGE_TAG} \
+    CODERAI_LONGCAT_VENV=/opt/coderai/longcat_venv \
+    CODERAI_LONGCAT_SRC=/opt/coderai/LongCat-Video
+
+EXPOSE 8000
+HEALTHCHECK --interval=30s --timeout=5s --start-period=120s --retries=5 \
+  CMD python -c "import urllib.request;urllib.request.urlopen('http://127.0.0.1:8000/healthz',timeout=4)"
+CMD ["/opt/coderai/boot.sh"]

packaging/runpod/profiles/video-longcat.apt

+ffmpeg
+git

packaging/runpod/profiles/video-longcat.dockerfile

+Dockerfile.capability-video-longcat

packaging/runpod/profiles/video-longcat.txt

+# LongCat-Video as a coderai pod.
+#
+# Nothing goes in the MAIN venv: LongCat runs in its own Python 3.10 venv with
+# torch 2.6+cu124, which the custom Dockerfile builds (see video-longcat.dockerfile).
+# The main venv only drives the subprocess and serves the API, so this profile is the
+# LIGHT core plus what coderai itself needs for /v1/video — nothing that pulls torch.
+#
+# The generic Dockerfile.capability cannot express this image: it is python:3.12 on the
+# cu128 torch index, and LongCat needs 3.10 on cu124. Hence the .dockerfile marker, the
+# same escape hatch the engines profile uses.
+#
+# Weights are NOT baked (the base repo is ~83 GB across variants); they download at
+# runtime into the pod's volume, like every other model.
+imageio[ffmpeg]>=2.33.0

tests/test_longcat_scale.py

+"""Step 7: context-parallel multi-GPU, and the pod image that can serve LongCat.
+
+Two things make this more than a flag:
+
+1. Context parallelism is N PROCESSES under torchrun, not N threads, and every rank must
+   enter the same pipeline call or the collectives deadlock. So rank 0 owns the socket and
+   broadcasts each job; the other ranks sit in a loop waiting for one. A rank-0-only
+   server that simply called generate() would hang on the first collective.
+2. The generic RunPod capability image is python:3.12 on the cu128 torch index. LongCat
+   needs 3.10 on cu124, so the 'video' image cannot serve it at all — it would boot, pass
+   its health check, accept a request and have nothing to run it with.
+"""
+
+import re
+import sys
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from codai.api import longcat_worker as W
+from codai.api.runpod_worker import (PUBLISHED_CAPABILITY_IMAGES,
+                                     default_capability_image, model_capability)
+
+SERVICE = (ROOT / "tools/longcat_service.py").read_text()
+WORKER = (ROOT / "codai/api/longcat_worker.py").read_text()
+PODFILE = (ROOT / "packaging/runpod/Dockerfile.capability-video-longcat").read_text()
+
+
+# ---------------------------------------------------------------- context parallel
+def test_the_service_accepts_a_context_parallel_size():
+    assert "--context-parallel-size" in SERVICE
+
+
+def test_nccl_and_context_parallel_are_initialised_before_the_model_loads():
+    """init_context_parallel has to run before the DiT is built — it is what makes the
+    DiT split its spatial dims — so load_pipeline must not be called first."""
+    main = SERVICE[SERVICE.index("def main(argv=None):"):]
+    assert main.index("cp_init(cp_size)") < main.index("load_pipeline(")
+    assert "init_context_parallel" in SERVICE
+    assert 'backend="nccl"' in SERVICE
+
+
+def test_only_rank_zero_binds_the_socket():
+    """Every rank binding the port would collide; every rank serving would mean N
+    different answers to one request."""
+    main = SERVICE[SERVICE.index("def main(argv=None):"):]
+    assert "if rank != 0:" in main
+    assert main.index("if rank != 0:") < main.index("ThreadingHTTPServer")
+
+
+def test_other_ranks_wait_for_a_broadcast_job():
+    assert "def cp_worker_loop" in SERVICE
+    assert "broadcast_object_list" in SERVICE
+
+
+def test_rank_zero_broadcasts_before_generating():
+    """This is the deadlock: rank 0 must publish the job BEFORE entering the pipeline,
+    because the other ranks have to be in the same call."""
+    handler = SERVICE[SERVICE.index('if self.path.startswith("/generate")'):]
+    handler = handler[:handler.index("elif")]
+    assert handler.index("cp_broadcast(job)") < handler.index("generate(job)")
+
+
+def test_the_group_is_told_to_shut_down():
+    """Otherwise the non-zero ranks outlive the server and the process never exits."""
+    assert "_CP_SHUTDOWN" in SERVICE
+    main = SERVICE[SERVICE.index("def main(argv=None):"):]
+    assert "cp_broadcast(dict(_CP_SHUTDOWN))" in main
+
+
+def test_a_single_gpu_does_not_initialise_distributed():
+    """cp_size 1 is the default path and must not need NCCL at all."""
+    main = SERVICE[SERVICE.index("def main(argv=None):"):]
+    assert "if cp_size > 1:" in main
+    cp = SERVICE[SERVICE.index("def cp_init"):SERVICE.index("def cp_broadcast")]
+    assert "if cp_size <= 1:" in cp and "return rank, local_rank" in cp
+
+
+def test_the_worker_launches_under_torchrun_only_when_asked():
+    launch = WORKER[WORKER.index("cp = 0"):WORKER.index('"--model", model_path')]
+    assert "torch.distributed.run" in launch
+    assert "nproc_per_node" in launch
+    assert "if cp > 1:" in launch and "launcher = [str(py)]" in launch
+
+
+def test_asking_for_more_ranks_than_gpus_is_refused_early():
+    """NCCL's own failure for this names nothing useful."""
+    launch = WORKER[WORKER.index("cp = 0"):WORKER.index('"--model", model_path')]
+    assert "_visible_gpu_count" in launch
+    assert "lower it" in launch
+
+
+@pytest.mark.parametrize("env,expected", [
+    ({"CUDA_VISIBLE_DEVICES": "0,1,2"}, 3),
+    ({"CUDA_VISIBLE_DEVICES": "0"}, 1),
+    ({"CUDA_VISIBLE_DEVICES": ""}, None),      # falls back to a real query
+])
+def test_the_visible_gpu_count_follows_the_selector(env, expected):
+    got = W._visible_gpu_count(env)
+    if expected is None:
+        assert got >= 0
+    else:
+        assert got == expected
+
+
+def test_cp_size_comes_from_the_model_entry():
+    """It is a per-model setting with a UI control, saved and whitelisted in step 5."""
+    html = (ROOT / "codai/admin/templates/models.html").read_text()
+    assert "cfg-longcat-cp" in html
+    routes = (ROOT / "codai/admin/routes.py").read_text()
+    block = routes[routes.index('for key in ("alias", "config_name", "backend"'):]
+    assert "cp_size" in block[:block.index("):")]
+
+
+# ---------------------------------------------------------------- the pod image
+def test_the_image_is_registered():
+    """A capability absent from this tuple has no image and demands an explicit one."""
+    assert "video-longcat" in PUBLISHED_CAPABILITY_IMAGES
+    assert default_capability_image("video-longcat").endswith(
+        "coderai-video-longcat:latest")
+
+
+def test_a_longcat_model_is_not_placed_on_the_plain_video_image():
+    """The 'video' image is python:3.12/cu128 with diffusers — it cannot run LongCat at
+    all. Placing it there gives a pod that boots, passes health and then has nothing to
+    generate with."""
+    assert model_capability({"model_type": "video_models",
+                             "backend": "longcat"}) == "video-longcat"
+    assert model_capability({"model_type": "video_models",
+                             "backend": "auto"}) == "video"
+
+
+def test_an_explicit_capability_still_wins():
+    """The operator's override must stay authoritative."""
+    assert model_capability({"model_type": "video_models", "backend": "longcat",
+                             "capability": "video"}) == "video"
+
+
+def test_the_profile_uses_the_custom_dockerfile_escape_hatch():
+    base = ROOT / "packaging/runpod/profiles"
+    assert (base / "video-longcat.txt").is_file()
+    assert (base / "video-longcat.light").is_file()       # light core: no torch in main
+    assert (base / "video-longcat.dockerfile").read_text().strip() == \
+        "Dockerfile.capability-video-longcat"
+
+
+def test_the_image_builds_the_310_venv_from_the_same_requirements_file():
+    """One source of truth: the pod and a local install must not drift."""
+    assert "requirements-longcat.txt" in PODFILE
+    assert "python3.10 -m venv" in PODFILE
+
+
+def test_the_image_verifies_the_venv_rather_than_hoping():
+    """A pip install that resolved differently would otherwise surface on a rented pod."""
+    assert "(3,10)" in PODFILE and "'2.6'" in PODFILE and "'4.41'" in PODFILE
+
+
+def test_the_image_clones_the_pipeline_source_and_checks_it_imports():
+    """`longcat_video` is in the repo, not on PyPI."""
+    assert "meituan-longcat/LongCat-Video" in PODFILE
+    assert "import longcat_video.pipeline_longcat_video" in PODFILE
+    assert "source_problems" in PODFILE
+
+
+def test_the_image_bakes_no_weights():
+    """~83 GB across variants, and the rule is that released images carry no models."""
+    assert "--depth 1" in PODFILE
+    assert "rm -rf /opt/coderai/LongCat-Video/.git" in PODFILE
+    assert not re.search(r"huggingface-cli download|hf_hub_download", PODFILE)
+
+
+def test_the_image_points_the_worker_at_what_it_baked():
+    assert "CODERAI_LONGCAT_VENV=/opt/coderai/longcat_venv" in PODFILE
+    assert "CODERAI_LONGCAT_SRC=/opt/coderai/LongCat-Video" in PODFILE
+
+
+def test_those_env_vars_are_the_ones_the_worker_reads():
+    """Seeded and read must agree, or the pod silently builds its own venv."""
+    assert "CODERAI_LONGCAT_VENV" in WORKER
+    assert "CODERAI_LONGCAT_SRC" in WORKER
+
+
+def test_the_image_answers_a_health_check():
+    assert "HEALTHCHECK" in PODFILE and "/healthz" in PODFILE

tools/longcat_service.py

@@ -670,7 +670,12 @@ def make_handler():
         def do_POST(self):
             try:
                 if self.path.startswith("/generate"):
-                    self._json(generate(self._read_json()))
+                    job = self._read_json()
+                    if int(_state.get("cp_size") or 1) > 1:
+                        # Every rank must enter the same pipeline call, or the
+                        # collectives deadlock.
+                        cp_broadcast(job)
+                    self._json(generate(job))
                 elif self.path.startswith("/yield"):
                     # Stop at the next segment boundary. The in-flight request still
                     # returns — shorter, with a warning — so a swap costs one segment
@@ -693,6 +698,70 @@ def make_handler():
     return Handler
 
 
+# ── context parallelism ───────────────────────────────────────────────────────
+#
+# Upstream scales by splitting the DiT's spatial dimensions across GPUs: NCCL is
+# initialised BEFORE the model loads, `init_context_parallel(cp_size)` is called, the DiT
+# gets `cp_split_hw`, and every rank runs `pipe.to(local_rank)`. Only rank 0 writes the
+# output.
+#
+# That is awkward behind an HTTP service, and the awkwardness is the point: all ranks must
+# enter the same pipeline call together or the collectives deadlock. So rank 0 owns the
+# socket and BROADCASTS each job to the others, which sit in a loop waiting for one. A
+# rank-0-only server that just called generate() would hang on the first collective.
+
+_CP_SHUTDOWN = {"__shutdown__": True}
+
+
+def cp_init(cp_size: int):
+    """Initialise NCCL + context parallelism. Returns (rank, local_rank)."""
+    import torch
+    import torch.distributed as dist
+
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if cp_size <= 1:
+        return rank, local_rank
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl")
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+    # Must happen before the DiT is built, which is why load_pipeline is not called
+    # until after this returns.
+    from longcat_video.context_parallel.context_parallel_util import (
+        init_context_parallel)
+    init_context_parallel(cp_size)
+    log(f"context parallel: rank {rank}/{dist.get_world_size()} "
+        f"(local_rank {local_rank}, cp_size {cp_size})")
+    return rank, local_rank
+
+
+def cp_broadcast(job):
+    """Send a job from rank 0 to every other rank, or receive one. Returns the job."""
+    import torch.distributed as dist
+    box = [job]
+    dist.broadcast_object_list(box, src=0)
+    return box[0]
+
+
+def cp_worker_loop():
+    """Ranks other than 0: wait for a job, run it, discard the result, repeat.
+
+    The result is discarded deliberately — rank 0's copy is the one that gets muxed and
+    returned. What matters here is entering the same collective."""
+    while True:
+        job = cp_broadcast(None)
+        if not job or job.get("__shutdown__"):
+            log("context-parallel worker shutting down")
+            return
+        try:
+            generate(job)
+        except Exception as exc:
+            # A failure on rank 0 raises there too; log and keep the loop alive so one
+            # bad request does not strand the group.
+            log(f"context-parallel worker error: {type(exc).__name__}: {exc}")
+
+
 def main(argv=None):
     ap = argparse.ArgumentParser(description="LongCat-Video generation service")
     ap.add_argument("--model", required=True, help="checkpoint directory")
@@ -702,24 +771,51 @@ def main(argv=None):
     ap.add_argument("--dtype", default="bfloat16")
     ap.add_argument("--offload", default="", help="'' | model | sequential")
     ap.add_argument("--attention", default="xformers", help="xformers | flash | sdpa")
+    ap.add_argument("--context-parallel-size", type=int, default=1,
+                    help="GPUs to split the DiT across (launch under torchrun)")
     ap.add_argument("--preload", action="store_true",
                     help="load the checkpoint at startup instead of on first request")
     args = ap.parse_args(argv)
 
+    cp_size = max(1, int(args.context_parallel_size))
     _state.update(model=os.path.expanduser(args.model),
                   source=os.path.expanduser(args.source),
-                  dtype=args.dtype)
+                  dtype=args.dtype, cp_size=cp_size)
     # Upstream reads the attention backend from the model config; the env var is what its
     # modules honour, and leaving it unset means FlashAttention-2, which is not installed
     # by default (requirements-longcat.txt installs xformers instead).
     os.environ.setdefault("LONGCAT_ATTENTION", args.attention)
 
+    rank, _local = (0, 0)
+    if cp_size > 1:
+        # The repo checkout has to be importable before init_context_parallel.
+        src = _state["source"]
+        if src and src not in sys.path:
+            sys.path.insert(0, src)
+        rank, _local = cp_init(cp_size)
+        # Splitting height and width across the group is what cp_size buys; the DiT
+        # takes it at construction.
+        _state["cp_split_hw"] = cp_size
+
     if args.preload:
         load_pipeline(_state["model"], _state["source"], args.dtype, args.offload)
 
+    if rank != 0:
+        # No socket on these ranks: one job, one broadcast, every rank in the same call.
+        cp_worker_loop()
+        return
+
     server = ThreadingHTTPServer((args.host, args.port), make_handler())
-    log(f"serving on http://{args.host}:{args.port} (model={_state['model']})")
-    server.serve_forever()
+    log(f"serving on http://{args.host}:{args.port} (model={_state['model']}"
+        + (f", cp_size={cp_size}" if cp_size > 1 else "") + ")")
+    try:
+        server.serve_forever()
+    finally:
+        if cp_size > 1:
+            try:
+                cp_broadcast(dict(_CP_SHUTDOWN))
+            except Exception:
+                pass
 
 
 if __name__ == "__main__":