Trace H3 final AdaLN parity

PARITY.md

@@ -0,0 +1,227 @@
+# H3 Parity Evidence Ledger
+
+ComfyUI is the oracle. This ledger records completed work and measured facts so
+we do not repeat investigations. A source review is never recorded as a passed
+numeric gate.
+
+## Reference Runs And Immutable Artifacts
+
+| Asset | Location on Spark | Contents | Status |
+| --- | --- | --- | --- |
+| Qwen trace | `/tmp/fl2va-qwen-trace` | `qwen_input_embeds.pt`, 50 decoder outputs in `qwen_layers/00.pt` through `49.pt`, and `qwen_layer50.pt` | Complete, immutable Comfy capture |
+| Qwen layer-0 trace | `/tmp/fl2va-qwen0-capture` | `qwen0_norm1.pt`, attention, post-attention, norm2, MLP, output | Complete, immutable Comfy capture |
+| H3 full trace | `/tmp/fl2va-full-capture` | assembled DiT input/output and `blocks/00.pt` through `49.pt` | Complete, immutable Comfy capture |
+| Working-copy capture mirror | `/home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/capture` | Qwen layer-0 and layer-50 assets plus H3 input/output | Present; do not treat as immutable if the working checkout is modified |
+| Sampler reference | `/home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts/fl2va-sampler-reference` | `initial.pt`, exact `sigmas`, `step_00.pt` through `step_11.pt` | Complete, immutable Comfy capture |
+
+All captures above use the established prompt-only FL2VA reference:
+
+```text
+prompt: A brass-and-paper dragon flies above a rain-washed old city at blue hour.
+size: 320x192, length: 22 frames
+seed: 440204
+sampler: beta / res_multistep / 12 steps
+Comfy attention baseline: --use-sage-attention (SageAttention2)
+```
+
+This reference is **not** the 72-frame cat preview. Do not overwrite it to
+diagnose preview output.
+
+## Completed Findings
+
+| Boundary | Evidence | Result | Conclusion |
+| --- | --- | --- | --- |
+| Direct runtime build | GB10 CUDA 13 image, Torch `2.9.1+cu130`, committed SageAttention3 wheel | Repeated direct runs complete | Environment is operational; PyTorch's `sm_121` warning is stale metadata in this environment |
+| H3 attention baseline | Comfy `--use-sage-attention` versus direct backend investigation | SageAttention2 (`sage2`) was established as the H3 correctness baseline | Use `sage2` for all H3 parity gates; Sage3 remains experimental |
+| Qwen tokenizer | Direct raw-prompt tokenizer versus MiniMax H3 source contract | Raw text, no chat template, `add_special_tokens=False` | Correct for nonempty text-only prompts; empty prompt differs: direct substitutes a space while Comfy injects token `151643` |
+| Qwen embedding loading, before fix | Direct versus Comfy input embedding capture | Example direct values were approximately `11` where Comfy was approximately `0.017`; downstream refined-text mean absolute delta `1.79263` | Root cause was omitted per-row `model.embed_tokens.weight_scale` |
+| Qwen embedding loading, after fix | Direct rows multiplied by `weight_scale` in FP32 | Input embeddings matched the Comfy capture exactly | Passed for the captured nonempty text-only reference |
+| Qwen layer-50, before embedding fix | Direct versus Comfy layer-50 capture | Mean absolute delta `23.50596`, max `38977.09375` | Failed; invalidated direct text conditioning |
+| Qwen layer-50, after embedding fix but before projection fix | Direct versus Comfy layer-50 capture | Mean absolute delta improved to about `5.63` | Embedding issue fixed; remaining error was material |
+| Qwen all-layer replay, after FP32 fix | `compare_qwen_layer_trace.py` replayed `/tmp/fl2va-qwen-trace` without new Comfy inference | Layers `00` through `49` and `qwen_layer50.pt` all reported mean/max `0` / `0` | **Passed: prompt-only Qwen conditioning is bit-exact for the immutable 17-token FL2VA reference.** |
+| Qwen layer-0 normalization | Direct `F.rms_norm` versus `qwen0_norm1.pt` | Mean absolute delta about `8.96e-06` | Passed to expected BF16-level tolerance |
+| Qwen layer-0 attention, before projection fix | Direct versus `qwen0_attention.pt` | Mean absolute delta about `0.0357`; layer output about `0.1496` | First material Qwen divergence was at projection/attention boundary, not tokenization, embeddings, RMSNorm, or RoPE |
+| Qwen attention backend | Comfy source inspection | Qwen chooses Comfy small-input attention, i.e. SDPA path, not Sage2 | Direct Qwen was changed to BF16 causal-mask SDPA with GQA |
+| Qwen SDPA dispatch audit and fix | Actual Spark Comfy block audit found direct raw SDPA omitted Comfy's four-dimensional mask wrapper, backend priority, and masked-GQA fallback. Direct ported that policy without importing Comfy. | Replaying immutable layer-0 capture after the change: attention mean/max `0.000244372` / `0.0161562`; post-attention `0.00026845` / `0.0213833`; layer output `0.000808998` / `0.0526199`. Previous attention mean was about `0.0357`; previous layer-output mean about `0.1496`. | Active code mismatch identified and substantially repaired. The remaining layer output error is now downstream of projection/elementwise BF16 boundaries, not an unreviewed SDPA dispatcher difference. |
+| Qwen NVFP4 metadata | Qwen layer-0 checkpoint sidecars inspected | Every Q/K/V/O and MLP projection has `full_precision_matrix_mult: true`; `o_proj` and `down_proj` also have `pre_quant_scale` | Direct was incorrectly quantizing activations for all Qwen projections |
+| Qwen NVFP4 fix | Commit `0cad1db` | Direct honors `full_precision_matrix_mult` and loads/applies `pre_quant_scale`; Spark smoke passed finite expected-shape Q and O projections | Structural fix passed; post-fix numeric QKV/MLP/layer-50 comparisons are still required |
+| Qwen layer-0 projections, after NVFP4 fix | Fresh Comfy projection capture at `/tmp/fl2va-qwen0-projections`; direct compared from captured input embedding and post-attention state | Q mean/max `6.6547e-05` / `0.00195485`; K `9.11704e-05` / `0.00241077`; V `6.0298e-05` / `0.00102112`; gate `0.000365695` / `0.0082469`; up `0.000335494` / `0.00792789`; activated `4.76882e-05` / `0.031621`; down `0.000557594` / `0.0749016` | The full-precision metadata fix removed the previous large projection failure. Projection parity is close but not strict; `down_proj` is the largest remaining layer-0 projection boundary. Preserve these numbers and do not recapture this case. |
+| Qwen projection code audit, post-fix | Statement-level comparison against actual Spark Comfy commit `43cb4fffc89bba20ab7bd61467a36d0339338dab`, Torch `2.9.1+cu130`, and `/h3-models/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` | All 350 Qwen projection sidecars set `full_precision_matrix_mult: true`; exactly 100 (`o_proj` and `down_proj` in each layer) have BF16 `pre_quant_scale`; no Qwen projection has bias. Packed weight is contiguous uint8, block scale is native E4M3, tensor scale is scalar FP32. | Direct and Comfy agree on the active full-precision branch, NVFP4 scale association, logical shapes, no-bias behavior, pre-scale order, BF16 compute dtype, and use the same installed Comfy Kitchen NVFP4 dequantizer. Do not change projection math based on speculation. |
+| H3 DiT block trace | `/tmp/fl2va-full-capture` | Block 0 was close, around mean absolute delta `0.00018`; small independent differences accumulated across blocks | Full 50-block strict parity remains failed/unresolved |
+| Sampler initial state | Comfy `initial.pt` inspected against direct preview | Comfy carries joint AV state; direct preview initializes video noise only and leaves audio zero | Known direct mismatch; identical integer seed is not sampler parity |
+| Preview output, before Qwen fixes | Dragon/locomotive prompt previews | Flower-like output | Confirmed incorrect conditioning had visible semantic impact |
+| Preview output, after Qwen metadata fix | 72-frame cat, 12-step direct preview | Cat semantic structure appears, with persistent 32-pixel-like square artifacts | Conditioning improved; denoiser/sampler/VAE boundary remains non-parity |
+| Video encoder | Direct ffmpeg raw RGB to H.264/yuv420p, no audio | Different from Comfy `CreateVideo`/`SaveVideo` | Not a valid latent-quality gate; cannot explain model-space corruption by itself |
+
+## Qwen Projection Code Audit: Remaining Non-Active Differences
+
+These are real implementation differences, but the audit established they are
+not active numerical branches for the captured Qwen checkpoint. They must not
+be presented as an explanation for the current `down_proj` delta without new
+evidence.
+
+| Direct location | Difference from Comfy | Active for current capture? |
+| --- | --- | --- |
+| `nvfp4.py:47-51` | Calls `.contiguous()` while Comfy preserves loaded storage/view identity. In particular Comfy treats E4M3 block scales as a dtype view, not a numeric cast. | No: checkpoint tensors are already contiguous and direct also uses a dtype view. |
+| `nvfp4.py:75-76` | Always flattens and contiguates activation rows; Comfy full-precision `F.linear` retains N-D input/stride. | No known value effect: elements and BF16 dtype are unchanged. |
+| `nvfp4.py:82` | Uses `packed_weight.dequantize().to(flat_x)` rather than an explicitly named Comfy compute-dtype dequantization call. | No: input and output compute dtype are BF16, and both call the same installed Kitchen layout/dequantizer. |
+| `nvfp4.py:46,81` | Does not implement Comfy global full-precision overrides, disabled-format controls, forced casts, weight functions, or LoRA/offload dispatch. | No: every Qwen sidecar itself forces full precision, no projection has bias, and no patches/LoRAs are attached. |
+
+**Audit conclusion:** projection source code contains cleanup/generalization gaps,
+but no identified active branch mismatch explains the current captured values.
+The next code audit must move outward to the exact Qwen block boundary:
+Comfy input dtype, full-precision weight construction, SDPA backend selection,
+and residual/BF16 rounding boundaries. Do not run another full inference merely
+to rediscover this projection result.
+
+## Qwen Layer-0 Arithmetic Audit After `4050afe`
+
+The follow-up statement audit covered every remaining active layer-0 arithmetic
+boundary outside the already-fixed SDPA wrapper: full-precision NVFP4 dispatch,
+weight/scales, `pre_quant_scale`, BF16 GEMM, activation layout, Q/K RMSNorm,
+O projection/residual, post-attention norm, SiLU/gate/up product, down
+projection, and final residual.
+
+**Result: no further source-proven active mismatch was found.** For the actual
+checkpoint, Comfy and direct both use BF16 input, BF16 NVFP4 dequantization,
+BF16 `F.linear`, no projection biases, identical pre-scale ordering, identical
+MLP evaluation order, and ordinary BF16 residual additions. Explicit flattening
+and `.contiguous()` in direct are not a demonstrated numerical difference for
+the captured contiguous batch-one input.
+
+Do not modify `nvfp4.py` or Qwen MLP/residual code to chase the remaining
+layer-0 delta. The next justified work is a bounded component experiment that
+identifies the actual selected SDPA kernel/GQA fallback and verifies capture
+consistency; it is not another full denoising or video run.
+
+### Bounded SDPA/Layer-0 Result
+
+The bounded replay was completed with the fresh projection-capture tensors.
+For the reference's 17-token layer-0 attention, both paths have native cuDNN
+GQA available (`flash=False`, `cudnn=True`, `efficient=False`). Therefore the
+direct SDPA wrapper takes the same native-GQA branch Comfy takes.
+
+| Boundary | Mean absolute error | Maximum absolute error |
+| --- | ---: | ---: |
+| O-projected attention | `0.000244372` | `0.0161562` |
+| First residual | `0.00026845` | `0.0213833` |
+| Post-attention RMSNorm | `3.87714e-05` | `0.000426032` |
+| MLP output | `0.000674194` | `0.0500984` |
+| Layer-0 output | `0.000808998` | `0.0526199` |
+
+This confirms the residual error is already present in the captured projection
+outputs and is amplified by the MLP. The NVFP4 and MLP code audits found no
+remaining active source-level branch difference. Further edits to flattening,
+contiguity, dequantization dtype, `pre_quant_scale`, SiLU ordering, or residual
+adds would be speculative and are prohibited until a new code-path difference
+is demonstrated.
+
+## Loaded-Comfy Projection Module Probe
+
+The source audit was followed by a direct replay against tensors consumed and
+produced by the actual loaded Comfy layer-0 modules. This proves the remaining
+difference is present **before attention** and is not inferred from source.
+
+| Item | Evidence |
+| --- | --- |
+| Capture location | `/tmp/fl2va-qwen0-loaded-modules` on Spark |
+| QKV capture SHA-256 | `8998722e73ba24c40c09a86133ccdb22b708998a3e8b8fb4c4293227345f19db` |
+| MLP capture SHA-256 | `e52ed70f4aa76f1ba26302b9955631d8f490fb6dc269d96a0ce3e1ea981e7ac2` |
+| Loaded Comfy class | `comfy.ops.mixed_precision_ops.<locals>.MixedPrecisionOps.Linear` |
+| Loaded Comfy weight | `comfy_kitchen.tensor.base.QuantizedTensor`, logical dtype BF16, `TensorCoreNVFP4Layout`, full precision enabled |
+| Q/K/V dimensions | `(8192,5120)`, `(1024,5120)`, `(1024,5120)`; no pre-scale |
+| Gate/up/down dimensions | `(25600,5120)`, `(25600,5120)`, `(5120,25600)`; only down has pre-scale |
+
+Direct replay on the **identical loaded-Comfy module inputs**:
+
+| Projection | Mean absolute error | Maximum absolute error |
+| --- | ---: | ---: |
+| Q | `6.11124e-05` | `0.0019514` |
+| K | `8.47002e-05` | `0.00167805` |
+| V | `5.4459e-05` | `0.00098893` |
+| Gate | `0.000308973` | `0.0082469` |
+| Up | `0.00028138` | `0.00792789` |
+| Down, using Comfy's captured activated input | `0.000424157` | `0.0500984` |
+
+**Conclusion:** same checkpoint tensors, same BF16 inputs, and the same named
+Kitchen layout do not currently produce identical projection outputs through
+direct `Nvfp4Linear` and Comfy's loaded `MixedPrecisionOps.Linear`. The next
+investigation must capture/compare the effective `QuantizedTensor` parameters
+and direct packed-weight wrapper at dispatch time. Do not revisit attention,
+sampler, or MLP ordering until that exact module-dispatch difference is found.
+
+### Effective Weight Root Cause And Resolution
+
+The storage-safe dispatch probe found the active difference. Comfy's text
+encoder calls Qwen with `dtype=torch.float32`; its full-precision NVFP4 branch
+casts each projection to an effective FP32 matrix before `F.linear`. Direct had
+constructed the Qwen encoder with BF16, causing its full-precision NVFP4 branch
+to dequantize the same packed data into BF16.
+
+The probe proved that packed qdata, tensor scale, and block scale are exact for
+all Q/K/V/gate/up/down projections. With direct Qwen switched to FP32:
+
+| Projection check | Result |
+| --- | --- |
+| Effective Q/K/V/gate/up/down FP32 weights | Bit-exact to Comfy |
+| Q/K/V/gate/up linear outputs | Bit-exact to Comfy |
+| Layer-0 O-projected attention | Mean/max `0` / `0` |
+| Layer-0 post-attention residual | Mean/max `0` / `0` |
+| Layer-0 post-attention RMSNorm | Mean/max `0` / `0` |
+| Layer-0 MLP output | Mean/max `0` / `0` |
+| Layer-0 final output | Mean/max `0` / `0` |
+
+The special `down` dispatch sub-probe that supplied Comfy's already-activated
+input directly to `F.linear` is intentionally invalid because it bypasses the
+module's required `pre_quant_scale`; the full layer replay above is the valid
+gate and is exact.
+
+## Current Runtime Scope
+
+| Component | Implemented | Known limitation |
+| --- | --- | --- |
+| 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 |
+| Audio VAE/decode/mux | No | Preview intentionally produces video only |
+| Full Ref2VA | No | References, vision conditioning, VAE encode, audio, and muxing are unimplemented |
+
+## Remaining Gates, In Dependency Order
+
+Only these are outstanding. Do not recapture or revisit rows marked complete
+unless the checkpoint, Comfy version, prompt, or backend changes.
+
+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.
+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.
+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.
+4. **End-to-end FL2VA preview.** Generate the reference dragon only after
+   gates 1-8 pass. Compare its raw final latent first, then video.
+5. **Feature/performance work.** Only then add audio, Ref2VA/reference paths,
+    Sage3, CUDA graphs, and multi-GPU execution.
+
+## Existing Tools And Their Intended Gate
+
+| Gate | Tools |
+| --- | --- |
+| Qwen captures | `patch_comfy_qwen_output_capture.py`, `patch_comfy_qwen_layer_trace.py`, `patch_comfy_qwen_layer0_sublayers.py`, `patch_comfy_qwen_layer0_projections.py` |
+| Qwen projection gate | `compare_qwen0_projections.py` |
+| Qwen local diagnosis | `trace_qwen0_attention.py` needs repair before reuse; it currently passes an obsolete encoder argument |
+| H3 block capture | `patch_comfy_h3_capture.py`, `patch_comfy_h3_block_capture.py`, `patch_comfy_h3_block0_sublayers.py` |
+| H3 local diagnosis | `compare_block0_qkv.py`, `compare_attention_backends.py`, `compare_block0_mlp_projections.py`, `localize_block_sublayers.py`, `localize_block_mismatch.py` |
+| Sampler | `patch_comfy_sampler_capture.py`, `compare_fl2va_steps.py` |
+| VAE | `smoke_h3_vae_decoder.py`; a same-latent direct-versus-Comfy pixel comparator still needs to be added |
+
+The `patch_comfy_*` tools modify `/opt/ComfyUI` and are only safe against the
+restorable Comfy reference image. Capture provenance must be appended here at
+the time a new capture is made.

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)

