"""Offline Qwen layer-0 attention trace against an existing Comfy capture."""

import torch
import torch.nn.functional as F

from h3_blackwell_runtime.attention import run_attention
from h3_blackwell_runtime.conditioning import H3PromptTokenizer
from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder, _rope


prompt = "A brass-and-paper dragon flies above a rain-washed old city at blue hour."
encoder = Qwen3VL32BTextEncoder("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", attention_backend="sage2")
ids = H3PromptTokenizer("/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer")(prompt)
x = (F.embedding(ids, encoder.embed_tokens).float() * F.embedding(ids, encoder.embed_scale)).to(encoder.dtype)
layer = encoder.layers[0]

with torch.inference_mode():
    norm = layer.input_layernorm(x)
    query = layer.q_proj(norm).view(1, 17, 64, 128).transpose(1, 2)
    key = layer.k_proj(norm).view(1, 17, 8, 128).transpose(1, 2)
    value = layer.v_proj(norm).view(1, 17, 8, 128).transpose(1, 2)
    query = layer.q_norm(query)
    key = layer.k_norm(key)
    query, key = _rope(query, key, layer.config.rope_theta)
    key = key.repeat_interleave(8, dim=1)
    value = value.repeat_interleave(8, dim=1)
    attention = layer.o_proj(run_attention(query, key, value, backend="sage2", is_causal=True).transpose(1, 2).reshape(1, 17, -1))

for name, actual in (("norm1", norm), ("attention", attention)):
    expected = torch.load(f"/capture/qwen0_{name}.pt", map_location="cuda", weights_only=False)
    delta = (actual.float() - expected.float()).abs()
    print(f"{name} mean_abs={delta.mean().item():.6g} max_abs={delta.max().item():.6g}")