Verify token refiner trace

PARITY.md

@@ -181,7 +181,7 @@ gate and is exact.
 | Component | Implemented | Known limitation |
 | --- | --- | --- |
 | Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags |
-| Token refiner | Yes | Architecture ported; no completed direct-versus-Comfy refiner tensor comparison |
+| 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 | No keyframe/reference condition rows; all preview work is text-only FL2VA |
 | 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 |
@@ -194,27 +194,24 @@ 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. **Token-refiner capture and comparison.** Capture exact Comfy pre-refiner,
+1. **H3 assembled input comparison.** Use the existing FL2VA input capture to
-   block 0, block 1, and final refined states once, then compare direct using
-   the matched Qwen layer-50 tensor.
-2. **H3 assembled input comparison.** Use the existing FL2VA input capture to
    compare text rows, audio/video rows, positions, times, and segments. This
    isolates packing from DiT error.
-3. **H3 first-divergence repair.** Reuse the existing block trace. Start at
+2. **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.
-4. **Final rows and video latent.** Compare final packed video rows and
+3. **Final rows and video latent.** Compare final packed video rows and
    `unpatchify_video` output with Comfy before entering the sampler.
-5. **Sampler replay using Comfy state.** Feed `initial.pt` and captured
+4. **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.
-6. **VAE pixel gate.** Decode the identical captured final video latent in
+5. **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.
-7. **End-to-end FL2VA preview.** Generate the reference dragon only after
+6. **End-to-end FL2VA preview.** Generate the reference dragon only after
    gates 1-8 pass. Compare its raw final latent first, then video.
-8. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
+7. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
     Sage3, CUDA graphs, and multi-GPU execution.
 
 ## Existing Tools And Their Intended Gate

tools/compare_refiner_trace.py

@@ -0,0 +1,29 @@
+"""Replay the direct token refiner against a Comfy boundary trace."""
+
+import argparse
+from pathlib import Path
+
+import torch
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.token_refiner import H3TokenRefiner
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--trace-dir", type=Path, required=True)
+parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+args = parser.parse_args()
+
+refiner = H3TokenRefiner(H3Checkpoint(args.checkpoint), attention_backend="sage2").eval()
+hidden = torch.load(args.trace_dir / "refiner_input.pt", map_location="cuda", weights_only=False)
+
+with torch.inference_mode():
+    for index, block in enumerate(refiner.blocks):
+        hidden = block(hidden)
+        expected = torch.load(args.trace_dir / f"refiner_block{index}.pt", map_location="cuda", weights_only=False)
+        delta = (hidden.float() - expected.float()).abs()
+        print(f"block={index} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+    output = refiner.final_norm
+    output = torch.nn.functional.rms_norm(hidden, output.shape, weight=output.to(hidden), eps=1e-5).unsqueeze(0)
+    expected = torch.load(args.trace_dir / "refiner_output.pt", map_location="cuda", weights_only=False)
+    delta = (output.float() - expected.float()).abs()
+    print(f"output max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/patch_comfy_h3_refiner_trace.py

@@ -0,0 +1,34 @@
+"""Capture Comfy H3 token-refiner boundaries once for direct parity replay."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = path.read_text(encoding="utf-8")
+if "import os\n" not in source:
+    source = source.replace("import math\n", "import math\nimport os\n", 1)
+
+old = (
+    "    def forward(self, x, transformer_options={}):\n"
+    "        for block in self.blocks:\n"
+    "            x = block(x, transformer_options=transformer_options)\n"
+    "        return self.final_norm(x)\n"
+)
+new = (
+    "    def forward(self, x, transformer_options={}):\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "        if capture_dir:\n"
+    "            torch.save(x.detach().cpu(), os.path.join(capture_dir, \"refiner_input.pt\"))\n"
+    "        for index, block in enumerate(self.blocks):\n"
+    "            x = block(x, transformer_options=transformer_options)\n"
+    "            if capture_dir:\n"
+    "                torch.save(x.detach().cpu(), os.path.join(capture_dir, f\"refiner_block{index}.pt\"))\n"
+    "        output = self.final_norm(x)\n"
+    "        if capture_dir:\n"
+    "            torch.save(output.detach().cpu(), os.path.join(capture_dir, \"refiner_output.pt\"))\n"
+    "        return output\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate TokenRefiner.forward.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied H3 token-refiner trace patch.")