src/h3_blackwell_runtime/final.py

@@ -42,11 +42,11 @@ class H3FinalLayer(nn.Module):
         return cls(
             checkpoint.tensor("adaln_t_table", dtype=torch.float32),
             checkpoint.tensor("final_layer.norm.weight", dtype=output_dtype),
-            checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.bfloat16),
-            checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float32),
+            checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float16),
-            checkpoint.tensor("final_layer.video_out.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.video_out.weight", dtype=torch.bfloat16),
             checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32),
-            checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.float32),
+            checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.bfloat16),
             checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32),
         )
 
@@ -64,10 +64,9 @@ class H3FinalLayer(nn.Module):
 
         video_start, video_stop, video_row = video_segment
         audio_start, audio_stop, audio_row = audio_segment
-        normalized = rms_norm(hidden, self.norm_weight, self.eps)
+        video_hidden = (rms_norm(hidden[video_start:video_stop], self.norm_weight, self.eps) * (1.0 + scale[video_row]) + shift[video_row]).to(torch.float32)
-        video_hidden = normalized[video_start:video_stop] * (1 + scale[video_row].to(hidden.dtype)) + shift[video_row].to(hidden.dtype)
+        audio_hidden = (rms_norm(hidden[audio_start:audio_stop], self.norm_weight, self.eps) * (1.0 + scale[audio_row]) + shift[audio_row]).to(torch.float32)
-        audio_hidden = normalized[audio_start:audio_stop] * (1 + scale[audio_row].to(hidden.dtype)) + shift[audio_row].to(hidden.dtype)
         return (
-            functional.linear(video_hidden.float(), self.video_weight, self.video_bias),
+            functional.linear(video_hidden, self.video_weight, self.video_bias),
-            functional.linear(audio_hidden.float(), self.audio_weight, self.audio_bias),
+            functional.linear(audio_hidden, self.audio_weight, self.audio_bias),
         )

