#!/usr/bin/env python3 """Algebra gate: H3 block-0 residual skip identity. Ask: is residual_0(step t) a function of residual_0(step t-1) plus a cheap exact predicate (timestep, modulation, or residual delta) that is byte-equal when the cache hits and fail-closed when it misses? Lab-only. Does not patch serving. Does not flip cache default. """ from __future__ import annotations import json import os import sys from pathlib import Path import torch HERE = Path(__file__).resolve().parent REPO = HERE.parents[1] COMFY = Path(os.environ.get("COMFYUI_ROOT", "$HOME/comfyui-h3-current")) CACHE_SRC = COMFY / "comfy/ldm/minimax/first_block_cache.py" SERVING_ENV = REPO / "config" / "h3_spark_serving.env" LIVE_NODE = COMFY / "custom_nodes/h3_sage_attention.py" SEED = 26081343 EPS = 1e-5 class TorchOperations: Linear = torch.nn.Linear RMSNorm = torch.nn.RMSNorm def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int: if left.shape != right.shape or left.dtype != right.dtype: raise RuntimeError("byte_diff shape/dtype mismatch") return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item()) def bits_equal(left: torch.Tensor, right: torch.Tensor) -> bool: return byte_diff(left, right) == 0 def rel_mean_abs(current: torch.Tensor, previous: torch.Tensor) -> float: numerator = (current.float() - previous.float()).abs().mean() denominator = previous.float().abs().mean().clamp(min=1e-8) return float((numerator / denominator).item()) def make_row_ids(rows: int, n_mod: int, device: torch.device) -> torch.Tensor: return ( torch.arange(rows, device=device, dtype=torch.int64) .mul_(n_mod) .div_(rows, rounding_mode="floor") .to(torch.int32) .contiguous() ) def segments_from_rows(row_ids: torch.Tensor) -> list[tuple[int, int, int]]: rows = row_ids.tolist() out: list[tuple[int, int, int]] = [] start = 0 for i in range(1, len(rows) + 1): if i == len(rows) or rows[i] != rows[start]: out.append((start, i, int(rows[start]))) start = i return out def serving_cache_default() -> tuple[str, str]: env_mode = "missing" for line in SERVING_ENV.read_text().splitlines(): if line.startswith("H3_CROSS_STEP_CACHE="): env_mode = line.split("=", 1)[1].strip() node_text = LIVE_NODE.read_text() node_default = "missing" if '"default": "off"' in node_text and "cross_step_cache" in node_text: node_default = "off" return env_mode, node_default def build_stack(hidden, heads, head_dim, ffn, t_dim, device, dtype, n_blocks): import comfy.ldm.minimax.model as minimax_model blocks = torch.nn.ModuleList() for _ in range(n_blocks): block = minimax_model.DiTBlock( hidden=hidden, heads=heads, head_dim=head_dim, ffn=ffn, t_dim=t_dim, eps=EPS, qk_eps=EPS, dtype=dtype, device=device, operations=TorchOperations, ) blocks.append(block) blocks.eval() for parameter in blocks.parameters(): parameter.requires_grad_(False) embedder = minimax_model.TimeEmbedder( freq_dim=t_dim, hidden=hidden, out=t_dim, dtype=torch.float32, device=device, operations=TorchOperations, ) embedder.eval() for parameter in embedder.parameters(): parameter.requires_grad_(False) return blocks, embedder def modulation(block, t_emb): parts = block.adaln_proj(t_emb) return { "shift_msa": parts[0], "scale_msa": parts[1], "gate_msa": parts[2], "shift_mlp": parts[3], "scale_mlp": parts[4], "gate_mlp": parts[5], } def run_one(block, hidden, t_emb, segments, row_ids): # DiTBlock _mod_gate is in-place on the residual stream. x = hidden.clone() y = block( x, t_emb, segments, None, transformer_options={"minimax_h3_modulation_fusion": "stock"}, mod_row_ids=row_ids, ) return y, y - hidden def run_stack(blocks, hidden, t_emb, segments, row_ids): h = hidden.clone() first_out = None first_res = None for index, block in enumerate(blocks): out, res = run_one(block, h, t_emb, segments, row_ids) if index == 0: first_out = out first_res = res h = out return first_out, first_res, h, h - first_out def cheap_scale_copy(prev, scale): return (prev.float() * float(scale)).to(prev.dtype) def per_row_gate_scale(prev, gate_now, gate_prev, row_ids): g0 = gate_prev[row_ids].to(torch.float32) g1 = gate_now[row_ids].to(torch.float32) scale = g1 / g0.clamp(min=1e-8) return (prev.float() * scale).to(prev.dtype) def evaluate_pair(name, blocks, x0, t0, x1, t1, segments, row_ids): first0, res0, full0, tail0 = run_stack(blocks, x0, t0, segments, row_ids) first1, res1, full1, tail1 = run_stack(blocks, x1, t1, segments, row_ids) mod0 = modulation(blocks[0], t0) mod1 = modulation(blocks[0], t1) pred = { "input_equal": bits_equal(x0, x1), "t_emb_equal": bits_equal(t0, t1), "modulation_equal": all(bits_equal(mod0[k], mod1[k]) for k in mod0), "residual_delta_zero": bits_equal(res0, res1), "first_out_equal": bits_equal(first0, first1), "state_equal": bits_equal(x0, x1) and bits_equal(t0, t1), } pred["cache_threshold"] = rel_mean_abs(res1, res0) <= 0.08 gate_mean0 = float(mod0["gate_msa"].float().abs().mean().item()) gate_mean1 = float(mod1["gate_msa"].float().abs().mean().item()) gate_ratio = gate_mean1 / max(gate_mean0, 1e-8) recon_res = { "copy_prev": byte_diff(res1, res0), "gate_mean_scale": byte_diff(res1, cheap_scale_copy(res0, gate_ratio)), "per_row_gate_msa": byte_diff( res1, per_row_gate_scale(res0, mod1["gate_msa"], mod0["gate_msa"], row_ids) ), "add_input_delta": byte_diff(res1, (res0.float() + (x1 - x0).float()).to(res1.dtype)), } cache_tail = first1 + tail0 # Serving cache stores tail in activation dtype (bf16). # a + (b - a) is not an involution there. fp32_tail = first1.float() + (full0.float() - first0.float()) recon_full = { "cache_tail_add": byte_diff(full1, cache_tail), "cache_tail_add_fp32": byte_diff(full1, fp32_tail.to(full1.dtype)), "copy_prev_full": byte_diff(full1, full0), "first_out_only": byte_diff(full1, first1), "bf16_involution": byte_diff(full0, first0 + (full0 - first0)), "fp32_involution": byte_diff(full0, (first0.float() + (full0.float() - first0.float())).to(full0.dtype)), } # Sound hit: predicate true implies the named reconstruction is byte-exact. sound = { "state_equal_copy_residual": (not pred["state_equal"]) or recon_res["copy_prev"] == 0, "state_equal_cache_tail": (not pred["state_equal"]) or recon_full["cache_tail_add"] == 0, "residual_delta_zero_copy": (not pred["residual_delta_zero"]) or recon_res["copy_prev"] == 0, "t_emb_only_copy": (not pred["t_emb_equal"]) or recon_res["copy_prev"] == 0, "modulation_only_copy": (not pred["modulation_equal"]) or recon_res["copy_prev"] == 0, "input_only_copy": (not pred["input_equal"]) or recon_res["copy_prev"] == 0, "residual_delta_zero_cache_tail": ( (not pred["residual_delta_zero"]) or recon_full["cache_tail_add"] == 0 ), "first_out_and_t_cache_tail": ( (not (pred["first_out_equal"] and pred["t_emb_equal"])) or recon_full["cache_tail_add"] == 0 ), "first_out_and_t_fp32_tail": ( (not (pred["first_out_equal"] and pred["t_emb_equal"])) or recon_full["cache_tail_add_fp32"] == 0 ), "cache_threshold_cache_tail": ( (not pred["cache_threshold"]) or recon_full["cache_tail_add"] == 0 ), } return { "pair": name, "predicates": pred, "rel_mean_abs": rel_mean_abs(res1, res0), "rel_mean_abs_full": rel_mean_abs(full1, full0), "gate_msa_mean_ratio": gate_ratio, "residual_recon_byte_diff": recon_res, "full_recon_byte_diff": recon_full, "sound": sound, "residual_elems": int(res1.numel()), "full_elems": int(full1.numel()), } def run_geometry(label, rows, hidden, heads, head_dim, ffn, t_dim, n_mod, device): dtype = torch.bfloat16 blocks, embedder = build_stack(hidden, heads, head_dim, ffn, t_dim, device, dtype, 2) row_ids = make_row_ids(rows, n_mod, device) segments = segments_from_rows(row_ids) n_t = max(1, n_mod // 3) def t_emb_from_sigma(video_sigma: float) -> torch.Tensor: # Seven packed timestep classes: cond pinned, audio shifted, video live. vals = [] for i in range(n_t): if i == 0: vals.append(0.999) elif i == 1: vals.append(min(1.0, video_sigma * 0.4 + 0.05)) else: vals.append(video_sigma) t = torch.tensor(vals, dtype=torch.float32, device=device) return embedder(t).to(dtype).contiguous() x0 = torch.randn(rows, hidden, device=device, dtype=dtype) t0 = t_emb_from_sigma(0.85) t1 = t_emb_from_sigma(0.70) x_ulp = x0.clone() x_ulp[0, 0] = torch.nextafter( x0[0, 0], torch.tensor(float("inf"), device=device, dtype=x0.dtype) ) pairs = { "replay": (x0, t0, x0.clone(), t0.clone()), "same_t_moved_x": (x0, t0, (x0 + 0.05 * torch.randn_like(x0)).contiguous(), t0.clone()), "same_x_moved_t": (x0, t0, x0.clone(), t1), "moved_both": (x0, t0, (x0 + 0.05 * torch.randn_like(x0)).contiguous(), t1), "tiny_x_same_t": (x0, t0, x_ulp, t0.clone()), "tiny_x_moved_t": (x0, t0, x_ulp.clone(), t1), } reports = [ evaluate_pair(name, blocks, a, ta, b, tb, segments, row_ids) for name, (a, ta, b, tb) in pairs.items() ] # Cheap predicates that never inspect residual(t). cheap_false_hits = [] for report in reports: if report["pair"] == "replay": continue if report["predicates"]["t_emb_equal"] and report["residual_recon_byte_diff"]["copy_prev"] != 0: cheap_false_hits.append((report["pair"], "t_emb_equal")) if report["predicates"]["modulation_equal"] and report["residual_recon_byte_diff"]["copy_prev"] != 0: cheap_false_hits.append((report["pair"], "modulation_equal")) if report["predicates"]["cache_threshold"] and report["full_recon_byte_diff"]["cache_tail_add"] != 0: cheap_false_hits.append((report["pair"], "cache_threshold")) replay = next(r for r in reports if r["pair"] == "replay") moved = next(r for r in reports if r["pair"] == "moved_both") same_t = next(r for r in reports if r["pair"] == "same_t_moved_x") same_x = next(r for r in reports if r["pair"] == "same_x_moved_t") tiny = next(r for r in reports if r["pair"] == "tiny_x_same_t") return { "label": label, "rows": rows, "hidden": hidden, "heads": heads, "blocks": 2, "mod_rows": n_mod, "pairs": reports, "replay_residual_exact": replay["residual_recon_byte_diff"]["copy_prev"] == 0, "replay_full_copy_exact": replay["full_recon_byte_diff"]["copy_prev_full"] == 0, "replay_bf16_tail_exact": replay["full_recon_byte_diff"]["cache_tail_add"] == 0, "replay_fp32_tail_exact": replay["full_recon_byte_diff"]["cache_tail_add_fp32"] == 0, "replay_bf16_involution_diff": replay["full_recon_byte_diff"]["bf16_involution"], "replay_fp32_involution_diff": replay["full_recon_byte_diff"]["fp32_involution"], "cheap_predicate_false_hits": [ {"pair": pair, "predicate": pred} for pair, pred in cheap_false_hits ], "nontrivial_residual_delta_zero": any( r["predicates"]["residual_delta_zero"] and r["pair"] != "replay" for r in reports ), "tiny_x_actually_moved": not tiny["predicates"]["input_equal"], "moved_both_copy_diff": moved["residual_recon_byte_diff"]["copy_prev"], "same_t_moved_x_copy_diff": same_t["residual_recon_byte_diff"]["copy_prev"], "same_x_moved_t_copy_diff": same_x["residual_recon_byte_diff"]["copy_prev"], "threshold_is_exact": all(r["sound"]["cache_threshold_cache_tail"] for r in reports) and any(r["predicates"]["cache_threshold"] and r["pair"] != "replay" for r in reports), } def main() -> int: torch.set_grad_enabled(False) if str(COMFY) not in sys.path: sys.path.insert(0, str(COMFY)) if not torch.cuda.is_available(): raise RuntimeError("CUDA required") device = torch.device("cuda") torch.manual_seed(SEED) torch.cuda.manual_seed_all(SEED) env_mode, node_default = serving_cache_default() geometries = [ run_geometry("small_s129", 129, 64, 4, 16, 128, 64, 21, device), run_geometry("small_s20423", 20423, 64, 4, 16, 128, 64, 21, device), run_geometry("prod_width_s129", 129, 5376, 56, 128, 14336, 2688, 21, device), ] all_replay_res = all(g["replay_residual_exact"] and g["replay_full_copy_exact"] for g in geometries) all_bf16_broken = all(not g["replay_bf16_tail_exact"] and g["replay_bf16_involution_diff"] > 0 for g in geometries) all_fp32_involution = all(g["replay_fp32_involution_diff"] == 0 for g in geometries) any_cheap_false = any( any(hit["predicate"] in ("t_emb_equal", "modulation_equal", "cache_threshold") for hit in g["cheap_predicate_false_hits"]) for g in geometries ) any_moved_copy_zero = any(g["moved_both_copy_diff"] == 0 for g in geometries) identity = ( "H3 block-0 residual(t) is F(x_t, t_emb_t), not a function of " "residual(t-1) plus timestep, modulation, or a cheap residual delta. " "Copy / gate-scale / input-delta reconstructions are not byte-exact " "once x or t_emb moves. residual_delta==0 is tautological and requires " "residual(t). Serving cache tail add first+(full-first) is not " "byte-exact even on replay: bf16 a+(b-a) is not an involution. " "fp32 tail add is involutive on replay and still misses once x or t " "moves. The 0.08 threshold is not an exact predicate." ) verdict = ( all_replay_res and all_bf16_broken and all_fp32_involution and any_cheap_false and not any_moved_copy_zero and env_mode == "off" and node_default == "off" ) receipt = { "identity": identity, "device": torch.cuda.get_device_name(0), "seed": SEED, "dtype": "bfloat16", "cache_source": str(CACHE_SRC), "serving_default": { "H3_CROSS_STEP_CACHE": env_mode, "node_cross_step_cache": node_default, "unchanged": env_mode == "off" and node_default == "off", }, "geometries": [ { "label": g["label"], "rows": g["rows"], "hidden": g["hidden"], "replay_residual_exact": g["replay_residual_exact"], "replay_full_copy_exact": g["replay_full_copy_exact"], "replay_bf16_tail_exact": g["replay_bf16_tail_exact"], "replay_fp32_tail_exact": g["replay_fp32_tail_exact"], "replay_bf16_involution_diff": g["replay_bf16_involution_diff"], "replay_fp32_involution_diff": g["replay_fp32_involution_diff"], "tiny_x_actually_moved": g["tiny_x_actually_moved"], "moved_both_copy_diff": g["moved_both_copy_diff"], "same_t_moved_x_copy_diff": g["same_t_moved_x_copy_diff"], "same_x_moved_t_copy_diff": g["same_x_moved_t_copy_diff"], "cheap_predicate_false_hits": g["cheap_predicate_false_hits"], "nontrivial_residual_delta_zero": g["nontrivial_residual_delta_zero"], "pairs": g["pairs"], } for g in geometries ], "algebra": { "block0": "residual = gate_msa*Attn(AdaLN(x,t)) + gate_mlp*MLP(AdaLN(x+attn_res,t))", "depends_on": ["x", "t_emb"], "not_a_function_of": [ "residual(t-1) + timestep", "residual(t-1) + modulation", "residual(t-1) * gate_ratio", "residual(t-1) + (x_t - x_{t-1})", ], "exact_fail_closed_predicate": "x bits equal AND t_emb bits equal (duplicate forward); then copy the previous residual/full, do not add a bf16 tail", "residual_delta_zero": "tautology; requires residual(t)", "serving_bf16_tail": "first+(full-first) is not involutive; not byte-exact on replay", "fp32_tail": "involutive on replay; still not a skip once x or t_emb moves", "serving_threshold_0p08": "approximation; not byte-exact", }, "decision": "no_exact_skip", "promoted": False, "serving_patched": False, "default_changed": False, "pass": bool(verdict), } out = HERE / "gate_skip_algebra.json" out.write_text(json.dumps(receipt, indent=2) + "\n") print(json.dumps({"pass": receipt["pass"], "receipt": str(out), "decision": receipt["decision"]}, indent=2)) return 0 if receipt["pass"] else 1 if __name__ == "__main__": raise SystemExit(main())