From b69903b7e6c05736b11a265f9a6d1aaba2d65a84 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 01:30:32 +0700 Subject: [PATCH] Match H3 DiT backbone trace --- PARITY.md | 17 ++++++------- src/h3_blackwell_runtime/attention.py | 27 +++++++++++++++++--- tools/patch_comfy_h3_final_capture.py | 36 +++++++++++++++++++++++++++ tools/trace_block0_exact.py | 21 +++++++++++----- 4 files changed, 82 insertions(+), 19 deletions(-) create mode 100644 tools/patch_comfy_h3_final_capture.py diff --git a/PARITY.md b/PARITY.md index c763a34..8f090f4 100644 --- a/PARITY.md +++ b/PARITY.md @@ -183,6 +183,7 @@ gate and is exact. | Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags | | Token refiner | Yes | Bit-exact from captured 5376-wide refiner input through both blocks and final RMSNorm; Qwen-to-refiner projection boundary is still not separately captured | | Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows | +| H3 DiT backbone | Yes | Bit-exact through all 50 blocks from the coherent assembled FL2VA input; requires the standalone Comfy Kitchen fused Q/K RMSNorm + split-half RoPE operator | | H3 DiT | Yes | Strict all-block numeric parity not achieved | | Beta/RES sampler | Yes | Exact sigma/state/update parity not demonstrated; direct preview has wrong initial audio state | | Video VAE decoder | Yes | No direct-versus-Comfy same-latent pixel comparison yet | @@ -194,21 +195,19 @@ gate and is exact. Only these are outstanding. Do not recapture or revisit rows marked complete unless the checkpoint, Comfy version, prompt, or backend changes. -1. **H3 first-divergence repair.** Reuse the existing block trace. Start at - block 0 sublayers, repair the first mismatch, then use Comfy block output as - the next direct input to distinguish local error from accumulation. Existing - evidence says block 0 is close but accumulation is unresolved. -2. **Final rows and video latent.** Compare final packed video rows and +1. **Final rows and video latent.** Final DiT hidden state, final AdaLN, and + final RMSNorm are exact. Repair the remaining final target-segment + modulation/FP32-head boundary (`video_hidden` max `0.00598395`) before `unpatchify_video` output with Comfy before entering the sampler. -3. **Sampler replay using Comfy state.** Feed `initial.pt` and captured +2. **Sampler replay using Comfy state.** Feed `initial.pt` and captured `sigmas` directly to the sampler. Compare video and audio denoised/update states for all 12 saved steps. Do not test seed equality until this passes. -4. **VAE pixel gate.** Decode the identical captured final video latent in +3. **VAE pixel gate.** Decode the identical captured final video latent in direct and Comfy and compare pixels before ffmpeg. This decides whether any residual grid comes from decoder behavior. -5. **End-to-end FL2VA preview.** Generate the reference dragon only after +4. **End-to-end FL2VA preview.** Generate the reference dragon only after gates 1-8 pass. Compare its raw final latent first, then video. -6. **Feature/performance work.** Only then add audio, Ref2VA/reference paths, +5. **Feature/performance work.** Only then add audio, Ref2VA/reference paths, Sage3, CUDA graphs, and multi-GPU execution. ## Existing Tools And Their Intended Gate diff --git a/src/h3_blackwell_runtime/attention.py b/src/h3_blackwell_runtime/attention.py index 19ed6d7..84de387 100644 --- a/src/h3_blackwell_runtime/attention.py +++ b/src/h3_blackwell_runtime/attention.py @@ -56,6 +56,24 @@ def apply_split_half_rope(x: torch.Tensor, rotation: torch.Tensor) -> torch.Tens return torch.cat((pair[..., 0], pair[..., 1], x[..., rotated_width:]), dim=-1) +def rms_rope_split_half_( + q: torch.Tensor, + k: torch.Tensor, + rotation: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run Comfy Kitchen's standalone fused H3 Q/K normalization and RoPE.""" + import comfy_kitchen # Registers the independent CUDA extension operators. + + del comfy_kitchen + torch.ops.comfy_kitchen.rms_rope_split_half_( + q, k, rotation, q_weight, k_weight, eps, rotation.shape[-3] * 2 + ) + return q, k + + class H3SageAttention(nn.Module): """One MiniMax H3 attention module, independent of ComfyUI and Raylight.""" @@ -101,12 +119,13 @@ class H3SageAttention(nn.Module): sequence = x.shape[0] inner = self.heads * self.head_dim q, k, v = self.qkv_proj(x).split(inner, dim=-1) - q = rms_norm(q.view(1, sequence, self.heads, self.head_dim), self.q_norm_weight, self.eps) - k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps) + q = q.view(1, sequence, self.heads, self.head_dim) + k = k.view(1, sequence, self.heads, self.head_dim) v = v.view(1, sequence, self.heads, self.head_dim) - q = apply_split_half_rope(q, rope_rotation).transpose(1, 2).contiguous() - k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous() + q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() v = v.transpose(1, 2).contiguous() out = run_attention(q, k, v, backend=self.backend, is_causal=False) diff --git a/tools/patch_comfy_h3_final_capture.py b/tools/patch_comfy_h3_final_capture.py new file mode 100644 index 0000000..aec2e74 --- /dev/null +++ b/tools/patch_comfy_h3_final_capture.py @@ -0,0 +1,36 @@ +"""Capture H3 final-layer intermediates from one reference inference.""" + +from pathlib import Path + + +model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py") +source = model.read_text(encoding="utf-8") +old = ( + " shift, scale = self.adaln_proj(t_emb)\n" + " va, vb, vrow = video_seg\n" + " aa, ab, arow = audio_seg\n" + " hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n" + " ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n" + " return self.video_out(hv), self.audio_out(ha)\n" +) +new = ( + " shift, scale = self.adaln_proj(t_emb)\n" + " va, vb, vrow = video_seg\n" + " aa, ab, arow = audio_seg\n" + " norm_v = self.norm(x[va:vb])\n" + " norm_a = self.norm(x[aa:ab])\n" + " hv = (norm_v * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n" + " ha = (norm_a * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n" + " video = self.video_out(hv)\n" + " audio = self.audio_out(ha)\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if H3_CAPTURE_ACTIVE else None\n" + " if capture_dir:\n" + " torch.save({\"hidden\": x.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu()}, os.path.join(capture_dir, \"final.pt\"))\n" + " return video, audio\n" +) +if source.count(old) == 1: + source = source.replace(old, new) +elif new not in source: + raise RuntimeError("Unable to locate H3 FinalLayer.forward.") +model.write_text(source, encoding="utf-8") +print("Applied H3 final-layer capture patch.") diff --git a/tools/trace_block0_exact.py b/tools/trace_block0_exact.py index 4d1aea9..8beeb69 100644 --- a/tools/trace_block0_exact.py +++ b/tools/trace_block0_exact.py @@ -4,7 +4,7 @@ import argparse import torch -from h3_blackwell_runtime.attention import apply_split_half_rope, rms_norm +from h3_blackwell_runtime.attention import rms_norm, rms_rope_split_half_ from h3_blackwell_runtime.block import gate_segments, modulate_segments from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.denoiser import H3PackedDenoiser @@ -30,8 +30,17 @@ with torch.inference_mode(): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln(inputs["timesteps"]) norm1 = modulate_segments(rms_norm(inputs["hidden"], block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"]) q, k, v = block.attention.qkv_proj(norm1).split(7168, dim=-1) - q_prepared = apply_split_half_rope(rms_norm(q.view(1, -1, 56, 128), block.attention.q_norm_weight, 1e-5), rotation).transpose(1, 2).contiguous() - k_prepared = apply_split_half_rope(rms_norm(k.view(1, -1, 56, 128), block.attention.k_norm_weight, 1e-5), rotation).transpose(1, 2).contiguous() + raw_q, raw_k, raw_v = q.clone(), k.clone(), v.clone() + q_prepared, k_prepared = rms_rope_split_half_( + q.view(1, -1, 56, 128), + k.view(1, -1, 56, 128), + rotation, + block.attention.q_norm_weight, + block.attention.k_norm_weight, + 1e-5, + ) + q_prepared = q_prepared.transpose(1, 2).contiguous() + k_prepared = k_prepared.transpose(1, 2).contiguous() v_prepared = v.view(1, -1, 56, 128).transpose(1, 2).contiguous() from sageattention import sageattn attention = block.attention.out_proj(sageattn(q_prepared, k_prepared, v_prepared, is_causal=False, tensor_layout="HND", smooth_k=False).transpose(1, 2).reshape(norm1.shape[0], -1)) @@ -42,9 +51,9 @@ with torch.inference_mode(): for name, actual, expected in ( ("norm1", norm1, capture["norm1"]), - ("raw_q", q, capture["qkv_raw"]["q"]), - ("raw_k", k, capture["qkv_raw"]["k"]), - ("raw_v", v, capture["qkv_raw"]["v"]), + ("raw_q", raw_q, capture["qkv_raw"]["q"]), + ("raw_k", raw_k, capture["qkv_raw"]["k"]), + ("raw_v", raw_v, capture["qkv_raw"]["v"]), ("q", q_prepared, capture["qkv_prepared"]["q"]), ("k", k_prepared, capture["qkv_prepared"]["k"]), ("v", v_prepared, capture["qkv_prepared"]["v"]),