src/h3_blackwell_runtime/nvfp4.py

@@ -67,8 +67,8 @@ class Nvfp4Linear(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if x.shape[-1] != self.in_features:
             raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.")
-        if x.dtype not in (torch.float16, torch.bfloat16):
+        if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
-            raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.")
+            raise ValueError("NVFP4 linear accepts FP16, BF16, or FP32 activations.")
 
         from comfy_kitchen.tensor import QuantizedTensor
 
@@ -82,6 +82,8 @@ class Nvfp4Linear(nn.Module):
             weight = packed_weight.dequantize().to(flat_x)
             output = functional.linear(flat_x, weight, bias)
             return output.reshape(*original_shape, self.out_features)
+        if x.dtype == torch.float32:
+            raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.")
         packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
         output = functional.linear(packed_x, packed_weight, bias)
         return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features)

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_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
-        self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=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.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.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])

src/h3_blackwell_runtime/qwen3vl_text.py

@@ -66,6 +66,29 @@ def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.T
     return query_output.to(query.dtype), key_output.to(key.dtype)
 
 
+def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Match Comfy's small-input SDPA wrapper for masked Qwen GQA."""
+    kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True}
+    if query.numel() >= 1024 * 128:
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]
+        params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True)
+        supports_native_gqa = (
+            torch.backends.cuda.can_use_flash_attention(params)
+            or torch.backends.cuda.can_use_cudnn_attention(params)
+            or torch.backends.cuda.can_use_efficient_attention(params)
+        )
+        if not supports_native_gqa:
+            repeats = query.shape[-3] // key.shape[-3]
+            key = key.repeat_interleave(repeats, dim=-3)
+            value = value.repeat_interleave(repeats, dim=-3)
+            kwargs["enable_gqa"] = False
+        with sdpa_kernel(priority, set_priority=True):
+            return F.scaled_dot_product_attention(query, key, value, **kwargs)
+    return F.scaled_dot_product_attention(query, key, value, **kwargs)
+
+
 class _Qwen3VLBlock(nn.Module):
     def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype):
         super().__init__()
