nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(h3): our own chunk-causal cache, geometry and DiT subclasses

Browse Source 
First half of owning the causal lane outright, so the feature needs no
third-party package installed.

h3_causal.py: a ChunkKVCache with sink/window eviction and host offload, the
chunk cut, and CausalAttention / CausalDiTBlock / CausalMiniMaxH3Model as
subclasses of ComfyUI's own H3 classes — same module names, same state dict, so
weights load with no remapping and a PEFT LoRA attaches through ComfyUI's
ordinary adapter (which already reads lora_A/lora_B keys).

The safety property: with cache=None every forward defers to its parent, so the
dense lane is bit-identical to stock H3 and nothing here can degrade an
ordinary generation.

Geometry is BORROWED, not re-derived. A chunk's RoPE positions must be the ones
the dense layout would have given those rows, so the full clip's grid is built
with upstream's own _frame_grid / _video_grid / _audio_grid and sliced per
chunk. An upstream grid change is inherited rather than silently disagreed with.

Two traps handled explicitly:
  - pack_audio is channel-major over the WHOLE clip, so a chunk's audio rows are
    TWO spans, not one. A contiguous slice would hand the model the wrong half
    of the stereo field, quietly.
  - blocks are RE-CLASSED in place, adopting submodules by reference, never
    wrapped. Wrapping renames the state-dict keys of everything inside, which
    has broken unpatching in this project before.

