Match H3 DiT backbone trace

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
+1. **Final rows and video latent.** Final DiT hidden state, final AdaLN, and
-   block 0 sublayers, repair the first mismatch, then use Comfy block output as
+   final RMSNorm are exact. Repair the remaining final target-segment
-   the next direct input to distinguish local error from accumulation. Existing
+   modulation/FP32-head boundary (`video_hidden` max `0.00598395`) before
-   evidence says block 0 is close but accumulation is unresolved.
-2. **Final rows and video latent.** Compare final packed video rows and
    `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

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)
+        q = q.view(1, sequence, self.heads, self.head_dim)
-        k = rms_norm(k.view(1, sequence, self.heads, self.head_dim), self.k_norm_weight, self.eps)
+        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()
+        q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps)
-        k = apply_split_half_rope(k, rope_rotation).transpose(1, 2).contiguous()
+        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)

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.")

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()
+    raw_q, raw_k, raw_v = q.clone(), k.clone(), v.clone()
-    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()
+    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_q", raw_q, capture["qkv_raw"]["q"]),
-    ("raw_k", k, capture["qkv_raw"]["k"]),
+    ("raw_k", raw_k, capture["qkv_raw"]["k"]),
-    ("raw_v", v, capture["qkv_raw"]["v"]),
+    ("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"]),