@@ -99,14 +122,7 @@ class _Qwen3VLBlock(nn.Module):
             dtype=query.dtype,
             device=query.device,
         ).triu_(1)
-        attention = F.scaled_dot_product_attention(
+        attention = _qwen_attention(query, key, value, causal_mask)
-            query,
-            key,
-            value,
-            attn_mask=causal_mask,
-            is_causal=False,
-            enable_gqa=True,
-        )
         hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1))
         residual = hidden_states
         x = self.post_attention_layernorm(hidden_states)
@@ -117,7 +133,7 @@ class Qwen3VL32BTextEncoder(nn.Module):
     """Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states."""
     config = Qwen3VL32BTextConfig()
 
-    def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
+    def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
         super().__init__()
         self.checkpoint_path = str(checkpoint_path)
         self.device_name = str(device)
@@ -179,7 +195,7 @@ class Qwen3VL32BTextEncoder(nn.Module):
 
 class Qwen3VLPromptConditioner:
     """Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning."""
-    def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16):
+    def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32):
         from .conditioning import H3PromptTokenizer
 
         tokenizer_dir = tokenizer_dir or Path(__file__).with_name("qwen25_tokenizer")

src/h3_blackwell_runtime/sampler.py

@@ -1,5 +1,7 @@
 """Minimal direct prompt-only H3 video sampler for visual smoke previews."""
 