Tests: 1751 passed (+21). prefill_text/forward_chunk still to come.
This commit is contained in:
digital-garbage
2026-08-26 18:53:31 +03:00
parent e0044425ca
commit 840098869c
2 changed files with 506 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
h3_causal.py
+296
View File
@@ -0,0 +1,296 @@
"""FunPack's own chunk-causal MiniMax H3: a KV cache, and a DiT that reads it.
WHY THIS EXISTS
---------------
H3's stock DiT does ONE dense forward over the whole packed sequence. Chunk-causal generation
needs something it does not have: a forward over a slice of the clip that attends to a cache
of the slices already finished. That is the machinery here.
Every class is a SUBCLASS of its ComfyUI counterpart, with the same module names and the same
state dict, so:
* the weights load with no remapping and no conversion;
* a LoRA trained for the chunk-causal pattern attaches through ComfyUI's ordinary LoRA path
(its adapter already reads PEFT `lora_A`/`lora_B` keys);
* with no cache passed, every forward defers to its parent and is bit-identical to stock H3.
That last point is the safety property. `cache=None` is the dense lane, unchanged, so nothing
here can degrade an ordinary generation.
GEOMETRY IS BORROWED, NOT REDERIVED
-----------------------------------
A chunk's RoPE positions must be the ones the dense layout would have given those same rows,
or the model is being told the clip is a different shape than it is. So the FULL clip's grid
is built with upstream's own `_frame_grid` / `_video_grid` / `_audio_grid` and then SLICED per
chunk. An upstream change to the grid is inherited rather than silently disagreed with.
THE LORA IS NOT OPTIONAL
------------------------
Chunk-causal attention is a pattern the base H3 weights have never seen — they were trained
with every row attending every other row. Running this without an adapter trained for it
produces output, and the output is out of distribution. The loader says so.
"""
from __future__ import annotations
import torch
#: video latents per chunk, and the pixel frames that many latents cover (1+4+4+4+4)
VIDEO_LATENTS_PER_CHUNK = 5
FRAMES_PER_CHUNK = 17
#: audio latents per pixel frame — 40/s against the model's fixed 24 fps
AUDIO_LATENTS_PER_FRAME = 40.0 / 24.0
#: AdaLN modality tags, from the dense model's own seg_tag
VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2
class CacheError(RuntimeError):
"""The cache was asked for something it cannot honestly answer."""
class ChunkKVCache:
"""Per-layer K/V of the chunks already finished, with sink + window eviction.
Two kinds of memory, and they do different jobs:
`sink` chunks are pinned from the very start and never evicted. Chunk 0 is the text, so a
sink of 2 means the prompt plus the opening shot stay visible to every later chunk — this
is what holds a character together across a long clip.
`window` is the most recent chunks besides those, which is short-term continuity: motion
and lighting carrying from one moment to the next.
Anything in between is dropped, which is what keeps memory flat instead of growing with
clip length. `window=None` keeps everything.
"""
def __init__(self, num_layers, sink=2, window=2, device=None, offload=True):
if int(num_layers) <= 0:
raise CacheError(f"num_layers must be positive, got {num_layers!r}")
if int(sink) < 0:
raise CacheError(f"sink must be >= 0, got {sink!r}")
if window is not None and int(window) < 0:
raise CacheError(f"window must be None or >= 0, got {window!r}")
self.num_layers = int(num_layers)
self.sink = int(sink)
self.window = None if window is None else int(window)
self.device = device
# The retained K/V of a long clip is measured in tens of GiB on the card. Holding it
# on the host and copying one chunk back per layer trades bandwidth for the ability to
# run at all.
self.offload = bool(offload)
self.committed_chunks = 0
self._store = {} # (layer, chunk_index) -> (k, v)
def _keep(self, chunk_index, next_index):
if chunk_index < self.sink:
return True
if self.window is None:
return True
return chunk_index > next_index - 1 - self.window
def retained_indices(self, next_index):
"""Which committed chunks are visible to the chunk about to be generated."""
return [i for i in range(self.committed_chunks) if self._keep(i, next_index)]
def read(self, layer, next_index):
"""The retained K and V for one layer, in chunk order, or (None, None)."""
parts = [self._store.get((layer, i)) for i in self.retained_indices(next_index)]
parts = [p for p in parts if p is not None]
if not parts:
return None, None
keys = torch.cat([p[0] for p in parts], dim=-2)
values = torch.cat([p[1] for p in parts], dim=-2)
if self.device is not None:
keys, values = keys.to(self.device), values.to(self.device)
return keys, values
def write(self, layer, chunk_index, keys, values):
"""Store one layer's K/V for a finished chunk."""
if self.offload:
keys, values = keys.detach().to("cpu"), values.detach().to("cpu")
else:
keys, values = keys.detach(), values.detach()
self._store[(layer, int(chunk_index))] = (keys, values)
def finish_chunk(self, chunk_index):
"""Mark a chunk committed, and drop whatever is now out of both sink and window."""
self.committed_chunks = max(self.committed_chunks, int(chunk_index) + 1)
self.prune()
def prune(self):
keep = set(self.retained_indices(self.committed_chunks))
for key in [k for k in self._store if k[1] not in keep]:
del self._store[key]
def clear(self):
self._store.clear()
self.committed_chunks = 0
def chunk_bounds(latent_t, audio_t):
"""(video_start, video_stop, audio_start, audio_stop) per chunk.
Video cuts every ``VIDEO_LATENTS_PER_CHUNK`` latents. Audio is derived from the SHARED
CLOCK — 40 latents per second against a fixed 24 fps, 17 pixel frames per chunk — so the
familiar 28/29/28/28 cadence comes out as a consequence and stays correct if the grid ever
moves, instead of being a remembered pattern that silently drifts.
"""
latent_t, audio_t = int(latent_t), int(audio_t)
if latent_t <= 0:
return []
bounds, cursor = [], 0
starts = list(range(0, latent_t, VIDEO_LATENTS_PER_CHUNK))
for index, video_start in enumerate(starts):
video_stop = min(latent_t, video_start + VIDEO_LATENTS_PER_CHUNK)
if index == len(starts) - 1:
audio_stop = audio_t # the tail owns whatever is left
else:
done = FRAMES_PER_CHUNK * (index + 1)
audio_stop = max(cursor, min(audio_t, int(round(done * AUDIO_LATENTS_PER_FRAME))))
bounds.append((video_start, video_stop, cursor, audio_stop))
cursor = audio_stop
return bounds
def clip_grids(text_len, latent_t, latent_h, latent_w, audio_t):
"""The whole clip's video and audio position grids, and the rows per latent frame.
Built with upstream's own helpers and sliced per chunk afterwards, so a chunk's RoPE
angles are exactly the ones the dense layout would have given those rows. Re-deriving the
geometry here instead would be a second source of truth for the clip's shape.
"""
from comfy.ldm.minimax.model import _audio_grid, _frame_grid, _video_grid
frame, w_grid = _frame_grid(int(latent_h), int(latent_w))
video = _video_grid(int(latent_t), frame, float(text_len))
audio = _audio_grid(float(text_len), int(audio_t),
float(w_grid[0]), float(w_grid[-1]))
return video, audio, int(frame.shape[0])
def chunk_rows(audio_t, audio_start, audio_stop):
"""Which packed audio ROWS a chunk owns.
`pack_audio` is channel-major over the WHOLE clip — all of the left channel, then all of
the right — so a chunk's rows are two separate spans, not one. Treating them as one
contiguous slice silently hands the model the wrong half of the stereo field.
"""
left = list(range(int(audio_start), int(audio_stop)))
right = [int(audio_t) + i for i in left]
return left + right
def _causal_classes():
"""Build the three subclasses against whatever H3 this ComfyUI has.
Deferred into a function because these subclass ``comfy.ldm.minimax.model``, which is not
importable on a ComfyUI without H3 — and a node pack must still register there.
"""
import comfy.model_management
import comfy.quant_ops
from comfy.ldm.minimax.model import (Attention, DiTBlock, MiniMaxH3Model,
_mod_gate, _mod_scale_shift, optimized_attention)
class CausalAttention(Attention):
"""Upstream attention, with the retained K/V prepended to this chunk's own.
Q stays current-only: this chunk is what is being generated. K and V are the cache
followed by this chunk, so the chunk sees all of itself and all of what was retained.
There is no mask — that visibility IS the causal pattern, expressed by what is in the
buffer rather than by hiding parts of a bigger one.
Cached keys keep the absolute positions they were rotated with, so nothing is ever
re-based. That is what makes the cache reusable across steps.
"""
def forward(self, x, rope_freqs=None, transformer_options={},
cache=None, layer_idx=0, chunk_index=0, update_cache=False):
if cache is None:
# the dense lane, untouched
return super().forward(x, rope_freqs=rope_freqs,
transformer_options=transformer_options)
s = x.shape[0]
q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)
v = v.view(s, self.heads, self.head_dim)
if rope_freqs is not None:
q = q.view(1, s, self.heads, self.head_dim)
k = k.view(1, s, self.heads, self.head_dim)
qw = comfy.model_management.cast_to(self.q_norm.weight, device=x.device)
kw = comfy.model_management.cast_to(self.k_norm.weight, device=x.device)
rot = rope_freqs.shape[-3] * 2
if comfy.model_management.in_training:
q, k = comfy.quant_ops.ck.rms_rope_split_half(
q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
else:
comfy.quant_ops.ck.rms_rope_split_half_(
q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
q, k = q[0], k[0]
else:
q = self.q_norm(q.view(s, self.heads, self.head_dim))
k = self.k_norm(k.view(s, self.heads, self.head_dim))
q = q.transpose(0, 1).unsqueeze(0)
k = k.transpose(0, 1).unsqueeze(0)
v = v.transpose(0, 1).unsqueeze(0)
# Written BEFORE the cache is prepended: what this chunk contributes to later
# chunks is its own K/V, never the concatenation it happened to attend over.
if update_cache:
cache.write(layer_idx, chunk_index, k, v)
past_k, past_v = cache.read(layer_idx, chunk_index)
if past_k is not None:
k = torch.cat((past_k.to(dtype=k.dtype, device=k.device), k), dim=-2)
v = torch.cat((past_v.to(dtype=v.dtype, device=v.device), v), dim=-2)
out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True,
transformer_options=transformer_options)
return self.out_proj(out.squeeze(0))
class CausalDiTBlock(DiTBlock):
"""Upstream block, with the cache threaded to its attention.
A verbatim fork of the parent body — the norm / AdaLN / gate / MLP order is the
checkpoint's contract and must not drift — differing only in the kwargs passed on.
"""
def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={},
cache=None, layer_idx=0, chunk_index=0, update_cache=False):
if cache is None:
return super().forward(x, t_emb, mod_segments, rope_freqs,
transformer_options=transformer_options)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)
attended = self.attn(h, rope_freqs=rope_freqs,
transformer_options=transformer_options,
cache=cache, layer_idx=layer_idx,
chunk_index=chunk_index, update_cache=update_cache)
x = _mod_gate(x, gate_msa, attended, mod_segments)
h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)
return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)
class CausalMiniMaxH3Model(MiniMaxH3Model):
"""Upstream H3 whose blocks can read a chunk cache. Same state dict, key for key.
The blocks are rebuilt as causal ones after construction, ONE AT A TIME so peak memory
stays at a single extra block rather than a second copy of all fifty. Each rebuilt
block adopts the original's submodules by reference, so no weight is copied and no key
is renamed — the hard rule this project already learned about wrapper modules.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
for index in range(len(self.blocks)):
self.blocks[index] = _to_causal_block(self.blocks[index],
CausalDiTBlock, CausalAttention)
return CausalAttention, CausalDiTBlock, CausalMiniMaxH3Model
def _to_causal_block(block, block_cls, attention_cls):
"""Re-class one block and its attention in place, adopting every submodule by reference.
NOT a wrapper module. Wrapping renames the state-dict keys of everything inside it, which
has broken unpatching here before; re-classing changes only which `forward` runs.
"""
block.attn.__class__ = attention_cls
block.__class__ = block_cls
return block
tests/test_h3_causal.py
+210
View File
@@ -0,0 +1,210 @@
"""FunPack's own chunk-causal H3: the cache, the chunk geometry, and the re-classed DiT.
H3's stock DiT does one dense forward over the whole packed sequence. Chunk-causal generation
needs a forward over a SLICE that attends to a cache of the slices already finished. This is
that machinery, owned rather than borrowed — no third-party package has to be installed for
the feature to exist.
The safety property everything rests on: with no cache passed, every forward defers to its
parent, so an ordinary generation cannot be degraded by any of this.
"""
import sys
from pathlib import Path
import pytest
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import h3_causal as hc
# ── the cache: what a chunk is allowed to remember ──────────────────────────
def _cache(**kw):
kw.setdefault("num_layers", 2)
return hc.ChunkKVCache(**kw)
def _kv(rows, value):
return torch.full((1, 2, rows, 4), float(value)), torch.full((1, 2, rows, 4), float(value))
def test_an_empty_cache_has_nothing_to_read():
assert _cache().read(0, 0) == (None, None)
def test_a_committed_chunk_is_readable():
cache = _cache()
k, v = _kv(3, 1.0)
cache.write(0, 0, k, v)
cache.finish_chunk(0)
got_k, got_v = cache.read(0, 1)
assert got_k.shape[-2] == 3 and got_v.shape[-2] == 3
def test_layers_do_not_share_their_cache():
"""Layer 7 reading layer 3's keys would be silent nonsense, not a crash."""
cache = _cache(num_layers=4)
cache.write(0, 0, *_kv(3, 1.0))
cache.finish_chunk(0)
assert cache.read(1, 1) == (None, None)
def test_chunks_are_read_back_in_order():
cache = _cache()
for index, value in enumerate((1.0, 2.0)):
cache.write(0, index, *_kv(2, value))
cache.finish_chunk(index)
keys, _v = cache.read(0, 2)
assert keys[0, 0, 0, 0] == 1.0 and keys[0, 0, 2, 0] == 2.0
def test_sink_chunks_are_never_evicted():
"""The prompt and the opening shot are what hold a character together over a long clip."""
cache = _cache(sink=1, window=1)
for index in range(6):
cache.write(0, index, *_kv(2, index))
cache.finish_chunk(index)
assert 0 in cache.retained_indices(6)
def test_the_window_keeps_only_the_most_recent():
cache = _cache(sink=1, window=2)
for index in range(6):
cache.write(0, index, *_kv(2, index))
cache.finish_chunk(index)
assert cache.retained_indices(6) == [0, 4, 5]
def test_a_zero_window_leans_entirely_on_the_sink():
cache = _cache(sink=2, window=0)
for index in range(5):
cache.write(0, index, *_kv(2, index))
cache.finish_chunk(index)
assert cache.retained_indices(5) == [0, 1]
def test_no_window_remembers_everything():
cache = _cache(sink=1, window=None)
for index in range(5):
cache.write(0, index, *_kv(2, index))
cache.finish_chunk(index)
assert cache.retained_indices(5) == [0, 1, 2, 3, 4]
def test_evicted_chunks_are_actually_dropped():
"""Retention that only hides rows would grow memory with clip length, which is the thing
chunking exists to avoid."""
cache = _cache(sink=1, window=1)
for index in range(8):
cache.write(0, index, *_kv(64, index))
cache.finish_chunk(index)
assert len(cache._store) <= 3
def test_a_negative_sink_is_refused():
with pytest.raises(hc.CacheError):
_cache(sink=-1)
def test_zero_layers_is_refused():
with pytest.raises(hc.CacheError):
hc.ChunkKVCache(0)
def test_clearing_forgets_everything():
cache = _cache()
cache.write(0, 0, *_kv(2, 1.0))
cache.finish_chunk(0)
cache.clear()
assert cache.committed_chunks == 0 and cache.read(0, 1) == (None, None)
def test_offloaded_entries_leave_the_compute_device():
"""The retained K/V of a long clip is tens of GiB; holding it on the card is what makes a
long clip impossible rather than merely slow."""
cache = _cache(offload=True)
cache.write(0, 0, *_kv(2, 1.0))
assert cache._store[(0, 0)][0].device.type == "cpu"
# ── the chunk cut ───────────────────────────────────────────────────────────
def test_video_cuts_every_five_latents():
bounds = hc.chunk_bounds(latent_t=12, audio_t=100)
assert [(a, b) for a, b, _c, _d in bounds] == [(0, 5), (5, 10), (10, 12)]
def test_every_latent_lands_in_exactly_one_chunk():
bounds = hc.chunk_bounds(latent_t=22, audio_t=180)
video, audio = [], []
for vs, ve, as_, ae in bounds:
video.extend(range(vs, ve))
audio.extend(range(as_, ae))
assert video == list(range(22)) and audio == list(range(180))
def test_the_audio_cadence_falls_out_of_the_clock():
"""40 latents/s against a fixed 24 fps, 17 frames per chunk. The 28/29 alternation is a
CONSEQUENCE — hard-coding it would drift silently if the grid ever moved."""
frames = 17 * 11 + 5
bounds = hc.chunk_bounds(latent_t=5 * 11 + 2, audio_t=int(round(frames / 24 * 40)))
per_chunk = [ae - as_ for _a, _b, as_, ae in bounds]
assert per_chunk[:4] == [28, 29, 28, 28]
def test_the_tail_chunk_is_kept_not_dropped():
"""H3's grid is 5k+2 latents, so there is ALWAYS a two-latent tail. Dropping it would cut
the last five frames off every clip."""
bounds = hc.chunk_bounds(latent_t=17, audio_t=140)
assert bounds[-1][1] - bounds[-1][0] == 2
# ── stereo rows: the trap in packed audio ───────────────────────────────────
def test_a_chunks_audio_rows_are_two_spans_not_one():
"""pack_audio is channel-major over the WHOLE clip — all of the left channel, then all of
the right. Treating a chunk's rows as one contiguous slice hands the model the wrong half
of the stereo field, quietly."""
rows = hc.chunk_rows(audio_t=100, audio_start=28, audio_stop=57)
assert rows[:29] == list(range(28, 57))
assert rows[29:] == list(range(128, 157))
def test_the_row_count_is_two_per_latent():
rows = hc.chunk_rows(audio_t=100, audio_start=0, audio_stop=28)
assert len(rows) == 56
def test_the_first_chunk_starts_at_both_channel_origins():
rows = hc.chunk_rows(audio_t=50, audio_start=0, audio_stop=4)
assert rows == [0, 1, 2, 3, 50, 51, 52, 53]
# ── re-classing, not wrapping ───────────────────────────────────────────────
def test_a_block_is_re_classed_in_place_and_keeps_its_submodules():
"""Wrapping renames the state-dict keys of everything inside, which has broken unpatching
here before. Re-classing changes only which forward runs."""
class FakeAttention:
pass
class FakeBlock:
def __init__(self):
self.attn = FakeAttention()
self.norm1 = "the-real-norm"
class CausalAttn(FakeAttention):
pass
class CausalBlock(FakeBlock):
def __init__(self):
raise AssertionError("re-classing must not re-run __init__")
block = FakeBlock()
attn, norm = block.attn, block.norm1
out = hc._to_causal_block(block, CausalBlock, CausalAttn)
assert out is block
assert out.attn is attn and out.norm1 is norm
assert isinstance(out, CausalBlock) and isinstance(out.attn, CausalAttn)