longcat: a conditioning pass tiles twice (v0.3.3)

An i2v generation died in the first denoising step with a bare AssertionError out
of flash_attn_bsa_3d, on a geometry the guard had just declared valid: 93 frames,
sides divisible by 64.

bsa_problems only ever checked the noise branch. A conditioning pass does not tile
the latent once — Attention.forward splits the sequence into a conditioning block
and a noise block and calls flash_attn_bsa_3d on each (modules/attention.py:123-134),
so the cond depth and the remaining noise depth must divide by the chunk as well.
generate_i2v hardcodes one conditioning latent (pipeline:796) and 1 % 4 != 0, which
means block-sparse attention has never once been able to run an i2v pass, at any
frame count or resolution. The avatar pipeline hardcodes the same 1 in four places,
and generate_vc computes a count it never rounds. Only generate_refine gets it
right, by padding both counts to the granularity itself (pipeline:1245-1250).

(a) The guard now knows about the split. _cond_latent_passes enumerates every
conditioning count the request can produce — segment 0 uses the task's own entry
point and every later segment uses generate_vc, so even t2v runs a conditioning
pass once num_segments > 1, and BSA is one flag for the whole request — and a
count that cannot tile puts the pass on dense attention with a line saying which
condition failed.

(b) bsa_pad_cond, OFF by default, rounds the count up to the granularity on the way
into the DiT, which is what generate_refine does for itself, so an i2v or
continuation pass can keep BSA. Patched onto the DiT instance rather than the
vendored source, the same way the denoising bars are rebound: the count is born
inside generate_i2v's loop, so there is no outer argument to correct. Off by
default because the pipeline marks only the real conditioning frames clean
(timestep[:, :1] = 0), so rounding 1 up to 4 hands the conditioning block three
latent frames that are still noise — well-defined attention, but not what the model
was trained on, and on a video model that kind of error shows up as temporal drift.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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.3.2"
+__version__ = "0.3.3"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/admin/routes.py