+import time
+
 import torch
 
 from .packing import H3PromptPacker, unpatchify_video
@@ -52,8 +54,12 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor
     audio_carried = audio * 4.0
     video_history = audio_history = None
     video_history_sigma = audio_history_sigma = None
-    for index, sigma in enumerate(sigmas[:-1]):
+    total_steps = len(sigmas) - 1
-        sigma_down = sigmas[index + 1]
+    started = time.perf_counter()
+    for index, sigma in enumerate(sigmas[:-1], start=1):
+        step_started = time.perf_counter()
+        previous_index = index - 1
+        sigma_down = sigmas[index]
         sigma_audio = _audio_sigma(sigma)
         native_audio = audio_carried * (sigma_audio / sigma)
         hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
@@ -64,11 +70,18 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor
         velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * sigma_audio) * velocity_audio_native
         video_denoised = video - sigma * velocity_video
         audio_denoised = audio_carried - sigma * velocity_audio
-        previous_sigma = sigmas[index - 1] if index else None
+        previous_sigma = sigmas[previous_index - 1] if previous_index else None
         video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
         audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
         video_history, audio_history = video_denoised, audio_denoised
         video_history_sigma = audio_history_sigma = sigma_down
+        elapsed = time.perf_counter() - started
+        eta = elapsed / index * (total_steps - index)
+        print(
+            f"sampling step {index}/{total_steps}: "
+            f"{time.perf_counter() - step_started:.1f}s, elapsed {elapsed:.1f}s, eta {eta:.1f}s",
+            flush=True,
+        )
     return video
 
 

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/compare_qwen0_loaded_modules.py

