File size: 1,797 Bytes
ffc54ec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | """Split prefill vs upstream prefill on identical step-0 inputs: output and prefix-KV agreement."""
import sys, types, torch
from PIL import Image
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache
from common import ROOT, BASE
from prompts import EVALUATION
precision = sys.argv[1]
if precision == 'fp8':
from fp8_runtime import load_pipeline
else:
from nvfp4_runtime import load_pipeline
from acceleration import accelerate_pipeline
pipe = accelerate_pipeline(load_pipeline(str(BASE), str(ROOT / precision / 'release' / 'transformer')))
tf = pipe.transformer; split = tf.forward; original = tf._image21_original_forward
e = lambda a, b: ((a.float() - b.float()).square().mean() / b.float().square().mean()).sqrt().item()
def check(*args, **kw):
if kw.get('kv_cache_mode') == 'extract':
scratch = QwenImage21KVCache(len(tf.transformer_blocks))
ref = original(*args, **{**kw, 'kv_cache': scratch})[0]
out = split(*args, **kw)[0]
n = out.shape[1]
kv = max(e(a.k, b.k) for a, b in zip(kw['kv_cache'].layer_caches, scratch.layer_caches))
print({'target_output_nrmse': e(out, ref[:, -n:]), 'max_prefix_k_nrmse': kv, 'tokens': ref.shape[1], 'target': n}, flush=True)
return (out,)
return split(*args, **kw)
tf.forward = check
with torch.inference_mode():
pipe(prompt=EVALUATION[8], width=1024, height=1024, generator=torch.Generator('cuda').manual_seed(1))
pipe(prompt=EVALUATION[1], width=2048, height=2048, generator=torch.Generator('cuda').manual_seed(1))
ref = Image.open(ROOT / 'evaluation' / 'bf16' / '03.png').resize((1024, 1024))
pipe(prompt='Turn this room into a warm evening scene', image=ref, width=1024, height=1024, generator=torch.Generator('cuda').manual_seed(1))
|