Match H3 assembled FL2VA input

PARITY.md

@@ -182,7 +182,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 | No keyframe/reference condition rows; all preview work is text-only FL2VA |
+| Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows |
 | 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,24 +194,21 @@ 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 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.
-2. **H3 first-divergence repair.** Reuse the existing block trace. Start at
+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.
-3. **Final rows and video latent.** Compare final packed video rows and
+2. **Final rows and video latent.** Compare final packed video rows and
    `unpatchify_video` output with Comfy before entering the sampler.
-4. **Sampler replay using Comfy state.** Feed `initial.pt` and captured
+3. **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.
-5. **VAE pixel gate.** Decode the identical captured final video latent in
+4. **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.
-6. **End-to-end FL2VA preview.** Generate the reference dragon only after
+5. **End-to-end FL2VA preview.** Generate the reference dragon only after
    gates 1-8 pass. Compare its raw final latent first, then video.
-7. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
+6. **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/packing.py

@@ -20,7 +20,7 @@ def patchify_video(latent: torch.Tensor) -> torch.Tensor:
 
 
 def pack_audio(latent: torch.Tensor) -> torch.Tensor:
-    return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1])
+    return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1)
 
 
 def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
@@ -58,10 +58,10 @@ class H3PromptPacker:
     """Build `[text | audio | video]` tokens for prompt-only H3 T2V."""
 
     def __init__(self, checkpoint):
-        self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.float32)
-        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.float32)
-        self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.float32)
-        self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.float32)
+        self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
+        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
+        self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
+        self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
         self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
         self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)
 
@@ -84,7 +84,7 @@ class H3PromptPacker:
         row = {value: index for index, value in enumerate(unique_times)}
         video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
         times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)
-        positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length) + audio.shape[-1])))
+        positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length))))
         block_video_segment = (text_length + audio_length, hidden.shape[0], video_row)
         block_audio_segment = (text_length, text_length + audio_length, audio_row + 2)
         final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time])

tools/compare_h3_assembled_input.py

@@ -0,0 +1,49 @@
+"""Compare direct prompt-only FL2VA packing with an immutable Comfy DiT input."""
+
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.packing import H3PromptPacker, pack_audio, patchify_video
+from h3_blackwell_runtime.sampler import _audio_sigma
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture", type=Path, required=True)
+parser.add_argument("--refiner-trace", type=Path, required=True)
+parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+args = parser.parse_args()
+
+captured = torch.load(args.capture, map_location="cuda", weights_only=False)
+text = torch.load(args.refiner_trace / "refiner_output.pt", map_location="cuda", weights_only=False)
+if text.ndim == 2:
+    text = text.unsqueeze(0)
+checkpoint = H3Checkpoint(args.checkpoint)
+packer = H3PromptPacker(checkpoint)
+
+video_shape = (1, 24, 7, 12, 20)
+audio_shape = (1, 32, 2, 37)
+video = captured["video_x"].to("cuda")
+audio = captured["audio_x"].to("cuda")
+sigma = 1 - captured["timesteps"].min()
+
+with torch.inference_mode():
+    hidden, times, segments, positions, _, _ = packer(text, video, audio, float(sigma))
+    text_rows = text[0].to(torch.bfloat16)
+    audio_rows = functional.linear(pack_audio(audio).float(), packer.audio_weight, packer.audio_bias).to(torch.bfloat16)
+    video_rows = functional.linear(patchify_video(video).float(), packer.video_weight, packer.video_bias).to(torch.bfloat16)
+
+for name, actual, expected in (
+    ("text", text_rows, captured["hidden"][:17]),
+    ("audio", audio_rows, captured["hidden"][17:91]),
+    ("video", video_rows, captured["hidden"][91:]),
+    ("hidden", hidden, captured["hidden"]),
+    ("times", times, captured["timesteps"]),
+    ("positions", positions, captured["position_ids"]),
+):
+    delta = (actual.float() - expected.to(actual.device).float()).abs()
+    print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+print(f"segments direct={segments} comfy={captured['segments']}")

tools/compare_h3_loaded_patch_projection.py

@@ -0,0 +1,30 @@
+"""Compare direct patch GEMMs with saved outputs from loaded Comfy modules."""
+
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture", type=Path, required=True)
+parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+args = parser.parse_args()
+
+captured = torch.load(args.capture, map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint(args.checkpoint)
+for name, data in captured.items():
+    rows = data["rows"].to("cuda")
+    weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
+    bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.float32)
+    with torch.inference_mode():
+        direct = functional.linear(rows, weight, bias)
+    row_delta = (rows.float() - data["rows"].to(rows.device).float()).abs()
+    output_delta = (direct.float() - data["output"].to(direct.device).float()).abs()
+    forward_delta = (data["forward_cast"].float() - data["output"].float()).abs()
+    print(f"{name}.rows stride={tuple(rows.stride())} comfy_stride={data['rows_stride']} max_abs={row_delta.max().item():.6g}")
+    print(f"{name}.direct max_abs={output_delta.max().item():.6g} mean_abs={output_delta.mean().item():.6g}")
+    print(f"{name}.comfy_forward_cast max_abs={forward_delta.max().item():.6g} mean_abs={forward_delta.mean().item():.6g}")

tools/compare_h3_patch_projections.py

