28 lines
1.4 KiB
Python
28 lines
1.4 KiB
Python
"""Compare K RMSNorm weight variants against captured fused K preparation."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.attention import apply_split_half_rope, rms_norm
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
from h3_blackwell_runtime.rope import h3_rope_rotation
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture-dir", required=True)
|
|
parser.add_argument("--model", required=True)
|
|
args = parser.parse_args()
|
|
|
|
inputs = torch.load(f"{args.capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
|
raw = torch.load(f"{args.capture_dir}/block0_qkv_raw.pt", map_location="cuda", weights_only=False)
|
|
expected = torch.load(f"{args.capture_dir}/block0_qkv_prepared.pt", map_location="cuda", weights_only=False)["k"]
|
|
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
|
|
attention = model.backbone.blocks[0].attention
|
|
rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, raw["k"].dtype)
|
|
|
|
for name, weight in (("q_norm", attention.q_norm_weight), ("k_norm", attention.k_norm_weight)):
|
|
actual = apply_split_half_rope(rms_norm(raw["k"].view(1, -1, 56, 128), weight, attention.eps), rotation).transpose(1, 2)
|
|
delta = (actual.float() - expected.float()).abs()
|
|
print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|