| |
| """VRAM Lab (alpha, 2026-08-18): exact low-memory execution of MiniMax-H3 blocks. |
| |
| The problem, measured 2026-08-17/18 on an RTX PRO 6000 (numbers below): a full-length H3 forward materialises, per DiT block, the fused |
| QKV projection ``[N, 3*7168]`` and the SwiGLU pre-activation ``[N, 2*14336]`` |
| for the whole packed sequence. At ~217k tokens (an 8 s clip de-roped at |
| d_max 4) those two tensors are 8.7 GiB and 15.4 GiB, and the MLP-phase peak is |
| ~24 GiB before weights. That is what OOMs 24 GB cards on the de-rope. |
| |
| What this module does instead, per block, with the SAME math: |
| |
| phase 1 for each token chunk: h = mod(norm1(x_c)); qkv = qkv_proj(h); |
| keep K, V (RMS-normed + RoPE'd) into full-sequence buffers, |
| discard Q. |
| phase 2 for each token chunk: recompute h and qkv, keep Q only; |
| attention(Q_c, K, V) is exact per query row (flash / SDPA); |
| out_proj; gated residual written in place into x[c]. |
| phase 3 for each token chunk: mod(norm2(x_c)) -> mlp -> gated residual. |
| |
| Peak per block falls from ~N x 118 KB (MLP phase) to ~N x 39 KB (x + K + V) |
| plus chunk-sized transients. The projection is run twice (phase 1 and 2); at |
| 217k tokens that is ~4% of the block's FLOPs, because attention is N^2 and |
| projections are N. This costs no exactness: the int8 kernel quantises |
| activations per row, so chunking rows is bit-identical (measured |
| 2026-08-18 on comfy_kitchen TensorWiseINT8), and every other weight format |
| goes through the module's own forward, weight streaming included. |
| |
| ``kv_block > 0`` (experimental): phase 2 attends K/V in blocks and combines |
| with the flash kernel's log-sum-exp (online softmax), measured to one bf16 |
| ulp against a single call. AS BUILT it does NOT lower memory: K and V are |
| still materialised in full by phase 1, and the blockwise combine adds |
| buffers (measured +13.0 GiB forward peak vs +11.9 plain at 216k tokens, |
| ~7% slower). A real "turtle" needs K/V staged in host RAM and streamed per |
| block; until then leave it at 0. |
| |
| Composition: MLP chunking here is optional (``mlp_chunk = 0`` leaves the |
| model's own ``mlp.forward`` alone, so KJNodes' MiniMax H3 Chunk FeedForward |
| keeps working). Attention goes through ``optimized_attention`` so Sol-Attn |
| style overrides still apply; head-group chunking from KJNodes' Low VRAM |
| Attention is read from ``transformer_options["minimax_head_chunks"]``. |
| |
| Not a quality dial. If it changes the picture, it is a bug; the gate is a |
| same-seed difference clip against the stock block. |
| """ |
| import logging |
| import math |
| import os |
| import sys |
| import time |
|
|
| import torch |
|
|
| import comfy.model_management |
| import comfy.quant_ops |
| import comfy.ldm.common_dit |
| import comfy.model_prefetch |
| import comfy.ldm.minimax.model as _h3m |
| from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention |
|
|
| try: |
| from sageattention import quant as _sage_quant |
| from sageattention import sm89_compile as _sage_sm89 |
| _SAGE_OK = True |
| except Exception: |
| _sage_quant = _sage_sm89 = None |
| _SAGE_OK = False |
|
|
| log = logging.getLogger("MAINodes.vram_lab") |
|
|
| |
| |
| |
| |
| _SAGE3_DIR = "/mnt/work/ai/venvs/sage3-test/SageAttention/sageattention3_blackwell" |
| _SAGE3 = None |
|
|
|
|
| def _sage3_api(): |
| """The sageattn3.api module, or None. Import is attempted once per process.""" |
| global _SAGE3 |
| if _SAGE3 is not None: |
| return _SAGE3 or None |
| try: |
| import fp4attn_cuda |
| except ImportError: |
| if not os.path.isdir(_SAGE3_DIR): |
| _SAGE3 = False |
| return None |
| if _SAGE3_DIR not in sys.path: |
| sys.path.insert(0, _SAGE3_DIR) |
| try: |
| import fp4attn_cuda |
| except Exception: |
| _SAGE3 = False |
| return None |
| try: |
| from sageattn3 import api as _sage3_api_mod |
| except Exception: |
| _SAGE3 = False |
| return None |
| _SAGE3 = _sage3_api_mod |
| return _sage3_api_mod |
|
|
|
|
| |
|
|
| def _ranges(n, size, min_tail=512): |
| """Row chunks of `size`; a tail shorter than `min_tail` is folded into the |
| previous chunk. Row-wise ops are exact at any chunk size, but comfy_kitchen |
| picks a different int8 GEMM kernel for m <= 128 rows (`_prefer_turing_fused_int8`) |
| and its fp32 epilogue is not proven identical across kernels, so never emit |
| a tiny tail (2026-08-18).""" |
| if size <= 0 or size >= n: |
| return [(0, n)] |
| out = [(a, min(a + size, n)) for a in range(0, n, size)] |
| if len(out) > 1 and out[-1][1] - out[-1][0] < min_tail: |
| a, _ = out.pop() |
| out[-1] = (out[-1][0], n) |
| return out |
|
|
|
|
| _SM_CACHE = {} |
|
|
|
|
| def _sm_count(device): |
| key = str(device) |
| if key not in _SM_CACHE: |
| try: |
| _SM_CACHE[key] = torch.cuda.get_device_properties(device).multi_processor_count |
| except Exception: |
| _SM_CACHE[key] = 128 |
| return _SM_CACHE[key] |
|
|
|
|
| def _min_query_chunk(device, heads_per_call, q_block=128, margin=2): |
| """Smallest query chunk that keeps the flash kernel on its non-split path. |
| |
| Measured 2026-08-18 on an RTX PRO 6000 (188 SMs, 56 heads): a 259-query |
| tail (168 query blocks < 188 SMs) made SDPA take its split-KV path and the |
| result stopped being bit-equal to the full-length call (43% of elements |
| off by ~2e-4, which diffusion amplified into a visibly different clip); |
| 512 queries (224 blocks) was bit-equal. Query chunking is exact only while |
| every chunk has at least ~SMs query blocks, so we require margin x SMs. |
| |
| Measured boundary (same card, 215k K/V, 2026-08-18): flash leaves its split-KV path exactly when |
| heads_per_call x ceil(L/64) >= 0.8 x 2 x SMs (PyTorch flash_api.cpp |
| set_params_splitkv / num_splits_heuristic): L >= 321 / 641 / 1345 for |
| 56 / 28 / 14 heads. This function returns 896 / 1792 / 3456 there, a |
| 2.6-2.8x margin over the measured line, so it is not the constraint on |
| chunk size in practice. The heuristic constants are PyTorch's and can |
| move; self_check stays the authority. cuDNN and mem-efficient SDPA are |
| chunk-invariant at every length (and 1.1-2x slower); no two backends are |
| bit-equal to each other, and stock renders run flash, so we stay on flash. |
| """ |
| return q_block * math.ceil(margin * _sm_count(device) / max(1, heads_per_call)) |
|
|
|
|
| def _balanced_ranges(n, size, min_size): |
| """Split [0, n) into chunks of at most `size` tokens, all of length >= min_size |
| (the tail is folded into its neighbour rather than left short). One chunk if |
| n < 2 * min_size.""" |
| if size <= 0 or size >= n or n < 2 * min_size: |
| return [(0, n)] |
| parts = max(1, math.ceil(n / size)) |
| while parts > 1 and n / parts < min_size: |
| parts -= 1 |
| base, extra = divmod(n, parts) |
| out, a = [], 0 |
| for i in range(parts): |
| b = a + base + (1 if i < extra else 0) |
| out.append((a, b)) |
| a = b |
| return out |
|
|
|
|
| def _mod_scale_shift_range(h, shift, scale, segments, c0, c1): |
| """h is norm(x[c0:c1]); apply the per-segment affine restricted to [c0, c1).""" |
| for a, b, row in segments: |
| lo, hi = max(a, c0), min(b, c1) |
| if lo < hi: |
| h[lo - c0:hi - c0].mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype)) |
| return h |
|
|
|
|
| def _mod_gate_range(x, gate, other, segments, c0, c1): |
| """x[c0:c1] += gate[row] * other, per segment, in place on the full residual.""" |
| for a, b, row in segments: |
| lo, hi = max(a, c0), min(b, c1) |
| if lo < hi: |
| x[lo:hi].addcmul_(other[lo - c0:hi - c0], gate[row].to(x.dtype)) |
| return x |
|
|
|
|
| def _norm_rope(attn, q, k, rope_c, s): |
| """Mirror of Attention.forward's fused RMSNorm + split-half RoPE on a chunk. |
| |
| q, k: [s, heads*hd] views into the chunk's qkv buffer (split on last dim, |
| viewable as [1, s, heads, hd] exactly as core does). Returns [s, heads, hd]. |
| """ |
| heads, hd = attn.heads, attn.head_dim |
| if rope_c is not None: |
| q = q.view(1, s, heads, hd) |
| k = k.view(1, s, heads, hd) |
| qw = comfy.model_management.cast_to(attn.q_norm.weight, device=q.device) |
| kw = comfy.model_management.cast_to(attn.k_norm.weight, device=q.device) |
| rot = rope_c.shape[-3] * 2 |
| comfy.quant_ops.ck.rms_rope_split_half_(q, k, rope_c, qw, kw, epsilon=attn.q_norm.eps, rot_dim=rot) |
| return q[0], k[0] |
| return attn.q_norm(q.view(s, heads, hd)), attn.k_norm(k.view(s, heads, hd)) |
|
|
|
|
| def _blockwise_attention(qc, K, V, block, out_dtype): |
| """Exact attention of qc [1,H,c,hd] against K,V [1,H,S,hd] in K/V blocks. |
| |
| Uses the flash kernel's per-row logsumexp for the online-softmax combine. |
| fp32 accumulator; result cast to out_dtype. Slow on purpose. |
| """ |
| op = torch.ops.aten._scaled_dot_product_flash_attention.default |
| S = K.shape[2] |
| acc = None |
| lse_acc = None |
| for a, b in _ranges(S, block): |
| r = op(qc, K[:, :, a:b].contiguous(), V[:, :, a:b].contiguous(), 0.0, False, False) |
| acc, lse_acc = _online_combine(acc, lse_acc, r[0], r[1]) |
| del r |
| return acc.to(out_dtype) |
|
|
|
|
| def _attend(qc, K, V, heads, head_chunks, transformer_options): |
| """qc [1,H,c,hd]; K,V [1,H,S,hd] -> [1, c, H*hd]. Honors KJ-style head groups.""" |
| hd = qc.shape[-1] |
| n = max(1, min(int(head_chunks or 1), heads)) |
| if n <= 1: |
| return optimized_attention(AttentionTensorContainer(qc), AttentionTensorContainer(K), |
| AttentionTensorContainer(V), heads, mask=None, skip_reshape=True, |
| transformer_options=transformer_options) |
| c = qc.shape[2] |
| out = torch.empty((1, c, heads * hd), dtype=qc.dtype, device=qc.device) |
| hs = 0 |
| sizes = [heads // n + (1 if i < heads % n else 0) for i in range(n)] |
| for size in sizes: |
| he = hs + size |
| o = optimized_attention(AttentionTensorContainer(qc[:, hs:he]), AttentionTensorContainer(K[:, hs:he]), |
| AttentionTensorContainer(V[:, hs:he]), size, mask=None, skip_reshape=True, |
| transformer_options=transformer_options) |
| out[:, :, hs * hd:he * hd] = o |
| hs = he |
| return out |
|
|
|
|
|
|
| |
|
|
| _HAD = {} |
|
|
|
|
| def _hadamard(n, device, dtype): |
| """Orthonormal Hadamard n x n (n power of two), cached per device/dtype.""" |
| key = (n, str(device), dtype) |
| if key not in _HAD: |
| h = torch.ones((1, 1), dtype=torch.float32) |
| while h.shape[0] < n: |
| h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) |
| _HAD[key] = (h / (n ** 0.5)).to(device=device, dtype=dtype) |
| return _HAD[key] |
|
|
|
|
| def _online_combine(acc, lse_acc, o, lse): |
| """Online-softmax merge of a new flash partial (o bf16, lse fp32) into the |
| fp32 accumulator, in place; returns (acc, lse_acc).""" |
| if acc is None: |
| return o.float(), lse |
| new = torch.logaddexp(lse_acc, lse) |
| |
| acc.mul_(torch.exp(lse_acc - new).unsqueeze(-1)).addcmul_(o, torch.exp(lse - new).unsqueeze(-1)) |
| return acc, new |
|
|
|
|
| class _Int8KV: |
| """kvi8r: K/V buffers stored as int8 rows with one bf16 scale per row (per |
| token, per head) in a fixed orthonormal Hadamard-rotated basis of the head |
| dim. Half the bytes of bf16 (S x heads x 128 + scales). |
| |
| Second cut (2026-08-18 evening): the store keeps K and V in the ROTATED basis and |
| never un-rotates them. Because the rotation is orthogonal, q.k = (qH).(kH) |
| and P.V = (P.(VH)).H^T, so the query chunk is rotated once (small) and the |
| attention output is un-rotated once (small); dequant of a K/V block is a |
| single fused int8->fp16 cast * row scale, no GEMM, no fp32 copy of the |
| block. Attention runs in fp16 in the rotated basis: bf16's 8-bit mantissa |
| costs a 4.8e-3 rel-rms floor there (measured, random data) because the |
| rotation spreads every value to the same magnitude, fp16's 11 bits cost |
| 1.7e-3; the int8 rounding (absmax/127 per 128-wide row) is ~1e-2 either |
| way. NOT bit-equal to stock: approximation tier, judged by the operator's |
| eyes and the sensor bank.""" |
|
|
| def __init__(self, heads, S, hd, device, dtype=torch.float16): |
| self.heads, self.S, self.hd, self.dtype = heads, S, hd, dtype |
| self.q = torch.empty((1, heads, S, hd), dtype=torch.int8, device=device) |
| self.s = torch.empty((1, heads, S, 1), dtype=dtype, device=device) |
| self.H = _hadamard(hd, device, dtype) |
|
|
| def store(self, a, b, x): |
| xr = x.to(self.dtype) @ self.H |
| |
| sc = (xr.abs().amax(dim=-1, keepdim=True).float().clamp_min_(1e-8) * (1.0 / 127.0 * (1.0 + 2.0 ** -7))).to(self.dtype) |
| self.q[:, :, a:b] = torch.round(xr.float() / sc.float()).clamp_(-127, 127).to(torch.int8) |
| self.s[:, :, a:b] = sc |
| del xr, sc |
|
|
| def load_rot(self, a, b): |
| return self.q[:, :, a:b].to(self.dtype).mul_(self.s[:, :, a:b]) |
|
|
| def load(self, a, b, dtype): |
| return (self.load_rot(a, b).float() @ self.H.float().T).to(dtype).contiguous() |
|
|
| def bytes(self): |
| return self.q.numel() + self.s.numel() * self.s.element_size() |
|
|
|
|
| def _blockwise_attention_q(qc, Kq, Vq, block, out_dtype): |
| """Attention of qc [1,H,c,hd] against kvi8r stores: rotate the query chunk |
| once, attend blockwise against rotated-dequantised K/V (transient = one |
| bf16 K block + one V block), online-softmax combine in fp32, un-rotate the |
| output once.""" |
| op = torch.ops.aten._scaled_dot_product_flash_attention.default |
| S = Kq.S |
| qr = (qc.to(Kq.dtype) @ Kq.H).contiguous() |
| acc = None |
| lse_acc = None |
| for a, b in _ranges(S, block): |
| kb = Kq.load_rot(a, b) |
| vb = Vq.load_rot(a, b) |
| r = op(qr, kb, vb, 0.0, False, False) |
| del kb, vb |
| acc, lse_acc = _online_combine(acc, lse_acc, r[0], r[1]) |
| del r |
| del qr |
| return (acc.to(Kq.dtype) @ Kq.H.T).to(out_dtype) |
|
|
|
|
| |
|
|
| class _SageKV: |
| """kvi8s: K/V kept in SageAttention's kernel layout, block by block: K int8 |
| with one fp32 scale per (head, 64-token block) after subtracting the block's |
| per-head mean (softmax-invariant, restored in the LSE), V fp8 e4m3 in Sage's |
| transposed/permuted layout with one fp32 scale per (head, channel, block). |
| Same bytes as kvi8r; attention runs on int8 QK^T + fp8 PV tensor cores |
| straight from the store, no dequant. Q and K are Hadamard-rotated first |
| (scores invariant; spreads outliers before Sage's per-block int8 rounding). |
| NOT bit-equal to stock: approximation tier, one rung below kvi8r on the |
| synthetic proxy, ~1.6x faster attention than the exact path.""" |
|
|
| def __init__(self, heads, S, hd, device): |
| if not _SAGE_OK: |
| raise RuntimeError("kvi8s needs the sageattention package (2.x) importable in this venv") |
| self.heads, self.S, self.hd, self.device = heads, S, hd, device |
| self.H = _hadamard(hd, device, torch.bfloat16) |
| self.blocks = [] |
|
|
| def store(self, a, b, k, v): |
| k = (k.to(torch.bfloat16) @ self.H).contiguous() |
| km = k.mean(dim=2, keepdim=True) |
| k_int8 = torch.empty(k.shape, dtype=torch.int8, device=k.device) |
| k_scale = torch.empty((1, self.heads, (b - a + 63) // 64), dtype=torch.float32, device=k.device) |
| _sage_quant._fused.quant_per_block_int8_fuse_sub_mean_cuda(k, km.squeeze(2), k_int8, k_scale, 64, 1) |
| v_fp8, v_scale, _ = _sage_quant.per_channel_fp8(v.to(torch.bfloat16).contiguous(), tensor_layout="HND", scale_max=2.25, smooth_v=False) |
| self.blocks.append((a, b, k_int8, k_scale, km, v_fp8, v_scale)) |
| del k |
|
|
| def bytes(self): |
| return sum(k.numel() + ks.numel() * 4 + v.numel() + vs.numel() * 4 for _, _, k, ks, _, v, vs in self.blocks) |
|
|
|
|
| def _sage_attention_q(qc, st, out_dtype): |
| """qc [1,H,c,hd] bf16 against a _SageKV store -> [1,H,c,hd] out_dtype.""" |
| sm_scale = qc.shape[-1] ** -0.5 |
| qr = (qc.to(torch.bfloat16) @ st.H).contiguous() |
| q_int8 = torch.empty(qr.shape, dtype=torch.int8, device=qr.device) |
| c = qr.shape[2] |
| q_scale = torch.empty((1, st.heads, ((c + 127) // 128) * 4), dtype=torch.float32, device=qr.device) |
| _sage_quant._fused.quant_per_warp_int8_cuda(qr, q_int8, q_scale, 128, 32, 1) |
| acc = None |
| lse_acc = None |
| for a, b, k_int8, k_scale, km, v_fp8, v_scale in st.blocks: |
| o = torch.empty(qr.shape, dtype=torch.bfloat16, device=qr.device) |
| lse = _sage_sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( |
| q_int8, k_int8, v_fp8, o, q_scale, k_scale, v_scale, 1, 0, 2, sm_scale, 1) |
| |
| lse = lse / 1.44269504 + torch.matmul(qr, km.transpose(2, 3)).squeeze(-1).float() * sm_scale |
| acc, lse_acc = _online_combine(acc, lse_acc, o, lse) |
| del o, lse |
| del qr, q_int8 |
| return acc.to(out_dtype) |
|
|
|
|
| |
|
|
| _KVFP4_LOGGED = False |
|
|
|
|
| class _Sage3FP4KV: |
| """kvfp4s: K/V kept as NVFP4 in SageAttention3's kernel layouts -- E2M1 packed |
| two per byte plus one e4m3 scale per 16 elements = 0.5625 byte/elem, i.e. |
| 0.28x bf16 (1.63 GiB vs 5.79 at 216k tokens, measured) -- and attended on |
| Blackwell fp4 tensor cores straight from the store. Q and K are Hadamard- |
| rotated first (scores invariant; measured 2026-08-19 to cut rel-rms from |
| 0.394 to 0.293 on the outlier proxy, no change on plain gaussian). V is not |
| rotated, so the output needs no un-rotation. |
| |
| Design notes (reports/overnight_2026-08-20/sage3.md s.4/s.6): |
| - ``per_block_mean`` MUST be False: at H3 lengths the True form's delta_s is |
| [B,H,ceil(L/128),KL] fp32 ~= 82 TB. With per_block_mean False, dropping |
| the Q-centering and delta_s entirely is not worse than computing them |
| (0.27104 vs 0.27167 rel-rms), so the store keeps ONLY the fp4 tensors: |
| no bf16 K, no delta_s, un-centred Q, and a zeros delta_s allocated once. |
| - ``preprocess_qkv`` is never used: it mutates the caller's k in place and |
| recomputes the K mean every call. The per-head K mean is computed once |
| over the full K and frozen (subtracting a constant from every key is |
| softmax-invariant per query row, so it needs no correction term). |
| - The kernel's softmax_scale is derived from the PACKED last dim |
| ((D/2)*2)**-0.5; correct at D=128, a trap if anything reshapes. |
| - The LSE the kernel returns is uninitialised memory (the epilogue store is |
| commented out upstream), so a K-blocked online-softmax combine is |
| impossible -- the store is one contiguous span and Q chunking alone |
| carries the memory win. That is why phase 1 buffers K/V in bf16 and |
| quantises once at finalize(): during that finalize the block holds both |
| (bf16 + fp4) before the bf16 is dropped. |
| |
| NOT bit-equal to stock: approximation tier, one rung below kvi8s (measured |
| at 217k tokens: rel-rms 0.293 rotated vs kvi8s-family Sage 2.2's 0.055 on the |
| outlier proxy), ~2.4x faster attention than the exact path. Judged by eyes.""" |
|
|
| def __init__(self, heads, S, hd, device, dtype=torch.bfloat16, rotate=True, |
| keep_exact=False): |
| api = _sage3_api() |
| if api is None: |
| raise RuntimeError("kvfp4s needs SageAttention3's fp4 extension (fp4attn_cuda) importable") |
| self.api = api |
| self.heads, self.S, self.hd, self.device = heads, S, hd, device |
| self.R = _hadamard(hd, device, torch.bfloat16) if rotate else None |
| pad = (128 - S % 128) % 128 |
| self.SP = S + pad |
| self.kbuf = torch.empty((1, heads, self.SP, hd), dtype=torch.bfloat16, device=device) |
| self.vbuf = torch.empty((1, heads, self.SP, hd), dtype=torch.bfloat16, device=device) |
| if pad: |
| self.kbuf[:, :, S:].zero_() |
| self.vbuf[:, :, S:].zero_() |
| self.k_fp4 = self.k_sf = self.v_fp4 = self.v_sf = self.ds = None |
| |
| |
| |
| self.keep_exact = bool(keep_exact) |
| self.k_bf = self.v_bf = None |
| self.t0 = time.time() |
|
|
| def store(self, a, b, k, v): |
| k = k.to(torch.bfloat16) |
| self.kbuf[:, :, a:b] = (k @ self.R) if self.R is not None else k |
| self.vbuf[:, :, a:b] = v.to(torch.bfloat16) |
|
|
| def finalize(self): |
| """Quantise the whole span once: K centred on its frozen per-head mean and |
| permuted, V transposed. The bf16 buffers are dropped here.""" |
| global _KVFP4_LOGGED |
| kmean = self.kbuf[:, :, :self.S].mean(dim=-2, keepdim=True) |
| self.kbuf[:, :, :self.S] -= kmean |
| if self.keep_exact: |
| |
| |
| self.k_bf = self.kbuf[:, :, :self.S].clone() |
| self.v_bf = self.vbuf[:, :, :self.S].clone() |
| self.k_fp4, self.k_sf = self.api.scale_and_quant_fp4_permute(self.kbuf) |
| self.v_fp4, self.v_sf = self.api.scale_and_quant_fp4_transpose(self.vbuf) |
| self.kbuf = self.vbuf = None |
| del kmean |
| self.ds = torch.zeros((1, self.heads, 1, self.v_fp4.size(-1) * 2), |
| dtype=torch.float32, device=self.device) |
| if not _KVFP4_LOGGED: |
| _KVFP4_LOGGED = True |
| torch.cuda.synchronize(self.device) |
| log.info("H3StreamedBlocks kvfp4s: store built, %d tokens, %.3f GiB, %.0f ms, rotated %s", |
| self.S, self.bytes() / 2 ** 30, (time.time() - self.t0) * 1e3, |
| "yes" if self.R is not None else "no") |
|
|
| def bytes(self): |
| if self.k_fp4 is None: |
| return 0 |
| n = (self.k_fp4.numel() + self.k_sf.numel() + self.v_fp4.numel() + self.v_sf.numel() |
| + self.ds.numel() * 4) |
| if self.k_bf is not None: |
| n += (self.k_bf.numel() + self.v_bf.numel()) * 2 |
| return n |
|
|
|
|
| def _sage3_attention_q(qc, st, out_dtype): |
| """qc [1,H,c,hd] bf16 against a _Sage3FP4KV store -> [1,H,c,hd] out_dtype.""" |
| QL = qc.shape[2] |
| q = qc.to(torch.bfloat16) |
| if st.R is not None: |
| q = q @ st.R |
| pad = (128 - QL % 128) % 128 |
| if pad: |
| q = torch.nn.functional.pad(q, (0, 0, 0, pad)) |
| ql = st.api.scale_and_quant_fp4(q.contiguous()) |
| del q |
| o = st.api.blockscaled_fp4_attn(ql, (st.k_fp4, st.k_sf), (st.v_fp4, st.v_sf), |
| st.ds, st.S, False, False, True) |
| del ql |
| return o[0][:, :, :QL, :].to(out_dtype) |
|
|
|
|
| _EXACT_AV_LOGGED = False |
|
|
|
|
| def _exact_av_rows(o, qc, st, segments, c0, c1, out_dtype): |
| """kvfp4s audio fix: re-run the TEXT and AUDIO query rows on the exact bf16 |
| K/V, leaving video rows on the fp4 store. |
| |
| Why: audio is 300 of 28,931 packed tokens, so a ~4% attention error that |
| video's redundancy hides as slight ghosting lands audibly on a soundtrack |
| carried by 1% of the sequence. Measured in |
| reports/sol_sa3/18_fp4_audio_diagnosis.md: audio is not quantised WORSE |
| than video (its per-token error is the lowest of the three segments), it is |
| merely far more sensitive, so the fix is routing, not a better quantiser. |
| mod_segments rows carry the modality in row % 3 (0 video, 1 text, 2 audio). |
| """ |
| global _EXACT_AV_LOGGED |
| n = 0 |
| for sa, sb, row in segments: |
| if row % 3 == 0: |
| continue |
| lo, hi = max(sa, c0), min(sb, c1) |
| if lo >= hi: |
| continue |
| qe = qc[:, :, lo - c0:hi - c0].to(torch.bfloat16) |
| if st.R is not None: |
| qe = qe @ st.R |
| o[:, :, lo - c0:hi - c0] = torch.nn.functional.scaled_dot_product_attention( |
| qe, st.k_bf, st.v_bf).to(out_dtype) |
| n += hi - lo |
| if n and not _EXACT_AV_LOGGED: |
| _EXACT_AV_LOGGED = True |
| log.info("H3StreamedBlocks kvfp4s: exact A/V rows ON (%d text+audio rows " |
| "in the first chunk take the bf16 path)", n) |
| return o |
|
|
|
|
| |
|
|
| |
| |
| |
| |
| _SOL_DIR = "/mnt/work/ai/venvs/sol-lab/comfy-kitchen-sol-lab/comfy_kitchen/backends/cuda" |
| _SOL = None |
|
|
|
|
| def _sol_ext(): |
| """The PR #117 CUDA extension, or None. Loaded once per process.""" |
| global _SOL |
| if _SOL is not None: |
| return _SOL or None |
| so = os.path.join(_SOL_DIR, "_C.abi3.so") |
| if not os.path.exists(so): |
| _SOL = False |
| return None |
| try: |
| import importlib.util |
| spec = importlib.util.spec_from_file_location("_C", so) |
| m = importlib.util.module_from_spec(spec) |
| sys.modules.setdefault("_C", m) |
| spec.loader.exec_module(m) |
| if not hasattr(m, "sol_attn"): |
| _SOL = False |
| return None |
| _SOL = m |
| except Exception as e: |
| log.warning("H3SolAttention: could not load the PR #117 extension (%s)", e) |
| _SOL = False |
| return None |
| return _SOL |
|
|
|
|
| _SOL_WS = {} |
| _SOL_LOGGED = False |
|
|
|
|
| def _sol_attention(q, k, v, heads, mask=None, skip_reshape=False, |
| transformer_options=None, **kw): |
| """optimized_attention-compatible shim backed by Sol's CUDA kernel. |
| |
| H3 hands us [1, heads, S, D] (skip_reshape). Sol wants [B, T, H, D] bf16 and |
| is self-attention only (q_len == kv_len), which is why this replaces the whole |
| op rather than riding inside the streamed executor's Q chunking. |
| """ |
| global _SOL_LOGGED, _SOL_AV_LOGGED |
| ext = _sol_ext() |
| qq = q.tensor if hasattr(q, "tensor") else q |
| kk = k.tensor if hasattr(k, "tensor") else k |
| vv = v.tensor if hasattr(v, "tensor") else v |
| B, H, S, D = qq.shape |
| if ext is None or D != 128 or qq.shape != kk.shape: |
| return _SOL_FALLBACK(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, |
| transformer_options=transformer_options, **kw) |
| qb = qq.transpose(1, 2).contiguous().to(torch.bfloat16) |
| kb = kk.transpose(1, 2).contiguous().to(torch.bfloat16) |
| vb = vv.transpose(1, 2).contiguous().to(torch.bfloat16) |
| key = (B, S, H, qq.device.index) |
| need = ext.sol_attn_workspace(B, S, H, 0) |
| ws = _SOL_WS.get(key) |
| if ws is None or ws.numel() < need: |
| ws = torch.empty(need, dtype=torch.uint8, device=qq.device) |
| _SOL_WS[key] = ws |
| |
| |
| out = torch.empty_like(qb) |
| dl = lambda t: t.detach().__dlpack__(stream=-1) |
| ext.sol_attn(dl(qb), dl(kb), dl(vb), dl(out), dl(ws), |
| B, S, H, D, 0, float(_SOL_TAU), float(D ** -0.5), |
| 0, 0, 0, 0, |
| list(qb.stride()[:3]), list(kb.stride()[:3]), list(vb.stride()[:3]), |
| torch.cuda.current_stream(qq.device).cuda_stream, |
| centroid_tail=True, key_bias=None) |
| o = out |
| if _SOL_EXACT_AV and (_SOL_AV_SLICES is None or _SOL_AV_SLICES[0] != S) and S > 4096 and not _SOL_AV_LOGGED: |
| _SOL_AV_LOGGED = True |
| log.warning("H3SolAttention: exact_av_rows requested but INACTIVE on the main " |
| "forward (S=%d, stashed=%s) - layout capture missed; AV rows are " |
| "riding Sol", S, _SOL_AV_SLICES[0] if _SOL_AV_SLICES else None) |
| if _SOL_EXACT_AV and _SOL_AV_SLICES is not None and _SOL_AV_SLICES[0] == S: |
| |
| |
| |
| |
| kbh = kb.transpose(1, 2) |
| vbh = vb.transpose(1, 2) |
| n_av = 0 |
| for a, b in _SOL_AV_SLICES[1]: |
| qe = qb[:, a:b].transpose(1, 2) |
| o[:, a:b] = torch.nn.functional.scaled_dot_product_attention( |
| qe, kbh, vbh).transpose(1, 2) |
| n_av += b - a |
| if n_av and not _SOL_AV_LOGGED: |
| _SOL_AV_LOGGED = True |
| log.info("H3SolAttention: exact A/V rows ON (%d of %d query rows dense " |
| "bf16; video rows ride Sol)", n_av, S) |
| if not _SOL_LOGGED: |
| _SOL_LOGGED = True |
| log.info("H3SolAttention: PR #117 CUDA kernel live (S=%d, heads=%d, tau=%.2f, " |
| "workspace %.2f GiB)", S, H, _SOL_TAU, need / 2 ** 30) |
| out = o.transpose(1, 2) |
| ret = out.transpose(1, 2).reshape(B, S, H * D) |
| return ret.to(qq.dtype) |
|
|
|
|
| _SOL_FALLBACK = None |
| _SOL_TAU = 1.0 |
| _SOL_EXACT_AV = False |
| _SOL_AV_SLICES = None |
| _SOL_AV_LOGGED = False |
| _SOL_FWD_ORIG = None |
|
|
|
|
| def _sol_capture_forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs): |
| """Stash the packed layout's text+audio row spans for the Sol exact-AV path. |
| |
| Mirrors the signature derivation at the top of MiniMaxH3Model._forward; the |
| layout is the payload's prebuilt one where present, else rebuilt (cheap, |
| cached on signature). cond/ref_img rows carry the video modality and stay on |
| Sol, matching the streamed executor's exact_av_rows split.""" |
| global _SOL_AV_SLICES |
| try: |
| video_x, audio_x = x[0], x[1] |
| vx = comfy.ldm.common_dit.pad_to_patch_size(video_x, self.patch_size) |
| payload = minimax_payload or {} |
| sig = (context.shape[1], vx.shape[2], vx.shape[3], vx.shape[4], audio_x.shape[-1]) |
| if _SOL_AV_SLICES is None or _SOL_AV_SLICES[2] != sig: |
| layout = payload.get("layout") |
| if layout is None or layout.signature != sig: |
| layout = _h3m.PackedLayout(*sig, keyframes=payload.get("keyframes"), |
| refs=payload.get("refs")) |
| av = [(a, b) for a, b, kind in layout.segments |
| if kind in ("text", "audio", "cond_audio", "ref_audio")] |
| total = max(b for _, b, _ in layout.segments) |
| _SOL_AV_SLICES = (total, av, sig) |
| log.info("H3SolAttention: exact-AV layout stashed (%d spans, %d rows, " |
| "total %d, kinds %s)", len(av), sum(b - a for a, b in av), total, |
| [k for _, _, k in layout.segments]) |
| except Exception as e: |
| log.warning("H3SolAttention: exact-AV layout capture failed (%s); " |
| "AV rows ride Sol this run", e) |
| return _SOL_FWD_ORIG(self, x, timestep, context, |
| transformer_options=transformer_options, |
| minimax_payload=minimax_payload, **kwargs) |
|
|
|
|
| class H3SolAttention: |
| """Replace H3's attention with Sol-Attn (comfy-kitchen PR #117, CUDA). |
| |
| NOT a quantiser: Sol is SPARSE -- it routes each query block to a subset of |
| key blocks. Measured on real captured H3 tensors it runs ~7x dense flash at |
| 20.8-22.2% routing density, with attention-output rel-rms 0.099-0.254, which |
| is 2-3x LARGER than the fp4 store at four of five depths |
| (reports/sol_sa3/23_sol_on_the_judged_ladder.md). Ships OFF; judged by eye.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": { |
| "model": ("MODEL",), |
| "tau": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0, "step": 0.05, |
| "tooltip": "Sol's routing threshold. Lower routes MORE blocks " |
| "(denser, slower, more accurate); 1.0 is the paper default."}), |
| }, "optional": { |
| "exact_av_rows": ("BOOLEAN", {"default": False, |
| "tooltip": "Route TEXT and AUDIO query rows through exact dense bf16 " |
| "attention; only video rows ride Sol's sparse routing. This is " |
| "the scoping MiniMax describe for native H3 sparsity (video-only " |
| "sparse, non-video dense). The identical split fixed kvfp4s " |
| "audio outright. Measured motivation: the tau sweep (0.5-1.4, " |
| "bakery 20260821) left speech pacing rough at every density, so " |
| "the loss is scoping, not density."}), |
| }} |
|
|
| RETURN_TYPES = ("MODEL",) |
| FUNCTION = "patch" |
| CATEGORY = "MAINodes/VRAM Lab" |
| DESCRIPTION = ("Sol-Attn sparse attention for MiniMax-H3, from the unmerged " |
| "comfy-kitchen PR #117 CUDA kernel built out-of-tree for sm120.") |
|
|
| def patch(self, model, tau, exact_av_rows=False): |
| try: |
| from . import h3_capabilities as _caps |
| for w in _caps.collision_warnings(_caps.block_patch_report(model)): |
| log.warning("H3SolAttention: %s", w) |
| except Exception as _e: |
| log.info("H3SolAttention: collision report skipped (%s)", _e) |
| global _SOL_FALLBACK, _SOL_TAU, _SOL_EXACT_AV, _SOL_FWD_ORIG, _SOL_AV_SLICES, _SOL_AV_LOGGED |
| if _sol_ext() is None: |
| log.warning("H3SolAttention: PR #117 extension not available; model unchanged") |
| return (model,) |
| import comfy.ldm.minimax.model as mm |
| if _SOL_FALLBACK is None: |
| _SOL_FALLBACK = mm.optimized_attention |
| _SOL_TAU = float(tau) |
| _SOL_EXACT_AV = bool(exact_av_rows) |
| _SOL_AV_SLICES = None |
| _SOL_AV_LOGGED = False |
| if _SOL_EXACT_AV and _SOL_FWD_ORIG is None: |
| _SOL_FWD_ORIG = _h3m.MiniMaxH3Model._forward |
| _h3m.MiniMaxH3Model._forward = _sol_capture_forward |
| mm.optimized_attention = _sol_attention |
| log.info("H3SolAttention: bound Sol to comfy.ldm.minimax.model.optimized_attention " |
| "(tau %.2f, exact_av_rows %s)", _SOL_TAU, _SOL_EXACT_AV) |
| return (model,) |
|
|
|
|
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| _KVMIX_DEFAULT_HI = (0, 2, 3, 4, 5, 11, 12, 17, 18, 19, 21, 23, 24, 26, 27, 30, |
| 31, 32, 35, 39, 41, 42, 43, 46, 48, 49, 51, 54) |
| _KVMIX_LOGGED = False |
|
|
|
|
| class _MixedKV: |
| """kvmix: `hi` heads on Sage 2.2 int8/fp8 (clean), the rest on SA3 fp4 (fast).""" |
|
|
| def __init__(self, heads, S, hd, device, hi_heads=None, keep_exact=False): |
| hi = sorted({h for h in (hi_heads if hi_heads is not None else _KVMIX_DEFAULT_HI) |
| if 0 <= h < heads}) |
| self.hi = torch.tensor(hi, dtype=torch.long, device=device) |
| lo = [h for h in range(heads) if h not in set(hi)] |
| self.lo = torch.tensor(lo, dtype=torch.long, device=device) |
| self.n_hi, self.n_lo = len(hi), len(lo) |
| self.fp4 = _Sage3FP4KV(self.n_lo, S, hd, device) if self.n_lo else None |
| self.i8 = _SageKV(self.n_hi, S, hd, device) if self.n_hi else None |
| self.S = S |
| |
| |
| |
| self.k_bf = self.v_bf = None |
| if keep_exact: |
| self.k_bf = torch.empty((1, heads, S, hd), dtype=torch.bfloat16, device=device) |
| self.v_bf = torch.empty((1, heads, S, hd), dtype=torch.bfloat16, device=device) |
|
|
| def store(self, a, b, k, v): |
| if self.k_bf is not None: |
| self.k_bf[:, :, a:b] = k.to(torch.bfloat16) |
| self.v_bf[:, :, a:b] = v.to(torch.bfloat16) |
| if self.fp4 is not None: |
| self.fp4.store(a, b, k[:, self.lo].contiguous(), v[:, self.lo].contiguous()) |
| if self.i8 is not None: |
| self.i8.store(a, b, k[:, self.hi].contiguous(), v[:, self.hi].contiguous()) |
|
|
| def finalize(self): |
| global _KVMIX_LOGGED |
| if self.fp4 is not None: |
| self.fp4.finalize() |
| if not _KVMIX_LOGGED: |
| _KVMIX_LOGGED = True |
| log.info("H3StreamedBlocks kvmix: %d heads int8/fp8 + %d heads fp4 " |
| "(avg %.1f bits)", self.n_hi, self.n_lo, |
| (8 * self.n_hi + 4 * self.n_lo) / max(1, self.n_hi + self.n_lo)) |
|
|
| def bytes(self): |
| return ((self.fp4.bytes() if self.fp4 is not None else 0) |
| + (self.i8.bytes() if self.i8 is not None and hasattr(self.i8, "bytes") else 0)) |
|
|
|
|
| def _exact_av_rows_mixed(o, qc, st, segments, c0, c1, out_dtype): |
| """Same audio fix as the fp4 path, over the mixed store's own bf16 K/V so it |
| covers all heads. Unrotated on both sides, so it is plain exact attention.""" |
| global _EXACT_AV_LOGGED |
| n = 0 |
| for sa, sb, row in segments: |
| if row % 3 == 0: |
| continue |
| lo, hi = max(sa, c0), min(sb, c1) |
| if lo >= hi: |
| continue |
| qe = qc[:, :, lo - c0:hi - c0].to(torch.bfloat16) |
| o[:, :, lo - c0:hi - c0] = torch.nn.functional.scaled_dot_product_attention( |
| qe, st.k_bf, st.v_bf).to(out_dtype) |
| n += hi - lo |
| if n and not _EXACT_AV_LOGGED: |
| _EXACT_AV_LOGGED = True |
| log.info("H3StreamedBlocks kvmix: exact A/V rows ON (%d text+audio rows " |
| "in the first chunk take the bf16 path, all heads)", n) |
| return o |
|
|
|
|
| def _mixed_attention_q(qc, st, out_dtype): |
| """qc [1,H,c,hd] -> [1,H,c,hd], each head through its own kernel.""" |
| o = torch.empty(qc.shape, dtype=out_dtype, device=qc.device) |
| if st.fp4 is not None: |
| o[:, st.lo] = _sage3_attention_q(qc[:, st.lo].contiguous(), st.fp4, out_dtype) |
| if st.i8 is not None: |
| o[:, st.hi] = _sage_attention_q(qc[:, st.hi].contiguous(), st.i8, out_dtype) |
| return o |
|
|
|
|
| |
|
|
| _STOCK_FORWARD_SHA = "f40e52b23fb2f9c7" |
|
|
|
|
| def _stock_forward_sha(): |
| import hashlib |
| import inspect |
| try: |
| return hashlib.sha256(inspect.getsource(_h3m.MiniMaxH3Model._forward).encode()).hexdigest()[:16] |
| except Exception: |
| return None |
|
|
|
|
| def _trimmed_forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs): |
| """Stock MiniMaxH3Model._forward (ComfyUI 0.33.0, source sha256 f40e52b2...) with the |
| embed/row buffers released after the sequence is assembled (F5). Same math.""" |
| video_x, audio_x = x[0], x[1] |
| orig_t, orig_h, orig_w = video_x.shape[2], video_x.shape[3], video_x.shape[4] |
| video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, self.patch_size) |
| if video_x.shape[0] != 1: |
| raise ValueError("MiniMax H3 supports batch size 1") |
| payload = minimax_payload or {} |
| device = video_x.device |
| dtype = context.dtype |
|
|
| latent_t, lat_h, lat_w = video_x.shape[2], video_x.shape[3], video_x.shape[4] |
| audio_t = audio_x.shape[-1] |
| text_len = context.shape[1] |
| |
| layout = payload.get("layout") |
| if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t): |
| layout = _h3m.PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t, |
| keyframes=payload.get("keyframes"), |
| refs=payload.get("refs")) |
|
|
| |
| shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video)) |
| shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio)) |
| sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6) |
| t_v = float(1.0 - sigma_v) |
| t_a = float(1.0 - _h3m.time_shift_sigma(sigma_v, shift_v, shift_a)) |
|
|
| |
| vis_aug = float(payload.get("visual_cond_noise_aug", _h3m.VISUAL_COND_TIMESTEP)) |
| aud_aug = float(payload.get("audio_cond_noise_aug", _h3m.AUDIO_COND_TIMESTEP)) |
| has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments) |
| has_aud_cond = any(k in ("cond_audio", "ref_audio") for _, _, k in layout.segments) |
| seg_t = {"text": t_v, "video": t_v, "audio": t_a, |
| "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug), |
| "cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)} |
| unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set()) |
| | ({seg_t["ref_audio"]} if has_aud_cond else set())) |
| t_row = {t: i for i, t in enumerate(unique_t)} |
| seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2} |
|
|
| text_tags = payload.get("text_token_tags") |
| mod_segments = [] |
| for a, b, kind in layout.segments: |
| row_base = t_row[seg_t[kind]] * 3 |
| if kind == "text" and text_tags is not None: |
| |
| tags = text_tags.view(-1).tolist() |
| run_start = 0 |
| for i in range(1, b - a + 1): |
| if i == b - a or tags[i] != tags[run_start]: |
| mod_segments.append((a + run_start, a + i, row_base + int(tags[run_start]))) |
| run_start = i |
| else: |
| mod_segments.append((a, b, row_base + seg_tag[kind])) |
|
|
| |
| img_update = layout.img_update.to(device) |
| audio_update = layout.audio_update.to(device) |
| video_rows = _h3m.patchify_video(video_x.to(torch.float32), self.patch_size) |
| audio_rows = _h3m.pack_audio(audio_x.to(torch.float32)) |
| cond_video_rows = self._cond_video_rows(payload, device) |
| cond_audio_rows = self._cond_audio_rows(payload, device) |
|
|
| all_video_rows = video_rows |
| if cond_video_rows is not None: |
| all_video_rows = torch.empty(img_update.shape[0], video_rows.shape[1], dtype=torch.float32, device=device) |
| all_video_rows[~img_update] = cond_video_rows |
| all_video_rows[img_update] = video_rows |
| all_audio_rows = audio_rows |
| if cond_audio_rows is not None: |
| all_audio_rows = torch.empty(audio_update.shape[0], audio_rows.shape[1], dtype=torch.float32, device=device) |
| all_audio_rows[~audio_update] = cond_audio_rows |
| all_audio_rows[audio_update] = audio_rows |
|
|
| video_embed = self.video_patch_proj(all_video_rows).to(dtype) |
| audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype) |
| text_states = context[0] |
| if text_states.shape[-1] != self.hidden_size: |
| text_states = self.token_refiner(self.condition_proj(text_states), |
| transformer_options=transformer_options) |
|
|
| |
| h = torch.empty(layout.seq_len, self.hidden_size, dtype=dtype, device=device) |
| voff = aoff = 0 |
| for a, b, kind in layout.segments: |
| n = b - a |
| if kind == "text": |
| h[a:b] = text_states |
| elif kind in ("cond", "ref_img", "video"): |
| h[a:b] = video_embed[voff:voff + n] |
| voff += n |
| else: |
| h[a:b] = audio_embed[aoff:aoff + n] |
| aoff += n |
|
|
| |
| |
| del video_embed, audio_embed, video_rows, audio_rows, all_video_rows, all_audio_rows, cond_video_rows, cond_audio_rows |
|
|
| t_vals = torch.tensor(unique_t, dtype=torch.float32, device=device) |
| if self.use_adaln_curves: |
| |
| table = comfy.model_management.cast_to(self.adaln_t_table, device=device) |
| pos = t_vals.clamp(0.0, 1.0) * (table.shape[0] - 1) |
| i0 = pos.floor().long().clamp(max=table.shape[0] - 2) |
| t_emb = torch.lerp(table[i0], table[i0 + 1], (pos - i0).unsqueeze(1)) |
| else: |
| t_emb = self.time_embedder(t_vals).to(dtype) |
|
|
| |
| rope_freqs = _h3m.rope_rotation_table(self.rope_freqs(layout.position_ids, device), dtype) |
|
|
| |
| patches_replace = transformer_options.get("patches_replace", {}) |
| blocks_replace = patches_replace.get("dit", {}) |
| prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options) |
| for i, block in enumerate(self.blocks): |
| comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block) |
| if ("double_block", i) in blocks_replace: |
| def block_wrap(args): |
| return {"img": block(args["img"], args["t_emb"], args["mod_segments"], args["rope_freqs"], |
| transformer_options=args["transformer_options"])} |
| h = blocks_replace[("double_block", i)]( |
| {"img": h, "t_emb": t_emb, "mod_segments": mod_segments, "rope_freqs": rope_freqs, |
| "transformer_options": transformer_options}, |
| {"original_block": block_wrap})["img"] |
| else: |
| h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options) |
| if prefetch_queue is not None: |
| comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None) |
|
|
| |
| video_seg = next((a, b, t_row[seg_t["video"]]) for a, b, k in layout.segments if k == "video") |
| audio_seg = next((a, b, t_row[seg_t["audio"]]) for a, b, k in layout.segments if k == "audio") |
| v, a = self.final_layer(h, t_emb, video_seg, audio_seg) |
|
|
| video_out = _h3m.unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, self.latents_dim, self.patch_size) |
| video_out = video_out[:, :, :orig_t, :orig_h, :orig_w] |
| audio_out = _h3m.unpack_audio(a) |
|
|
| return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)] |
|
|
|
|
| |
|
|
| class _PrecProbe: |
| """Per (forward, block, projection, segment) record of the real input activation of |
| each block linear and the extra output error three cheaper activation quantisers |
| would add on top of what the layer does today (the shipped layers quantise the |
| activation to int8 rowwise-convrot internally): NVFP4 with one activation scale |
| (kitchen quantize/dequantize), NVFP4 after a 128-wide Hadamard, FP8 e4m3 per tensor. |
| Plus per (forward, block) the block's residual change ||x_out - x_in|| / ||x_in|| on |
| a strided row sample (report 7's router signal). Sampled 1 in `every` chunks per |
| phase. Writes jsonl; the map is drawn offline (benchmarks/scripts/prec_map.py).""" |
|
|
| def __init__(self, path=None, every=8, row_stride=64): |
| self.path, self.every, self.row_stride = path, int(every), int(row_stride) |
| self.fwd = -1 |
| self.rows = [] |
| self.counter = {} |
| self.sigma = None |
| self.H = None |
|
|
| def begin_forward(self, sigma): |
| self.fwd += 1 |
| self.sigma = sigma |
| self.counter = {} |
|
|
| def _take(self, key): |
| c = self.counter.get(key, 0) |
| self.counter[key] = c + 1 |
| return c % self.every == 0 |
|
|
| @staticmethod |
| def _seg_kind(row): |
| return ("cond_" if row >= 3 else "") + ("video", "text", "audio")[row % 3] |
|
|
| def _fq_nvfp4(self, x2d): |
| from comfy_kitchen.tensor.nvfp4 import TensorCoreNVFP4Layout |
| q, prm = TensorCoreNVFP4Layout.quantize(x2d) |
| return TensorCoreNVFP4Layout.dequantize(q, prm)[: x2d.shape[0], : x2d.shape[1]] |
|
|
| def _fq_fp8(self, x2d): |
| amax = x2d.abs().amax().float().clamp_min(1e-12) |
| sc = amax / 448.0 |
| return (x2d.float() / sc).clamp_(-448, 448).to(torch.float8_e4m3fn).to(x2d.dtype) * sc.to(x2d.dtype) |
|
|
| def _fq_nvfp4_had(self, x2d): |
| n = x2d.shape[1] |
| if self.H is None or self.H.shape[0] != 128 or self.H.device != x2d.device: |
| self.H = _hadamard(128, x2d.device, torch.float32) |
| xr = (x2d.float().view(-1, n // 128, 128) @ self.H).view(-1, n).to(x2d.dtype) |
| yr = self._fq_nvfp4(xr) |
| return (yr.float().view(-1, n // 128, 128) @ self.H.T).view(-1, n).to(x2d.dtype) |
|
|
| def _fq_int4(self, x2d): |
| """A UNIFORM 4-bit lattice at NVFP4's block granularity: 16 evenly spaced codes with an |
| exact per-16-value scale. Not a deployable format on sm120 (the 4-bit tensor cores decode |
| E2M1 only) - this is the diagnostic arm that says whether E2M1's non-uniform spacing is |
| what the activation error is paying for, as the weight-side study found it is.""" |
| f = x2d.float() |
| g = f.reshape(-1, 16) |
| sc = g.abs().amax(1, keepdim=True).clamp_min(1e-12) / 7.0 |
| q = (g / sc).round_().clamp_(-8, 7) |
| return (q * sc).reshape(f.shape).to(x2d.dtype) |
|
|
| def _fq_int4_had(self, x2d): |
| n = x2d.shape[1] |
| if self.H is None or self.H.shape[0] != 128 or self.H.device != x2d.device: |
| self.H = _hadamard(128, x2d.device, torch.float32) |
| xr = (x2d.float().view(-1, n // 128, 128) @ self.H).view(-1, n).to(x2d.dtype) |
| yr = self._fq_int4(xr) |
| return (yr.float().view(-1, n // 128, 128) @ self.H.T).view(-1, n).to(x2d.dtype) |
|
|
| @staticmethod |
| def _group_crest(xs): |
| """Per-16-value crest factor (peak/rms), the quantity MixFP4's 2.224 threshold is defined |
| on. The recorded whole-tensor crest is a different and much coarser statistic.""" |
| n = xs.numel() - xs.numel() % 16 |
| if n < 16: |
| return None |
| g = xs.reshape(-1)[:n].reshape(-1, 16) |
| rms = g.pow(2).mean(1).sqrt().clamp_min(1e-12) |
| c = g.abs().amax(1) / rms |
| return {"g_crest_med": c.median().item(), "g_crest_p95": c.quantile(0.95).item(), |
| "g_crest_frac_lt_2224": (c < 2.224).float().mean().item()} |
|
|
| def amax(self, block, proj, x): |
| """Every chunk: running per-(block, proj) activation amax for this forward (calibrated |
| input_scale for a chunk-invariant NVFP4/FP8 activation path needs the whole-layer amax).""" |
| k = ("amax", block, proj) |
| v = x.abs().amax() |
| cur = self.counter.get(k) |
| self.counter[k] = v if cur is None else torch.maximum(cur, v) |
|
|
| def flush_amax(self): |
| for k, v in list(self.counter.items()): |
| if isinstance(k, tuple) and k and k[0] == "amax": |
| self.rows.append({"fwd": self.fwd, "sigma": self.sigma, "block": k[1], "proj": k[2], |
| "seg": "all", "amax_layer": float(v)}) |
| del self.counter[k] |
|
|
| @staticmethod |
| def _bf16_weight(layer, device): |
| """The layer's weight dequantised to bf16 ON the compute device (under dynamic VRAM the |
| stored weight may live on the CPU); None if that is not possible.""" |
| w = getattr(layer, "weight", None) |
| if w is None: |
| return None |
| try: |
| if getattr(w, "device", None) is not None and w.device != device: |
| w = w.to(device) |
| if hasattr(w, "dequantize"): |
| w = w.dequantize() |
| return w.to(device=device, dtype=torch.bfloat16) |
| except Exception: |
| return None |
|
|
| @torch.no_grad() |
| def projection(self, block, proj, layer, x, a, b, mod_segments): |
| """x: the layer's real input for chunk [a:b] (rows x K), layer: the module. Cheap sample.""" |
| self.amax(block, proj, x) |
| if not self._take((block, proj)): |
| return |
| if x.shape[-1] % 128: |
| return |
| y = layer(x) |
| fq = {"nvfp4": self._fq_nvfp4, "nvfp4_had": self._fq_nvfp4_had, "fp8": self._fq_fp8} |
| if os.environ.get("MAINODES_PREC_INT4") == "1": |
| |
| fq["int4"] = self._fq_int4 |
| fq["int4_had"] = self._fq_int4_had |
| outs = {k: layer(f(x)) for k, f in fq.items()} |
| wb = self._bf16_weight(layer, x.device) |
| y_ref = None |
| if wb is not None and wb.shape[-1] == x.shape[-1]: |
| bias = getattr(layer, "bias", None) |
| bias = bias.to(device=x.device, dtype=x.dtype) if bias is not None else None |
| y_ref = torch.nn.functional.linear(x, wb, bias) |
| del wb |
| for sa, sb, row in mod_segments: |
| lo, hi = max(sa, a), min(sb, b) |
| if hi - lo < 16: |
| continue |
| xs = x[lo - a:hi - a].float() |
| ys = y[lo - a:hi - a].float() |
| rms = xs.pow(2).mean().sqrt() |
| rec = {"fwd": self.fwd, "sigma": self.sigma, "block": block, "proj": proj, "seg": self._seg_kind(row), |
| "rows": int(hi - lo), "x_rms": rms.item(), "x_amax": xs.abs().amax().item(), |
| "x_outlier_frac": (xs.abs() > 6 * rms).float().mean().item(), |
| "x_kurt": ((xs / rms) ** 4).mean().item()} |
| if os.environ.get("MAINODES_PREC_INT4") == "1": |
| gc = self._group_crest(xs) |
| if gc: |
| rec.update(gc) |
| yn = ys.pow(2).sum().sqrt().clamp_min(1e-12) |
| for k, o in outs.items(): |
| rec["err_" + k] = ((o[lo - a:hi - a].float() - ys).pow(2).sum().sqrt() / yn).item() |
| if y_ref is not None: |
| yr = y_ref[lo - a:hi - a].float() |
| rn = yr.pow(2).sum().sqrt().clamp_min(1e-12) |
| rec["err_today_vs_bf16"] = ((ys - yr).pow(2).sum().sqrt() / rn).item() |
| for k, o in outs.items(): |
| rec["err_" + k + "_vs_bf16"] = ((o[lo - a:hi - a].float() - yr).pow(2).sum().sqrt() / rn).item() |
| self.rows.append(rec) |
| del y, outs, y_ref |
|
|
| @torch.no_grad() |
| def block_in(self, x): |
| return x[:: self.row_stride].detach().clone() |
|
|
| @torch.no_grad() |
| def block_out(self, block, x_in_sample, x): |
| s = x[:: self.row_stride].float() |
| d = (s - x_in_sample.float()).pow(2).sum().sqrt() / x_in_sample.float().pow(2).sum().sqrt().clamp_min(1e-12) |
| self.rows.append({"fwd": self.fwd, "sigma": self.sigma, "block": block, "proj": "block", "seg": "all", |
| "rel_change": d.item()}) |
|
|
| def flush(self): |
| if not self.rows or self.path is None: |
| return |
| import json |
| with open(self.path, "a") as f: |
| for r in self.rows: |
| f.write(json.dumps(r) + "\n") |
| self.rows = [] |
|
|
|
|
| class _FakeQuant: |
| """A6 phase 0d: simulate a lower activation precision on chosen projections / blocks / |
| segments inside the streamed block (fake-quant: quantise -> dequantise the layer input, |
| the layer then runs as shipped). For the end-to-end sensitivity sweep: same seed, one |
| region at a time, latent divergence + eyes. Not a speed path.""" |
|
|
| def __init__(self, fmt, projs, blocks, segs): |
| self.fmt, self.projs, self.blocks, self.segs = fmt, set(projs), blocks, (set(segs) if segs else None) |
| self.p = _PrecProbe(path=None) |
|
|
| def apply(self, block, proj, x, a, b, mod_segments): |
| if proj not in self.projs or not (self.blocks[0] <= block <= self.blocks[1]) or x.shape[-1] % 128: |
| return x |
| f = {"nvfp4": self.p._fq_nvfp4, "nvfp4_had": self.p._fq_nvfp4_had, "fp8": self.p._fq_fp8}[self.fmt] |
| if self.segs is None: |
| return f(x) |
| out = x.clone() |
| for sa, sb, row in mod_segments: |
| lo, hi = max(sa, a), min(sb, b) |
| if hi > lo and _PrecProbe._seg_kind(row) in self.segs: |
| out[lo - a:hi - a] = f(x[lo - a:hi - a]) |
| return out |
|
|
|
|
| |
|
|
| def _phase1_kv(block, x, shift_msa, scale_msa, mod_segments, rope_freqs, kv_chunk, heads, hd, inner, kv_int8=False, kv_sage=False, kv_fp4=False, kv_fp4_exact_av=False, kv_mix=False, prec=None, block_index=0, fq=None): |
| """K, V for the whole sequence, chunk by chunk (Q computed and dropped). |
| kv_int8: store them as _Int8KV (kvi8r, half the bytes; approximation tier). |
| kv_sage: store them as one _SageKV (kvi8s, half the bytes, Sage kernels; approximation tier). |
| kv_fp4: store them as one _Sage3FP4KV (kvfp4s, 0.28x the bytes, sage3 fp4 kernels; |
| approximation tier) -- buffered bf16 while streaming, quantised once at finalize().""" |
| S = x.shape[0] |
| attn = block.attn |
| if kv_mix: |
| K = _MixedKV(heads, S, hd, x.device, keep_exact=kv_fp4_exact_av) |
| V = K |
| elif kv_fp4: |
| K = _Sage3FP4KV(heads, S, hd, x.device, keep_exact=kv_fp4_exact_av) |
| V = K |
| elif kv_sage: |
| K = _SageKV(heads, S, hd, x.device) |
| V = K |
| elif kv_int8: |
| K = _Int8KV(heads, S, hd, x.device) |
| V = _Int8KV(heads, S, hd, x.device) |
| else: |
| K = torch.empty((1, heads, S, hd), dtype=x.dtype, device=x.device) |
| V = torch.empty((1, heads, S, hd), dtype=x.dtype, device=x.device) |
| for a, b in _ranges(S, kv_chunk): |
| s = b - a |
| h = _mod_scale_shift_range(block.norm1(x[a:b]), shift_msa, scale_msa, mod_segments, a, b) |
| if prec is not None: |
| prec.projection(block_index, "qkv", attn.qkv_proj, h, a, b, mod_segments) |
| if fq is not None: |
| h = fq.apply(block_index, "qkv", h, a, b, mod_segments) |
| qkv = attn.qkv_proj(h) |
| q, k, v = qkv.split(inner, dim=-1) |
| rope_c = rope_freqs[:, a:b].contiguous() if rope_freqs is not None else None |
| _, k = _norm_rope(attn, q, k, rope_c, s) |
| if kv_sage or kv_fp4 or kv_mix: |
| K.store(a, b, k.transpose(0, 1).unsqueeze(0), v.view(s, heads, hd).transpose(0, 1).unsqueeze(0)) |
| elif kv_int8: |
| K.store(a, b, k.transpose(0, 1).unsqueeze(0)) |
| V.store(a, b, v.view(s, heads, hd).transpose(0, 1).unsqueeze(0)) |
| else: |
| K[0, :, a:b] = k.transpose(0, 1) |
| V[0, :, a:b] = v.view(s, heads, hd).transpose(0, 1) |
| del h, qkv, q, k, v, rope_c |
| if kv_fp4 or kv_mix: |
| K.finalize() |
| return K, V |
|
|
|
|
| def _phase2_q_attn(block, x, K, V, shift_msa, scale_msa, gate_msa, mod_segments, rope_freqs, |
| transformer_options, q_chunk, kv_block, heads, hd, inner, head_chunks, prec=None, block_index=0, fq=None): |
| """Q per chunk, attention against full K/V, out_proj, gated residual in place.""" |
| S = x.shape[0] |
| attn = block.attn |
| heads_per_call = max(1, heads // max(1, min(int(head_chunks or 1), heads))) |
| q_ranges = _balanced_ranges(S, q_chunk, _min_query_chunk(x.device, heads_per_call)) |
| for a, b in q_ranges: |
| s = b - a |
| h = _mod_scale_shift_range(block.norm1(x[a:b]), shift_msa, scale_msa, mod_segments, a, b) |
| if fq is not None: |
| h = fq.apply(block_index, "qkv", h, a, b, mod_segments) |
| qkv = attn.qkv_proj(h) |
| q, k, _v = qkv.split(inner, dim=-1) |
| rope_c = rope_freqs[:, a:b].contiguous() if rope_freqs is not None else None |
| q, _ = _norm_rope(attn, q, k, rope_c, s) |
| qc = q.transpose(0, 1).unsqueeze(0).contiguous() |
| del h, qkv, q, k, _v, rope_c |
| if isinstance(K, _MixedKV): |
| o = _mixed_attention_q(qc, K, x.dtype) |
| if K.k_bf is not None: |
| _exact_av_rows_mixed(o, qc, K, mod_segments, a, b, x.dtype) |
| o = o.transpose(1, 2).reshape(1, s, inner) |
| elif isinstance(K, _Sage3FP4KV): |
| o = _sage3_attention_q(qc, K, x.dtype) |
| if K.k_bf is not None: |
| _exact_av_rows(o, qc, K, mod_segments, a, b, x.dtype) |
| o = o.transpose(1, 2).reshape(1, s, inner) |
| elif isinstance(K, _SageKV): |
| o = _sage_attention_q(qc, K, x.dtype) |
| o = o.transpose(1, 2).reshape(1, s, inner) |
| elif isinstance(K, _Int8KV): |
| o = _blockwise_attention_q(qc, K, V, kv_block if kv_block and kv_block > 0 else 16384, x.dtype) |
| o = o.transpose(1, 2).reshape(1, s, inner) |
| elif kv_block and kv_block > 0: |
| o = _blockwise_attention(qc, K, V, kv_block, x.dtype) |
| o = o.transpose(1, 2).reshape(1, s, inner) |
| else: |
| o = _attend(qc, K, V, heads, head_chunks, transformer_options) |
| del qc |
| o = o.squeeze(0) |
| if prec is not None: |
| prec.projection(block_index, "out", attn.out_proj, o, a, b, mod_segments) |
| if fq is not None: |
| o = fq.apply(block_index, "out", o, a, b, mod_segments) |
| o = attn.out_proj(o) |
| _mod_gate_range(x, gate_msa, o, mod_segments, a, b) |
| del o |
|
|
|
|
| def _phase3_mlp(block, x, shift_mlp, scale_mlp, gate_mlp, mod_segments, mlp_chunk, prec=None, block_index=0, fq=None): |
| """MLP per chunk, gated residual in place.""" |
| S = x.shape[0] |
| for a, b in _ranges(S, mlp_chunk): |
| h = _mod_scale_shift_range(block.norm2(x[a:b]), shift_mlp, scale_mlp, mod_segments, a, b) |
| if prec is not None and hasattr(block.mlp, "fc1") and hasattr(block.mlp, "fc2"): |
| prec.projection(block_index, "fc1", block.mlp.fc1, h, a, b, mod_segments) |
| gate, up = block.mlp.fc1(h).chunk(2, dim=-1) |
| act = torch.nn.functional.silu(gate).mul_(up) |
| prec.projection(block_index, "fc2", block.mlp.fc2, act, a, b, mod_segments) |
| del gate, up, act |
| if fq is not None and hasattr(block.mlp, "fc1") and hasattr(block.mlp, "fc2"): |
| h = fq.apply(block_index, "fc1", h, a, b, mod_segments) |
| gate, up = block.mlp.fc1(h).chunk(2, dim=-1) |
| act = torch.nn.functional.silu(gate).mul_(up) |
| act = fq.apply(block_index, "fc2", act, a, b, mod_segments) |
| o = block.mlp.fc2(act) |
| del gate, up, act |
| else: |
| o = block.mlp(h) |
| _mod_gate_range(x, gate_mlp, o, mod_segments, a, b) |
| del h, o |
|
|
|
|
| def streamed_block_forward(block, x, t_emb, mod_segments, rope_freqs, transformer_options, |
| q_chunk=16384, kv_chunk=16384, mlp_chunk=16384, kv_block=0, probe=None, index=None, kv_int8=False, kv_sage=False, kv_fp4=False, kv_fp4_exact_av=False, kv_mix=False): |
| """Exact replacement for DiTBlock.forward with chunk-bounded transients. |
| |
| The three phases are separate named functions so an allocator trace |
| (torch.cuda.memory._record_memory_history / memory_viz, see H3MemoryProbe) |
| labels every band by phase from the Python stack alone; `probe` (an |
| H3MemoryProbe ledger, read from transformer_options["h3_memprobe"]) gets a |
| zero-sync mark after each phase. |
| """ |
| attn = block.attn |
| heads, hd = attn.heads, attn.head_dim |
| inner = heads * hd |
| head_chunks = transformer_options.get("minimax_head_chunks", 1) if isinstance(transformer_options, dict) else 1 |
|
|
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaln_proj(t_emb) |
| prec = transformer_options.get("h3_precprobe") if isinstance(transformer_options, dict) else None |
| fq = transformer_options.get("h3_fakequant") if isinstance(transformer_options, dict) else None |
| x_in_sample = prec.block_in(x) if prec is not None else None |
|
|
| K, V = _phase1_kv(block, x, shift_msa, scale_msa, mod_segments, rope_freqs, kv_chunk, heads, hd, inner, kv_int8=kv_int8, kv_sage=kv_sage, |
| kv_fp4=kv_fp4, kv_fp4_exact_av=kv_fp4_exact_av, kv_mix=kv_mix, |
| prec=prec, block_index=index if index is not None else 0, fq=fq) |
| if probe is not None: |
| probe.mark(index, "kv") |
| _phase2_q_attn(block, x, K, V, shift_msa, scale_msa, gate_msa, mod_segments, rope_freqs, |
| transformer_options, q_chunk, kv_block, heads, hd, inner, head_chunks, |
| prec=prec, block_index=index if index is not None else 0, fq=fq) |
| del K, V |
| if probe is not None: |
| probe.mark(index, "attn") |
| _phase3_mlp(block, x, shift_mlp, scale_mlp, gate_mlp, mod_segments, mlp_chunk, prec=prec, block_index=index if index is not None else 0, fq=fq) |
| if probe is not None: |
| probe.mark(index, "mlp") |
| if prec is not None: |
| prec.block_out(index if index is not None else 0, x_in_sample, x) |
| if index == 49 or index is None: |
| prec.flush() |
| return x |
|
|
|
|
| def _self_check(block, args, extra, cfg, tag): |
| """Run stock and streamed on the same input and log per-phase divergence. Diagnostic only.""" |
| x = args["img"] |
| S = x.shape[0] |
| ref = extra["original_block"](dict(args, img=x.clone()))["img"] |
| scale = ref.float().abs().max().item() |
| ulp = 2.0 ** (math.floor(math.log2(max(scale, 1e-30))) - 7) |
| variants = {"full": cfg, |
| "q_only": dict(cfg, kv_chunk=S, mlp_chunk=S), |
| "kv_only": dict(cfg, q_chunk=S, mlp_chunk=S), |
| "mlp_only": dict(cfg, q_chunk=S, kv_chunk=S), |
| "none": dict(cfg, q_chunk=S, kv_chunk=S, mlp_chunk=S)} |
| parts = [] |
| for name, c in variants.items(): |
| out = streamed_block_forward(block, x.clone(), args["t_emb"], args["mod_segments"], args["rope_freqs"], |
| args["transformer_options"], q_chunk=c["q_chunk"], kv_chunk=c["kv_chunk"], |
| mlp_chunk=c["mlp_chunk"], kv_block=c["kv_block"], kv_int8=cfg.get("kv_int8", False), kv_sage=cfg.get("kv_sage", False), |
| kv_fp4=cfg.get("kv_fp4", False), |
| kv_fp4_exact_av=cfg.get("kv_fp4_exact_av", False), |
| kv_mix=cfg.get("kv_mix", False)) |
| d = (out.float() - ref.float()).abs() |
| parts.append(f"{name}: max {d.max().item():.3e} ({d.max().item() / ulp:.2f} ulp) mean {d.mean().item():.2e}") |
| del out, d |
| to = args["transformer_options"] |
| keys = sorted(k for k in to.keys()) if isinstance(to, dict) else type(to).__name__ |
| log.warning("H3StreamedBlocks self_check[%s] S=%d segs=%d rope=%s x=%s/%s to_keys=%s | %s", |
| tag, S, len(args["mod_segments"]), tuple(args["rope_freqs"].shape) if args["rope_freqs"] is not None else None, |
| x.dtype, x.is_contiguous(), keys, " | ".join(parts)) |
| del ref |
|
|
|
|
| def _make_replacement(block, cfg, index=0): |
| state = {"checked": False} |
|
|
| def fn(args, extra): |
| x = args["img"] |
| if x.shape[0] < cfg["min_tokens"]: |
| return extra["original_block"](args) |
| if cfg.get("self_check") and index == 0 and not state["checked"]: |
| state["checked"] = True |
| _self_check(block, args, extra, cfg, f"block{index}") |
| to = args["transformer_options"] |
| probe = to.get("h3_memprobe") if isinstance(to, dict) else None |
| x = streamed_block_forward(block, x, args["t_emb"], args["mod_segments"], args["rope_freqs"], |
| to, q_chunk=cfg["q_chunk"], kv_chunk=cfg["kv_chunk"], |
| mlp_chunk=cfg["mlp_chunk"], kv_block=cfg["kv_block"], probe=probe, index=index, |
| kv_int8=cfg.get("kv_int8", False), kv_sage=cfg.get("kv_sage", False), |
| kv_fp4=cfg.get("kv_fp4", False), |
| kv_fp4_exact_av=cfg.get("kv_fp4_exact_av", False), |
| kv_mix=cfg.get("kv_mix", False)) |
| return {"img": x} |
| return _named(fn, f"block{index:02d}") |
|
|
|
|
| def _named(fn, name): |
| """Return `fn` under a new code-object name so it shows as `name` in Python |
| stacks (allocator traces label bands by frame name; there is no other |
| per-call label channel).""" |
| import types |
| code = fn.__code__.replace(co_name=name) |
| g = types.FunctionType(code, fn.__globals__, name, fn.__defaults__, fn.__closure__) |
| g.__qualname__ = name |
| return g |
|
|
|
|
|
|
| |
|
|
| def streamed_final_layer_forward(fl, x, t_emb, video_seg, audio_seg, chunk=16384, probe=None, exact_gemm=True): |
| """Row-chunked FinalLayer.forward. Stock (comfy/ldm/minimax/model.py:295) |
| builds `norm(x[span]) * (1 + scale) + shift` for the whole target span and |
| the fp32 modulation promotes it to fp32 twice: 2 x 4.28 GiB at 216k tokens, |
| measured as the forward's peak on a 216k-token de-rope (2026-08-18). Every op |
| is per row (RMSNorm, per-element mod, per-row fp32 linear), so the same |
| math per chunk yields the same [rows, out] result with ~chunk-sized |
| transients. Whether the fp32 cuBLAS GEMM stays bit-equal under a different |
| M is a kernel property: gated, not assumed.""" |
| shift, scale = fl.adaln_proj(t_emb) |
|
|
| def head(a, b, row, out_mod): |
| n = b - a |
| if exact_gemm: |
| |
| |
| |
| hbuf = torch.empty((n, x.shape[1]), dtype=torch.float32, device=x.device) |
| for c0, c1 in _ranges(n, chunk): |
| hbuf[c0:c1] = fl.norm(x[a + c0:a + c1]) * (1.0 + scale[row]) + shift[row] |
| out = out_mod(hbuf) |
| del hbuf |
| return out |
| |
| |
| parts = [] |
| for c0, c1 in _ranges(n, chunk): |
| h = (fl.norm(x[a + c0:a + c1]) * (1.0 + scale[row]) + shift[row]).to(torch.float32) |
| parts.append(out_mod(h)) |
| del h |
| return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) |
|
|
| va, vb, vrow = video_seg |
| aa, ab, arow = audio_seg |
| v = head(va, vb, vrow, fl.video_out) |
| a = head(aa, ab, arow, fl.audio_out) |
| if probe is not None: |
| probe.mark(None, "final") |
| return v, a |
|
|
|
|
| |
|
|
| class H3StreamedBlocks: |
| """Run every H3 DiT block in token chunks: exact, chunk-bounded VRAM.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "model": ("MODEL",), |
| "q_chunk": ("INT", {"default": 16384, "min": 1024, "max": 262144, "step": 1024, |
| "tooltip": "Query tokens per attention call. Smaller = lower transient VRAM, more calls. Exact at any value."}), |
| "kv_chunk": ("INT", {"default": 16384, "min": 1024, "max": 262144, "step": 1024, |
| "tooltip": "Tokens per K/V projection chunk while building the full K/V buffers."}), |
| "mlp_chunk": ("INT", {"default": 16384, "min": 0, "max": 262144, "step": 1024, |
| "tooltip": "Tokens per MLP chunk. 0 = leave the model's mlp.forward alone (e.g. KJNodes' chunk node)."}), |
| "min_tokens": ("INT", {"default": 32768, "min": 0, "max": 1048576, "step": 1024, |
| "tooltip": "Below this packed sequence length the stock block runs (short clips gain nothing)."}), |
| "kv_block": ("INT", {"default": 0, "min": 0, "max": 262144, "step": 1024, |
| "tooltip": "EXPERIMENTAL, leave at 0. Attends K/V in blocks with a log-sum-exp combine (1 bf16 ulp). As built it does not lower memory (K/V are still fully built; measured +1.1 GiB and ~7% slower at 216k tokens); kept for the host-staged K/V design to come."}), |
| }, |
| "optional": { |
| "final_layer_chunk": ("INT", {"default": 16384, "min": 0, "max": 262144, "step": 1024, |
| "tooltip": "Rows per chunk through the output head's norm -> mod -> fp32 promotion. Stock promotes the whole span to fp32 twice (~10 GiB at 216k tokens, the forward's peak). 0 = stock."}), |
| "final_layer_gemm": (["exact (whole GEMM, one fp32 buffer)", "streamed (chunked GEMM, ~1e-6 fp32 diffs)"], |
| {"default": "exact (whole GEMM, one fp32 buffer)", |
| "tooltip": "exact: same head GEMM as stock, transient = one fp32 [rows, hidden] (bit-equal). streamed: GEMM per chunk, transient ~chunk-sized, but fp32 cuBLAS is not chunk-invariant (numerically-equivalent tier)."}), |
| "kv_store": (["bf16 (exact)", "kvi8r: rotated int8 K/V (approximate)", "kvi8s: Sage int8/fp8 K/V, rotated (approximate, faster attention)", |
| "kvfp4s: Sage3 FP4 K/V, rotated (approximate, fastest attention)", "kvmix: per-head 4/8-bit, fp4 + Sage int8 (the ~6-bit tier)"], |
| {"default": "bf16 (exact)", |
| "tooltip": "kvi8r = rotated int8 K/V. K and V held as int8 with one fp16 scale per row (per token, per head) in a fixed orthonormal 128-wide Hadamard-rotated basis of the head dim; the query chunk is rotated to match and the output un-rotated once, so a K/V block dequant is a single int8->fp16 cast * scale and attention runs in fp16 blockwise (kv_block, default 16384) with an online-softmax combine. Halves the K/V bytes. NOT bit-equal to stock: a same-seed render is a sibling take (first cut 2026-08-18: operator judged the de-rope side by side 'almost perfect'). Second cut (2026-08-18 evening): standalone at 217k tokens the attention costs +16% over the exact path (first cut +24%) with a ~1 GiB transient (first cut ~3), so the K/V saving now shows in the forward peak; the live numbers are in LOWVRAM.md. kvi8s = the same K/V bytes kept in SageAttention's kernel layout (int8 K per 64-token block, fp8 V per channel, Q/K Hadamard-rotated first) and attended on int8/fp8 tensor cores straight from the store, no dequant: standalone ~1.6x faster attention than the exact path at 217k tokens, one rung more approximate than kvi8r; needs the sageattention package (2.x). kvfp4s = the same idea one precision rung further down, on SageAttention3's Blackwell fp4 kernels: K/V kept as NVFP4 (E2M1 + one e4m3 scale per 16 elements = 0.5625 byte/elem, 0.28x bf16 -- 1.63 GiB vs 5.79 at 216k tokens), Q/K Hadamard-rotated, no delta_s and no Q-centering (measured to cost nothing at per_block_mean False, which is the only affordable setting at H3 lengths). Standalone at 217k tokens: ~2.4x faster attention than the exact path and ~1.4x faster than kvi8s, at rel-rms 0.293 vs Sage 2.2's 0.055 on the outlier proxy (0.192 vs 0.039 plain gaussian) -- clearly the most approximate rung, gated on the operator's eyes. Because the fp4 kernel's log-sum-exp is not retrievable there is no K-blocked combine: phase 1 buffers K/V in bf16 and quantises once, so the block's peak during that one finalize is bf16 + fp4 before the bf16 is dropped. Needs SageAttention3's fp4 extension (fp4attn_cuda) importable; falls back to bf16 (exact) with a warning if it is not."}), |
| "trim_forward": ("BOOLEAN", {"default": True, |
| "tooltip": "F5: run a copy of the stock model forward that releases the patch-embed and row buffers once the packed sequence is assembled (stock keeps them for the whole forward: 2.14 GiB at 216k tokens). Same math, exact. Applied only if the installed ComfyUI's forward matches the copy (source hash); otherwise skipped with a log line."}), |
| "self_check": ("BOOLEAN", {"default": False, |
| "tooltip": "Diagnostic: on block 0's first call, run stock and streamed on the same input and log per-phase divergence. Costs one extra block forward."}), |
| "exact_av_rows": ("BOOLEAN", {"default": False, |
| "tooltip": "kvfp4s only. Route the TEXT and AUDIO query rows through an exact bf16 attention over a retained bf16 K/V, leaving the video rows (96.3% of the packed sequence, where the speed win lives) on the fp4 store. Audio is ~300 of 28,931 tokens, so the ~4% attention error that video's redundancy hides as slight ghosting lands audibly on a soundtrack carried by 1% of the sequence -- measured: audio is not quantised worse than video, it is merely far more sensitive, so the fix is routing rather than a better quantiser. Costs the bf16 K/V of the block in flight (~0.8 GiB at 29k tokens, 5.8 at 216k) and ~5% more attention time. No effect unless kv_store is kvfp4s."}), |
| } |
| } |
|
|
| RETURN_TYPES = ("MODEL",) |
| FUNCTION = "patch" |
| CATEGORY = "MAINodes/VRAM Lab" |
| DESCRIPTION = ("Exact low-VRAM execution of MiniMax-H3: never materialises the full-sequence " |
| "fused QKV or SwiGLU tensors. Same math as the stock block; costs ~4% extra " |
| "projection work at long lengths. See vram_lab.py for the ledger.") |
|
|
| def patch(self, model, q_chunk, kv_chunk, mlp_chunk, min_tokens, kv_block, self_check=False, final_layer_chunk=16384, |
| final_layer_gemm="exact (whole GEMM, one fp32 buffer)", kv_store="bf16 (exact)", trim_forward=True, |
| exact_av_rows=False): |
| dm = getattr(getattr(model, "model", None), "diffusion_model", None) |
| blocks = getattr(dm, "blocks", None) |
| if not blocks or not hasattr(blocks[0], "attn") or not hasattr(blocks[0].attn, "qkv_proj"): |
| log.warning("H3StreamedBlocks: model does not look like MiniMax H3 (no blocks[*].attn.qkv_proj); unchanged") |
| return (model,) |
| cfg = {"q_chunk": q_chunk, "kv_chunk": kv_chunk, "mlp_chunk": mlp_chunk, |
| "min_tokens": min_tokens, "kv_block": kv_block, "self_check": bool(self_check), |
| "kv_int8": str(kv_store).startswith("kvi8r"), "kv_sage": str(kv_store).startswith("kvi8s"), |
| "kv_fp4": str(kv_store).startswith("kvfp4s"), |
| "kv_fp4_exact_av": bool(exact_av_rows), |
| "kv_mix": str(kv_store).startswith("kvmix")} |
| if cfg["kv_sage"] and not _SAGE_OK: |
| log.warning("H3StreamedBlocks: kv_store kvi8s needs the sageattention package; falling back to bf16 (exact)") |
| cfg["kv_sage"] = False |
| if cfg["kv_fp4"] and _sage3_api() is None: |
| log.warning("H3StreamedBlocks: kv_store kvfp4s needs SageAttention3's fp4 extension (fp4attn_cuda, tried %s); falling back to bf16 (exact)", _SAGE3_DIR) |
| cfg["kv_fp4"] = False |
| if cfg["kv_mix"] and (not _SAGE_OK or _sage3_api() is None): |
| log.warning("H3StreamedBlocks: kv_store kvmix needs BOTH the sageattention package (2.x: %s) and SageAttention3's fp4 extension (fp4attn_cuda, tried %s: %s); falling back to bf16 (exact)", |
| "present" if _SAGE_OK else "missing", _SAGE3_DIR, |
| "present" if _sage3_api() is not None else "missing") |
| cfg["kv_mix"] = False |
| try: |
| from . import h3_capabilities as _caps |
| for w in _caps.collision_warnings(_caps.block_patch_report(model)): |
| log.warning("H3StreamedBlocks: %s", w) |
| except Exception as _e: |
| log.info("H3StreamedBlocks: collision report skipped (%s)", _e) |
| m = model.clone() |
| for i, block in enumerate(blocks): |
| m.set_model_patch_replace(_make_replacement(block, cfg, i), "dit", "double_block", i) |
| fl = getattr(dm, "final_layer", None) |
| if final_layer_chunk and fl is not None and hasattr(fl, "video_out") and hasattr(fl, "audio_out"): |
| _exact = str(final_layer_gemm).startswith("exact") |
|
|
| def _fl_forward(x, t_emb, video_seg, audio_seg, _fl=fl, _c=int(final_layer_chunk), _e=_exact): |
| if x.shape[0] < cfg["min_tokens"]: |
| return type(_fl).forward(_fl, x, t_emb, video_seg, audio_seg) |
| return streamed_final_layer_forward(_fl, x, t_emb, video_seg, audio_seg, chunk=_c, |
| probe=getattr(dm, "_h3_memprobe", None), exact_gemm=_e) |
| m.add_object_patch("diffusion_model.final_layer.forward", _fl_forward) |
| if trim_forward: |
| sha = _stock_forward_sha() |
| if sha == _STOCK_FORWARD_SHA and hasattr(dm, "_forward"): |
| import types |
| m.add_object_patch("diffusion_model._forward", types.MethodType(_trimmed_forward, dm)) |
| log.info("H3StreamedBlocks: trim_forward on (stock _forward %s matches the copy)", sha) |
| else: |
| log.warning("H3StreamedBlocks: trim_forward skipped: installed _forward source hash %s != %s (ComfyUI changed; the copy needs refreshing)", sha, _STOCK_FORWARD_SHA) |
| log.info("H3StreamedBlocks: %d blocks patched (q %d, kv %d, mlp %d, min %d, kv_block %d, final_layer_chunk %d, %s, kv_store %s)", |
| len(blocks), q_chunk, kv_chunk, mlp_chunk, min_tokens, kv_block, final_layer_chunk, final_layer_gemm, kv_store) |
| return (m,) |
|
|
|
|
|
|
| |
|
|
| def _rss(): |
| """(RssAnon, RssFile) of this process in bytes from /proc/self/status; ~µs. |
| RssAnon is the host RAM the process really holds (mirrors, retained |
| allocator arenas, pinned buffers); RssFile is mmap'd model files.""" |
| anon = file = 0 |
| try: |
| with open("/proc/self/status") as f: |
| for line in f: |
| if line.startswith("RssAnon:"): |
| anon = int(line.split()[1]) * 1024 |
| elif line.startswith("RssFile:"): |
| file = int(line.split()[1]) * 1024 |
| except OSError: |
| pass |
| return anon, file |
|
|
|
|
| class _MemLedger: |
| """Zero-sync memory ledger: reads the caching allocator's host-side counters |
| (allocated, peak since last mark, reserved) at phase boundaries. No device |
| sync: allocations happen at launch time, so the counters are exact on the |
| host timeline; the timestamps are launch times, not completion times.""" |
|
|
| def __init__(self, path, dev): |
| self.path = path |
| self.dev = dev |
| self.fwd = -1 |
| self.rows = [] |
| self.t0 = time.perf_counter() |
|
|
| def begin_forward(self, shape): |
| self.fwd += 1 |
| torch.cuda.reset_peak_memory_stats(self.dev) |
| self._base = torch.cuda.memory_allocated(self.dev) |
| self.rows.append({"fwd": self.fwd, "block": None, "phase": "start", "t": time.perf_counter() - self.t0, |
| "alloc": self._base, "peak": self._base, |
| "reserved": torch.cuda.memory_reserved(self.dev), "shape": list(shape), |
| "rss_anon": _rss()[0], "rss_file": _rss()[1]}) |
|
|
| def mark(self, block, phase): |
| st = torch.cuda.memory_stats(self.dev) |
| self.rows.append({"fwd": self.fwd, "block": block, "phase": phase, "t": time.perf_counter() - self.t0, |
| "alloc": st.get("allocated_bytes.all.current", 0), |
| "peak": st.get("allocated_bytes.all.peak", 0), |
| "reserved": st.get("reserved_bytes.all.current", 0), |
| "rss_anon": _rss()[0], "rss_file": _rss()[1]}) |
| torch.cuda.reset_peak_memory_stats(self.dev) |
|
|
| def end_forward(self): |
| st = torch.cuda.memory_stats(self.dev) |
| self.rows.append({"fwd": self.fwd, "block": None, "phase": "end", "t": time.perf_counter() - self.t0, |
| "alloc": st.get("allocated_bytes.all.current", 0), |
| "peak": st.get("allocated_bytes.all.peak", 0), |
| "reserved": st.get("reserved_bytes.all.current", 0), |
| "rss_anon": _rss()[0], "rss_file": _rss()[1]}) |
| self.flush() |
|
|
| def flush(self): |
| import json |
| os.makedirs(os.path.dirname(self.path), exist_ok=True) |
| with open(self.path, "w") as f: |
| for r in self.rows: |
| f.write(json.dumps(r) + "\n") |
| try: |
| with open(os.path.splitext(self.path)[0] + ".html", "w") as f: |
| f.write(_ledger_html(self.rows)) |
| except Exception as e: |
| log.debug("ledger html failed: %s", e) |
|
|
| def summary(self): |
| rows = [r for r in self.rows if r["fwd"] == self.fwd] |
| if not rows: |
| return "" |
| peak = max(r["peak"] for r in rows) |
| top = max(rows, key=lambda r: r["peak"]) |
| return (f"fwd {self.fwd}: base {rows[0]['alloc'] / 2**30:.1f} GiB, peak {peak / 2**30:.1f} GiB " |
| f"(at block {top['block']} {top['phase']}), reserved {rows[-1]['reserved'] / 2**30:.1f} GiB, " |
| f"{rows[-1]['t'] - rows[0]['t']:.1f} s; RSS anon {rows[0]['rss_anon'] / 2**30:.1f} -> {rows[-1]['rss_anon'] / 2**30:.1f} GiB " |
| f"(max {max(r['rss_anon'] for r in rows) / 2**30:.1f})") |
|
|
|
|
|
|
| _PHASE_COLOR = {"start": "#888", "kv": "#4c78a8", "attn": "#f58518", "mlp": "#54a24b", "end": "#888"} |
|
|
|
|
| def _ledger_html(rows): |
| """Self-contained SVG timeline of the ledger (no external assets): allocated, |
| per-phase peak, reserved and process RSS in GiB against wall time, one |
| forward per band; hover a mark for its numbers. Deep dive = trace.html.""" |
| if not rows: |
| return "<p>empty ledger</p>" |
| G = 2.0 ** 30 |
| W, H, L, T, B = 1400, 520, 70, 30, 60 |
| t0 = rows[0]["t"] |
| tmax = max(r["t"] for r in rows) - t0 or 1.0 |
| ymax = max(max(r["peak"], r["reserved"], r.get("rss_anon", 0)) for r in rows) / G * 1.05 or 1.0 |
| xs = lambda t: L + (t - t0) / tmax * (W - L - 20) |
| ys = lambda v: T + (H - T - B) * (1 - v / ymax) |
|
|
| def path(key): |
| return " ".join(f"{'M' if i == 0 else 'L'}{xs(r['t']):.1f},{ys(r[key] / G):.1f}" for i, r in enumerate(rows)) |
|
|
| out = [f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {W} {H}" style="max-width:100%;font:12px sans-serif;background:#fff">'] |
| for g in range(0, int(ymax) + 1, max(1, int(ymax // 10) or 1)): |
| out.append(f'<line x1="{L}" x2="{W-20}" y1="{ys(g):.1f}" y2="{ys(g):.1f}" stroke="#eee"/>' |
| f'<text x="{L-6}" y="{ys(g)+4:.1f}" text-anchor="end" fill="#666">{g} GiB</text>') |
| |
| fwds = sorted({r["fwd"] for r in rows}) |
| for fw in fwds: |
| rr = [r for r in rows if r["fwd"] == fw] |
| x0, x1 = xs(rr[0]["t"]), xs(rr[-1]["t"]) |
| out.append(f'<rect x="{x0:.1f}" y="{T}" width="{max(1.0, x1-x0):.1f}" height="{H-T-B}" fill="{"#f7f7f7" if fw % 2 else "#fff"}"/>' |
| f'<text x="{x0+3:.1f}" y="{T+12}" fill="#999">fwd {fw}</text>') |
| out.append(f'<path d="{path("reserved")}" fill="none" stroke="#bbb" stroke-width="1.5"/>') |
| out.append(f'<path d="{path("peak")}" fill="none" stroke="#e45756" stroke-width="1" stroke-dasharray="3,2"/>') |
| out.append(f'<path d="{path("alloc")}" fill="none" stroke="#222" stroke-width="1.5"/>') |
| if any(r.get("rss_anon") for r in rows): |
| out.append(f'<path d="{path("rss_anon")}" fill="none" stroke="#9467bd" stroke-width="1.5"/>') |
| for r in rows: |
| c = _PHASE_COLOR.get(r["phase"], "#333") |
| tip = (f"fwd {r['fwd']} block {r['block']} {r['phase']}\\nt = {r['t']-t0:.1f} s\\nallocated {r['alloc']/G:.2f} GiB\\n" |
| f"peak since last mark {r['peak']/G:.2f} GiB\\nreserved {r['reserved']/G:.2f} GiB\\nRSS anon {r.get('rss_anon',0)/G:.2f} GiB, file {r.get('rss_file',0)/G:.2f} GiB") |
| out.append(f'<circle cx="{xs(r["t"]):.1f}" cy="{ys(r["peak"]/G):.1f}" r="3" fill="{c}"><title>{tip}</title></circle>') |
| lg = [("#222", "allocated"), ("#e45756", "peak since last mark"), ("#bbb", "reserved"), ("#9467bd", "process RSS (anon)"), |
| ("#4c78a8", "mark: kv"), ("#f58518", "mark: attn"), ("#54a24b", "mark: mlp")] |
| for i, (col, name) in enumerate(lg): |
| x = L + i * 190 |
| out.append(f'<rect x="{x}" y="{H-28}" width="14" height="10" fill="{col}"/><text x="{x+18}" y="{H-19}" fill="#333">{name}</text>') |
| out.append(f'<text x="{W/2:.0f}" y="{H-2}" text-anchor="middle" fill="#666">wall time, {tmax:.0f} s span; hover a mark</text></svg>') |
| return ("<!doctype html><meta charset=utf-8><title>H3 memory ledger</title>" |
| "<style>body{margin:12px;font-family:sans-serif}</style><h3>H3MemoryProbe ledger</h3>" + "".join(out)) |
|
|
| class H3MemoryProbe: |
| """See what holds VRAM, per block and phase, and optionally record the |
| allocator trace for a hoverable timeline (torch memory_viz).""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "model": ("MODEL",), |
| "tag": ("STRING", {"default": "probe", "tooltip": "Run label; files land in out_dir/<tag>_<time>/"}), |
| "ledger": ("BOOLEAN", {"default": True, |
| "tooltip": "Per-forward JSONL of allocator counters at every H3StreamedBlocks phase boundary (start, block i kv/attn/mlp, end). No device sync, negligible cost. Stock blocks contribute start/end only."}), |
| "record_history_forwards": ("INT", {"default": 0, "min": 0, "max": 64, |
| "tooltip": "Record the caching allocator's alloc/free trace (with Python stacks) for this many model forwards, then dump snapshot.pickle and trace.html (standalone; hover a band for the stack that allocated it). 0 = off. ~20k events per H3 forward at 200k tokens; a few percent while recording."}), |
| "max_entries": ("INT", {"default": 300000, "min": 10000, "max": 5000000, "step": 10000, |
| "tooltip": "Ring size for the allocator trace."}), |
| "out_dir": ("STRING", {"default": "output/h3_memprobe", |
| "tooltip": "Relative to the ComfyUI working directory. Not /tmp."}), |
| } |
| } |
|
|
| RETURN_TYPES = ("MODEL", "STRING") |
| RETURN_NAMES = ("model", "report") |
| FUNCTION = "patch" |
| CATEGORY = "MAINodes/VRAM Lab" |
| DESCRIPTION = ("Memory instrument for the H3 diffusion model: a per-block/per-phase ledger of PyTorch's " |
| "allocator counters (with H3StreamedBlocks upstream), and an optional allocator trace " |
| "rendered to a hoverable HTML timeline. Off = no cost; not installed at all.") |
|
|
| def patch(self, model, tag, ledger, record_history_forwards, max_entries, out_dir): |
| import folder_paths |
| import comfy.patcher_extension as pe |
| dev = comfy.model_management.get_torch_device() |
| base = out_dir if os.path.isabs(out_dir) else os.path.join(os.path.dirname(folder_paths.get_output_directory()), out_dir) |
| run_dir = os.path.join(base, f"{tag}_{time.strftime('%Y%m%d_%H%M%S')}") |
| os.makedirs(run_dir, exist_ok=True) |
| led = _MemLedger(os.path.join(run_dir, "ledger.jsonl"), dev) if ledger else None |
| state = {"fwd": 0, "recording": False, "done": False, "n": int(record_history_forwards)} |
|
|
| m = model.clone() |
| dm = getattr(getattr(m, "model", None), "diffusion_model", None) |
| if led is not None: |
| to = m.model_options.setdefault("transformer_options", {}) |
| to["h3_memprobe"] = led |
|
|
| def wrapper(executor, *args, **kwargs): |
| if led is not None and dm is not None: |
| dm._h3_memprobe = led |
| x = args[0] if args else None |
| shape = tuple(x.shape) if hasattr(x, "shape") else () |
| if state["n"] > 0 and not state["done"] and not state["recording"]: |
| try: |
| torch.cuda.memory._record_memory_history(enabled="all", context="all", stacks="python", |
| max_entries=int(max_entries), device=dev, |
| record_pinned_host_memory=True) |
| state["recording"] = True |
| log.info("H3MemoryProbe[%s]: allocator trace ON (fwd %d)", tag, state["fwd"]) |
| except RuntimeError as e: |
| |
| |
| state["done"] = True |
| log.warning("H3MemoryProbe[%s]: allocator trace unavailable (%s). Start ComfyUI with " |
| "--disable-cuda-malloc to record it; the ledger still runs.", tag, str(e)[:120]) |
| if led is not None: |
| led.begin_forward(shape) |
| try: |
| return executor(*args, **kwargs) |
| finally: |
| if dm is not None and hasattr(dm, "_h3_memprobe"): |
| del dm._h3_memprobe |
| if led is not None: |
| led.end_forward() |
| log.info("H3MemoryProbe[%s]: %s", tag, led.summary()) |
| state["fwd"] += 1 |
| if state["recording"] and state["fwd"] >= state["n"]: |
| _dump_trace(run_dir, dev, tag) |
| torch.cuda.memory._record_memory_history(enabled=None, device=dev) |
| state["recording"] = False |
| state["done"] = True |
|
|
| m.add_wrapper_with_key(pe.WrappersMP.DIFFUSION_MODEL, "h3_memprobe", wrapper) |
| rep = f"H3MemoryProbe: {run_dir} (ledger {'on' if ledger else 'off'}, trace forwards {record_history_forwards})" |
| log.info(rep) |
| return (m, rep) |
|
|
|
|
| def _dump_trace(run_dir, dev, tag): |
| try: |
| os.makedirs(run_dir, exist_ok=True) |
| snap = torch.cuda.memory._snapshot(device=dev) |
| import pickle |
| with open(os.path.join(run_dir, "snapshot.pickle"), "wb") as f: |
| pickle.dump(snap, f) |
| try: |
| from torch.cuda._memory_viz import trace_plot |
| html = trace_plot(snap, device=None) |
| with open(os.path.join(run_dir, "trace.html"), "w") as f: |
| f.write(html) |
| except Exception as e: |
| log.warning("H3MemoryProbe[%s]: trace_plot failed (%s); snapshot.pickle kept for pytorch.org/memory_viz", tag, e) |
| log.info("H3MemoryProbe[%s]: allocator trace written to %s", tag, run_dir) |
| except Exception as e: |
| log.warning("H3MemoryProbe[%s]: snapshot failed: %s", tag, e) |
|
|
|
|
|
|
|
|
| class H3FreeCache: |
| """Passthrough that returns the allocator's cached-but-free VRAM to the |
| driver before the next stage. Measured motivation (2026-08-18): the |
| VAE decode after a long H3 pass grew the pool 69.6 -> 77.9 GiB while live |
| tensors were LOWER than during sampling; decode-shaped blocks could not |
| reuse sampling's freed ones. Costs a few ms; changes no math.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": {"samples": ("LATENT",)}, |
| "optional": {"also_gc": ("BOOLEAN", {"default": True, "tooltip": "gc.collect() first so dead Python refs release their tensors too."})}} |
|
|
| RETURN_TYPES = ("LATENT", "STRING") |
| RETURN_NAMES = ("samples", "report") |
| FUNCTION = "free" |
| CATEGORY = "MAINodes/VRAM Lab" |
| DESCRIPTION = "Empty the CUDA caching allocator (torch.cuda.empty_cache) between stages; passthrough for a LATENT so it can sit right before VAE Decode." |
|
|
| def free(self, samples, also_gc=True): |
| import gc |
| dev = comfy.model_management.get_torch_device() |
| before = torch.cuda.memory_reserved(dev) |
| if also_gc: |
| gc.collect() |
| comfy.model_management.soft_empty_cache(force=True) |
| after = torch.cuda.memory_reserved(dev) |
| rep = f"H3FreeCache: reserved {before / 2**30:.2f} -> {after / 2**30:.2f} GiB (live {torch.cuda.memory_allocated(dev) / 2**30:.2f})" |
| log.info(rep) |
| return (samples, rep) |
|
|
|
|
|
|
| class H3EvictTextEncoder: |
| """Passthrough for CONDITIONING that unloads the text encoder (the CLIP |
| patcher and its clones) the moment encoding is done - the same call |
| ckinpdx's MMH3Tools makes inside its nodes. Under --gpu-only it is a no-op |
| (offload device is the GPU); in normal mode it frees the TE's VRAM before |
| the DiT loads instead of letting the planner evict it on demand. Exists to |
| MEASURE whether explicit eviction matters on a small card (2026-08-18).""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": {"conditioning": ("CONDITIONING",), "clip": ("CLIP",)}} |
|
|
| RETURN_TYPES = ("CONDITIONING", "STRING") |
| RETURN_NAMES = ("conditioning", "report") |
| FUNCTION = "evict" |
| CATEGORY = "MAINodes/VRAM Lab" |
| DESCRIPTION = "Unload the text encoder right after encoding (passthrough for the conditioning)." |
|
|
| def evict(self, conditioning, clip): |
| dev = comfy.model_management.get_torch_device() |
| before = comfy.model_management.get_free_memory(dev) |
| try: |
| comfy.model_management.unload_model_and_clones(clip.patcher) |
| except Exception as e: |
| log.warning("H3EvictTextEncoder: unload failed: %s", e) |
| comfy.model_management.soft_empty_cache() |
| after = comfy.model_management.get_free_memory(dev) |
| rep = f"H3EvictTextEncoder: device free {before / 2**30:.1f} -> {after / 2**30:.1f} GiB" |
| log.info(rep) |
| return (conditioning, rep) |
|
|
| class H3PrecisionProbe: |
| """A6 phase 0: record, per (forward, block, projection, segment), the real input |
| activation of every block linear and the extra output error NVFP4 / NVFP4+Hadamard / |
| FP8 activation quantisation would add over the shipped int8 path, plus the per-block |
| residual change. Needs H3StreamedBlocks downstream (it isolates the projections). |
| Cost ~1.2x at every=8. Output jsonl under out_dir; draw with benchmarks/scripts/prec_map.py.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": { |
| "model": ("MODEL",), |
| "tag": ("STRING", {"default": "precprobe"}), |
| "every": ("INT", {"default": 8, "min": 1, "max": 256, "tooltip": "sample one chunk in N per (block, projection)"}), |
| "out_dir": ("STRING", {"default": "output/h3_precprobe"}), |
| }} |
|
|
| RETURN_TYPES = ("MODEL",) |
| FUNCTION = "patch" |
| CATEGORY = "MAINodes/VRAM Lab" |
|
|
| def patch(self, model, tag, every, out_dir): |
| import folder_paths |
| import comfy.patcher_extension as pe |
| base = out_dir if os.path.isabs(out_dir) else os.path.join(os.path.dirname(folder_paths.get_output_directory()), out_dir) |
| run_dir = os.path.join(base, f"{tag}_{time.strftime('%Y%m%d_%H%M%S')}") |
| os.makedirs(run_dir, exist_ok=True) |
| pr = _PrecProbe(os.path.join(run_dir, "prec.jsonl"), every=every) |
| m = model.clone() |
| to = m.model_options.setdefault("transformer_options", {}) |
| to["h3_precprobe"] = pr |
|
|
| def wrapper(executor, *args, **kwargs): |
| sig = None |
| try: |
| sig = float(args[1].flatten()[0].item()) / 1000.0 |
| except Exception: |
| pass |
| pr.begin_forward(sig) |
| out = executor(*args, **kwargs) |
| pr.flush_amax() |
| pr.flush() |
| return out |
| m.add_wrapper_with_key(pe.WrappersMP.DIFFUSION_MODEL, "h3_precprobe", wrapper) |
| log.info("H3PrecisionProbe[%s]: -> %s (every %d)", tag, run_dir, every) |
| return (m,) |
|
|
|
|
| class H3FakeQuant: |
| """A6 phase 0d: simulate NVFP4 / FP8 activation precision on a chosen set of |
| projections, block range and token segments (fake-quant of the layer input; the |
| layer runs as shipped). For the same-seed sensitivity sweep that decides which |
| layers can go to 4 bits. Needs H3StreamedBlocks downstream. Not a speed path.""" |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return {"required": { |
| "model": ("MODEL",), |
| "format": (["nvfp4", "nvfp4_had", "fp8"], {"default": "nvfp4"}), |
| "projections": ("STRING", {"default": "qkv,out,fc1,fc2", "tooltip": "comma list of qkv,out,fc1,fc2"}), |
| "block_lo": ("INT", {"default": 0, "min": 0, "max": 63}), |
| "block_hi": ("INT", {"default": 49, "min": 0, "max": 63}), |
| "segments": ("STRING", {"default": "", "tooltip": "empty = all rows; else comma list of video,text,audio,cond_video,cond_audio"}), |
| }} |
|
|
| RETURN_TYPES = ("MODEL",) |
| FUNCTION = "patch" |
| CATEGORY = "MAINodes/VRAM Lab" |
|
|
| def patch(self, model, format, projections, block_lo, block_hi, segments): |
| projs = [p.strip() for p in projections.split(",") if p.strip()] |
| segs = [p.strip() for p in segments.split(",") if p.strip()] |
| m = model.clone() |
| to = m.model_options.setdefault("transformer_options", {}) |
| to["h3_fakequant"] = _FakeQuant(format, projs, (int(block_lo), int(block_hi)), segs) |
| log.info("H3FakeQuant: %s on %s blocks %d..%d segs %s", format, projs, block_lo, block_hi, segs or "all") |
| return (m,) |
|
|
|
|
| |
| |
| |
| NODE_CLASS_MAPPINGS = {"H3StreamedBlocks": H3StreamedBlocks, "H3MemoryProbe": H3MemoryProbe, "H3FreeCache": H3FreeCache, "H3EvictTextEncoder": H3EvictTextEncoder, |
| "H3PrecisionProbe": H3PrecisionProbe, "H3FakeQuant": H3FakeQuant} |
| NODE_DISPLAY_NAME_MAPPINGS = {"H3StreamedBlocks": "H3 Streamed Blocks (exact low-VRAM, alpha)", |
| "H3MemoryProbe": "H3 Memory Probe (ledger + allocator trace, alpha)", |
| "H3FreeCache": "H3 Free Cache (empty allocator between stages)", |
| "H3EvictTextEncoder": "H3 Evict Text Encoder (unload after encode)", |
| "H3PrecisionProbe": "H3 Precision Probe (A6: 4-bit vs 8-bit activation map, alpha)", |
| "H3FakeQuant": "H3 Fake Quant (A6: simulate NVFP4/FP8 activations per region, alpha)"} |
|
|