@@ -0,0 +1,30 @@
+"""Replay direct layer-0 projections from actual loaded-Comfy module inputs."""
+
+import argparse
+
+import torch
+
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture-dir", required=True)
+parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
+args = parser.parse_args()
+
+encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
+layer = encoder.layers[0]
+qkv = torch.load(f"{args.capture_dir}/qwen0_loaded_qkv.pt", map_location="cuda", weights_only=False)
+mlp = torch.load(f"{args.capture_dir}/qwen0_loaded_mlp.pt", map_location="cuda", weights_only=False)
+
+with torch.inference_mode():
+    qkv_actual = {"q": layer.q_proj(qkv["input"].to(encoder.dtype)), "k": layer.k_proj(qkv["input"].to(encoder.dtype)), "v": layer.v_proj(qkv["input"].to(encoder.dtype))}
+    gate = layer.gate_proj(mlp["input"].to(encoder.dtype))
+    up = layer.up_proj(mlp["input"].to(encoder.dtype))
+    mlp_actual = {"gate": gate, "up": up, "down": layer.down_proj(mlp["output"]["activated"].to(encoder.dtype))}
+
+for group, actual, captured in (("qkv", qkv_actual, qkv), ("mlp", mlp_actual, mlp)):
+    print(f"{group}.loaded_metadata={captured['metadata']}")
+    for name, value in actual.items():
+        delta = (value.float() - captured["output"][name].float()).abs()
+        print(f"{group}.{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_qwen0_projections.py

@@ -0,0 +1,35 @@
+"""Compare direct Qwen layer-0 projections with a Comfy capture."""
+
+import argparse
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture-dir", required=True)
+parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
+args = parser.parse_args()
+
+encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
+layer = encoder.layers[0]
+embeds = torch.load(f"{args.capture_dir}/qwen_input_embeds.pt", map_location="cuda", weights_only=False)
+qkv_expected = torch.load(f"{args.capture_dir}/qwen0_qkv.pt", map_location="cuda", weights_only=False)
+mlp_expected = torch.load(f"{args.capture_dir}/qwen0_mlp_projections.pt", map_location="cuda", weights_only=False)
+
+with torch.inference_mode():
+    norm1 = layer.input_layernorm(embeds.to(encoder.dtype))
+    qkv_actual = {"q": layer.q_proj(norm1), "k": layer.k_proj(norm1), "v": layer.v_proj(norm1)}
+    post_attention = torch.load(f"{args.capture_dir}/qwen0_post_attention.pt", map_location="cuda", weights_only=False)
+    norm2 = layer.post_attention_layernorm(post_attention.to(encoder.dtype))
+    gate = layer.gate_proj(norm2)
+    up = layer.up_proj(norm2)
+    activated = functional.silu(gate) * up
+    mlp_actual = {"gate": gate, "up": up, "activated": activated, "down": layer.down_proj(activated)}
+
+for group, actual, expected in (("qkv", qkv_actual, qkv_expected), ("mlp", mlp_actual, mlp_expected)):
+    for name, value in actual.items():
+        delta = (value.float() - expected[name].float()).abs()
+        print(f"{group}.{name} shape={tuple(value.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_qwen0_quantized_dispatch.py

@@ -0,0 +1,46 @@
+"""Compare direct NVFP4 wrapper and effective Comfy BF16 weights."""
+
+import argparse
+
+import torch
+import torch.nn.functional as functional
+
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--capture-dir", required=True)
+parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
+parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
+args = parser.parse_args()
+
+encoder = Qwen3VL32BTextEncoder(args.checkpoint, dtype=getattr(torch, args.dtype)).eval()
+layer = encoder.layers[0]
+qkv = torch.load(f"{args.capture_dir}/qwen0_quantized_qkv.pt", map_location="cuda", weights_only=False)
+mlp = torch.load(f"{args.capture_dir}/qwen0_quantized_mlp.pt", map_location="cuda", weights_only=False)
+
+for group, capture, modules, inputs in (
+    ("qkv", qkv, {"q": layer.q_proj, "k": layer.k_proj, "v": layer.v_proj}, {"q": qkv["input"], "k": qkv["input"], "v": qkv["input"]}),
+    ("mlp", mlp, {"gate": layer.gate_proj, "up": layer.up_proj, "down": layer.down_proj}, {"gate": mlp["input"], "up": mlp["input"], "down": mlp["activated"]}),
+):
+    for name, module in modules.items():
+        expected = capture["modules"][name]
+        packed = module._packed_weight()
+        params = packed._params
+        qdata, scale, block_scale = packed.layout_cls.get_plain_tensors(packed)
+        reference_qdata = expected["qdata"].as_subclass(torch.Tensor).cpu()
+        print(f"{group}.{name}.qdata exact={torch.equal(qdata.cpu(), reference_qdata)} shape={tuple(qdata.shape)}")
+        for field, actual in (("scale", scale), ("block_scale", block_scale)):
+            reference = expected[field].to(actual.device)
+            print(f"{group}.{name}.{field} exact={torch.equal(actual, reference)} shape={tuple(actual.shape)} dtype={actual.dtype}")
+        value = inputs[name].to(encoder.dtype)
+        with torch.inference_mode():
+            dequantized = packed.dequantize()
+            direct = module(value)
+            effective = expected["effective_weight"].to("cuda")
+            comfy_effective = functional.linear(value.to(effective.dtype), effective, None)
+        print(f"{group}.{name}.effective dtype={expected['effective_dtype']} shape={expected['effective_shape']} direct_dtype={encoder.dtype}")
+        for variant, actual in (("weight", dequantized), ("linear", direct), ("effective_linear", comfy_effective)):
+            reference = expected["effective_weight"].to(actual.device) if variant == "weight" else expected["output"].to(actual.device)
+            delta = (actual.float() - reference.float()).abs()
+            print(f"{group}.{name}.{variant}_vs_comfy max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

tools/compare_qwen_layer_trace.py

@@ -0,0 +1,27 @@
+"""Replay direct Qwen against an immutable per-layer Comfy trace."""
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--trace-dir", type=Path, required=True)
+parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
+args = parser.parse_args()
+
+encoder = Qwen3VL32BTextEncoder(args.checkpoint).eval()
+hidden = torch.load(args.trace_dir / "qwen_input_embeds.pt", map_location="cuda", weights_only=False).to(encoder.dtype)
+
+with torch.inference_mode():
+    for index, layer in enumerate(encoder.layers):
+        hidden = layer(hidden)
+        expected = torch.load(args.trace_dir / "qwen_layers" / f"{index:02d}.pt", map_location="cuda", weights_only=False)
+        delta = (hidden.float() - expected.float()).abs()
+        print(f"layer={index:02d} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
+    expected = torch.load(args.trace_dir / "qwen_layer50.pt", map_location="cuda", weights_only=False)
+    delta = (hidden.float() - expected.float()).abs()
+    print(f"layer50 max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")

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_capture.py

@@ -24,14 +24,14 @@ 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(), \"t_emb\": t_emb.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(
     model,
     "        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
     "        if capture_dir and H3_CAPTURE_ACTIVE:\n"
-    "            torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg}, os.path.join(capture_dir, \"output.pt\"))\n"
+    "            torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg, \"adaln_t_table\": self.adaln_t_table.detach().cpu()}, os.path.join(capture_dir, \"output.pt\"))\n"
     "            H3_CAPTURE_ACTIVE = False\n"
     "        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
 )

tools/patch_comfy_h3_final_capture.py

@@ -0,0 +1,42 @@
+"""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 = (
+    "        capture_input = t_emb.clone() if H3_CAPTURE_ACTIVE else None\n"
+    "        adaln_output = self.adaln_proj(t_emb)\n"
+    "        capture_replay = self.adaln_proj.linear(capture_input.clone()) if H3_CAPTURE_ACTIVE else None\n"
+    "        capture_forward = self.adaln_proj.linear._forward(capture_input.clone(), self.adaln_proj.linear.weight, self.adaln_proj.linear.bias) if H3_CAPTURE_ACTIVE else None\n"
+    "        shift, scale = adaln_output\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.cuda.synchronize(x.device)\n"
+    "            linear = self.adaln_proj.linear\n"
+    "            torch.save({\"hidden\": x.detach().cpu(), \"adaln_input\": capture_input.detach().cpu(), \"adaln_output\": torch.cat((shift, scale), dim=-1).detach().cpu(), \"adaln_replay\": capture_replay.detach().cpu(), \"adaln_forward\": capture_forward.detach().cpu(), \"t_emb\": t_emb.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"adaln_weight\": linear.weight.detach().cpu(), \"adaln_bias\": linear.bias.detach().cpu(), \"adaln_metadata\": {\"class\": f\"{type(linear).__module__}.{type(linear).__qualname__}\", \"weight_class\": f\"{type(linear.weight).__module__}.{type(linear.weight).__qualname__}\", \"quant_format\": getattr(linear, \"quant_format\", None), \"layout_type\": getattr(linear, \"layout_type\", None), \"full_precision_mm\": getattr(linear, \"_full_precision_mm\", None), \"input_scale\": getattr(linear, \"input_scale\", None)}, \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video_weight\": self.video_out.weight.detach().cpu(), \"video_bias\": self.video_out.bias.detach().cpu(), \"audio_weight\": self.audio_out.weight.detach().cpu(), \"audio_bias\": self.audio_out.bias.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, 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/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.")

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

tools/patch_comfy_h3_temb_trace.py

@@ -0,0 +1,29 @@
+"""Capture H3 curve embeddings after every transformer block."""
+
+from pathlib import Path
+
+
+model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
+source = model.read_text(encoding="utf-8")
+old = (
+    "            else:\n"
+    "                h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
+    "        if prefetch_queue is not None:\n"
+)
+new = (
+    "            else:\n"
+    "                if capture_dir and H3_CAPTURE_ACTIVE and i == 0:\n"
+    "                    torch.save({\"t_emb\": t_emb.clone().detach().cpu(), \"timesteps\": t_vals.clone().detach().cpu()}, os.path.join(capture_dir, \"t_emb_before_block0.pt\"))\n"
+    "                h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
+    "            if capture_dir and H3_CAPTURE_ACTIVE:\n"
+    "                temb_dir = os.path.join(capture_dir, \"t_emb\")\n"
+    "                os.makedirs(temb_dir, exist_ok=True)\n"
+    "                torch.save({\"value\": t_emb.detach().cpu(), \"class\": f\"{type(block.adaln_proj.linear).__module__}.{type(block.adaln_proj.linear).__qualname__}\", \"data_ptr\": t_emb.data_ptr(), \"stride\": t_emb.stride()}, os.path.join(temb_dir, f\"{i:02d}.pt\"))\n"
+    "        if prefetch_queue is not None:\n"
+)
+if source.count(old) == 1:
+    source = source.replace(old, new)
+elif new not in source:
+    raise RuntimeError("Unable to locate the H3 transformer block loop.")
+model.write_text(source, encoding="utf-8")
+print("Applied H3 t_emb trace patch.")

tools/patch_comfy_qwen_layer0_projections.py

@@ -0,0 +1,32 @@
+"""Capture Comfy Qwen layer-0 projection boundaries for direct parity checks."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
+source = path.read_text(encoding="utf-8")
+old = "        xq = self.q_proj(hidden_states)\n        xk = self.k_proj(hidden_states)\n        xv = self.v_proj(hidden_states)\n"
+new = old + (
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            torch.save({\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, os.path.join(capture_dir, \"qwen0_qkv.pt\"))\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen QKV projections.")
+source = source.replace(old, new)
+
+old = "        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
+new = (
+    "        gate = self.gate_proj(x)\n"
+    "        up = self.up_proj(x)\n"
+    "        activated = self.activation(gate) * up\n"
+    "        output = self.down_proj(activated)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            torch.save({\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, os.path.join(capture_dir, \"qwen0_mlp_projections.pt\"))\n"
+    "        return output\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen MLP projections.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied Qwen layer-0 projection capture patch.")

tools/patch_comfy_qwen_layer0_sublayers.py

@@ -57,7 +57,12 @@ if source.count(old) != 1:
 source = source.replace(old, new)
 
 old = "        for i, layer in enumerate(self.layers):\n"
-new = "        for i, layer in enumerate(self.layers):\n            layer._h3_trace_index = i\n"
+new = (
+    "        for i, layer in enumerate(self.layers):\n"
+    "            layer._h3_trace_index = i\n"
+    "            layer.self_attn._h3_trace_index = i\n"
+    "            layer.mlp._h3_trace_index = i\n"
+)
 if source.count(old) != 1:
     raise RuntimeError("Unable to locate Qwen decoder layer loop.")
 path.write_text(source.replace(old, new), encoding="utf-8")

tools/patch_comfy_qwen_loaded_projection_probe.py

@@ -0,0 +1,36 @@
+"""Capture actual loaded Comfy Qwen layer-0 projection module behavior."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
+source = path.read_text(encoding="utf-8")
+old = "        xq = self.q_proj(hidden_states)\n        xk = self.k_proj(hidden_states)\n        xv = self.v_proj(hidden_states)\n"
+new = old + (
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
+    "            metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
+    "            torch.save({\"input\": hidden_states.detach().cpu(), \"output\": {\"q\": xq.detach().cpu(), \"k\": xk.detach().cpu(), \"v\": xv.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_qkv.pt\"))\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen QKV projections.")
+source = source.replace(old, new)
+
+old = "        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
+new = (
+    "        gate = self.gate_proj(x)\n"
+    "        up = self.up_proj(x)\n"
+    "        activated = self.activation(gate) * up\n"
+    "        output = self.down_proj(activated)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
+    "            metadata = {name: {\"class\": f\"{type(module).__module__}.{type(module).__qualname__}\", \"weight_type\": f\"{type(module.weight).__module__}.{type(module.weight).__qualname__}\", \"weight_dtype\": str(module.weight.dtype), \"weight_shape\": tuple(module.weight.shape), \"layout_type\": getattr(module, \"layout_type\", None), \"full_precision\": getattr(module, \"_full_precision_mm\", None), \"pre_quant_scale\": getattr(module, \"pre_quant_scale\", None) is not None} for name, module in modules.items()}\n"
+    "            torch.save({\"input\": x.detach().cpu(), \"output\": {\"gate\": gate.detach().cpu(), \"up\": up.detach().cpu(), \"activated\": activated.detach().cpu(), \"down\": output.detach().cpu()}, \"metadata\": metadata}, os.path.join(capture_dir, \"qwen0_loaded_mlp.pt\"))\n"
+    "        return output\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen MLP projections.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied loaded Qwen layer-0 projection probe.")

tools/patch_comfy_qwen_quantized_tensor_probe.py

@@ -0,0 +1,58 @@
+"""Capture loaded Comfy Qwen QuantizedTensor fields and dispatch variants."""
+
+from pathlib import Path
+
+
+path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
+source = path.read_text(encoding="utf-8")
+helper = (
+    "def _h3_owned_cpu(tensor):\n"
+    "    return tensor.detach().to(device=\"cpu\", copy=True).contiguous().clone()\n\n"
+)
+if "def _h3_owned_cpu(tensor):\n" not in source:
+    source = helper + source
+old = "        xq = self.q_proj(hidden_states)\n        xk = self.k_proj(hidden_states)\n        xv = self.v_proj(hidden_states)\n"
+new = old + (
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n"
+    "            def probe(module):\n"
+    "                stored = module.weight\n"
+    "                qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
+    "                weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, hidden_states, offloadable=True)\n"
+    "                try:\n"
+    "                    output = torch.nn.functional.linear(hidden_states, weight, bias)\n"
+    "                    return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
+    "                finally:\n"
+    "                    comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
+    "            torch.save({\"input\": hidden_states.detach().cpu(), \"modules\": {name: probe(module) for name, module in modules.items()}}, os.path.join(capture_dir, \"qwen0_quantized_qkv.pt\"))\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen QKV projections.")
+source = source.replace(old, new)
+
+old = "        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n"
+new = (
+    "        gate = self.gate_proj(x)\n"
+    "        up = self.up_proj(x)\n"
+    "        activated = self.activation(gate) * up\n"
+    "        output = self.down_proj(activated)\n"
+    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n"
+    "        if capture_dir:\n"
+    "            modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n"
+    "            def probe(module, value):\n"
+    "                stored = module.weight\n"
+    "                qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n"
+    "                weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, value, offloadable=True)\n"
+    "                try:\n"
+    "                    output = torch.nn.functional.linear(value, weight, bias)\n"
+    "                    return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n"
+    "                finally:\n"
+    "                    comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n"
+    "            torch.save({\"input\": x.detach().cpu(), \"activated\": activated.detach().cpu(), \"modules\": {\"gate\": probe(self.gate_proj, x), \"up\": probe(self.up_proj, x), \"down\": probe(self.down_proj, activated)}}, os.path.join(capture_dir, \"qwen0_quantized_mlp.pt\"))\n"
+    "        return output\n"
+)
+if source.count(old) != 1:
+    raise RuntimeError("Unable to locate Qwen MLP projections.")
+path.write_text(source.replace(old, new), encoding="utf-8")
+print("Applied Qwen QuantizedTensor dispatch probe.")

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"]),