@@ -2959,7 +2959,8 @@ async def api_model_configure(request: Request, username: str = Depends(require_
                 
                     # LongCat component quantisation, block-sparse attention, the base
                     # checkpoint an avatar family borrows from, and the CP tile.
-                    "text_encoder_quant", "bsa", "base_model", "cp_split_hw",
+                    "text_encoder_quant", "bsa", "bsa_pad_cond", "base_model",
+                    "cp_split_hw",
                     # The prompt-embedding cache: an encode outlives the request that
                     # paid for it, and the text encoder is then built only on a miss.
                     "prompt_cache", "prompt_cache_dir", "prompt_cache_max_gb",):

codai/admin/templates/models.html

@@ -718,6 +718,11 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
               <option value="on">on — needs sides /64 and 16n+13 frames</option>
               <option value="auto">auto — on when triton is available</option>
             </select></div>
+          <div class="form-row" style="margin:0"><label class="form-label">Pad conditioning for BSA</label>
+            <select id="cfg-longcat-bsa-pad-cond" class="form-input">
+              <option value="">off — i2v/continuation run dense (safe)</option>
+              <option value="on">on — round cond latents up so BSA applies (experimental)</option>
+            </select></div>
         </div>
         <div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:.6rem;margin-top:.5rem">
           <div class="form-row" style="margin:0"><label class="form-label">Prompt cache</label>
@@ -4642,6 +4647,7 @@ function _populateCluster(m, s){
   set('cfg-longcat-offload', s.offload_strategy || '');
   set('cfg-longcat-te-quant', s.text_encoder_quant || '');
   set('cfg-longcat-bsa', s.bsa || '');
+  set('cfg-longcat-bsa-pad-cond', s.bsa_pad_cond || '');
   set('cfg-longcat-prompt-cache', s.prompt_cache || '');
   set('cfg-longcat-prompt-cache-dir', s.prompt_cache_dir || '');
   set('cfg-longcat-prompt-cache-max', s.prompt_cache_max_gb || '');
@@ -4692,6 +4698,7 @@ function _collectCluster(){
     out.offload_strategy = val('cfg-longcat-offload') || '';
     out.text_encoder_quant = val('cfg-longcat-te-quant') || '';
     out.bsa = val('cfg-longcat-bsa') || '';
+    out.bsa_pad_cond = val('cfg-longcat-bsa-pad-cond') || '';
     out.prompt_cache = val('cfg-longcat-prompt-cache') || '';
     out.prompt_cache_dir = val('cfg-longcat-prompt-cache-dir') || '';
     out.prompt_cache_max_gb = num('cfg-longcat-prompt-cache-max');

codai/api/longcat_worker.py

@@ -45,6 +45,10 @@ Per-model config keys (models.json entry, all optional)::
     "quality": "fast",             # draft | fast | best
     "offload_strategy": "",        # '' | model | sequential
     "bsa": "off",                  # block-sparse attention: on | off | auto
+    "bsa_pad_cond": "off",         # on | off — let i2v/continuation passes use BSA by
+                                   #   rounding the conditioning latents up to the
+                                   #   granularity; off, because it widens the
+                                   #   conditioning block past what is actually clean
     "base_model": "",              # avatar families borrow tokenizer/text_encoder/vae
     "text_encoder_quant": "none",   # none | int8 | nf4 | fp4 (keeps it resident)
     "prompt_cache": "off",         # off | disk — cache prompt embeddings between
@@ -498,6 +502,12 @@ def ensure_service(model_path: str, config: dict = None,
                   (getattr(sec, "bsa", "") if sec is not None else "") or "").strip()
         if bsa:
             cmd += ["--bsa", bsa]
+        # Only meaningful with bsa on, and the service validates it either way.
+        pad = str(config.get("bsa_pad_cond") or
+                  (getattr(sec, "bsa_pad_cond", "") if sec is not None else "")
+                  or "").strip()
+        if pad:
+            cmd += ["--bsa-pad-cond", pad]
         split = str(config.get("cp_split_hw") or "").strip()
         if split:
             cmd += ["--cp-split-hw", split]

docs/longcat-video.md

@@ -74,6 +74,45 @@ checkout and venv, and the avatar fields. Global runtime settings (venv, the 3.1
 interpreter, `auto_build`, attention backend, timeouts, the eviction drain) live in
 **Settings → LongCat**.
 
+## Block-sparse attention
+
+Attention dominates at video token counts — ~37k tokens for 93 frames at 480x832, and it
+is O(N^2) in that — so the checkpoint carries `bsa_params` with sparsity 0.9375, roughly a
+16x cut, implemented in pure Triton. **Block-sparse attention** on the model page turns it
+on. It ships off: it is an approximation, and on a video model the artifacts show up as
+temporal flicker rather than one soft frame.
+
+It also cannot run every geometry. `flash_attn_bsa_3d` tiles the latent into 3D chunks and
+asserts the latent divides evenly by the chunk on every axis, which with the shipped
+`(4, 4, 4)` means:
+
+- **sides divisible by 64** — the VAE's 16-pixel latent cell times 4. Upstream's own
+  default 480x832 is *not* one (480/16 = 30).
+- **frame counts of 16n+13** — 13, 29, 45, 61, 77, 93, … not the 4n+1 the generation rule
+  suggests, because the latent depth is `(F-1)/4 + 1` and *that* is what must divide by 4.
+
+And there is a third condition that is easy to miss. A **conditioning** pass — i2v, or any
+continuation — does not tile the latent once: `Attention.forward` splits the sequence into
+a conditioning block and a noise block and tiles each separately, so both depths must
+divide by the chunk too. `generate_i2v` hardcodes a single conditioning latent, and
+`1 % 4 != 0`, so block-sparse attention can never run an i2v pass at any frame count or
+resolution. coderai detects all three before the first denoising step and runs that pass
+**dense**, with a log line saying which condition failed — rather than letting a bare
+`AssertionError` surface from inside the DiT.
+
+### Padding the conditioning (experimental, off)
+
+**Pad conditioning for BSA** rounds the conditioning latent count up to the chunk
+granularity on the way into the DiT, which is exactly what upstream's own `generate_refine`
+does for itself. With it on, an i2v or continuation pass keeps block-sparse attention
+instead of falling back.
+
+It is off by default, and it is not free: the pipeline marks only the real conditioning
+frames clean (`timestep[:, :1] = 0`), so rounding 1 up to 4 hands the conditioning block
+three latent frames that are still noise. The attention is well-defined but it is not what
+the model was trained on. Measure a known-good clip against the dense path before trusting
+it on anything you care about.
+
 ## Generating
 
 ```jsonc

tests/test_longcat_bsa.py

@@ -10,6 +10,7 @@ argument whose checkpoint value cannot be used as shipped.
 """
 import importlib.util
 import pathlib
+import re
 import sys
 
 import pytest
@@ -217,4 +218,266 @@ def test_the_service_falls_back_instead_of_asserting(service):
     assert "enable_bsa = True" in body, "it must be restored for the next request"
     assert "WARNING" in body, "a silent fallback hides that nothing got faster"
     # and it is actually applied around the generation
-    assert "with _bsa_for(pipe, ctx):" in src
+    assert "with _bsa_for(pipe, ctx, task):" in src
+
+
+# ──────────────────────────────── the conditioning branch ─────────────────────
+# A geometry can be perfectly tileable and still assert, because a conditioning pass
+# tiles the cond and noise blocks SEPARATELY (modules/attention.py:123-134). This is
+# what killed a 93-frame i2v run with 64-divisible sides: the guard checked the noise
+# depth, passed, and flash_attn_bsa_3d asserted on Tq=1 in the first step.
+
+def test_the_i2v_hardcode_can_never_tile(LC):
+    """generate_i2v passes num_cond_latents=1, and 1 % 4 != 0 under any frame count."""
+    for frames in (13, 29, 45, 61, 77, 93, 173):
+        problems = LC.bsa_problems(frames, 512, 832,
+                                   num_cond_latents=LC.I2V_COND_LATENTS)
+        assert any("conditioning" in p for p in problems), frames
+
+
+def test_a_pass_with_no_conditioning_is_unaffected(LC):
+    """t2v has no cond branch, so the old answer must not change."""
+    assert LC.bsa_problems(93, 512, 832, num_cond_latents=0) == []
+    assert LC.bsa_problems(93, 512, 832, num_cond_latents=None) == []
+    assert LC.bsa_problems(93, 512, 832) == []
+
+
+def test_both_halves_of_the_split_must_tile(LC):
+    """Latent depth 24: 4 cond + 20 noise tiles; 8 cond + 16 noise tiles; 6 does not."""
+    frames = 93                                   # (93-1)/4 + 1 = 24 latent frames
+    assert LC.bsa_problems(frames, 512, 832, num_cond_latents=4) == []
+    assert LC.bsa_problems(frames, 512, 832, num_cond_latents=8) == []
+    assert LC.bsa_problems(frames, 512, 832, num_cond_latents=6) != []
+
+
+def test_a_conditioning_block_that_swallows_the_noise_is_rejected(LC):
+    """Padding must not be allowed to leave zero noise latents to denoise."""
+    assert LC.bsa_problems(13, 512, 832, num_cond_latents=4) != []   # depth 4, noise 0
+
+
+def test_the_message_names_the_way_out(LC):
+    problems = LC.bsa_problems(93, 512, 832, num_cond_latents=1)
+    joined = " ".join(problems)
+    assert "bsa_pad_cond" in joined
+    assert "4" in joined
+
+
+def test_the_continuation_count_matches_the_pipeline(LC):
+    """pipeline_longcat_video.py:1016 — 1 + (num_cond_frames - 1) // 4."""
+    assert LC.cond_latents_for(1) == 1
+    assert LC.cond_latents_for(13) == 4
+    assert LC.cond_latents_for(0) == 0
+    assert LC.cond_latents_for(None) == 0
+
+
+def test_rounding_matches_what_generate_refine_does(LC):
+    """pipeline_longcat_video.py:1245 — ceil(n / granularity) * granularity."""
+    assert LC.round_to_granularity(1) == 4
+    assert LC.round_to_granularity(4) == 4
+    assert LC.round_to_granularity(5) == 8
+    assert LC.round_to_granularity(3, chunk=(2, 4, 4)) == 4
+
+
+# ───────────────────────── (a) the guard refuses a cond pass ──────────────────
+
+def test_an_i2v_request_counts_one_conditioning_latent(service):
+    assert service._cond_latent_passes("i2v", {"num_frames": 93}) == [1]
+
+
+def test_a_t2v_request_counts_none(service):
+    assert service._cond_latent_passes("t2v", {"num_frames": 93}) == [0]
+
+
+def test_multiple_segments_add_a_continuation_pass(service):
+    """Segment 0 is t2v, every later one is generate_vc — one DiT, one BSA flag."""
+    counts = service._cond_latent_passes(
+        "t2v", {"num_frames": 93, "num_segments": 3, "num_cond_frames": 13})
+    assert 0 in counts and 4 in counts
+
+
+def test_a_continuation_request_uses_its_cond_frames(service):
+    counts = service._cond_latent_passes(
+        "vc", {"num_frames": 93, "num_cond_frames": 13})
+    assert counts == [4]
+
+
+def test_an_avatar_continuation_allows_for_the_reference_image(service):
+    """pipeline_longcat_video_avatar.py:1375 adds one for a ref image."""
+    counts = service._cond_latent_passes(
+        "ai2v", {"num_frames": 93, "num_cond_frames": 13})
+    assert 5 in counts
+
+
+def test_the_guard_checks_every_pass_the_request_will_run(service):
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def _bsa_for"):src.index("def _place_pipeline")]
+    assert "_cond_latent_passes" in body
+    assert "num_cond_latents=cond" in body
+    assert "with _bsa_for(pipe, ctx, task):" in src, "the pass kind must reach the guard"
+
+
+# ───────────────────────── (b) padding, off by default ────────────────────────
+
+def test_padding_is_off_by_default(service, monkeypatch):
+    monkeypatch.delenv("LONGCAT_BSA_PAD_COND", raising=False)
+    assert service._bsa_pad_cond() is False
+
+
+@pytest.mark.parametrize("value", ["", "off", "no", "false", "0", "none"])
+def test_every_spelling_of_off_is_off(service, value):
+    assert service._bsa_pad_cond(value) is False
+
+
+@pytest.mark.parametrize("value", ["on", "yes", "true", "1"])
+def test_it_can_be_turned_on_explicitly(service, value):
+    assert service._bsa_pad_cond(value) is True
+
+
+def test_a_nonsense_setting_is_rejected_before_the_load(service):
+    """A bad value must not surface after a multi-minute checkpoint load."""
+    with pytest.raises(ValueError):
+        service._bsa_pad_cond("maybe")
+    src = SERVICE.read_text(encoding="utf-8")
+    assert "_bsa_pad_cond()" in src[src.index("_bsa_enabled()\n    _bsa_pad_cond()"):]
+
+
+def test_the_env_var_is_read(service, monkeypatch):
+    monkeypatch.setenv("LONGCAT_BSA_PAD_COND", "on")
+    assert service._bsa_pad_cond() is True
+
+
+def test_with_padding_on_the_guard_checks_the_padded_count(service, monkeypatch):
+    """1 cond latent cannot tile; padded to 4 against a depth-24 latent, it can."""
+    monkeypatch.setenv("LONGCAT_BSA_PAD_COND", "on")
+    pipe, dit = _fake_pipe()
+    with service._bsa_for(pipe, {"num_frames": 93, "height": 512, "width": 832},
+                          "i2v") as on:
+        assert on is True, "padding is the whole point: BSA must survive an i2v pass"
+
+
+def test_with_padding_off_an_i2v_pass_falls_back_to_dense(service, monkeypatch):
+    monkeypatch.delenv("LONGCAT_BSA_PAD_COND", raising=False)
+    pipe, dit = _fake_pipe()
+    with service._bsa_for(pipe, {"num_frames": 93, "height": 512, "width": 832},
+                          "i2v") as on:
+        assert on is False
+        assert dit.block.enable_bsa is False, "it must actually be turned off"
+    assert dit.block.enable_bsa is True, "and restored for the next request"
+
+
+def test_padding_is_not_installed_unless_both_flags_are_on(service):
+    """bsa off + padding on must not patch the DiT: there is nothing to pad for."""
+    src = SERVICE.read_text(encoding="utf-8")
+    assert "_bsa_enabled() and _bsa_pad_cond() and _install_cond_padding(pipe)" in src
+
+
+def test_the_wrapper_rounds_the_count_on_the_way_into_the_dit(service):
+    pipe, dit = _fake_pipe()
+    assert service._install_cond_padding(pipe) is True
+    pipe.dit(hidden_states=_FakeLatent(depth=24), num_cond_latents=1)
+    assert dit.seen[-1] == 4
+
+
+def test_the_wrapper_is_idempotent(service):
+    """A reload must not stack wrappers, each rounding the previous one's answer."""
+    pipe, dit = _fake_pipe()
+    assert service._install_cond_padding(pipe) is True
+    assert service._install_cond_padding(pipe) is False
+
+
+def test_the_wrapper_leaves_a_pass_with_no_conditioning_alone(service):
+    pipe, dit = _fake_pipe()
+    service._install_cond_padding(pipe)
+    pipe.dit(hidden_states=_FakeLatent(depth=24), num_cond_latents=0)
+    assert dit.seen[-1] == 0
+    pipe.dit(hidden_states=_FakeLatent(depth=24))
+    assert dit.seen[-1] is None
+
+
+def test_the_wrapper_does_not_swallow_the_noise_block(service):
+    """Depth 4 padded to 4 leaves nothing to denoise — hand the real count through."""
+    pipe, dit = _fake_pipe()
+    service._install_cond_padding(pipe)
+    pipe.dit(hidden_states=_FakeLatent(depth=4), num_cond_latents=1)
+    assert dit.seen[-1] == 1
+
+
+def test_the_wrapper_uses_the_checkpoints_own_chunk(service):
+    pipe, dit = _fake_pipe(chunk=[2, 4, 4])
+    service._install_cond_padding(pipe)
+    pipe.dit(hidden_states=_FakeLatent(depth=24), num_cond_latents=1)
+    assert dit.seen[-1] == 2
+
+
+def test_the_guard_wraps_only_the_segment_loop(service):
+    """generate_refine rounds both counts itself (pipeline:1245-1250), so it is safe
+    with BSA on and must not be forced down the dense path. The guard therefore wraps
+    _run_segments and nothing else."""
+    src = SERVICE.read_text(encoding="utf-8")
+    sites = [m.start() for m in re.finditer(r"with _bsa_for\(", src)]
+    assert len(sites) == 1, "one wrap point, or a pass can slip past the guard"
+    after = src[sites[0]:].split("\n", 2)[1]
+    assert "_run_segments(" in after
+
+
+# ────────────────────────────── config plumbing ───────────────────────────────
+
+def test_the_worker_passes_the_flag_through():
+    src = WORKER.read_text(encoding="utf-8")
+    assert '"bsa_pad_cond": "off"' in src, "it must be documented AND default off"
+    assert '"--bsa-pad-cond"' in src
+
+
+def test_the_admin_save_does_not_drop_it():
+    """A key the whitelist does not name is silently lost on every GUI save."""
+    routes = (ROOT / "codai" / "admin" / "routes.py").read_text(encoding="utf-8")
+    assert '"bsa_pad_cond"' in routes
+    html = (ROOT / "codai" / "admin" / "templates"
+            / "models.html").read_text(encoding="utf-8")
+    assert "cfg-longcat-bsa-pad-cond" in html
+    assert "out.bsa_pad_cond" in html
+    assert "s.bsa_pad_cond" in html
+
+
+# ───────────────────────────────── fakes ──────────────────────────────────────
+
+class _FakeLatent:
+    """Stands in for the [B, C, T, H, W] latent the DiT is handed."""
+
+    def __init__(self, depth):
+        self.shape = (1, 16, depth, 32, 52)
+        self.ndim = 5
+
+
+class _FakeBlock:
+    def __init__(self, chunk):
+        self.enable_bsa = True
+        self.bsa_params = {"chunk_3d_shape_q": list(chunk),
+                           "chunk_3d_shape_k": list(chunk)}
+
+
+class _FakeDit:
+    def __init__(self, chunk):
+        self.block = _FakeBlock(chunk)
+        self.seen = []
+
+    def modules(self):
+        return [self, self.block]
+
+    def forward(self, *args, **kwargs):
+        self.seen.append(kwargs.get("num_cond_latents"))
+        return None
+
+    def __call__(self, *args, **kwargs):
+        # nn.Module._call_impl reads self.forward, which is what the wrapper rebinds.
+        return self.forward(*args, **kwargs)
+
+
+class _FakePipe:
+    def __init__(self, dit):
+        self.dit = dit
+
+
+def _fake_pipe(chunk=(4, 4, 4)):
+    dit = _FakeDit(chunk)
+    return _FakePipe(dit), dit

tools/longcat_common.py

@@ -286,7 +286,35 @@ VAE_TEMPORAL_STRIDE = 4
 DEFAULT_BSA_CHUNK = (4, 4, 4)
 
 
-def bsa_problems(num_frames: int, height: int, width: int, chunk=None) -> list:
+# What generate_i2v passes, hardcoded, whatever num_frames is
+# (pipeline_longcat_video.py:796; the avatar pipeline does the same at 1076-1113).
+I2V_COND_LATENTS = 1
+
+
+def cond_latents_for(num_cond_frames) -> int:
+    """How many latent frames a continuation pass treats as conditioning.
+
+    Mirrors pipeline_longcat_video.py:1016 — and, like it, does NOT round up to the
+    BSA granularity the way generate_refine does at 1245.
+    """
+    try:
+        frames = int(num_cond_frames)
+    except (TypeError, ValueError):
+        return 0
+    if frames <= 0:
+        return 0
+    return 1 + (frames - 1) // VAE_TEMPORAL_STRIDE
+
+
+def round_to_granularity(count, chunk=None) -> int:
+    """``count`` rounded UP to the BSA temporal granularity, as generate_refine does."""
+    tq = int((tuple(chunk or DEFAULT_BSA_CHUNK))[0]) or 1
+    n = max(0, int(count))
+    return -(-n // tq) * tq
+
+
+def bsa_problems(num_frames: int, height: int, width: int, chunk=None,
+                 num_cond_latents=None) -> list:
     """Why block-sparse attention cannot run THIS geometry, or [].
 
     flash_attn_bsa_3d tiles the latent into 3D chunks and asserts the latent divides
@@ -299,6 +327,15 @@ def bsa_problems(num_frames: int, height: int, width: int, chunk=None) -> list:
     num_frames at 13, 29, 45, 61, 77, 93 … Upstream's own default 480x832 is NOT one of
     them (480/16 = 30), which is presumably why the flag ships off: an AssertionError
     out of the first denoising step says nothing about geometry.
+
+    ``num_cond_latents`` is the OTHER half of the same assert, and the half that cost a
+    production run. A conditioning pass does not tile the latent once: Attention.forward
+    slices the sequence into a conditioning block and a noise block and calls
+    flash_attn_bsa_3d on EACH (modules/attention.py:123-134), so the cond depth and the
+    remaining noise depth must both divide by ``tq`` as well. generate_i2v hardcodes one
+    cond latent and 1 % 4 != 0, which means block-sparse attention has never once been
+    able to run an i2v pass at any frame count or resolution. Pass 0 or None for a pass
+    with no conditioning branch.
     """
     tq, hq, wq = tuple(chunk or DEFAULT_BSA_CHUNK)
     problems = []
@@ -329,6 +366,18 @@ def bsa_problems(num_frames: int, height: int, width: int, chunk=None) -> list:
             f"num_frames must leave a latent depth divisible by {tq} for block-sparse "
             f"attention ({step}n+1: 13, 29, 45, 61, 77, 93 …); got {frames}"
             + (f", try {' or '.join(str(n) for n in options)}" if options else ""))
+
+    if num_cond_latents:
+        cond = int(num_cond_latents)
+        noise = latent_t - cond
+        if cond % tq or noise <= 0 or noise % tq:
+            problems.append(
+                f"a conditioning pass tiles {cond} cond + {noise} noise latent frames "
+                f"separately, so BOTH must divide by {tq} for block-sparse attention; "
+                f"generate_i2v hardcodes {I2V_COND_LATENTS}, which never can. Turn "
+                f"bsa_pad_cond on to round the conditioning up to "
+                f"{round_to_granularity(cond, chunk)}, or leave block-sparse attention "
+                f"off for conditioning passes")
     return problems
 
 

tools/longcat_service.py

@@ -288,6 +288,11 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
         # Both pipeline modules are imported by now (the avatar one only for an
         # avatar family), so this is the first point where the bars can be rebound.
         _install_progress_bars()
+        if _bsa_enabled() and _bsa_pad_cond() and _install_cond_padding(pipe):
+            log(f"block-sparse attention: conditioning latents padded up to "
+                f"{int(_bsa_chunk(pipe)[0])} (bsa_pad_cond on) — i2v and continuation "
+                f"passes keep BSA, at the cost of attending over latent frames the "
+                f"pipeline still marks as noise")
         _place_pipeline(pipe, offload)
         _state.update(pipe=pipe, model=root, source=src, dtype=dtype)
         log("ready")
@@ -353,6 +358,81 @@ def _bsa_enabled(setting=None):
     return True
 
 
+def _bsa_pad_cond(setting=None) -> bool:
+    """Whether to round a conditioning pass up to the block-sparse granularity. OFF.
+
+    This is the only way to get block-sparse attention on an i2v or continuation pass:
+    generate_i2v hardcodes one conditioning latent, the cond and noise blocks are tiled
+    separately, and 1 % 4 != 0, so the pass asserts out of the first denoising step.
+    generate_refine already solves it for itself (pipeline_longcat_video.py:1245) by
+    rounding both counts up to the granularity; this does the same thing for the paths
+    upstream left out.
+
+    It is off by default because it is NOT free. The pipeline marks only the real
+    conditioning frames clean (``timestep[:, :1] = 0``), so rounding 1 up to 4 hands the
+    conditioning block three latent frames that are still noise. Attention over them is
+    well-defined but it is not what the model was trained on, and the failure mode for a
+    video model is temporal drift — the same class of defect as the dual-expert grey
+    video. Measure a known-good clip against the dense path before trusting it.
+    """
+    want = str(setting if setting is not None
+               else os.environ.get("LONGCAT_BSA_PAD_COND") or "off").strip().lower()
+    if want in ("", "0", "off", "false", "no", "none"):
+        return False
+    if want not in ("1", "on", "true", "yes"):
+        raise ValueError(f"bsa_pad_cond must be on or off; got {want!r}")
+    return True
+
+
+def _bsa_chunk(pipe):
+    """The checkpoint's own 3D chunk, or the shipped default."""
+    blocks = [m for m in pipe.dit.modules() if getattr(m, "enable_bsa", False)] \
+        if getattr(pipe, "dit", None) is not None else []
+    if blocks:
+        chunk = (getattr(blocks[0], "bsa_params", None) or {}).get("chunk_3d_shape_q")
+        if chunk:
+            return tuple(chunk)
+    return LC.DEFAULT_BSA_CHUNK
+
+
+def _install_cond_padding(pipe):
+    """Round ``num_cond_latents`` up to the granularity on the way into the DiT.
+
+    Patched onto the DiT instance rather than the vendored source, the same way the
+    denoising bars are rebound: the count is born inside generate_i2v's loop, so there
+    is no outer argument to correct, but every path funnels through one ``self.dit(...)``
+    call. nn.Module._call_impl reads ``self.forward``, so an instance attribute is the
+    hook.
+
+    Returns False when there is nothing to do, so the caller can say so.
+    """
+    dit = getattr(pipe, "dit", None)
+    if dit is None or getattr(dit, "_lc_cond_padded", False):
+        return False
+    inner = dit.forward
+    tq = int(_bsa_chunk(pipe)[0]) or 1
+
+    def forward(*args, **kwargs):
+        n = kwargs.get("num_cond_latents")
+        if n:
+            hidden = kwargs.get("hidden_states")
+            if hidden is None and args:
+                hidden = args[0]
+            # [B, C, T, H, W]: the latent depth this pass is splitting.
+            depth = int(hidden.shape[2]) if getattr(hidden, "ndim", 0) == 5 else 0
+            padded = -(-int(n) // tq) * tq
+            # Leave it alone if padding would swallow the noise block — _bsa_for has
+            # already turned BSA off for that case, so dense attention wants the real
+            # count, not a widened one.
+            if padded != n and (not depth or padded < depth):
+                kwargs["num_cond_latents"] = padded
+        return inner(*args, **kwargs)
+
+    dit.forward = forward
+    dit._lc_cond_padded = True
+    return True
+
+
 def _dit_kwargs():
     """Everything the DiT has to be CONSTRUCTED with, for either variant.
 
@@ -798,8 +878,39 @@ def _wrap_encode_prompt(pipe, swap: bool, cache=None, lazy=None):
     # Set on the INSTANCE: all four generate_* methods call self.encode_prompt.
     pipe.encode_prompt = encode_prompt
     return pipe
+def _cond_latent_passes(task: str, ctx: dict) -> list:
+    """How many conditioning latents each pass of THIS request will attend over.
+
+    One request is not one pass. _run_segments renders segment 0 with the task's own
+    entry point and every later segment with generate_vc, so even a t2v request runs a
+    conditioning pass once num_segments > 1 — and block-sparse attention is a property
+    of the DiT, set for the whole request. So collect every count the request can
+    produce and let the caller refuse if ANY of them cannot tile.
+
+    0 is a pass with no conditioning branch.
+    """
+    counts = []
+    cond_frames = ctx.get("num_cond_frames")
+    continuing = ctx.get("cond_video") is not None
+    if continuing or task in ("vc", "avc"):
+        counts.append(LC.cond_latents_for(cond_frames))
+    elif task in ("i2v", "ai2v"):
+        counts.append(LC.I2V_COND_LATENTS)
+    else:
+        counts.append(0)
+    if max(1, int(ctx.get("num_segments") or 1)) > 1:
+        counts.append(LC.cond_latents_for(cond_frames))
+    if task in LC.AVATAR_TASKS or task == "avc":
+        # The avatar continuation adds one more for a reference image
+        # (pipeline_longcat_video_avatar.py:1375).
+        base = LC.cond_latents_for(cond_frames)
+        if base:
+            counts.append(base + 1)
+    return counts
+
+
 @contextlib.contextmanager
-def _bsa_for(pipe, ctx):
+def _bsa_for(pipe, ctx, task: str = ""):
     """Run this pass with block-sparse attention only if the geometry allows it.
 
     flash_attn_bsa_3d asserts the latent divides evenly by its 3D chunk, so a shape it
@@ -807,16 +918,30 @@ def _bsa_for(pipe, ctx):
     AssertionError out of the first denoising step. Turning the flag on in a config
     must not break requests that work; falling back silently must not leave someone
     wondering why nothing got faster. So: fall back, and say why.
+
+    The geometry is not only num_frames x height x width. A conditioning pass tiles the
+    cond and noise blocks separately, and generate_i2v's hardcoded single cond latent
+    tiles under no chunk at all — which is how an i2v request with a perfectly valid
+    93-frame, 64-divisible geometry still died in the first step. With bsa_pad_cond on,
+    the DiT wrapper rounds that count up before the split, so the counts checked here
+    are the padded ones.
     """
     blocks = [m for m in pipe.dit.modules() if getattr(m, "enable_bsa", False)] \
         if getattr(pipe, "dit", None) is not None else []
     problems = []
     if blocks:
         chunk = (blocks[0].bsa_params or {}).get("chunk_3d_shape_q")
-        problems = LC.bsa_problems(ctx.get("num_frames") or LC.DEFAULT_NUM_FRAMES,
+        conds = _cond_latent_passes(task, ctx)
+        if _bsa_pad_cond():
+            conds = [LC.round_to_granularity(c, chunk) if c else c for c in conds]
+        for cond in dict.fromkeys(conds):
+            for problem in LC.bsa_problems(
+                    ctx.get("num_frames") or LC.DEFAULT_NUM_FRAMES,
                     ctx.get("height") or LC.DEFAULT_BASE_SIZE[0],
                     ctx.get("width") or LC.DEFAULT_BASE_SIZE[1],
-                                   chunk)
+                    chunk, num_cond_latents=cond):
+                if problem not in problems:
+                    problems.append(problem)
     if not blocks or not problems:
         yield bool(blocks)
         return
@@ -1455,7 +1580,7 @@ def generate(body: dict) -> dict:
                 log(f"stage 'refinement': refining {len(frames)} frames")
                 frames = _frames_from_output(pipe.generate_refine(**kw)[0])
             else:
-                with _bsa_for(pipe, ctx):
+                with _bsa_for(pipe, ctx, task):
                     frames, yielded = _run_segments(pipe, task, stage, ctx, log=log)
             log(f"stage '{stage}': {len(frames)} frames")
     finally:
@@ -1648,6 +1773,11 @@ def main(argv=None):
     ap.add_argument("--bsa", default="off",
                     help="block-sparse attention: on | off | auto (auto stays off "
                          "when triton is missing instead of failing)")
+    ap.add_argument("--bsa-pad-cond", default="off",
+                    help="round a conditioning pass up to the block-sparse granularity "
+                         "so i2v/continuation can use BSA at all: on | off (default "
+                         "off; it widens the conditioning block past what the pipeline "
+                         "marks clean)")
     ap.add_argument("--cp-split-hw", default="",
                     help="context-parallel tile as HxW, e.g. 1x2; "
                          "defaults to 1x<cp_size>")
@@ -1666,6 +1796,7 @@ def main(argv=None):
     # by default (requirements-longcat.txt installs xformers instead).
     os.environ.setdefault("LONGCAT_ATTENTION", args.attention)
     os.environ.setdefault("LONGCAT_BSA", args.bsa)
+    os.environ.setdefault("LONGCAT_BSA_PAD_COND", args.bsa_pad_cond)
     _state["base_model"] = args.base_model or ""
     _problems = LC.component_quant_problems("text_encoder_quant",
                                             args.text_encoder_quant)
@@ -1691,6 +1822,7 @@ def main(argv=None):
     # Validate both before anything slow happens: a bad value should not surface after
     # a multi-minute load.
     _bsa_enabled()
+    _bsa_pad_cond()
     _state["cp_split_hw"] = _cp_split_hw(1, args.cp_split_hw)
 
     rank, _local = (0, 0)