@@ -0,0 +1,37 @@
+"""Compare direct FP32 patch GEMMs under Comfy's captured matmul policy."""
+
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture", type=Path, required=True)
+parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+args = parser.parse_args()
+
+captured = torch.load(args.capture, map_location="cuda", weights_only=False)
+checkpoint = H3Checkpoint(args.checkpoint)
+previous_tf32 = torch.backends.cuda.matmul.allow_tf32
+previous_precision = torch.get_float32_matmul_precision()
+torch.backends.cuda.matmul.allow_tf32 = captured["matmul"]["allow_tf32"]
+torch.set_float32_matmul_precision(captured["matmul"]["precision"])
+try:
+    for name in ("video", "audio"):
+        rows = captured[f"{name}_rows"].to("cuda")
+        weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
+        bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
+        output = functional.linear(rows, weight, bias)
+        for label, actual, expected in (("weight", weight, captured[f"{name}_weight"]), ("bias", bias, captured[f"{name}_bias"]), ("fp32", output, captured[f"{name}_embed_fp32"])):
+            delta = (actual.float() - expected.to(actual.device).float()).abs()
+            print(f"{name}.{label} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+        bf16_delta = (output.to(torch.bfloat16).float() - captured[f"{name}_embed_fp32"].to(torch.bfloat16).float()).abs()
+        print(f"{name}.bf16 max_abs={bf16_delta.max().item():.6g} mean_abs={bf16_delta.mean().item():.6g}")
+    print(f"matmul={captured['matmul']}")
+finally:
+    torch.backends.cuda.matmul.allow_tf32 = previous_tf32
+    torch.set_float32_matmul_precision(previous_precision)

tools/patch_comfy_h3_capture.py

@@ -24,7 +24,7 @@ replace_once(
     "        if capture_dir and not H3_CAPTURE_ACTIVE:\n"
     "            H3_CAPTURE_ACTIVE = True\n"
     "            os.makedirs(capture_dir, exist_ok=True)\n"
-    "            torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments}, os.path.join(capture_dir, \"input.pt\"))\n\n"
+    "            torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments, \"video_x\": video_x.detach().cpu(), \"audio_x\": audio_x.detach().cpu()}, os.path.join(capture_dir, \"input.pt\"))\n\n"
     "        # blocks\n        patches_replace = transformer_options.get(\"patches_replace\", {})\n",
 )
 replace_once(

tools/patch_comfy_h3_loaded_patch_projection_probe.py

@@ -0,0 +1,33 @@
+"""Probe loaded Comfy H3 patch projections on their captured FP32 rows."""
+
+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 = (
+    "        video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
+    "        audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
+)
+new = (
+    "        video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
+    "        audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "        if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\")):\n"
+    "            def probe(layer, rows, output):\n"
+    "                weight, bias, state = comfy.ops.cast_bias_weight(layer, rows, offloadable=True)\n"
+    "                try:\n"
+    "                    cast_output = torch.nn.functional.linear(rows, weight, bias)\n"
+    "                    return {\"rows\": rows.detach().cpu(), \"output\": output.detach().cpu(), \"forward_cast\": layer.forward_comfy_cast_weights(rows).detach().cpu(), \"effective_weight\": weight.detach().cpu(), \"effective_bias\": None if bias is None else bias.detach().cpu(), \"effective_weight_dtype\": str(weight.dtype), \"effective_bias_dtype\": None if bias is None else str(bias.dtype), \"cast_output\": cast_output.detach().cpu(), \"rows_dtype\": str(rows.dtype), \"rows_stride\": tuple(rows.stride()), \"weight_dtype\": str(layer.weight.dtype), \"weight_stride\": tuple(layer.weight.stride()), \"bias_dtype\": None if layer.bias is None else str(layer.bias.dtype), \"bias_stride\": None if layer.bias is None else tuple(layer.bias.stride()), \"autocast\": torch.is_autocast_enabled(), \"force_cast\": getattr(layer, \"comfy_force_cast_weights\", False), \"weight_functions\": len(layer.weight_function), \"bias_functions\": len(layer.bias_function)}\n"
+    "                finally:\n"
+    "                    comfy.ops.uncast_bias_weight(layer, weight, bias, state)\n"
+    "            torch.save({\"video\": probe(self.video_patch_proj, all_video_rows, video_embed_fp32), \"audio\": probe(self.audio_patch_proj, all_audio_rows, audio_embed_fp32)}, os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\"))\n"
+    "        video_embed = video_embed_fp32.to(dtype)\n"
+    "        audio_embed = audio_embed_fp32.to(dtype)\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate H3 patch-projection calls.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied loaded H3 patch-projection probe.")

tools/patch_comfy_h3_patch_projection_probe.py

@@ -0,0 +1,24 @@
+"""Capture Comfy H3 FP32 patch-projection inputs, outputs, and matmul policy."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = path.read_text(encoding="utf-8")
+old = (
+    "        video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
+    "        audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
+)
+new = (
+    "        video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
+    "        audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
+    "        if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_patch_projections.pt\")):\n"
+    "            torch.save({\"video_rows\": all_video_rows.detach().cpu(), \"audio_rows\": all_audio_rows.detach().cpu(), \"video_embed_fp32\": video_embed_fp32.detach().cpu(), \"audio_embed_fp32\": audio_embed_fp32.detach().cpu(), \"video_weight\": self.video_patch_proj.weight.detach().cpu(), \"video_bias\": self.video_patch_proj.bias.detach().cpu(), \"audio_weight\": self.audio_patch_proj.weight.detach().cpu(), \"audio_bias\": self.audio_patch_proj.bias.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"h3_patch_projections.pt\"))\n"
+    "        video_embed = video_embed_fp32.to(dtype)\n"
+    "        audio_embed = audio_embed_fp32.to(dtype)\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate H3 patch-projection calls.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied H3 patch-projection probe.")